#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#
# pylint: disable=invalid-name

"""Dataloader of WIDER_FACE dataset."""

import os
from collections import OrderedDict
from itertools import islice
from typing import Dict, Iterator, List, Union

from tensorbay.dataset import Data, Dataset
from tensorbay.label import Classification, LabeledBox2D

DATASET_NAME = "WIDER_FACE"
_SEGMENT_LIST = {
    "test": "wider_face_test_filelist.txt",
    "train": "wider_face_train_bbx_gt.txt",
    "val": "wider_face_val_bbx_gt.txt",
}
_ATTRIBUTE_MAP_TYPE = Dict[str, List[Union[bool, str]]]


def WIDER_FACE(path: str) -> Dataset:
    """`WIDER FACE <http://shuoyang1213.me/WIDERFACE/>`_ dataset.

    The file structure should be like::

        <path>
            WIDER_train/
                images/
                    0--Parade/
                        0_Parade_marchingband_1_100.jpg
                        0_Parade_marchingband_1_1015.jpg
                        0_Parade_marchingband_1_1030.jpg
                        ...
                    1--Handshaking/
                    ...
                    59--people--driving--car/
                    61--Street_Battle/
            WIDER_val/
                ...
            WIDER_test/
                ...
            wider_face_split/
                wider_face_train_bbx_gt.txt
                wider_face_val_bbx_gt.txt

    Arguments:
        path: The root directory of the dataset.

    Returns:
        Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance.

    """
    dataset = Dataset(DATASET_NAME)
    dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json"))
    attribute_map = _get_attribute_map(dataset)
    for segment_name, label_file in _SEGMENT_LIST.items():
        segment = dataset.create_segment(segment_name)
        for data in _load_data(
            os.path.join(path, "wider_face_split", label_file), attribute_map, segment_name
        ):
            segment.append(data)
        segment.sort()
    return dataset


def _get_attribute_map(dataset: Dataset) -> _ATTRIBUTE_MAP_TYPE:
    attribute_map: _ATTRIBUTE_MAP_TYPE = OrderedDict()
    for info in dataset.catalog.box2d.attributes:
        if getattr(info, "type", None) == "boolean":
            attribute_map[info.name] = [False, True]
        else:
            attribute_map[info.name] = info.enum  # type: ignore[assignment]
    return attribute_map


def _load_data(path: str, attribute_map: _ATTRIBUTE_MAP_TYPE, segment_name: str) -> Iterator[Data]:
    """Loads the box2d and classification label in to data.

    The train and val label file context should be like::

        0--Parade/0_Parade_marchingband_1_849.jpg
        1
        449 330 122 149 0 0 0 0 0 0
        0--Parade/0_Parade_Parade_0_452.jpg
        0
        0 0 0 0 0 0 0 0 0 0
        0--Parade/0_Parade_marchingband_1_799.jpg
        21
        78 221 7 8 2 0 0 0 0 0
        78 238 14 17 2 0 0 0 0 0
        113 212 11 15 2 0 0 0 0 0
        134 260 15 15 2 0 0 0 0 0

    Arguments:
        path: The path of label file.
        attribute_map: A attribute value enum table.
        segment_name: Name of the segment.

    Yields:
        Data with loaded lables.

    """
    is_test = segment_name == "test"
    with open(path, encoding="utf-8") as fp:
        for image_path in fp:
            event, file_name = image_path.split("/")
            # translate directory name to category. like 0--Parade -> Parade
            category = "_".join(event.split("--")[1:])
            root_path = path.rsplit(os.sep, 2)[0]
            data = Data(
                os.path.join(root_path, f"WIDER_{segment_name}", "images", event, file_name)
            )
            data.label.classification = Classification(category)
            if not is_test:
                label_num = int(fp.readline())

                # when the label num is 0, a line of "0 0 0 0 0 0 0 0 0 0" also given
                if label_num == 0:
                    fp.readline()
                data.label.box2d = []
                for line in islice(fp, label_num):
                    labels = line.strip().split()
                    attributes = {
                        key: mapping[int(value)]
                        for (key, mapping), value in zip(attribute_map.items(), labels[4:10])
                    }
                    data.label.box2d.append(
                        LabeledBox2D.from_xywh(
                            x=int(labels[0]),
                            y=int(labels[1]),
                            width=int(labels[2]),
                            height=int(labels[3]),
                            attributes=attributes,
                        )
                    )

            yield data
