"""Convert dataset: Public dataset (Jester[RawFrames], AVA) --> Datumaro dataset (CVAT).

This script contains a lot of hardcoding to create an.xml file that Datumaro can consume.

Current Datumaro format for video (CVAT)

root
|- video_0
    |- images
        |- frames_001.png
        |- frames_002.png
    |- annotations.xml
|- video_1
    |- images
    |- annotations.xml
|- video_2

"""

# Copyright (C) 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions
# and limitations under the License.

# pylint: disable=too-many-locals, c-extension-no-member, invalid-name, too-many-statements
import argparse
import csv
import os
import os.path as osp
import pathlib
import shutil
from typing import List

import cv2
from lxml import etree
from tqdm import tqdm


def generate_default_cvat_xml_fields(i, video_path, frame_list):
    """Generate default CVAT xml fields required to make multi-video CVAT format by using Jester, and AVA dataset."""
    n_frames = len(frame_list)
    annotations = etree.Element("annotations")

    version = etree.Element("version")
    version.text = "1.1"
    annotations.append(version)

    meta = etree.Element("meta")
    annotations.append(meta)

    task = etree.Element("task")
    meta.append(task)

    _id = etree.Element("id")
    _id.text = str(i)
    task.append(_id)

    name = etree.Element("name")
    name.text = f"v{i}"
    task.append(name)

    size = etree.Element("size")
    size.text = str(n_frames)
    task.append(size)

    mode = etree.Element("mode")
    mode.text = "interpolation"
    task.append(mode)

    overlap = etree.Element("overlap")
    overlap.text = "2"
    task.append(overlap)

    bugtracker = etree.Element("bugtracker")
    bugtracker.text = ""
    task.append(bugtracker)

    created = etree.Element("created")
    created.text = ""
    task.append(created)

    updated = etree.Element("updated")
    updated.text = ""
    task.append(updated)

    start_frame = etree.Element("start_frame")
    start_frame.text = "0"
    task.append(start_frame)

    stop_frame = etree.Element("stop_frame")
    stop_frame.text = str(n_frames - 1)
    task.append(stop_frame)

    frame_filter = etree.Element("frame_filter")
    frame_filter.text = "1"
    task.append(frame_filter)

    z_order = etree.Element("z_order")
    z_order.text = str(True)
    task.append(z_order)

    labels = etree.Element("labels")
    task.append(labels)

    segments = etree.Element("segments")
    segments.text = ""
    task.append(segments)

    original_size = etree.Element("original_size")
    task.append(original_size)

    sample_frame = cv2.imread(osp.join(video_path, frame_list[0]))
    original_size_width = etree.Element("width")
    original_size_width.text = str(sample_frame.shape[1])
    original_size.append(original_size_width)

    original_size_height = etree.Element("height")
    original_size_height.text = str(sample_frame.shape[0])
    original_size.append(original_size_height)

    return annotations, sample_frame.shape, labels


# classification
def convert_action_cls_dataset_to_datumaro(src_path: str, dst_path: str, ann_file: str, label_map=None):
    """Convert a public dataset to multi-video CVAT (Datumaro) format.

    Supported datasets are: Jester, HMDB51, UCF101

    Args:
        src_path (str): The path to the directory containing the video files as rawframe folder.
        dst_path (str): The path to the directory where the multi-video CVAT (Datumaro) format dataset will be saved.
        ann_file (str): The path to the file containing the annotations for the videos.
        label_map (optional): The path to the file containing the mapping between class IDs and class names.

    Returns:
        None

    Examples:
        src_path = "./data/hmdb51/rawframes"
        dst_path = "./data/hmdb51/CVAT/train"
        ann_file = "./data/hmdb51/hmdb51_train_split_1_rawframes.txt"
        label_map = "./data/hmdb51/label_map.txt"
        convert_public_dataset_to_datumaro(src_path, dst_path, ann_file, label_map=label_map)
    """

    # Load the annotations file. Annotation is supposed as whitespace separated format: Video_name num_frames class_idx
    # Ex) kiss/American_History_X_kiss_h_cm_np2_le_goo_40 69 22
    with open(ann_file, "r", encoding="utf-8") as anns:
        pathlib.Path(osp.join(dst_path)).mkdir(parents=True, exist_ok=True)

        lines = anns.readlines()
        for i, line in tqdm(enumerate(lines), total=len(lines)):
            if line[0] == "#":
                continue

            # Parse the video directory and class ID from the annotations file
            video_dir, _, class_idx = line[:-1].split(" ")

            if label_map is not None:
                # Load label mapping file : brush_hair\ncarwheel\ncatch, ...
                with open(label_map, "r", encoding="utf-8") as f:
                    label_names = f.read().splitlines()
                label_mapping_dict = {str(idx): label_name for idx, label_name in enumerate(label_names)}
                class_name = label_mapping_dict[class_idx]
            else:
                class_name = class_idx

            # Prepare the output directories and file names
            video_path = osp.join(src_path, video_dir)
            video_name = f"Video_{i}"
            images_dir = osp.join(dst_path, f"{video_name}/images")

            # List the frames in the video and sort them
            frame_list = os.listdir(video_path)
            frame_list.sort()

            # Generate default CVAT XML fields for the video annotation
            annotations, img_shape, labels = generate_default_cvat_xml_fields(i, video_path, frame_list)

            # Add the video label to the annotations
            label = etree.Element("label")
            labels.append(label)

            name = etree.Element("name")
            name.text = class_name
            label.append(name)

            attributes = etree.Element("attributes")
            attributes.text = ""
            label.append(attributes)

            # Copy the video frames to the output directory and create the image tags in the annotations
            for j, frame in enumerate(frame_list):
                if not osp.exists(images_dir):
                    os.makedirs(images_dir, exist_ok=True)

                image_name = f"{j+1:05d}.jpg"
                shutil.copy(osp.join(video_path, frame), osp.join(images_dir, image_name))
                image = etree.Element(
                    "image", id=str(j), name=image_name, width=str(img_shape[1]), height=str(img_shape[0])
                )
                tag = etree.Element("tag", label=class_name, source="manual")
                tag.text = ""
                image.append(tag)
                annotations.append(image)

            et = etree.ElementTree(annotations)
            et.write(
                osp.join(dst_path, f"{video_name}/annotations.xml"),
                pretty_print=True,
                xml_declaration=True,
                encoding="utf-8",
            )


def convert_ava_dataset_to_datumaro(src_path: str, dst_path: str, ann_file: str):
    """Convert AVA dataset to multi-video CVAT (Datumaro) format.

    Args:
        src_path (str): The path to the directory containing the video files as rawframe folder.
        dst_path (str): The path to the directory where the multi-video CVAT (Datumaro) format dataset will be saved.
        ann_file (str): The path to the file containing the annotations for the videos.

    Returns:
        None

    Examples:
        src_path = "./data/ava_dataset/frames"
        dst_path = "./data/ava_dataset/CVAT/train"
        ann_file = "./data/ava_dataset/annotations/train.csv"
        convert_ava_dataset_to_datumaro(src_path, dst_path, ann_file)
    """

    video_dict = {}
    video_idx = 0
    annot_info = read_ava_csv(ann_file)
    video_dirs = os.listdir(src_path)
    video_dirs.sort()
    for video_dir in video_dirs:
        if video_dir not in annot_info:
            continue
        video_path = osp.join(src_path, video_dir)
        frame_list = os.listdir(video_path)
        frame_list.sort()
        if video_dir not in video_dict:
            annotations, _, labels = generate_default_cvat_xml_fields(video_idx, video_path, frame_list)
            label_list: List[str] = []
            video_dict[video_dir] = {"annotations": annotations, "labels": labels, "label_list": label_list}
        for frame in frame_list:
            frame_idx = int(frame.split(".")[0].split("_")[-1])
            if frame_idx in annot_info[video_dir]:
                annots = annot_info[video_dir][frame_idx]
            else:
                annots = [[0, 0, 0, 0, "EmptyFrame"]]

            for annot in annots:
                bboxes = annot[:4]
                class_idx = annot[4]
                track = etree.Element("track", id=str(video_idx), label=str(class_idx))
                if class_idx not in video_dict[video_dir]["label_list"]:
                    video_dict[video_dir]["label_list"].append(class_idx)

                    label = etree.Element("label")
                    video_dict[video_dir]["labels"].append(label)

                    name = etree.Element("name")
                    name.text = str(class_idx)
                    label.append(name)

                    attributes = etree.Element("attributes")
                    attributes.text = ""
                    label.append(attributes)
                box = etree.Element(
                    "box",
                    frame=str(frame_idx),
                    outside="0",  # not used
                    occluded="0",  # not used
                    xtl=str(bboxes[0]),
                    ytl=str(bboxes[1]),
                    xbr=str(bboxes[2]),
                    ybr=str(bboxes[3]),
                    z_order="0",
                    label=str(class_idx),
                )
                box.text = ""
                track.append(box)
                video_dict[video_dir]["annotations"].append(track)

    for (video_dir, info) in video_dict.items():
        video_path = osp.join(src_path, video_dir)
        shutil.copytree(video_path, osp.join(dst_path, f"{video_dir}/images"), copy_function=rename_and_copy)
        et = etree.ElementTree(info["annotations"])
        et.write(
            osp.join(dst_path, f"{video_dir}/annotations.xml"),
            pretty_print=True,
            xml_declaration=True,
            encoding="utf-8",
        )


def rename_and_copy(_src, _dst):
    """Change frame name to cvat format."""
    img_name = _dst.rsplit("/", maxsplit=1)[-1]
    # FIXME This only support AVA dataset name
    if "_" in img_name:
        frame_index = int(img_name.split(".")[0].split("_")[-1])
    else:
        frame_index = int(img_name.split(".")[0])
    new_img_name = f"frame_{frame_index:06d}.png"
    _dst = _dst.replace(img_name, new_img_name)
    shutil.copy2(_src, _dst)


def read_ava_csv(csv_path):
    """Read ava format annotation csv file."""
    annot_info = {}  # {video_id: {frame_idx: [annot0, annot1, ..., annotN]}}
    with open(csv_path, "r", encoding="utf-8") as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=",")
        for line in csv_reader:
            video_id, frame_idx, bboxes, class_idx = line[0], line[1], line[2:6], line[6]
            frame_idx = int(frame_idx)
            bboxes.append(class_idx)
            if video_id in annot_info:
                if frame_idx in annot_info[video_id]:
                    annot_info[video_id][frame_idx].append(bboxes)
                else:
                    annot_info[video_id][frame_idx] = [bboxes]
            else:
                annot_info[video_id] = {frame_idx: [bboxes]}
    return annot_info


def parse_args():
    """Parses command line arguments."""
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--task", required=True, type=str)
    parser.add_argument("--src_path", required=True, type=str)
    parser.add_argument("--dst_path", required=True, type=str)
    parser.add_argument("--ann_file", required=True, type=str)
    parser.add_argument("--label_map", type=str, default=None)
    return parser.parse_args()


def main():
    """Main function."""
    args = parse_args()
    if args.task == "action_classification":
        convert_action_cls_dataset_to_datumaro(args.src_path, args.dst_path, args.ann_file, args.label_map)
    elif args.task == "action_detection":
        convert_ava_dataset_to_datumaro(args.src_path, args.dst_path, args.ann_file)


if __name__ == "__main__":
    main()
