#!/usr/bin/env python3
#
# Copytright 2020 Graviti. All Rights Reserved
#
# pylint: disable=invalid-name

"""This file handles with the Bosch Small Traffic Light Dataset"""

import os

from ...dataset import Data, Dataset
from ...label import LabeledBox2D, LabelType

DATASET_NAME = "Bosch Small Traffic Lights"

_LABEL_FILENAME_DICT = {
    "test": "test.yaml",
    "train": "train.yaml",
    "additional": "additional_train.yaml",
}


def BSTLD(path: str) -> Dataset:
    """
    Load the Bosch Small Traffic Light Dataset to TensorBay
    :param path: the root directory of the dataset
    The file structure should be like:
    <path>
        rgb/
            additional/
                2015-10-05-10-52-01_bag/
                    <image_name>.jpg
                    ...
                ...
            test/
                <image_name>.jpg
                ...
            train/
                2015-05-29-15-29-39_arastradero_traffic_light_loop_bag/
                    <image_name>.jpg
                    ...
                ...
        test.yaml
        train.yaml
        additional_train.yaml
    """
    import yaml  # pylint: disable=import-outside-toplevel

    root_path = os.path.abspath(os.path.expanduser(path))

    dataset = Dataset(DATASET_NAME)
    dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json"))

    for mode, label_file_name in _LABEL_FILENAME_DICT.items():
        segment = dataset.create_segment(mode)
        label_file_path = os.path.join(root_path, label_file_name)

        with open(label_file_path) as fp:
            labels = yaml.load(fp, yaml.FullLoader)

        for label in labels:
            if mode == "test":
                # the path in test label file looks like:
                # /absolute/path/to/<image_name>.png
                file_path = os.path.join(root_path, "rgb", "test", label["path"].rsplit("/", 1)[-1])
            else:
                # the path in label file looks like:
                # ./rgb/additional/2015-10-05-10-52-01_bag/<image_name>.png
                file_path = os.path.join(root_path, *label["path"][2:].split("/"))
            data = Data(file_path)
            data.register_label(LabelType.BOX2D)

            for box in label["boxes"]:
                label = LabeledBox2D(
                    box["x_min"],
                    box["y_min"],
                    box["x_max"],
                    box["y_max"],
                    category=box["label"],
                    attributes={"occluded": box["occluded"]},
                )
                data.append_label(label)
            segment.append(data)

    return dataset
