import numpy as np
import torch
from torchvision import datasets, transforms
from torch.utils.data.sampler import SubsetRandomSampler
from constants import *
from smallNorb import SmallNORB


def build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset):
    # Compute validation split
    train_size = len(train_dataset)
    indices = list(range(train_size))
    split = int(np.floor(valid_size * train_size))
    np.random.shuffle(indices)
    train_idx, valid_idx = indices[split:], indices[:split]
    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)

    # Create data loaders
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size,
                                               sampler=train_sampler)
    valid_loader = torch.utils.data.DataLoader(valid_dataset,
                                               batch_size=batch_size,
                                               sampler=valid_sampler)
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=batch_size,
                                              shuffle=False)
    return train_loader, valid_loader, test_loader


def load_mnist(batch_size, valid_size=0.1):
    train_transform = transforms.Compose([
        transforms.RandomAffine(0, translate=[0.08, 0.08]),
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    valid_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = datasets.MNIST('../data',
                                   train=True,
                                   download=True,
                                   transform=train_transform)
    valid_dataset = datasets.MNIST('../data',
                                   train=True,
                                   download=True,
                                   transform=valid_transform)
    test_dataset = datasets.MNIST('../data',
                                  train=False,
                                  download=True,
                                  transform=test_transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)


def load_smallnorb(batch_size, valid_size=0.1):
    path = SMALL_NORB_PATH
    train_transform = transforms.Compose([
        transforms.Resize(48),
        transforms.RandomCrop(32),
        transforms.ColorJitter(brightness=32. / 255, contrast=0.5),
        transforms.ToTensor(),
        transforms.Normalize((0.0,), (0.3081,))
    ])
    valid_transform = transforms.Compose([
        transforms.Resize(48),
        transforms.CenterCrop(32),
        transforms.ToTensor(),
        transforms.Normalize((0.,), (0.3081,))
    ])
    test_transform = transforms.Compose([
        transforms.Resize(48),
        transforms.CenterCrop(32),
        transforms.ToTensor(),
        transforms.Normalize((0.,), (0.3081,))
    ])

    train_dataset = SmallNORB(path, train=True, download=True, transform=train_transform)
    valid_dataset = SmallNORB(path, train=True, download=True, transform=valid_transform)
    test_dataset = SmallNORB(path, train=False, transform=test_transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)


def load_cifar10(batch_size, valid_size=0.1):
    train_transform = transforms.Compose([
        transforms.ColorJitter(brightness=63. / 255, contrast=0.8),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0, 0, 0), (0.5, 0.5, 0.5))
    ])
    valid_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0, 0, 0), (0.5, 0.5, 0.5))
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0, 0, 0), (0.5, 0.5, 0.5))
    ])
    train_dataset = datasets.CIFAR10('../data',
                                     train=True,
                                     download=True,
                                     transform=train_transform)
    valid_dataset = datasets.CIFAR10('../data',
                                     train=True,
                                     download=True,
                                     transform=valid_transform)
    test_dataset = datasets.CIFAR10('../data',
                                    train=False,
                                    download=True,
                                    transform=test_transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)


def load_trashnet(batch_size, valid_size=0.1):
    transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
    ])

    train_dataset = datasets.ImageFolder("../trashnet/augmented/", transform=transform)
    valid_dataset = datasets.ImageFolder("../trashnet/dataset-resized/", transform=transform)
    test_dataset = datasets.ImageFolder("../trashnet/augmented/", transform=transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)

import numpy as np
import torch
from torchvision import datasets, transforms
from torch.utils.data.sampler import SubsetRandomSampler
from constants import *
from smallNorb import SmallNORB


def build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset):
    # Compute validation split
    train_size = len(train_dataset)
    indices = list(range(train_size))
    split = int(np.floor(valid_size * train_size))
    np.random.shuffle(indices)
    train_idx, valid_idx = indices[split:], indices[:split]
    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)

    # Create data loaders
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size,
                                               sampler=train_sampler)
    valid_loader = torch.utils.data.DataLoader(valid_dataset,
                                               batch_size=batch_size,
                                               sampler=valid_sampler)
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=batch_size,
                                              shuffle=False)
    return train_loader, valid_loader, test_loader


def load_mnist(batch_size, valid_size=0.1):
    train_transform = transforms.Compose([
        transforms.RandomAffine(0, translate=[0.08, 0.08]),
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    valid_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = datasets.MNIST('../data',
                                   train=True,
                                   download=True,
                                   transform=train_transform)
    valid_dataset = datasets.MNIST('../data',
                                   train=True,
                                   download=True,
                                   transform=valid_transform)
    test_dataset = datasets.MNIST('../data',
                                  train=False,
                                  download=True,
                                  transform=test_transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)


def load_smallnorb(batch_size, valid_size=0.1):
    path = SMALL_NORB_PATH
    train_transform = transforms.Compose([
        transforms.Resize(48),
        transforms.RandomCrop(32),
        transforms.ColorJitter(brightness=32. / 255, contrast=0.5),
        transforms.ToTensor(),
        transforms.Normalize((0.0,), (0.3081,))
    ])
    valid_transform = transforms.Compose([
        transforms.Resize(48),
        transforms.CenterCrop(32),
        transforms.ToTensor(),
        transforms.Normalize((0.,), (0.3081,))
    ])
    test_transform = transforms.Compose([
        transforms.Resize(48),
        transforms.CenterCrop(32),
        transforms.ToTensor(),
        transforms.Normalize((0.,), (0.3081,))
    ])

    train_dataset = SmallNORB(path, train=True, download=True, transform=train_transform)
    valid_dataset = SmallNORB(path, train=True, download=True, transform=valid_transform)
    test_dataset = SmallNORB(path, train=False, transform=test_transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)


def load_cifar10(batch_size, valid_size=0.1):
    train_transform = transforms.Compose([
        transforms.ColorJitter(brightness=63. / 255, contrast=0.8),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0, 0, 0), (0.5, 0.5, 0.5))
    ])
    valid_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0, 0, 0), (0.5, 0.5, 0.5))
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0, 0, 0), (0.5, 0.5, 0.5))
    ])
    train_dataset = datasets.CIFAR10('../data',
                                     train=True,
                                     download=True,
                                     transform=train_transform)
    valid_dataset = datasets.CIFAR10('../data',
                                     train=True,
                                     download=True,
                                     transform=valid_transform)
    test_dataset = datasets.CIFAR10('../data',
                                    train=False,
                                    download=True,
                                    transform=test_transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)


def load_trashnet(batch_size, valid_size=0.1):
    transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
    ])

    train_dataset = datasets.ImageFolder("../trashnet/augmented/", transform=transform)
    valid_dataset = datasets.ImageFolder("../trashnet/dataset-resized/", transform=transform)
    test_dataset = datasets.ImageFolder("../trashnet/augmented/", transform=transform)

    return build_data_loaders(batch_size, valid_size, train_dataset, valid_dataset, test_dataset)

# Compute validation split # Create data loaders