import numpy as np
from random import sample 
from PIL import Image

import torch
import torchvision.datasets as datasets
from torch.distributions.multivariate_normal import MultivariateNormal
from torch.distributions.normal import Normal

from sklearn.utils import shuffle



class CIFAR10Noise(datasets.CIFAR10):
    def __init__(self, corrupt_input=0.0, corrupt_class=1, perct_original=1.0, **kwargs):
        """
        corrupt_input: Percentage of data we add in each class
        corrupt_class: Amount of classes (of cifar100) where we take data for each class
        """
        super(CIFAR10Noise, self).__init__(**kwargs)

        if perct_original < 1:
            ind_train = sample(list(range(len(self))), int(len(self)*perct_original))
            self.data = self.data[ind_train]
            print(self.data.shape)
            self.targets = list(torch.tensor(self.targets)[ind_train])

        if corrupt_input > 0:
            dataset = datasets.CIFAR100(**kwargs)
            d, t = self.get_data_cifar100(dataset.data, dataset.targets, corrupt_input, corrupt_class)
            self.corrupt_input(d, t, corrupt_input, corrupt_class)
        else:
            self.corrupted = torch.zeros(len(self.targets))

    def get_data_cifar100(self, data, targets, corrupt_input, corrupt_class):
        if corrupt_class > 10 or corrupt_class < 1:
            raise("Number of corrupted classes must be less than 11 and more than 0")

        sample_classes = sample(list(set(targets)), corrupt_class*10)

        d = []
        t = []
        for i in range(0,corrupt_class*10,corrupt_class):
            total = 0
            for j in range(corrupt_class):
                c = sample_classes[i+j]
                if j < (corrupt_class-1):
                    n = int(((6000*corrupt_input)/corrupt_class))
                else:
                    n = int(6000*corrupt_input) - total
                total += n

                sample_data = sample(list(torch.nonzero((torch.tensor(targets) == c), as_tuple=False)), n)
                d.extend(data[sample_data])
                for k in range(n):
                    t.append(int(i/corrupt_class))
                    self.targets.append(int(i/corrupt_class))
                #t.extend([int(i/corrupt_class)]*n)

        return np.array(d), t

    def corrupt_input(self, data, target, corrupt_input, corrupt_class):
        self.corrupted = torch.cat((torch.zeros(len(self.data)), torch.ones(len(target))))
        self.data = np.append(self.data, data, 0)
        #self.targets.extend(target)

        self.targets = torch.tensor(self.targets)

        print("Total amount of data: {}. Each class have {} new element of {} different classes of CIFAR10".format(len(self.targets), int(6000*corrupt_input), corrupt_class))

    def __getitem__(self, index):
        img, target, corrupt = self.data[index], self.targets[index], bool(self.corrupted[index])

        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, corrupt







class CIFAR10RandomLabels(datasets.CIFAR10):
    def __init__(self, corrupt_prob=0.0, perct_original=1.0, num_classes=10, **kwargs):
        super(CIFAR10RandomLabels, self).__init__(**kwargs)
        self.n_classes = num_classes

        if perct_original < 1:
            ind_train = sample(list(range(len(self))), int(len(self)*perct_original))
            self.data = self.data[ind_train]
            self.targets = list(torch.tensor(self.targets)[ind_train])

        if corrupt_prob > 0:
            self.corrupt_labels(corrupt_prob)
        else:
            self.corrupted = torch.zeros(len(self.targets))

        self.targets = torch.tensor(self.targets)

    def corrupt_labels(self, corrupt_prob):
        # self.targets = [ x if np.random.uniform()>corrupt_prob else np.random.randint(0,self.n_classes)
        #                               for x in self.targets]

        temp = []
        self.corrupted = torch.zeros(len(self.targets))
        for i,x in enumerate(self.targets):
            if np.random.uniform()>corrupt_prob:
                temp.append(x)
            else:
                temp.append(np.random.randint(0,self.n_classes))
                self.corrupted[i] = 1

        self.targets = temp

    def __getitem__(self, index):
        img, target, corrupt = self.data[index], self.targets[index], bool(self.corrupted[index])

        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, corrupt

class CIFAR10RandomInputs(datasets.CIFAR10):
    def __init__(self, corrupt_prob=0.0, num_classes=10, **kwargs):
        super(CIFAR10RandomInputs, self).__init__(**kwargs)
        self.n_classes = num_classes
        if corrupt_prob > 0:
            self.get_statistic()
            self.corrupt_inputs(corrupt_prob)
        else:
            self.corrupted = torch.zeros(len(self.targets))

    def corrupt_inputs(self, corrupt_prob):
        idx = list(range(32 * 32 * 3))
        #perm_pixel = shuffle(idx, random_state=42 * 100 + 1)
        #self.data = [ x if np.random.uniform()>corrupt_prob else x[perm_pixel]
        #                               for i,x in enumerate(self.data)]

        data = []
        for i,x in enumerate(self.data):
            if np.random.uniform()<corrupt_prob:
                perm_pixel = shuffle(idx, random_state=42 * 100 + i)
                data.append(x.reshape(-1)[perm_pixel].reshape(32,32,3))
                self.targets.append(self.targets[i])

        self.corrupted = torch.cat([torch.zeros(self.data.shape[0]),torch.ones(len(data))])
        self.data = np.concatenate([self.data, data])

        #sample_data = int(((torch.tensor(self.targets) == self.targets[0]).nonzero()).size(0)*corrupt_prob)
        #data = []
        #y = []
        #for i in list(set(self.targets)):
        #    data.extend(self.get_sample(i, sample_data))
        #    self.targets.extend([i]*sample_data)

        #self.corrupted = torch.cat([torch.zeros(self.data.shape[0]),torch.ones(sample_data*10)])
        #self.data = np.concatenate([self.data, data])

    def get_sample(self, x, t, num_elem):
        #return np.random.normal(self.normales[t]['mean'], self.normales[t]['std'], (32,32,3)).astype(np.uint8)
        #955return np.random.normal(self.normales[t]['mean'], self.normales[t]['std'], (num_elem,32,32,3)).astype(np.uint8)
        #956return ((self.normales[t]['image'] + np.random.normal(self.normales[t]['mean'], self.normales[t]['std'], (num_elem,32,32,3)))/2).astype(np.uint8)
        return ((x + np.random.normal(np.zeros_like(self.normales[t]['mean']), self.normales[t]['std'], (num_elem,32,32,3)))/2).astype(np.uint8)
        #957return np.random.normal(self.normales[t]['mean'], self.normales[t]['std'], (num_elem,32*32*3)).reshape(num_elem,32,32,3).astype(np.uint8)
        #958return (self.normales[t]['image'] + np.random.normal(np.zeros_like(self.normales[t]['mean']), self.normales[t]['std'], (num_elem,32,32,3))).astype(np.uint8)

    def get_statistic(self):
        self.normales = {}
        m = self.data.mean()
        std = self.data.std()
        for elem in list(set(self.targets)):
            m = self.data[torch.tensor(self.targets) == elem].mean(axis=(0,1,2))#.reshape(-1)
            std = self.data[torch.tensor(self.targets) == elem].std(axis=(0,1,2))#.reshape(-1)
            img = self.data[torch.tensor(self.targets) == elem].std(axis=(0))
            self.normales[elem] = {'mean': m, 'std': std, 'image': img.reshape(1,32,32,3)}

    def __getitem__(self, index):
        img, target, corrupt = self.data[index], self.targets[index], bool(self.corrupted[index])

        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, corrupt





class CIFAR100RandomLabels(datasets.CIFAR100):
    def __init__(self, corrupt_prob=0.0, num_classes=10, **kwargs):
        super(CIFAR100RandomLabels, self).__init__(**kwargs)
        self.n_classes = num_classes
        if corrupt_prob > 0:
            self.corrupt_labels(corrupt_prob)
        else:
            self.corrupted = torch.zeros(len(self.targets))

    def corrupt_labels(self, corrupt_prob):
        # self.targets = [ x if np.random.uniform()>corrupt_prob else np.random.randint(0,self.n_classes)
        #                               for x in self.targets]

        temp = []
        self.corrupted = torch.zeros(len(self.targets))
        for i,x in enumerate(self.targets):
            if np.random.uniform()>corrupt_prob:
                temp.append(x)
            else:
                temp.append(np.random.randint(0,self.n_classes))
                self.corrupted[i] = 1

        self.targets = temp

    def __getitem__(self, index):
        img, target, corrupt = self.data[index], self.targets[index], bool(self.corrupted[index])

        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, corrupt


class CIFAR100RandomInputs(datasets.CIFAR100):
    def __init__(self, corrupt_prob=0.0, num_classes=10, **kwargs):
        super(CIFAR100RandomInputs, self).__init__(**kwargs)
        self.n_classes = num_classes
        if corrupt_prob > 0:
            self.get_statistic()
            self.corrupt_inputs(corrupt_prob)
        else:
            self.corrupted = torch.zeros(len(self.targets))

    def corrupt_inputs(self, corrupt_prob):
        # self.data = [ x if np.random.uniform()>corrupt_prob else x + self.get_sample(self.targets[i])
        #                               for i,x in enumerate(self.data)]

        temp = []
        self.corrupted = torch.zeros(len(self.targets))
        for i,x in enumerate(self.data):
            if np.random.uniform()>corrupt_prob:
                temp.append(x)
            else:
                temp.append(x + self.get_sample(self.targets[i]))
                self.corrupted[i] = 1

        self.data = temp

    def get_sample(self, t):
        return np.random.normal(self.normales[t]['mean'], self.normales[t]['std'], (32,32,3)).astype(np.uint8)

    def get_statistic(self):
        self.normales = {}
        m = self.data.mean()
        std = self.data.std()
        for elem in list(set(self.targets)):
            # m = self.data[torch.tensor(self.targets) == elem].mean()
            # std = self.data[torch.tensor(self.targets) == elem].std()
            self.normales[elem] = {'mean': m, 'std': std}

    def __getitem__(self, index):
        img, target, corrupt = self.data[index], self.targets[index], bool(self.corrupted[index])

        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, corrupt




class MNISTRandomLabels(datasets.MNIST):
    def __init__(self, corrupt_prob=0.0, num_classes=10, **kwargs):
        super(MNISTRandomLabels, self).__init__(**kwargs)
        self.n_classes = num_classes
        self.corrupted = torch.zeros(len(self.targets))
        if corrupt_prob > 0:
            self.corrupt_labels(corrupt_prob)
            

    def corrupt_labels(self, corrupt_prob):
        # self.targets = [ x if np.random.uniform()>corrupt_prob else np.random.randint(0,self.n_classes)
        #                               for x in self.targets]

        temp = []
        self.corrupted = torch.zeros(len(self.targets))
        for i,x in enumerate(self.targets):
            if np.random.uniform()>corrupt_prob:
                temp.append(x)
            else:
                temp.append(np.random.randint(0,self.n_classes))
                self.corrupted[i] = 1

        self.targets = temp


    def __getitem__(self, index):
        img, target, corrupt = self.data[index], int(self.targets[index]), bool(self.corrupted[index])
        
        img = Image.fromarray(img.numpy(), mode='L')

        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, corrupt


class MNISTRandomInputs(datasets.MNIST):
    def __init__(self, corrupt_prob=0.0, num_classes=10, **kwargs):
        super(MNISTRandomInputs, self).__init__(**kwargs)
        self.n_classes = num_classes
        self.to_image = True
        self.corrupted = torch.zeros(len(self.targets))
        if corrupt_prob > 0:
            self.corrupt_inputs(corrupt_prob)
            

    def __getitem__(self, index):
        img, target, corrupt = self.data[index], int(self.targets[index]), bool(self.corrupted[index])

        if self.to_image:
            #if len(img.size(0)) == 3:
            img = img.squeeze()
            img = Image.fromarray(img.numpy(), mode='L')

        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, corrupt

    def corrupt_inputs(self, corrupt_prob):
        data = []
        t = []
        for i,x in enumerate(self.data):
            if np.random.uniform()<corrupt_prob:
                idx = list(range(28*28))
                perm_pixel = shuffle(idx, random_state=42 * 100 + i)
                data.append(x.reshape(-1)[perm_pixel].reshape(28,28))
                t.append(self.targets[i])

        if len(data) > 0:
            self.corrupted = torch.cat([torch.zeros(self.data.shape[0]),torch.ones(len(data))])
            self.data = torch.cat([self.data, torch.stack(data)])
            self.targets = torch.cat([self.targets, torch.tensor(t)])
