from __future__ import print_function

import os
import torch
import socket
import numpy as np
import torchvision
from helper.util import AugDataset
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from PIL import Image
from helper.util import plot_tensor

"""
mean = {
    'cifar100': (0.5071, 0.4867, 0.4408),
}

std = {
    'cifar100': (0.2675, 0.2565, 0.2761),
}
"""


class Datasubset(torch.utils.data.Dataset):
    def __init__(self, dataset, len):
        self.dataset = dataset
        self.len = len

    def __getitem__(self, i):
        return self.dataset[i % self.len]

    def __len__(self):
        return self.len  # max(len(d) for d in self.datasets)


class ConcatDataset(torch.utils.data.Dataset):
    def __init__(self, *datasets, len, opt):
        self.datasets = datasets
        self.len = len
        self.opt = opt

    def __getitem__(self, i):
        res = []
        for j, d in enumerate(self.datasets):
            l = min(len(d), self.len)
            # print(l)

            if (self.opt.aug_type == 'mixup' or self.opt.aug_type == 'cutmix') and j == 1:
                i += self.opt.batch_size * 10
            res.append(d[i % l])

        # return tuple(d[i % len(d)] for d in self.datasets)
        return tuple(res)

    def __len__(self):
        return self.len  # max(len(d) for d in self.datasets)


class DatasetMasked(torch.utils.data.Dataset):
    def __init__(self, dataset, opt):
        self.dataset = dataset
        self.len = len(dataset)
        self.opt = opt

    def __getitem__(self, i):
        res = self.dataset[i]  # 3x32x32
        mask = torch.zeros([32, 32]).type(torch.FloatTensor)

        # set a random square area in the mask to one
        lambda_aug = np.random.beta(self.opt.aug_alpha, self.opt.aug_alpha)

        s_w = int(32 * np.sqrt(1 - lambda_aug))
        if s_w == 32:
            s_w = 31



        rand = torch.randint(0, 32 - s_w, size=[2])
        mask[int(rand[0]):int(rand[0]) + s_w, int(rand[1]):int(rand[1]) + s_w] = 1
        mask = mask.view(1, 32, 32)
        # res.append(mask)
        return res + tuple(mask)  # append the mask to  output

    def __len__(self):
        return self.len  # max(len(d) for d in self.datasets)


def get_data_folder():
    """
    return server-dependent path to store the data
    """
    hostname = socket.gethostname()
    if hostname.startswith('visiongpu'):
        data_folder = '/data/vision/phillipi/rep-learn/datasets'
    elif hostname.startswith('yonglong-home'):
        data_folder = '/home/yonglong/Data/data'
    else:
        data_folder = './data/'

    if not os.path.isdir(data_folder):
        os.makedirs(data_folder)

    return data_folder


class CIFAR100Instance(datasets.CIFAR100):
    """CIFAR100Instance Dataset.
    """

    def __getitem__(self, index):

        # if torch.__version__[0] == '0':

        if self.train:
            img, target = self.train_data[index], self.train_labels[index]
        else:
            img, target = self.test_data[index], self.test_labels[index]
        # else:
        #     img, target = self.data[index], self.targets[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target, index


def get_cifar100_dataloaders(opt, is_instance=False):
    """
    cifar 100
    """
    data_folder = get_data_folder()

    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])

    if is_instance:
        train_set = CIFAR100Instance(root=data_folder,
                                     download=True,
                                     train=True,
                                     transform=train_transform)
        n_data = len(train_set)
    else:
        train_set = datasets.CIFAR100(root=data_folder,
                                      download=True,
                                      train=True,
                                      transform=train_transform)

    # prepare the augmentation dataset
    if opt.aug_type is not None:
        train_transform_aug = transforms.Compose([
            transforms.RandomCrop(32, padding=2),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
        ])

        if opt.aug_type == 'supermix':
            train_set_aug = torchvision.datasets.ImageFolder(
                root=opt.aug_dir,
                transform=train_transform_aug
            )
            if opt.aug_size == -1:
                # max(len(d) for d in self.datasets)
                opt.aug_size = max(len(train_set), len(train_set_aug))
                opt.aug_size -= opt.aug_size % 100
        elif opt.aug_type == 'mixup':
            train_set_aug = datasets.CIFAR100(root=data_folder,
                                              download=True,
                                              train=True,
                                              transform=train_transform)
            opt.aug_size = 50000
        elif opt.aug_type == 'cutmix':
            train_set_aug = datasets.CIFAR100(root=data_folder,
                                              download=True,
                                              train=True,
                                              transform=train_transform)
            # generate masks for the data
            train_set_aug = DatasetMasked(train_set_aug, opt=opt)
            opt.aug_size = 50000
        #
        # img = train_set.__getitem__(798)
        # plot_tensor([img[0]])
        # exit()

        train_loader = torch.utils.data.DataLoader(
            ConcatDataset(train_set, train_set_aug, len=opt.aug_size, opt=opt), batch_size=opt.batch_size, shuffle=True,
            num_workers=opt.num_workers, pin_memory=True)
    else:
        train_loader = DataLoader(train_set,
                                  batch_size=opt.batch_size,
                                  shuffle=True,
                                  num_workers=opt.num_workers)
        opt.aug_size = 50000

    test_set = datasets.CIFAR100(root=data_folder,
                                 download=True,
                                 train=False,
                                 transform=test_transform)
    test_loader = DataLoader(test_set,
                             batch_size=opt.batch_size,
                             shuffle=False,
                             num_workers=opt.num_workers)

    print("size of the augment set: ", opt.aug_size)

    if is_instance:
        return train_loader, test_loader, n_data
    else:
        return train_loader, test_loader


class CIFAR100InstanceSample(datasets.CIFAR100):
    """
    CIFAR100Instance+Sample Dataset
    """

    def __init__(self, root, train=True,
                 transform=None, target_transform=None,
                 download=False, k=4096, mode='exact', is_sample=True, percent=1.0):
        super().__init__(root=root, train=train, download=download,
                         transform=transform, target_transform=target_transform)
        self.k = k
        self.mode = mode
        self.is_sample = is_sample

        num_classes = 100
        if self.train:
            num_samples = len(self.train_data)
            label = self.train_labels
        else:
            num_samples = len(self.test_data)
            label = self.test_labels

        self.cls_positive = [[] for i in range(num_classes)]
        for i in range(num_samples):
            self.cls_positive[label[i]].append(i)

        self.cls_negative = [[] for i in range(num_classes)]
        for i in range(num_classes):
            for j in range(num_classes):
                if j == i:
                    continue
                self.cls_negative[i].extend(self.cls_positive[j])

        self.cls_positive = [np.asarray(self.cls_positive[i]) for i in range(num_classes)]
        self.cls_negative = [np.asarray(self.cls_negative[i]) for i in range(num_classes)]

        if 0 < percent < 1:
            n = int(len(self.cls_negative[0]) * percent)
            self.cls_negative = [np.random.permutation(self.cls_negative[i])[0:n]
                                 for i in range(num_classes)]

        self.cls_positive = np.asarray(self.cls_positive)
        self.cls_negative = np.asarray(self.cls_negative)

    def __getitem__(self, index):
        if self.train:
            img, target = self.train_data[index], self.train_labels[index]
        else:
            img, target = self.test_data[index], self.test_labels[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        if not self.is_sample:
            # directly return
            return img, target, index
        else:
            # sample contrastive examples
            if self.mode == 'exact':
                pos_idx = index
            elif self.mode == 'relax':
                pos_idx = np.random.choice(self.cls_positive[target], 1)
                pos_idx = pos_idx[0]
            else:
                raise NotImplementedError(self.mode)
            replace = True if self.k > len(self.cls_negative[target]) else False
            neg_idx = np.random.choice(self.cls_negative[target], self.k, replace=replace)
            sample_idx = np.hstack((np.asarray([pos_idx]), neg_idx))
            return img, target, index, sample_idx


def get_cifar100_dataloaders_sample(batch_size=128, num_workers=8, k=4096, mode='exact',
                                    is_sample=True, percent=1.0):
    """
    cifar 100
    """
    data_folder = get_data_folder()

    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])

    train_set = CIFAR100InstanceSample(root=data_folder,
                                       download=True,
                                       train=True,
                                       transform=train_transform,
                                       k=k,
                                       mode=mode,
                                       is_sample=is_sample,
                                       percent=percent)
    n_data = len(train_set)
    train_loader = DataLoader(train_set,
                              batch_size=batch_size,
                              shuffle=True,
                              num_workers=num_workers)

    test_set = datasets.CIFAR100(root=data_folder,
                                 download=True,
                                 train=False,
                                 transform=test_transform)
    test_loader = DataLoader(test_set,
                             batch_size=int(batch_size / 2),
                             shuffle=False,
                             num_workers=int(num_workers / 2))

    return train_loader, test_loader, n_data
