import torch
import torchvision
import torchvision.transforms as transforms
from os.path import join
from random import sample 
import numpy as np
from dataloaders.sample import MNISTRandomLabels, MNISTRandomInputs, CIFAR10RandomLabels, \
                                CIFAR10RandomInputs, CIFAR100RandomLabels, CIFAR100RandomInputs, \
                                CIFAR10Noise



def get_datasets(dataset, path, noise, noise_data, noise_input, noise_classes, perct_original, model=None, normalize=True):
    if dataset == 'cifar10':
        trans = [
            transforms.Resize(32),
            transforms.CenterCrop(28),
            transforms.ToTensor()]

        if normalize:
            trans += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
            
        transform_train = transforms.Compose(trans)
        transform_test = transforms.Compose(trans)

    elif dataset == 'mnist':
        trans = [ transforms.ToTensor()]
        if normalize:
            trans += [transforms.Normalize((0.1307,), (0.3081,))]
        transform_train = transforms.Compose(trans)
        transform_test = transforms.Compose(trans)

    elif model == 'resnet-18':
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])

        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])

    if noise == 'input':
        if dataset == 'cifar10':
            if noise_data == 'cifar100':
                trainset = CIFAR10Noise(noise_input, noise_classes, perct_original, root=path, train=True,
                                                download=True, transform=transform_train)
            else:
                trainset = CIFAR10RandomInputs(noise_input, perct_original, root=path, train=True,
                                                download=True, transform=transform_train)
        elif dataset == 'mnist':
            trainset = MNISTRandomInputs(corrupt_prob=noise_input, root='../data', train=True, download=True,
                       transform=transform_train)
    else:
        trainset = CIFAR10RandomLabels(noise_input, perct_original, root=path, train=True,
                                             download=True, transform=transform_train)

    if dataset == 'cifar10':
        testset = CIFAR10Noise(0.0, noise_classes, root=path, train=False,
                                           download=True, transform=transform_test)
    elif dataset == 'mnist':
        testset = MNISTRandomInputs(corrupt_prob=0.0, root='../data', train=False, transform=transform_test)

    return trainset, testset

def get_dataloader(dataset, path='./data', batch_size=64, noise_input=0.0, noise_classes=1, perct_original=1.0, noise='input', noise_data='cifar100', model = None, normalize=True):
    trainset, testset = get_datasets(dataset, path, noise, noise_data, noise_input, noise_classes, perct_original, normalize=normalize)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size[0],
                                              shuffle=True, num_workers=1)

    trainloader_test = {'noise': None, 'benign': None, 'total': None}
    trainloader_test['total'] = torch.utils.data.DataLoader(trainset, batch_size=max(batch_size[0],batch_size[1]),
                                              shuffle=False, num_workers=1)


    n_total = torch.nonzero((trainset.corrupted == 1), as_tuple=False).size(0)

    hard_metric = 'agr'
    filename = join(path, 'order_{}_{}_train.npy'.format(dataset, hard_metric))
    indices = np.load(filename)
    start_index = int(indices.shape[0]*0.2)
    ind_benign = sample(list(indices[start_index:]), min([1024, n_total]))
    #ind_benign = sample(list((trainset.corrupted == 0).nonzero()), 1024)
    ind_noise = sample(list(torch.nonzero((trainset.corrupted == 1), as_tuple=False)), min([1024, n_total]))
 
    if noise_input > 0:
        noise_set = torch.utils.data.Subset(trainset, ind_noise)
        trainloader_test['noise'] = torch.utils.data.DataLoader(noise_set, batch_size=batch_size[1],
                                                  shuffle=False, num_workers=1)
        benign_set = torch.utils.data.Subset(trainset, ind_benign)
        trainloader_test['benign'] = torch.utils.data.DataLoader(benign_set, batch_size=batch_size[1],
                                                  shuffle=False, num_workers=1)
    else:
        ind_benign = sample(list(torch.nonzero((trainset.corrupted == 0), as_tuple=False)), 1024)
        benign_set = torch.utils.data.Subset(trainset, ind_benign)
        trainloader_test['benign'] = torch.utils.data.DataLoader(benign_set, batch_size=batch_size[1],
                                                  shuffle=False, num_workers=1)

    testloader = {'total': torch.utils.data.DataLoader(testset, batch_size=max(batch_size[0],batch_size[1]),
                                             shuffle=False, num_workers=1)}

    grad_loaders = []
    for _ in range(10):
        if noise_input > 0:
            factor = (1-noise_input)
        else:
            factor = 1
            ind_noise = []

        ind_benign = sample(list(indices[start_index:]), int(1024*factor))
        # ind_benign = sample(list(indices[:start_index]), int(1024*factor))
        if noise_input > 0:
            ind_noise = sample(list(torch.nonzero((trainset.corrupted == 1), as_tuple=False)), int(1024*(noise_input)))
            ind_noise = [ elem.item() for elem in ind_noise ]

        grad_set = torch.utils.data.Subset(trainset, ind_benign + ind_noise)
        grad_loaders.append(torch.utils.data.DataLoader(grad_set, batch_size=batch_size[1],
                                                  shuffle=True, num_workers=1))

    return trainloader, trainloader_test, testloader, grad_loaders

def getDataset(dataset, args, normalize=True):
    if dataset == 'mnist':
        channels = 1

    elif dataset == 'cifar10':
        channels = 3

    return get_dataloader(dataset, batch_size=[args.batch_size_train,args.batch_size_test], 
                noise_input=args.noise_input, noise_classes=args.noise_classes, noise=args.noise, perct_original=args.perct_original_data,
                noise_data=args.noise_data, normalize=normalize, model=args.net), channels, 10

    #elif dataset == 'cifar100':
    #    return cifar100(batch_size=[args.batch_size_train,args.batch_size_test], 
    #            noise_input=args.noise_input, noise=args.noise), 3, 100

    #else:
    #    raise "Not Implemented"
