import os
from typing import Optional
from torchvision.datasets.folder import DatasetFolder, IMG_EXTENSIONS, default_loader
from torchvision.datasets.utils import download_and_extract_archive
from ._util import check_exits


class OfficeCaltech(DatasetFolder):
    """Office+Caltech Dataset.

    Parameters:
        - **root** (str): Root directory of dataset
        - **task** (str): The task (domain) to create dataset. Choices include ``'A'``: amazon, \
            ``'D'``: dslr, ``'W'``:webcam and ``'C'``: caltech.
        - **download** (bool, optional): If true, downloads the dataset from the internet and puts it \
            in root directory. If dataset is already downloaded, it is not downloaded again.
        - **transform** (callable, optional): A function/transform that  takes in an PIL image and returns a \
            transformed version. E.g, ``transforms.RandomCrop``.
        - **target_transform** (callable, optional): A function/transform that takes in the target and transforms it.

    .. note:: In `root`, there will exist following files after downloading.
        ::
            amazon/
                images/
                    backpack/
                        *.jpg
                        ...
            dslr/
            webcam/
            caltech/
            image_list/
                amazon.txt
                dslr.txt
                webcam.txt
                caltech.txt
    """

    directories = {"A": "amazon", "D": "dslr", "W": "webcam", "C": "caltech"}
    CLASSES = [
        "back_pack",
        "bike",
        "calculator",
        "headphones",
        "keyboard",
        "laptop_computer",
        "monitor",
        "mouse",
        "mug",
        "projector",
    ]

    def __init__(
        self, root: str, task: str, download: Optional[bool] = False, **kwargs
    ):
        if download:
            for dir in self.directories.values():
                if not os.path.exists(os.path.join(root, dir)):
                    download_and_extract_archive(
                        url="https://cloud.tsinghua.edu.cn/f/bc427c57459c4194baf0/?dl=1",
                        download_root=os.path.join(root, "download"),
                        filename="officecaltech.tgz",
                        remove_finished=False,
                        extract_root=root,
                    )
                    break
        else:
            list(map(lambda dir, _: check_exits(root, dir), self.directories.values()))

        super(OfficeCaltech, self).__init__(
            os.path.join(root, self.directories[task]),
            default_loader,
            extensions=IMG_EXTENSIONS,
            **kwargs
        )
        self.classes = OfficeCaltech.classes
        self.class_to_idx = {
            cls: idx for idx, clss in enumerate(self.classes) for cls in clss
        }

    @property
    def num_classes(self):
        """Number of classes"""
        return len(self.classes)
