import os
from typing import List

import numpy as np
from torchvision import datasets as torchdata

from continuum.datasets import ImageFolderDataset
from continuum import download
from continuum.tasks import TaskType


class DTD(ImageFolderDataset):
    """Describable Textures Dataset (DTD)

    Reference:
        * Describing Textures in the Wild
          M. Cimpoi and S. Maji and I. Kokkinos and S. Mohamed and and A. Vedaldi
          CVPR 2014
    """
    url = "https://www.robots.ox.ac.uk/~vgg/data/dtd/download/dtd-r1.0.1.tar.gz"

    def __init__(self, data_path: str, train: bool = True, download: bool = True, split: int = 1):
        super().__init__(data_path=data_path, train=train, download=download, data_type=TaskType.IMAGE_PATH)

        if not (1 <= int(split) <= 10):
            raise ValueError(f"Available splits are [1, ..., 10], not {split}")
        self.split = split

    def _download(self):
        archive_path = os.path.join(self.data_path, "dtd-r1.0.1.tar.gz")
        if not os.path.exists(archive_path):
            print("Downloading DTD dataset...")
            download.download(self.url, self.data_path)
        if not os.path.exists(os.path.join(self.data_path, "dtd")):
            print("Uncompressing images...")
            download.untar(archive_path)

    def get_data(self):
        x, y, t = self._format(torchdata.ImageFolder(os.path.join(self.data_path, "dtd", "images")).imgs)

        if self.train:
            index_files = [
                os.path.join(self.data_path, "dtd", "labels", f"train{str(self.split)}.txt"),
                os.path.join(self.data_path, "dtd", "labels", f"val{str(self.split)}.txt")
            ]
        else:
            index_files = [
                os.path.join(self.data_path, "dtd", "labels", f"test{str(self.split)}.txt")
            ]

        valid_paths = set()
        for index_file in index_files:
            with open(index_file) as f:
                valid_paths.update(
                    map(lambda p: os.path.join(self.data_path, "dtd", "images", p.strip()),
                        f.readlines()
                    )
                )
        valid_paths = np.array(list(valid_paths))
        indexes = np.isin(x, valid_paths)

        return x[indexes], y[indexes], None


[((30, 23, 30, 72), 'os.path.join', 'os.path.join', ({(30, 36, 30, 50): 'self.data_path', (30, 52, 30, 71): '"""dtd-r1.0.1.tar.gz"""'}, {}), "(self.data_path, 'dtd-r1.0.1.tar.gz')", False, 'import os\n'), ((60, 18, 60, 41), 'numpy.isin', 'np.isin', ({(60, 26, 60, 27): 'x', (60, 29, 60, 40): 'valid_paths'}, {}), '(x, valid_paths)', True, 'import numpy as np\n'), ((31, 15, 31, 43), 'os.path.exists', 'os.path.exists', ({(31, 30, 31, 42): 'archive_path'}, {}), '(archive_path)', False, 'import os\n'), ((33, 12, 33, 55), 'continuum.download.download', 'download.download', ({(33, 30, 33, 38): 'self.url', (33, 40, 33, 54): 'self.data_path'}, {}), '(self.url, self.data_path)', False, 'from continuum import download\n'), ((36, 12, 36, 40), 'continuum.download.untar', 'download.untar', ({(36, 27, 36, 39): 'archive_path'}, {}), '(archive_path)', False, 'from continuum import download\n'), ((34, 30, 34, 65), 'os.path.join', 'os.path.join', ({(34, 43, 34, 57): 'self.data_path', (34, 59, 34, 64): '"""dtd"""'}, {}), "(self.data_path, 'dtd')", False, 'import os\n'), ((39, 53, 39, 98), 'os.path.join', 'os.path.join', ({(39, 66, 39, 80): 'self.data_path', (39, 82, 39, 87): '"""dtd"""', (39, 89, 39, 97): '"""images"""'}, {}), "(self.data_path, 'dtd', 'images')", False, 'import os\n')]