import torch
import torchvision
import torch.utils.data as data
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import numpy as np


def compute_mean_std(dataset):
    """
    compute the mean and std of cifar100 dataset
    """
    data_r = np.dstack([dataset[i][0] for i in range(len(dataset))])
    data_g = np.dstack([dataset[i][1] for i in range(len(dataset))])
    data_b = np.dstack([dataset[i][2] for i in range(len(dataset))])
    '''
    data_r = numpy.dstack([cifar100_dataset[i][1][:, :, 0] for i in range(len(cifar100_dataset))])
    data_g = numpy.dstack([cifar100_dataset[i][1][:, :, 1] for i in range(len(cifar100_dataset))])
    data_b = numpy.dstack([cifar100_dataset[i][1][:, :, 2] for i in range(len(cifar100_dataset))])
    '''
    mean = np.mean(data_r), np.mean(data_g), np.mean(data_b)
    std = np.std(data_r), np.std(data_g), np.std(data_b)

    return mean, std


def get_mnist_dataloader(data_dir, batch_size, num_workers, spilt='train'):
    if spilt == 'train':
        train_dataset = torchvision.datasets.MNIST(
            root=data_dir,
            train=True,
            transform=transforms.ToTensor(),
            download=True
        )

        train_data_loader = data.DataLoader(
            dataset=train_dataset,
            batch_size=batch_size,
            shuffle=True,
            drop_last=True,
            num_workers=num_workers,
            pin_memory=True
        )
        return train_data_loader

    if spilt == 'test':
        test_dataset = torchvision.datasets.MNIST(
            root=data_dir,
            train=False,
            transform=transforms.ToTensor(),
            download=True
        )

        test_data_loader = data.DataLoader(
            dataset=test_dataset,
            batch_size=batch_size,
            shuffle=False,
            drop_last=False,
            num_workers=num_workers,
            pin_memory=True
        )
        return test_data_loader


def get_cifar10_dataloader(data_dir, batch_size, num_workers, mean=(0.49186853, 0.48265487, 0.44717678), std=(0.24697104, 0.24338877, 0.26159248), spilt='train'):
    if spilt == 'train':
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(15),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
        train_dataset = torchvision.datasets.CIFAR10(
            root=data_dir,
            train=True,
            download=False,
            transform=transform_train
        )
        train_data_loader = DataLoader(
            train_dataset,
            shuffle=True,
            num_workers=num_workers,
            batch_size=batch_size
        )
        return train_data_loader

    if spilt == 'test':
        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
        test_dataset = torchvision.datasets.CIFAR10(
            root=data_dir,
            train=False,
            download=False,
            transform=transform_test
        )
        test_data_loader = DataLoader(
            test_dataset,
            shuffle=False,
            num_workers=num_workers,
            batch_size=batch_size
        )

        return test_data_loader


def get_cifar100_dataloader(data_dir, batch_size, num_workers, shuffle=True, mean=(0.5070751592371323, 0.48654887331495095, 0.4409178433670343), std=(0.2673342858792401, 0.2564384629170883, 0.27615047132568404), spilt='train'):
    if spilt == 'train':
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(15),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
        train_dataset = torchvision.datasets.CIFAR100(
            root=data_dir,
            train=True,
            download=False,
            transform=transform_train
        )
        train_data_loader = DataLoader(
            train_dataset,
            shuffle=shuffle,
            num_workers=num_workers,
            batch_size=batch_size
        )
        return train_data_loader

    if spilt == 'test':
        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
        test_dataset = torchvision.datasets.CIFAR100(
            root=data_dir,
            train=False,
            download=False,
            transform=transform_test
        )
        test_data_loader = DataLoader(
            test_dataset,
            shuffle=shuffle,
            num_workers=num_workers,
            batch_size=batch_size
        )

        return test_data_loader


def get_imagenet_dataloader(data_dir, batch_size, num_workers, shuffle=True, spilt='train'):
    transform_train = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    transform_test = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    if spilt == 'train':
        train_dataset = torchvision.datasets.ImageNet(
            root=data_dir,
            split='train',
            # train=True,
            # download=False,
            transform=transform_train
        )
        train_data_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=num_workers
        )
        return train_data_loader

    if spilt == 'test':
        test_datatset = torchvision.datasets.ImageNet(
            root=data_dir,
            split='val',
            # train=False,
            # download=False,
            transform=transform_test
        )
        test_data_loader = torch.utils.data.DataLoader(
            test_datatset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers
        )

        return test_data_loader
