import os
from typing import Optional
from .imagelist import ImageList
from ._util import download as download_data, check_exits


class DomainNet(ImageList):
    """`DomainNet <http://ai.bu.edu/M3SDA/#dataset>`_ (cleaned version, recommended)

    See `Moment Matching for Multi-Source Domain Adaptation <https://arxiv.org/abs/1812.01754>`_ for details.

    Parameters:
        - **root** (str): Root directory of dataset
        - **task** (str): The task (domain) to create dataset. Choices include ``'c'``:clipart, \
            ``'i'``: infograph, ``'p'``: painting, ``'q'``: quickdraw, ``'r'``: real, ``'s'``: sketch
        - **evaluate** (bool, optional): If true, use the test set. Otherwise, use the training set. Default: False
        - **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.
        ::
            clipart/
            infograph/
            painting/
            quickdraw/
            real/
            sketch/
            image_list/
                clipart.txt
                ...
    """

    download_list = [
        (
            "image_list",
            "image_list.zip",
            "https://cloud.tsinghua.edu.cn/f/73489ae10aea45d58194/?dl=1",
        ),
        (
            "clipart",
            "clipart.zip",
            "http://csr.bu.edu/ftp/visda/2019/multi-source/groundtruth/clipart.zip",
        ),
        (
            "infograph",
            "infograph.zip",
            "http://csr.bu.edu/ftp/visda/2019/multi-source/infograph.zip",
        ),
        (
            "painting",
            "painting.zip",
            "http://csr.bu.edu/ftp/visda/2019/multi-source/groundtruth/painting.zip",
        ),
        (
            "quickdraw",
            "quickdraw.zip",
            "http://csr.bu.edu/ftp/visda/2019/multi-source/quickdraw.zip",
        ),
        ("real", "real.zip", "http://csr.bu.edu/ftp/visda/2019/multi-source/real.zip"),
        (
            "sketch",
            "sketch.zip",
            "http://csr.bu.edu/ftp/visda/2019/multi-source/sketch.zip",
        ),
    ]
    image_list = {
        "c": "clipart",
        "i": "infograph",
        "p": "painting",
        "q": "quickdraw",
        "r": "real",
        "s": "sketch",
    }
    CLASSES = [
        "aircraft_carrier",
        "airplane",
        "alarm_clock",
        "ambulance",
        "angel",
        "animal_migration",
        "ant",
        "anvil",
        "apple",
        "arm",
        "asparagus",
        "axe",
        "backpack",
        "banana",
        "bandage",
        "barn",
        "baseball",
        "baseball_bat",
        "basket",
        "basketball",
        "bat",
        "bathtub",
        "beach",
        "bear",
        "beard",
        "bed",
        "bee",
        "belt",
        "bench",
        "bicycle",
        "binoculars",
        "bird",
        "birthday_cake",
        "blackberry",
        "blueberry",
        "book",
        "boomerang",
        "bottlecap",
        "bowtie",
        "bracelet",
        "brain",
        "bread",
        "bridge",
        "broccoli",
        "broom",
        "bucket",
        "bulldozer",
        "bus",
        "bush",
        "butterfly",
        "cactus",
        "cake",
        "calculator",
        "calendar",
        "camel",
        "camera",
        "camouflage",
        "campfire",
        "candle",
        "cannon",
        "canoe",
        "car",
        "carrot",
        "castle",
        "cat",
        "ceiling_fan",
        "cello",
        "cell_phone",
        "chair",
        "chandelier",
        "church",
        "circle",
        "clarinet",
        "clock",
        "cloud",
        "coffee_cup",
        "compass",
        "computer",
        "cookie",
        "cooler",
        "couch",
        "cow",
        "crab",
        "crayon",
        "crocodile",
        "crown",
        "cruise_ship",
        "cup",
        "diamond",
        "dishwasher",
        "diving_board",
        "dog",
        "dolphin",
        "donut",
        "door",
        "dragon",
        "dresser",
        "drill",
        "drums",
        "duck",
        "dumbbell",
        "ear",
        "elbow",
        "elephant",
        "envelope",
        "eraser",
        "eye",
        "eyeglasses",
        "face",
        "fan",
        "feather",
        "fence",
        "finger",
        "fire_hydrant",
        "fireplace",
        "firetruck",
        "fish",
        "flamingo",
        "flashlight",
        "flip_flops",
        "floor_lamp",
        "flower",
        "flying_saucer",
        "foot",
        "fork",
        "frog",
        "frying_pan",
        "garden",
        "garden_hose",
        "giraffe",
        "goatee",
        "golf_club",
        "grapes",
        "grass",
        "guitar",
        "hamburger",
        "hammer",
        "hand",
        "harp",
        "hat",
        "headphones",
        "hedgehog",
        "helicopter",
        "helmet",
        "hexagon",
        "hockey_puck",
        "hockey_stick",
        "horse",
        "hospital",
        "hot_air_balloon",
        "hot_dog",
        "hot_tub",
        "hourglass",
        "house",
        "house_plant",
        "hurricane",
        "ice_cream",
        "jacket",
        "jail",
        "kangaroo",
        "key",
        "keyboard",
        "knee",
        "knife",
        "ladder",
        "lantern",
        "laptop",
        "leaf",
        "leg",
        "light_bulb",
        "lighter",
        "lighthouse",
        "lightning",
        "line",
        "lion",
        "lipstick",
        "lobster",
        "lollipop",
        "mailbox",
        "map",
        "marker",
        "matches",
        "megaphone",
        "mermaid",
        "microphone",
        "microwave",
        "monkey",
        "moon",
        "mosquito",
        "motorbike",
        "mountain",
        "mouse",
        "moustache",
        "mouth",
        "mug",
        "mushroom",
        "nail",
        "necklace",
        "nose",
        "ocean",
        "octagon",
        "octopus",
        "onion",
        "oven",
        "owl",
        "paintbrush",
        "paint_can",
        "palm_tree",
        "panda",
        "pants",
        "paper_clip",
        "parachute",
        "parrot",
        "passport",
        "peanut",
        "pear",
        "peas",
        "pencil",
        "penguin",
        "piano",
        "pickup_truck",
        "picture_frame",
        "pig",
        "pillow",
        "pineapple",
        "pizza",
        "pliers",
        "police_car",
        "pond",
        "pool",
        "popsicle",
        "postcard",
        "potato",
        "power_outlet",
        "purse",
        "rabbit",
        "raccoon",
        "radio",
        "rain",
        "rainbow",
        "rake",
        "remote_control",
        "rhinoceros",
        "rifle",
        "river",
        "roller_coaster",
        "rollerskates",
        "sailboat",
        "sandwich",
        "saw",
        "saxophone",
        "school_bus",
        "scissors",
        "scorpion",
        "screwdriver",
        "sea_turtle",
        "see_saw",
        "shark",
        "sheep",
        "shoe",
        "shorts",
        "shovel",
        "sink",
        "skateboard",
        "skull",
        "skyscraper",
        "sleeping_bag",
        "smiley_face",
        "snail",
        "snake",
        "snorkel",
        "snowflake",
        "snowman",
        "soccer_ball",
        "sock",
        "speedboat",
        "spider",
        "spoon",
        "spreadsheet",
        "square",
        "squiggle",
        "squirrel",
        "stairs",
        "star",
        "steak",
        "stereo",
        "stethoscope",
        "stitches",
        "stop_sign",
        "stove",
        "strawberry",
        "streetlight",
        "string_bean",
        "submarine",
        "suitcase",
        "sun",
        "swan",
        "sweater",
        "swing_set",
        "sword",
        "syringe",
        "table",
        "teapot",
        "teddy-bear",
        "telephone",
        "television",
        "tennis_racquet",
        "tent",
        "The_Eiffel_Tower",
        "The_Great_Wall_of_China",
        "The_Mona_Lisa",
        "tiger",
        "toaster",
        "toe",
        "toilet",
        "tooth",
        "toothbrush",
        "toothpaste",
        "tornado",
        "tractor",
        "traffic_light",
        "train",
        "tree",
        "triangle",
        "trombone",
        "truck",
        "trumpet",
        "t-shirt",
        "umbrella",
        "underwear",
        "van",
        "vase",
        "violin",
        "washing_machine",
        "watermelon",
        "waterslide",
        "whale",
        "wheel",
        "windmill",
        "wine_bottle",
        "wine_glass",
        "wristwatch",
        "yoga",
        "zebra",
        "zigzag",
    ]

    def __init__(
        self,
        root: str,
        task: str,
        evaluate: Optional[bool] = False,
        download: Optional[float] = False,
        **kwargs
    ):
        assert task in self.image_list
        phase = "test" if evaluate else "train"
        data_list_file = os.path.join(
            root, "image_list", "{}_{}.txt".format(self.image_list[task], phase)
        )
        print("loading {}".format(data_list_file))

        if download:
            list(map(lambda args: download_data(root, *args), self.download_list))
        else:
            list(map(lambda args: check_exits(root, args[0]), self.download_list))

        super(DomainNet, self).__init__(
            root, DomainNet.CLASSES, data_list_file=data_list_file, **kwargs
        )
