# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
import collections
import json
import os
import resource

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image

from .dataset.vtab import base
from .dataset.vtab.caltech import Caltech101
from .dataset.vtab.cars import CarsData
from .dataset.vtab.cifar import CifarData
from .dataset.vtab.clevr import CLEVRData
from .dataset.vtab.cub import CUB2011Data
from .dataset.vtab.diabetic_retinopathy import RetinopathyData
from .dataset.vtab.dmlab import DmlabData
from .dataset.vtab.dsprites import DSpritesData
from .dataset.vtab.dtd import DTDData
from .dataset.vtab.eurosat import EurosatData
from .dataset.vtab.food101 import Food101Data
from .dataset.vtab.inaturalist import INaturalistData
from .dataset.vtab.kitti import KittiData
from .dataset.vtab.oxford_flowers102 import OxfordFlowers102Data
from .dataset.vtab.oxford_iiit_pet import OxfordIIITPetData
from .dataset.vtab.patch_camelyon import PatchCamelyonData
from .dataset.vtab.resisc45 import Resisc45Data
from .dataset.vtab.smallnorb import SmallNORBData
from .dataset.vtab.sun397 import Sun397Data
from .dataset.vtab.svhn import SvhnData

low, high = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (high, high))


try:
    from IPython import embed
except:
    pass

_DATASETS = {}

Dataset = collections.namedtuple("Dataset", ["trainset", "testset"])


def _add_dataset(dataset_fn):
    _DATASETS[dataset_fn.__name__] = dataset_fn
    return dataset_fn


def _get_transforms(augment=True, normalize=None):
    if normalize is None:
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        )

    basic_transform = [transforms.ToTensor(), normalize]

    transform_train = []
    if augment:
        transform_train += [
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
        ]
    else:
        transform_train += [
            transforms.Resize(256),
            transforms.CenterCrop(224),
        ]
    transform_train += basic_transform
    transform_train = transforms.Compose(transform_train)

    transform_test = [
        transforms.Resize(256),
        transforms.CenterCrop(224),
    ]
    transform_test += basic_transform
    transform_test = transforms.Compose(transform_test)

    return transform_train, transform_test


def _get_mnist_transforms(augment=True, invert=False, transpose=False):
    transform = [
        transforms.ToTensor(),
    ]
    if invert:
        transform += [transforms.Lambda(lambda x: 1.0 - x)]
    if transpose:
        transform += [transforms.Lambda(lambda x: x.transpose(2, 1))]
    transform += [
        transforms.Normalize((0.5,), (0.5,)),
        transforms.Lambda(lambda x: x.expand(3, 32, 32)),
    ]

    transform_train = []
    transform_train += [transforms.Pad(padding=2)]
    if augment:
        transform_train += [transforms.RandomCrop(32, padding=4)]
    transform_train += transform
    transform_train = transforms.Compose(transform_train)

    transform_test = []
    transform_test += [transforms.Pad(padding=2)]
    transform_test += transform
    transform_test = transforms.Compose(transform_test)

    return transform_train, transform_test


def _get_cifar_transforms(augment=True):
    transform = [
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ]
    transform_train = []
    if augment:
        transform_train += [
            transforms.Pad(padding=4, fill=(125, 123, 113)),
            transforms.RandomCrop(32, padding=0),
            transforms.RandomHorizontalFlip(),
        ]
    transform_train += transform
    transform_train = transforms.Compose(transform_train)
    transform_test = []
    transform_test += transform
    transform_test = transforms.Compose(transform_test)
    return transform_train, transform_test


def set_metadata(trainset, testset, config, dataset_name):
    trainset.metadata = {
        "dataset": dataset_name,
        "task_id": config.task_id,
        "task_name": trainset.task_name,
    }
    testset.metadata = {
        "dataset": dataset_name,
        "task_id": config.task_id,
        "task_name": testset.task_name,
    }
    return trainset, testset


@_add_dataset
def inat2018(root, config):
    from .dataset.inat import iNat2018Dataset

    transform_train, transform_test = _get_transforms()
    trainset = iNat2018Dataset(
        root, split="train", transform=transform_train, task_id=config.task_id
    )
    testset = iNat2018Dataset(
        root, split="val", transform=transform_test, task_id=config.task_id
    )
    trainset, testset = set_metadata(trainset, testset, config, "inat2018")
    return trainset, testset


def load_tasks_map(tasks_map_file):
    assert os.path.exists(tasks_map_file), tasks_map_file
    with open(tasks_map_file, "r") as f:
        tasks_map = json.load(f)
        tasks_map = {int(k): int(v) for k, v in tasks_map.items()}
    return tasks_map


@_add_dataset
def cub_inat2018(root, config):
    """This meta-task is the concatenation of CUB-200 (first 25 tasks) and iNat (last 207 tasks).

    - The first 10 tasks are classification of the animal species inside one of 10 orders of birds in CUB-200
        (considering all orders except passeriformes).
    - The next 15 tasks are classification of species inside the 15 families of the order of passerifomes
    - The remaining 207 tasks are classification of the species inside each of 207 families in iNat

    As noted above, for CUB-200 10 taks are classification of species inside an order, rather than inside of a family
    as done in the iNat (recall order > family > species). This is done because CUB-200 has very few images
    in each family of bird (expect for the families of passeriformes). Hence, we go up step in the taxonomy and
    consider classification inside a orders and not families.
    """
    NUM_CUB = 25
    NUM_CUB_ORDERS = 10
    NUM_INAT = 207
    assert 0 <= config.task_id < NUM_CUB + NUM_INAT
    transform_train, transform_test = _get_transforms()
    if 0 <= config.task_id < NUM_CUB:
        # CUB
        from .dataset.cub import CUBDataset, CUBTasks

        tasks_map_file = os.path.join(root, "cub/CUB_200_2011", "final_tasks_map.json")
        tasks_map = load_tasks_map(tasks_map_file)
        task_id = tasks_map[config.task_id]

        if config.task_id < NUM_CUB_ORDERS:
            # CUB orders
            train_tasks = CUBTasks(CUBDataset(root, split="train"))
            trainset = train_tasks.generate(
                task_id=task_id, use_species_names=True, transform=transform_train
            )
            test_tasks = CUBTasks(CUBDataset(root, split="test"))
            testset = test_tasks.generate(
                task_id=task_id, use_species_names=True, transform=transform_test
            )
        else:
            # CUB passeriformes families
            train_tasks = CUBTasks(CUBDataset(root, split="train"))
            trainset = train_tasks.generate(
                task_id=task_id,
                task="family",
                taxonomy_file="passeriformes.txt",
                use_species_names=True,
                transform=transform_train,
            )
            test_tasks = CUBTasks(CUBDataset(root, split="test"))
            testset = test_tasks.generate(
                task_id=task_id,
                task="family",
                taxonomy_file="passeriformes.txt",
                use_species_names=True,
                transform=transform_test,
            )
    else:
        # iNat2018
        from .dataset.inat import iNat2018Dataset

        tasks_map_file = os.path.join(root, "inat2018", "final_tasks_map.json")
        tasks_map = load_tasks_map(tasks_map_file)
        task_id = tasks_map[config.task_id - NUM_CUB]

        trainset = iNat2018Dataset(
            root, split="train", transform=transform_train, task_id=task_id
        )
        testset = iNat2018Dataset(
            root, split="val", transform=transform_test, task_id=task_id
        )
    trainset, testset = set_metadata(trainset, testset, config, "cub_inat2018")
    return trainset, testset


@_add_dataset
def imat2018fashion(root, config):
    NUM_IMAT = 228
    assert 0 <= config.task_id < NUM_IMAT
    from .dataset.imat import iMat2018FashionDataset, iMat2018FashionTasks

    transform_train, transform_test = _get_transforms()
    train_tasks = iMat2018FashionTasks(iMat2018FashionDataset(root, split="train"))
    trainset = train_tasks.generate(task_id=config.task_id, transform=transform_train)
    test_tasks = iMat2018FashionTasks(iMat2018FashionDataset(root, split="validation"))
    testset = test_tasks.generate(task_id=config.task_id, transform=transform_test)
    trainset, testset = set_metadata(trainset, testset, config, "imat2018fashion")
    return trainset, testset


@_add_dataset
def split_mnist(root, config):
    assert isinstance(config.task_id, tuple)
    from .dataset.mnist import MNISTDataset, SplitMNISTTask

    transform_train, transform_test = _get_mnist_transforms()
    train_tasks = SplitMNISTTask(MNISTDataset(root, train=True))
    trainset = train_tasks.generate(classes=config.task_id, transform=transform_train)
    test_tasks = SplitMNISTTask(MNISTDataset(root, train=False))
    testset = test_tasks.generate(classes=config.task_id, transform=transform_test)
    trainset, testset = set_metadata(trainset, testset, config, "split_mnist")
    return trainset, testset


@_add_dataset
def split_cifar(root, config):
    assert 0 <= config.task_id < 11
    from .dataset.cifar import CIFAR10Dataset, CIFAR100Dataset, SplitCIFARTask

    transform_train, transform_test = _get_cifar_transforms()
    train_tasks = SplitCIFARTask(
        CIFAR10Dataset(root, train=True), CIFAR100Dataset(root, train=True)
    )
    trainset = train_tasks.generate(task_id=config.task_id, transform=transform_train)
    test_tasks = SplitCIFARTask(
        CIFAR10Dataset(root, train=False), CIFAR100Dataset(root, train=False)
    )
    testset = test_tasks.generate(task_id=config.task_id, transform=transform_test)
    trainset, testset = set_metadata(trainset, testset, config, "split_cifar")
    return trainset, testset


@_add_dataset
def cifar10_mnist(root, config):
    from .dataset.cifar import CIFAR10Dataset
    from .dataset.expansion import UnionClassificationTaskExpander
    from .dataset.mnist import MNISTDataset

    transform_train, transform_test = _get_cifar_transforms()
    trainset = UnionClassificationTaskExpander(merge_duplicate_images=False)(
        [CIFAR10Dataset(root, train=True), MNISTDataset(root, train=True, expand=True)],
        transform=transform_train,
    )
    testset = UnionClassificationTaskExpander(merge_duplicate_images=False)(
        [
            CIFAR10Dataset(root, train=False),
            MNISTDataset(root, train=False, expand=True),
        ],
        transform=transform_test,
    )
    return trainset, testset


@_add_dataset
def cifar10(root):
    from torchvision.datasets import CIFAR10

    transform = transforms.Compose(
        [
            transforms.Resize(224),
            transforms.ToTensor(),
            transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
        ]
    )
    trainset = CIFAR10(root, train=True, transform=transform, download=True)
    testset = CIFAR10(root, train=False, transform=transform)
    return trainset, testset


@_add_dataset
def cifar100(root):
    from torchvision.datasets import CIFAR100

    transform = transforms.Compose(
        [
            transforms.Resize(224),
            transforms.ToTensor(),
            transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
        ]
    )
    trainset = CIFAR100(root, train=True, transform=transform, download=True)
    testset = CIFAR100(root, train=False, transform=transform)
    return trainset, testset


@_add_dataset
def mnist(root):
    from torchvision.datasets import MNIST

    transform = transforms.Compose(
        [
            lambda x: x.convert("RGB"),
            transforms.Resize(224),
            transforms.ToTensor(),
            # transforms.Normalize((0.5, 0.5, 0.5), (1., 1., 1.)),
        ]
    )
    trainset = MNIST(root, train=True, transform=transform, download=True)
    testset = MNIST(root, train=False, transform=transform)
    return trainset, testset


class TfdsWrapper(torch.utils.data.Dataset):
    def __init__(self, name, split, transform=None, max_samples=-1, data_dir=None):
        self.transform = transform
        ds = tfds.load(name=name, split=split, data_dir=data_dir, batch_size=-1)
        if max_samples < 0:
            self.data = np.array([x for i, x in enumerate(ds["image"])])
            self.targets = np.array([x for i, x in enumerate(ds["label"])])
        else:
            self.data = np.array(
                [x for i, x in enumerate(ds["image"]) if i < max_samples]
            )
            self.targets = np.array(
                [x for i, x in enumerate(ds["label"]) if i < max_samples]
            )

    def __getitem__(self, index):
        if self.transform:
            img = self.transform(self.data[index])
            return (img, self.targets[index])
        return (self.data[index], self.targets[index])

    def __len__(self) -> int:
        return len(self.targets)


def vtab_preprocess_fn(data, size=224, input_range=(0.0, 1.0)):
    image = data["image"]
    image = tf.image.resize(image, [size, size])

    image = tf.cast(image, tf.float32) / 255.0
    image = image * (input_range[1] - input_range[0]) + input_range[0]

    data["image"] = tf.transpose(image, perm=[2, 0, 1])
    return data


class VtabWrapper(torch.utils.data.Dataset):
    def __init__(self, dataset, split):
        dummy_decoder = tfds.decode.SkipDecoding()
        preprocess_fn = base.compose_preprocess_fn(
            dataset._image_decoder, dataset._base_preprocess_fn, vtab_preprocess_fn
        )
        ds = dataset._dataset_builder.as_dataset(
            split=dataset._tfds_splits[split],
            decoders={dataset._image_key: dummy_decoder},
        )
        ds = ds.map(preprocess_fn)

        self.data = np.array([x["image"].numpy() for i, x in enumerate(ds)])
        self.targets = np.array([x["label"] for i, x in enumerate(ds)]).astype(np.int_)

    def __getitem__(self, index):
        img = self.data[index]
        return (img, self.targets[index])

    def __len__(self) -> int:
        return len(self.targets)


# class TfdsWrapper(torch.utils.data.IterableDataset):
#  def __init__(self, name, split):
#    ds = tfds.load(name=name, split=split, shuffle_files=False)
#    self.iter = ds.as_numpy_iterator()
#    self.targets = np.array([x['label'] for x in ds])
#
#  def __iter__(self):
#    return self.iter
#
#  def __len__(self) -> int:
#    return len(self.targets)


@_add_dataset
def mnist_tfds(root):
    transform = torchvision.transforms.Compose(
        [
            lambda x: Image.fromarray(x.squeeze(), mode="L"),
            lambda x: x.convert("RGB"),
            torchvision.transforms.Resize(224),
            torchvision.transforms.ToTensor(),
            # transforms.Normalize((0.5, 0.5, 0.5), (1., 1., 1.)),
        ]
    )
    trainset = TfdsWrapper(name="mnist", split="train", transform=transform)
    testset = TfdsWrapper(name="mnist", split="test", transform=transform)

    return trainset, testset


@_add_dataset
def cifar10_tfds(root):
    transform = torchvision.transforms.Compose(
        [
            lambda x: Image.fromarray(x),
            transforms.Resize(224),
            transforms.ToTensor(),
            transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
        ]
    )
    trainset = TfdsWrapper(name="cifar10", split="train", transform=transform)
    testset = TfdsWrapper(name="cifar10", split="test", transform=transform)

    return trainset, testset


@_add_dataset
def vtab_caltech(root):
    ds = Caltech101()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_cifar(root):
    ds = CifarData(num_classes=100)
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_clevrdist(root):
    ds = CLEVRData(task="closest_object_distance")
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_clevrcnt(root):
    ds = CLEVRData(task="count_all")
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_retinopathy(root):
    ds = RetinopathyData(config="btgraham-300")
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_dmlab(root):
    ds = DmlabData()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_dspritespos(root):
    ds = DSpritesData(predicted_attribute="label_x_position", num_classes=16)
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_dspritesori(root):
    ds = DSpritesData(predicted_attribute="label_orientation", num_classes=16)
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_dtd(root):
    ds = DTDData()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_eurosat(root):
    ds = EurosatData()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_kitti(root):
    ds = KittiData(task="closest_vehicle_distance")
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_flowers(root):
    ds = OxfordFlowers102Data()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_pets(root):
    ds = OxfordIIITPetData()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_camelyon(root):
    ds = PatchCamelyonData()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_resisc(root):
    ds = Resisc45Data()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_smallnorbazi(root):
    ds = SmallNORBData(predicted_attribute="label_azimuth")
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_smallnorbele(root):
    ds = SmallNORBData(predicted_attribute="label_elevation")
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_sun397(root):
    ds = Sun397Data()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def vtab_svhn(root):
    ds = SvhnData()
    trainset = VtabWrapper(dataset=ds, split="train800val200")
    return trainset, None


@_add_dataset
def letters(root):
    from torchvision.datasets import EMNIST

    transform = transforms.Compose(
        [
            lambda x: x.convert("RGB"),
            transforms.Resize(224),
            transforms.ToTensor(),
            # transforms.Normalize((0.5, 0.5, 0.5), (1., 1., 1.)),
        ]
    )
    trainset = EMNIST(
        root, train=True, split="letters", transform=transform, download=True
    )
    testset = EMNIST(root, train=False, split="letters", transform=transform)
    return trainset, testset


@_add_dataset
def kmnist(root):
    from torchvision.datasets import KMNIST

    transform = transforms.Compose(
        [
            lambda x: x.convert("RGB"),
            transforms.Resize(224),
            transforms.ToTensor(),
        ]
    )
    trainset = KMNIST(root, train=True, transform=transform, download=True)
    testset = KMNIST(root, train=False, transform=transform)
    return trainset, testset


@_add_dataset
def stl10(root):
    from torchvision.datasets import STL10

    transform = transforms.Compose(
        [
            transforms.Resize(224),
            transforms.ToTensor(),
            transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
        ]
    )
    trainset = STL10(root, split="train", transform=transform, download=True)
    testset = STL10(root, split="test", transform=transform)
    trainset.targets = trainset.labels
    testset.targets = testset.labels
    return trainset, testset


def get_dataset(root, config=None):
    return _DATASETS[config.name](os.path.expanduser(root), config)
