import torch
import numpy as np
import torch.utils.data as data
from torch.utils.data import Subset

from data.fast_mnist import create_MNIST_dataset
from data.ambiguous_mnist.ambiguous_mnist_dataset import AmbiguousMNIST


def get_train_valid_loader(root, batch_size, val_seed=1, val_size=0.1, **kwargs):

    error_msg = "[!] val_size should be in the range [0, 1]."
    assert (val_size >= 0) and (val_size <= 1), error_msg

    # load the dataset
    mnist_train_dataset, _ = create_MNIST_dataset()

    # AmbiguousMNIST does whiten the data itself
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_dataset = data.ConcatDataset(
        [mnist_train_dataset, AmbiguousMNIST(root=root, train=True, device=device),]
    )
    valid_dataset = data.ConcatDataset(
        [mnist_train_dataset, AmbiguousMNIST(root=root, train=True, device=device),]
    )

    num_train = len(train_dataset)
    indices = list(range(num_train))
    split = int(np.floor(val_size * num_train))

    np.random.seed(val_seed)
    np.random.shuffle(indices)

    train_idx, valid_idx = indices[split:], indices[:split]
    train_subset = Subset(train_dataset, train_idx)
    valid_subset = Subset(valid_dataset, valid_idx)

    train_loader = torch.utils.data.DataLoader(train_subset, batch_size=batch_size, num_workers=0, shuffle=True)

    valid_loader = torch.utils.data.DataLoader(valid_subset, batch_size=batch_size, num_workers=0, shuffle=False)

    return train_loader, valid_loader


def get_test_loader(root, batch_size, **kwargs):

    # load the dataset
    _, mnist_test_dataset = create_MNIST_dataset()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    test_dataset = data.ConcatDataset(
        [mnist_test_dataset, AmbiguousMNIST(root=root, train=False, device=device),]
    )

    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

    return test_loader

[((480, 502), 'data.fast_mnist.create_MNIST_dataset', 'create_MNIST_dataset', ([], {}), '()\n', (500, 502), False, 'from data.fast_mnist import create_MNIST_dataset\n'), ((1015, 1039), 'numpy.random.seed', 'np.random.seed', (['val_seed'], {}), '(val_seed)\n', (1029, 1039), True, 'import numpy as np\n'), ((1044, 1070), 'numpy.random.shuffle', 'np.random.shuffle', (['indices'], {}), '(indices)\n', (1061, 1070), True, 'import numpy as np\n'), ((1151, 1183), 'torch.utils.data.Subset', 'Subset', (['train_dataset', 'train_idx'], {}), '(train_dataset, train_idx)\n', (1157, 1183), False, 'from torch.utils.data import Subset\n'), ((1203, 1235), 'torch.utils.data.Subset', 'Subset', (['valid_dataset', 'valid_idx'], {}), '(valid_dataset, valid_idx)\n', (1209, 1235), False, 'from torch.utils.data import Subset\n'), ((1256, 1353), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', (['train_subset'], {'batch_size': 'batch_size', 'num_workers': '(0)', 'shuffle': '(True)'}), '(train_subset, batch_size=batch_size,\n    num_workers=0, shuffle=True)\n', (1283, 1353), False, 'import torch\n'), ((1370, 1468), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', (['valid_subset'], {'batch_size': 'batch_size', 'num_workers': '(0)', 'shuffle': '(False)'}), '(valid_subset, batch_size=batch_size,\n    num_workers=0, shuffle=False)\n', (1397, 1468), False, 'import torch\n'), ((1607, 1629), 'data.fast_mnist.create_MNIST_dataset', 'create_MNIST_dataset', ([], {}), '()\n', (1627, 1629), False, 'from data.fast_mnist import create_MNIST_dataset\n'), ((1854, 1953), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', (['test_dataset'], {'batch_size': 'batch_size', 'shuffle': '(False)', 'num_workers': '(0)'}), '(test_dataset, batch_size=batch_size, shuffle=\n    False, num_workers=0)\n', (1881, 1953), False, 'import torch\n'), ((978, 1008), 'numpy.floor', 'np.floor', (['(val_size * num_train)'], {}), '(val_size * num_train)\n', (986, 1008), True, 'import numpy as np\n'), ((589, 614), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (612, 614), False, 'import torch\n'), ((697, 749), 'data.ambiguous_mnist.ambiguous_mnist_dataset.AmbiguousMNIST', 'AmbiguousMNIST', ([], {'root': 'root', 'train': '(True)', 'device': 'device'}), '(root=root, train=True, device=device)\n', (711, 749), False, 'from data.ambiguous_mnist.ambiguous_mnist_dataset import AmbiguousMNIST\n'), ((828, 880), 'data.ambiguous_mnist.ambiguous_mnist_dataset.AmbiguousMNIST', 'AmbiguousMNIST', ([], {'root': 'root', 'train': '(True)', 'device': 'device'}), '(root=root, train=True, device=device)\n', (842, 880), False, 'from data.ambiguous_mnist.ambiguous_mnist_dataset import AmbiguousMNIST\n'), ((1667, 1692), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1690, 1692), False, 'import torch\n'), ((1773, 1826), 'data.ambiguous_mnist.ambiguous_mnist_dataset.AmbiguousMNIST', 'AmbiguousMNIST', ([], {'root': 'root', 'train': '(False)', 'device': 'device'}), '(root=root, train=False, device=device)\n', (1787, 1826), False, 'from data.ambiguous_mnist.ambiguous_mnist_dataset import AmbiguousMNIST\n')]