import csv
import functools
from typing import Any, Dict, List, Optional, Tuple, Iterator, Sequence, BinaryIO

from torchdata.datapipes.iter import (
    IterDataPipe,
    Mapper,
    Filter,
    Zipper,
    IterKeyZipper,
)
from torchvision.prototype.datasets.utils import (
    Dataset,
    DatasetConfig,
    DatasetInfo,
    GDriveResource,
    OnlineResource,
)
from torchvision.prototype.datasets.utils._internal import (
    INFINITE_BUFFER_SIZE,
    getitem,
    path_accessor,
    hint_sharding,
    hint_shuffling,
)
from torchvision.prototype.features import EncodedImage, _Feature, Label, BoundingBox


csv.register_dialect("celeba", delimiter=" ", skipinitialspace=True)


class CelebACSVParser(IterDataPipe[Tuple[str, Dict[str, str]]]):
    def __init__(
        self,
        datapipe: IterDataPipe[Tuple[Any, BinaryIO]],
        *,
        fieldnames: Optional[Sequence[str]] = None,
    ) -> None:
        self.datapipe = datapipe
        self.fieldnames = fieldnames

    def __iter__(self) -> Iterator[Tuple[str, Dict[str, str]]]:
        for _, file in self.datapipe:
            file = (line.decode() for line in file)

            if self.fieldnames:
                fieldnames = self.fieldnames
            else:
                # The first row is skipped, because it only contains the number of samples
                next(file)

                # Empty field names are filtered out, because some files have an extra white space after the header
                # line, which is recognized as extra column
                fieldnames = [name for name in next(csv.reader([next(file)], dialect="celeba")) if name]
                # Some files do not include a label for the image ID column
                if fieldnames[0] != "image_id":
                    fieldnames.insert(0, "image_id")

            for line in csv.DictReader(file, fieldnames=fieldnames, dialect="celeba"):
                yield line.pop("image_id"), line


class CelebA(Dataset):
    def _make_info(self) -> DatasetInfo:
        return DatasetInfo(
            "celeba",
            homepage="https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html",
            valid_options=dict(split=("train", "val", "test")),
        )

    def resources(self, config: DatasetConfig) -> List[OnlineResource]:
        splits = GDriveResource(
            "0B7EVK8r0v71pY0NSMzRuSXJEVkk",
            sha256="fc955bcb3ef8fbdf7d5640d9a8693a8431b5f2ee291a5c1449a1549e7e073fe7",
            file_name="list_eval_partition.txt",
        )
        images = GDriveResource(
            "0B7EVK8r0v71pZjFTYXZWM3FlRnM",
            sha256="46fb89443c578308acf364d7d379fe1b9efb793042c0af734b6112e4fd3a8c74",
            file_name="img_align_celeba.zip",
        )
        identities = GDriveResource(
            "1_ee_0u7vcNLOfNLegJRHmolfH5ICW-XS",
            sha256="c6143857c3e2630ac2da9f782e9c1232e5e59be993a9d44e8a7916c78a6158c0",
            file_name="identity_CelebA.txt",
        )
        attributes = GDriveResource(
            "0B7EVK8r0v71pblRyaVFSWGxPY0U",
            sha256="f0e5da289d5ccf75ffe8811132694922b60f2af59256ed362afa03fefba324d0",
            file_name="list_attr_celeba.txt",
        )
        bounding_boxes = GDriveResource(
            "0B7EVK8r0v71pbThiMVRxWXZ4dU0",
            sha256="7487a82e57c4bb956c5445ae2df4a91ffa717e903c5fa22874ede0820c8ec41b",
            file_name="list_bbox_celeba.txt",
        )
        landmarks = GDriveResource(
            "0B7EVK8r0v71pd0FJY3Blby1HUTQ",
            sha256="6c02a87569907f6db2ba99019085697596730e8129f67a3d61659f198c48d43b",
            file_name="list_landmarks_align_celeba.txt",
        )
        return [splits, images, identities, attributes, bounding_boxes, landmarks]

    _SPLIT_ID_TO_NAME = {
        "0": "train",
        "1": "val",
        "2": "test",
    }

    def _filter_split(self, data: Tuple[str, Dict[str, str]], *, split: str) -> bool:
        return self._SPLIT_ID_TO_NAME[data[1]["split_id"]] == split

    def _prepare_sample(
        self,
        data: Tuple[
            Tuple[str, Tuple[Tuple[str, List[str]], Tuple[str, BinaryIO]]],
            Tuple[
                Tuple[str, Dict[str, str]],
                Tuple[str, Dict[str, str]],
                Tuple[str, Dict[str, str]],
                Tuple[str, Dict[str, str]],
            ],
        ],
    ) -> Dict[str, Any]:
        split_and_image_data, ann_data = data
        _, (_, image_data) = split_and_image_data
        path, buffer = image_data

        image = EncodedImage.from_file(buffer)
        (_, identity), (_, attributes), (_, bounding_box), (_, landmarks) = ann_data

        return dict(
            path=path,
            image=image,
            identity=Label(int(identity["identity"])),
            attributes={attr: value == "1" for attr, value in attributes.items()},
            bounding_box=BoundingBox(
                [int(bounding_box[key]) for key in ("x_1", "y_1", "width", "height")],
                format="xywh",
                image_size=image.image_size,
            ),
            landmarks={
                landmark: _Feature((int(landmarks[f"{landmark}_x"]), int(landmarks[f"{landmark}_y"])))
                for landmark in {key[:-2] for key in landmarks.keys()}
            },
        )

    def _make_datapipe(
        self,
        resource_dps: List[IterDataPipe],
        *,
        config: DatasetConfig,
    ) -> IterDataPipe[Dict[str, Any]]:
        splits_dp, images_dp, identities_dp, attributes_dp, bounding_boxes_dp, landmarks_dp = resource_dps

        splits_dp = CelebACSVParser(splits_dp, fieldnames=("image_id", "split_id"))
        splits_dp = Filter(splits_dp, functools.partial(self._filter_split, split=config.split))
        splits_dp = hint_shuffling(splits_dp)
        splits_dp = hint_sharding(splits_dp)

        anns_dp = Zipper(
            *[
                CelebACSVParser(dp, fieldnames=fieldnames)
                for dp, fieldnames in (
                    (identities_dp, ("image_id", "identity")),
                    (attributes_dp, None),
                    (bounding_boxes_dp, None),
                    (landmarks_dp, None),
                )
            ]
        )

        dp = IterKeyZipper(
            splits_dp,
            images_dp,
            key_fn=getitem(0),
            ref_key_fn=path_accessor("name"),
            buffer_size=INFINITE_BUFFER_SIZE,
            keep_key=True,
        )
        dp = IterKeyZipper(
            dp,
            anns_dp,
            key_fn=getitem(0),
            ref_key_fn=getitem(0, 0),
            buffer_size=INFINITE_BUFFER_SIZE,
        )
        return Mapper(dp, self._prepare_sample)
