text:
import unittest

from os.path import expanduser
from avalanche.models import SimpleMLP
from torch.optim import SGD
from torch.nn import CrossEntropyLoss
from avalanche.training.strategies import Naive
from avalanche.benchmarks.generators import dataset_benchmark
import PIL
import torch
from PIL import ImageChops
from PIL.Image import Image
from torch import Tensor
from torch.utils.data import TensorDataset, Subset, ConcatDataset
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor, RandomCrop, ToPILImage, Compose, \
    Lambda, CenterCrop
from typing import List

from avalanche.benchmarks.scenarios.generic_benchmark_creation import \
    create_generic_benchmark_from_tensor_lists
from avalanche.benchmarks.utils import AvalancheDataset, \
    AvalancheSubset, AvalancheConcatDataset, AvalancheDatasetType, \
    AvalancheTensorDataset
from avalanche.benchmarks.utils.dataset_utils import ConstantSequence
from avalanche.training.utils import load_all_dataset
import random

import numpy as np

def pil_images_equal(img_a, img_b):
    diff = ImageChops.difference(img_a, img_b)

    return not diff.getbbox()

class AvalancheDatasetTests(unittest.TestCase):
    def test_mnist_no_transforms(self):
        dataset = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                        download=True)
        x, y = dataset[0]
        self.assertIsInstance(x, Image)
        self.assertEqual([x.width, x.height], [28, 28])
        self.assertIsInstance(y, int)

    def test_mnist_native_transforms(self):
        dataset = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                        download=True, transform=ToTensor())
        x, y = dataset[0]
        self.assertIsInstance(x, Tensor)
        self.assertEqual(x.shape, (1, 28, 28))
        self.assertIsInstance(y, int)

    def test_avalanche_dataset_transform(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[0]
        dataset = AvalancheDataset(dataset_mnist, transform=ToTensor())
        x2, y2, t2 = dataset[0]

        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, int)
        self.assertIsInstance(t2, int)
        self.assertEqual(0, t2)
        self.assertTrue(torch.equal(ToTensor()(x), x2))
        self.assertEqual(y, y2)

    def test_avalanche_dataset_slice(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x0, y0 = dataset_mnist[0]
        x1, y1 = dataset_mnist[1]
        dataset = AvalancheDataset(dataset_mnist, transform=ToTensor())
        x2, y2, t2 = dataset[:2]
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, Tensor)
        self.assertIsInstance(t2, Tensor)
        self.assertTrue(torch.equal(ToTensor()(x0), x2[0]))
        self.assertTrue(torch.equal(ToTensor()(x1), x2[1]))
        self.assertEqual(y0, y2[0].item())
        self.assertEqual(y1, y2[1].item())
        self.assertEqual(0, t2[0].item())
        self.assertEqual(0, t2[1].item())

    def test_avalanche_dataset_indexing(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x0, y0 = dataset_mnist[0]
        x1, y1 = dataset_mnist[5]
        dataset = AvalancheDataset(dataset_mnist, transform=ToTensor())
        x2, y2, t2 = dataset[0, 5]
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, Tensor)
        self.assertIsInstance(t2, Tensor)
        self.assertTrue(torch.equal(ToTensor()(x0), x2[0]))
        self.assertTrue(torch.equal(ToTensor()(x1), x2[1]))
        self.assertEqual(y0, y2[0].item())
        self.assertEqual(y1, y2[1].item())
        self.assertEqual(0, t2[0].item())
        self.assertEqual(0, t2[1].item())

    def test_avalanche_dataset_composition(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True, transform=RandomCrop(16))
        x, y = dataset_mnist[0]
        self.assertIsInstance(x, Image)
        self.assertEqual([x.width, x.height], [16, 16])
        self.assertIsInstance(y, int)

        dataset = AvalancheDataset(
            dataset_mnist, transform=ToTensor(),
            target_transform=lambda target: -1)

        x2, y2, t2 = dataset[0]
        self.assertIsInstance(x2, Tensor)
        self.assertEqual(x2.shape, (1, 16, 16))
        self.assertIsInstance(y2, int)
        self.assertEqual(y2, -1)
        self.assertIsInstance(t2, int)
        self.assertEqual(0, t2)

    def test_avalanche_dataset_add(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True, transform=CenterCrop(16))

        dataset1 = AvalancheDataset(
            dataset_mnist, transform=ToTensor(),
            target_transform=lambda target: -1)

        dataset2 = AvalancheDataset(
            dataset_mnist, target_transform=lambda target: -2,
            task_labels=ConstantSequence(2, len(dataset_mnist)))

        dataset3 = dataset1 + dataset2

        self.assertEqual(len(dataset_mnist)*2, len(dataset3))

        x1, y1, t1 = dataset1[0]
        x2, y2, t2 = dataset2[0]

        x3, y3, t3 = dataset3[0]
        x3_2, y3_2, t3_2 = dataset3[len(dataset_mnist)]

        self.assertIsInstance(x1, Tensor)
        self.assertEqual(x1.shape, (1, 16, 16))
        self.assertEqual(-1, y1)
        self.assertEqual(0, t1)

        self.assertIsInstance(x2, PIL.Image.Image)
        self.assertEqual(x2.size, (16, 16))
        self.assertEqual(-2, y2)
        self.assertEqual(2, t2)

        self.assertEqual((y1, t1), (y3, t3))
        self.assertEqual(16 * 16, torch.sum(torch.eq(x1, x3)).item())

        self.assertEqual((y2, t2), (y3_2, t3_2))
        self.assertTrue(pil_images_equal(x2, x3_2))

    def test_avalanche_dataset_radd(self):
        dataset_mnist = MNIST(
            expanduser("~") + "/.avalanche/data/mnist/",
            download=True,
            transform=CenterCrop(16))

        dataset1 = AvalancheDataset(
            dataset_mnist, transform=ToTensor(),
            target_transform=lambda target: -1)

        dataset2 = dataset_mnist + dataset1
        self.assertIsInstance(dataset2, AvalancheDataset)
        self.assertEqual(len(dataset_mnist) * 2, len(dataset2))

        dataset3 = dataset_mnist + dataset1 + dataset_mnist
        self.assertIsInstance(dataset3, AvalancheDataset)
        self.assertEqual(len(dataset_mnist) * 3, len(dataset3))

        dataset4 = dataset_mnist + dataset_mnist + dataset1
        self.assertIsInstance(dataset4, AvalancheDataset)
        self.assertEqual(len(dataset_mnist) * 3, len(dataset4))

    def test_dataset_add_monkey_patch_vanilla_behaviour(self):
        dataset_mnist = MNIST(
            expanduser("~") + "/.avalanche/data/mnist/",
            download=True,
            transform=CenterCrop(16))

        dataset_mnist2 = MNIST(
            expanduser("~") + "/.avalanche/data/mnist/",
            download=True,
            transform=CenterCrop(16))

        dataset = dataset_mnist + dataset_mnist2

        self.assertIsInstance(dataset, ConcatDataset)

        self.assertEqual(len(dataset_mnist) * 2, len(dataset))

    def test_avalanche_dataset_uniform_task_labels(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[0]
        dataset = AvalancheDataset(dataset_mnist, transform=ToTensor(),
                                   task_labels=[1] * len(dataset_mnist))
        x2, y2, t2 = dataset[0]

        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, int)
        self.assertIsInstance(t2, int)
        self.assertEqual(1, t2)
        self.assertTrue(torch.equal(ToTensor()(x), x2))
        self.assertEqual(y, y2)

        self.assertListEqual([1] * len(dataset_mnist),
                             list(dataset.targets_task_labels))

        subset_task1 = dataset.task_set[1]
        self.assertIsInstance(subset_task1, AvalancheDataset)
        self.assertEqual(len(dataset), len(subset_task1))

        with self.assertRaises(KeyError):
            subset_task0 = dataset.task_set[0]

    def test_avalanche_dataset_tensor_task_labels(self):
        x = torch.rand(32, 10)
        y = torch.rand(32, 10)
        t = torch.ones(32)  # Single task
        dataset = AvalancheTensorDataset(x, y, targets=1, task_labels=t)

        x2, y2, t2 = dataset[:]

        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, Tensor)
        self.assertIsInstance(t2, Tensor)
        self.assertTrue(torch.equal(x, x2))
        self.assertTrue(torch.equal(y, y2))
        self.assertTrue(torch.equal(t.to(int), t2))

        self.assertListEqual([1] * 32,
                             list(dataset.targets_task_labels))

        # Regression test for #654
        self.assertEqual(1, len(dataset.task_set))

        subset_task1 = dataset.task_set[1]
        self.assertIsInstance(subset_task1, AvalancheDataset)
        self.assertEqual(len(dataset), len(subset_task1))

        with self.assertRaises(KeyError):
            subset_task0 = dataset.task_set[0]

        with self.assertRaises(KeyError):
            subset_task0 = dataset.task_set[2]

        # Check single instance types
        x2, y2, t2 = dataset[0]

        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, Tensor)
        self.assertIsInstance(t2, int)

    def test_avalanche_dataset_uniform_task_labels_simple_def(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        dataset = AvalancheDataset(dataset_mnist, transform=ToTensor(),
                                   task_labels=1)
        _, _, t2 = dataset[0]

        self.assertIsInstance(t2, int)
        self.assertEqual(1, t2)

        self.assertListEqual([1] * len(dataset_mnist),
                             list(dataset.targets_task_labels))

        subset_task1 = dataset.task_set[1]
        self.assertIsInstance(subset_task1, AvalancheDataset)
        self.assertEqual(len(dataset), len(subset_task1))

        with self.assertRaises(KeyError):
            subset_task0 = dataset.task_set[0]

    def test_avalanche_dataset_mixed_task_labels(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[0]

        random_task_labels = [random.randint(0, 10)
                              for _ in range(len(dataset_mnist))]
        dataset = AvalancheDataset(dataset_mnist, transform=ToTensor(),
                                   task_labels=random_task_labels)
        x2, y2, t2 = dataset[0]

        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, int)
        self.assertIsInstance(t2, int)
        self.assertEqual(random_task_labels[0], t2)
        self.assertTrue(torch.equal(ToTensor()(x), x2))
        self.assertEqual(y, y2)

        self.assertListEqual(random_task_labels,
                             list(dataset.targets_task_labels))

        u_labels, counts = np.unique(random_task_labels, return_counts=True)
        for i, task_label in enumerate(u_labels.tolist()):
            subset_task = dataset.task_set[task_label]
            self.assertIsInstance(subset_task, AvalancheDataset)
            self.assertEqual(int(counts[i]), len(subset_task))

            unique_task_labels = list(subset_task.targets_task_labels)
            self.assertListEqual([task_label] * int(counts[i]),
                                 unique_task_labels)

        with self.assertRaises(KeyError):
            subset_task11 = dataset.task_set[11]

    def test_avalanche_tensor_dataset_task_labels_train(self):
        tr_ds = [AvalancheTensorDataset(
            torch.randn(10, 4),
            torch.randint(0, 3, (10,)),
            dataset_type=AvalancheDatasetType.CLASSIFICATION,
            task_labels=torch.randint(0, 5, (10,)).tolist()) for i in range(3)]
        ts_ds = [AvalancheTensorDataset(
            torch.randn(10, 4), torch.randint(0, 3, (10,)),
            dataset_type=AvalancheDatasetType.CLASSIFICATION,
            task_labels=torch.randint(0, 5, (10,)).tolist()) for i in range(3)]
        scenario = dataset_benchmark(train_datasets=tr_ds, test_datasets=ts_ds)
        model = SimpleMLP(input_size=4, num_classes=3)
        cl_strategy = Naive(model, SGD(model.parameters(), lr=0.001,
                                       momentum=0.9),
                            CrossEntropyLoss(), train_mb_size=5,
                            train_epochs=1, eval_mb_size=5,
                            device='cpu', evaluator=None)
        exp = []
        for i, experience in enumerate(scenario.train_stream):
            exp.append(i)
            cl_strategy.train(experience)
        self.assertEqual(len(exp), 3)

    def test_avalanche_dataset_task_labels_inheritance(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        random_task_labels = [random.randint(0, 10)
                              for _ in range(len(dataset_mnist))]
        dataset_orig = AvalancheDataset(dataset_mnist, transform=ToTensor(),
                                        task_labels=random_task_labels)

        dataset_child = AvalancheDataset(dataset_orig)
        x2, y2, t2 = dataset_orig[0]
        x3, y3, t3 = dataset_child[0]

        self.assertIsInstance(t2, int)
        self.assertEqual(random_task_labels[0], t2)

        self.assertIsInstance(t3, int)
        self.assertEqual(random_task_labels[0], t3)

        self.assertListEqual(random_task_labels,
                             list(dataset_orig.targets_task_labels))

        self.assertListEqual(random_task_labels,
                             list(dataset_child.targets_task_labels))

    def test_avalanche_dataset_tensor_dataset_input(self):
        train_x = torch.rand(500, 3, 28, 28)
        train_y = torch.zeros(500)
        test_x = torch.rand(200, 3, 28, 28)
        test_y = torch.ones(200)

        train = TensorDataset(train_x, train_y)
        test = TensorDataset(test_x, test_y)
        train_dataset = AvalancheDataset(train)
        test_dataset = AvalancheDataset(test)

        self.assertEqual(500, len(train_dataset))
        self.assertEqual(200, len(test_dataset))

        x, y, t = train_dataset[0]
        self.assertIsInstance(x, Tensor)
        self.assertEqual(0, y)
        self.assertEqual(0, t)

        x2, y2, t2 = test_dataset[0]
        self.assertIsInstance(x2, Tensor)
        self.assertEqual(1, y2)
        self.assertEqual(0, t2)

    def test_avalanche_dataset_multiple_outputs_and_float_y(self):
        train_x = torch.rand(500, 3, 28, 28)
        train_y = torch.zeros(500)
        train_z = torch.ones(500)
        test_x = torch.rand(200, 3, 28, 28)
        test_y = torch.ones(200)
        test_z = torch.full((200,), 5)

        train = TensorDataset(train_x, train_y, train_z)
        test = TensorDataset(test_x, test_y, test_z)
        train_dataset = AvalancheDataset(train)
        test_dataset = AvalancheDataset(test)

        self.assertEqual(500, len(train_dataset))
        self.assertEqual(200, len(test_dataset))

        x, y, z, t = train_dataset[0]
        self.assertIsInstance(x, Tensor)
        self.assertEqual(0, y)
        self.assertEqual(1, z)
        self.assertEqual(0, t)

        x2, y2, z2, t2 = test_dataset[0]
        self.assertIsInstance(x2, Tensor)
        self.assertEqual(1, y2)
        self.assertEqual(5, z2)
        self.assertEqual(0, t2)

    def test_avalanche_dataset_from_pytorch_subset(self):
        tensor_x = torch.rand(500, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (500,))

        whole_dataset = TensorDataset(tensor_x, tensor_y)

        train = Subset(whole_dataset, indices=list(range(400)))
        test = Subset(whole_dataset, indices=list(range(400, 500)))

        train_dataset = AvalancheDataset(train)
        test_dataset = AvalancheDataset(test)

        self.assertEqual(400, len(train_dataset))
        self.assertEqual(100, len(test_dataset))

        x, y, t = train_dataset[0]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x[0], x))
        self.assertTrue(torch.equal(tensor_y[0], y))
        self.assertEqual(0, t)

        self.assertTrue(torch.equal(torch.as_tensor(train_dataset.targets),
                                    tensor_y[:400]))

        x2, y2, t2 = test_dataset[0]
        self.assertIsInstance(x2, Tensor)
        self.assertTrue(torch.equal(tensor_x[400], x2))
        self.assertTrue(torch.equal(tensor_y[400], y2))
        self.assertEqual(0, t2)

        self.assertTrue(torch.equal(torch.as_tensor(test_dataset.targets),
                                    tensor_y[400:]))

    def test_avalanche_dataset_from_pytorch_concat_dataset(self):
        tensor_x = torch.rand(500, 3, 28, 28)
        tensor_x2 = torch.rand(300, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (500,))
        tensor_y2 = torch.randint(0, 100, (300,))

        dataset1 = TensorDataset(tensor_x, tensor_y)
        dataset2 = TensorDataset(tensor_x2, tensor_y2)

        concat_dataset = ConcatDataset((dataset1, dataset2))

        av_dataset = AvalancheDataset(concat_dataset)

        self.assertEqual(500, len(dataset1))
        self.assertEqual(300, len(dataset2))

        x, y, t = av_dataset[0]
        x2, y2, t2 = av_dataset[500]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x[0], x))
        self.assertTrue(torch.equal(tensor_y[0], y))
        self.assertEqual(0, t)

        self.assertIsInstance(x2, Tensor)
        self.assertTrue(torch.equal(tensor_x2[0], x2))
        self.assertTrue(torch.equal(tensor_y2[0], y2))
        self.assertEqual(0, t2)

        self.assertTrue(torch.equal(torch.as_tensor(av_dataset.targets),
                                    torch.cat((tensor_y, tensor_y2))))

    def test_avalanche_dataset_from_chained_pytorch_concat_dataset(self):
        tensor_x = torch.rand(500, 3, 28, 28)
        tensor_x2 = torch.rand(300, 3, 28, 28)
        tensor_x3 = torch.rand(200, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (500,))
        tensor_y2 = torch.randint(0, 100, (300,))
        tensor_y3 = torch.randint(0, 100, (200,))

        dataset1 = TensorDataset(tensor_x, tensor_y)
        dataset2 = TensorDataset(tensor_x2, tensor_y2)
        dataset3 = TensorDataset(tensor_x3, tensor_y3)

        concat_dataset = ConcatDataset((dataset1, dataset2))
        concat_dataset2 = ConcatDataset((concat_dataset, dataset3))

        av_dataset = AvalancheDataset(concat_dataset2)

        self.assertEqual(500, len(dataset1))
        self.assertEqual(300, len(dataset2))

        x, y, t = av_dataset[0]
        x2, y2, t2 = av_dataset[500]
        x3, y3, t3 = av_dataset[800]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x[0], x))
        self.assertTrue(torch.equal(tensor_y[0], y))
        self.assertEqual(0, t)

        self.assertIsInstance(x2, Tensor)
        self.assertTrue(torch.equal(tensor_x2[0], x2))
        self.assertTrue(torch.equal(tensor_y2[0], y2))
        self.assertEqual(0, t2)

        self.assertIsInstance(x3, Tensor)
        self.assertTrue(torch.equal(tensor_x3[0], x3))
        self.assertTrue(torch.equal(tensor_y3[0], y3))
        self.assertEqual(0, t3)

        self.assertTrue(torch.equal(
            torch.as_tensor(av_dataset.targets),
            torch.cat((tensor_y, tensor_y2, tensor_y3))))

    def test_avalanche_dataset_from_chained_pytorch_subsets(self):
        tensor_x = torch.rand(500, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (500,))

        whole_dataset = TensorDataset(tensor_x, tensor_y)

        subset1 = Subset(whole_dataset, indices=list(range(400, 500)))
        subset2 = Subset(subset1, indices=[5, 7, 0])

        dataset = AvalancheDataset(subset2)

        self.assertEqual(3, len(dataset))

        x, y, t = dataset[0]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x[405], x))
        self.assertTrue(torch.equal(tensor_y[405], y))
        self.assertEqual(0, t)

        self.assertTrue(
            torch.equal(
                torch.as_tensor(dataset.targets),
                torch.as_tensor([tensor_y[405], tensor_y[407], tensor_y[400]])
            )
        )

    def test_avalanche_dataset_from_chained_pytorch_concat_subset_dataset(self):
        tensor_x = torch.rand(200, 3, 28, 28)
        tensor_x2 = torch.rand(100, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (200,))
        tensor_y2 = torch.randint(0, 100, (100,))

        dataset1 = TensorDataset(tensor_x, tensor_y)
        dataset2 = TensorDataset(tensor_x2, tensor_y2)

        indices = [random.randint(0, 299) for _ in range(1000)]

        concat_dataset = ConcatDataset((dataset1, dataset2))
        subset = Subset(concat_dataset, indices)

        av_dataset = AvalancheDataset(subset)

        self.assertEqual(200, len(dataset1))
        self.assertEqual(100, len(dataset2))
        self.assertEqual(1000, len(av_dataset))

        for idx in range(1000):
            orig_idx = indices[idx]
            if orig_idx < 200:
                expected_x, expected_y = dataset1[orig_idx]
            else:
                expected_x, expected_y = dataset2[orig_idx-200]

            x, y, t = av_dataset[idx]
            self.assertIsInstance(x, Tensor)
            self.assertTrue(torch.equal(expected_x, x))
            self.assertTrue(torch.equal(expected_y, y))
            self.assertEqual(0, t)
            self.assertEqual(int(expected_y), int(av_dataset.targets[idx]))

    def test_avalanche_dataset_from_chained_pytorch_datasets(self):
        tensor_x = torch.rand(200, 3, 28, 28)
        tensor_x2 = torch.rand(100, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (200,))
        tensor_y2 = torch.randint(0, 100, (100,))

        dataset1 = TensorDataset(tensor_x, tensor_y)
        dataset1_sub = Subset(dataset1, range(199, -1, -1))
        dataset2 = TensorDataset(tensor_x2, tensor_y2)

        indices = [random.randint(0, 299) for _ in range(1000)]

        concat_dataset = ConcatDataset((dataset1_sub, dataset2))
        subset = Subset(concat_dataset, indices)

        av_dataset = AvalancheDataset(subset)

        self.assertEqual(200, len(dataset1_sub))
        self.assertEqual(100, len(dataset2))
        self.assertEqual(1000, len(av_dataset))

        for idx in range(1000):
            orig_idx = indices[idx]
            if orig_idx < 200:
                orig_idx = range(199, -1, -1)[orig_idx]
                expected_x, expected_y = dataset1[orig_idx]
            else:
                expected_x, expected_y = dataset2[orig_idx-200]

            x, y, t = av_dataset[idx]
            self.assertIsInstance(x, Tensor)
            self.assertTrue(torch.equal(expected_x, x))
            self.assertTrue(torch.equal(expected_y, y))
            self.assertEqual(0, t)
            self.assertEqual(int(expected_y), int(av_dataset.targets[idx]))

    def test_avalanche_dataset_from_chained_pytorch_datasets_task_labels(self):
        tensor_x = torch.rand(200, 3, 28, 28)
        tensor_x2 = torch.rand(100, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (200,))
        tensor_y2 = torch.randint(0, 100, (100,))
        tensor_t = torch.randint(0, 100, (200,))
        tensor_t2 = torch.randint(0, 100, (100,))

        dataset1 = AvalancheTensorDataset(
            tensor_x, tensor_y, task_labels=tensor_t)
        dataset1_sub = Subset(dataset1, range(199, -1, -1))
        dataset2 = AvalancheDataset(
            TensorDataset(tensor_x2, tensor_y2), task_labels=tensor_t2)

        indices = [random.randint(0, 299) for _ in range(1000)]

        concat_dataset = ConcatDataset((dataset1_sub, dataset2))
        subset = Subset(concat_dataset, indices)

        av_dataset = AvalancheDataset(subset)

        self.assertEqual(200, len(dataset1_sub))
        self.assertEqual(100, len(dataset2))
        self.assertEqual(1000, len(av_dataset))

        for idx in range(1000):
            orig_idx = indices[idx]
            if orig_idx < 200:
                orig_idx = range(199, -1, -1)[orig_idx]
                expected_x = tensor_x[orig_idx]
                expected_y = tensor_y[orig_idx]
                expected_t = tensor_t[orig_idx]
            else:
                orig_idx -= 200
                expected_x = tensor_x2[orig_idx]
                expected_y = tensor_y2[orig_idx]
                expected_t = tensor_t2[orig_idx]

            x, y, t = av_dataset[idx]
            self.assertIsInstance(x, Tensor)
            self.assertTrue(torch.equal(expected_x, x))
            self.assertTrue(torch.equal(expected_y, y))
            self.assertIsInstance(t, int)
            self.assertEqual(int(expected_t), int(t))
            self.assertEqual(int(expected_y), int(av_dataset.targets[idx]))

    def test_avalanche_dataset_collate_fn(self):
        tensor_x = torch.rand(500, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (500,))
        tensor_z = torch.randint(0, 100, (500,))

        def my_collate_fn(patterns):
            x_values = torch.stack([pat[0] for pat in patterns], 0)
            y_values = torch.tensor([pat[1] for pat in patterns]) + 1
            z_values = torch.tensor([-1 for _ in patterns])
            t_values = torch.tensor([pat[3] for pat in patterns])
            return x_values, y_values, z_values, t_values

        whole_dataset = TensorDataset(tensor_x, tensor_y, tensor_z)
        dataset = AvalancheDataset(whole_dataset, collate_fn=my_collate_fn)

        x, y, z, t = dataset[0]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x[0], x))
        self.assertTrue(torch.equal(tensor_y[0], y))
        self.assertEqual(0, t)

        x2, y2, z2, t2 = dataset[0:5]
        self.assertIsInstance(x2, Tensor)
        self.assertTrue(torch.equal(tensor_x[0:5], x2))
        self.assertTrue(torch.equal(tensor_y[0:5]+1, y2))
        self.assertTrue(torch.equal(torch.full((5,), -1, dtype=torch.long), z2))
        self.assertTrue(torch.equal(torch.zeros(5, dtype=torch.long), t2))

        inherited = AvalancheDataset(dataset)

        x3, y3, z3, t3 = inherited[0:5]
        self.assertIsInstance(x3, Tensor)
        self.assertTrue(torch.equal(tensor_x[0:5], x3))
        self.assertTrue(torch.equal(tensor_y[0:5] + 1, y3))
        self.assertTrue(torch.equal(torch.full((5,), -1, dtype=torch.long), z3))
        self.assertTrue(torch.equal(torch.zeros(5, dtype=torch.long), t3))

        with self.assertRaises(ValueError):
            # Can't define a custom collate when dataset_type != UNDEFINED
            bad_definition = AvalancheDataset(
                dataset, dataset_type=AvalancheDatasetType.CLASSIFICATION,
                collate_fn=my_collate_fn)

    def test_avalanche_dataset_collate_fn_inheritance(self):
        tensor_x = torch.rand(200, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (200,))
        tensor_z = torch.randint(0, 100, (200,))

        def my_collate_fn(patterns):
            x_values = torch.stack([pat[0] for pat in patterns], 0)
            y_values = torch.tensor([pat[1] for pat in patterns]) + 1
            z_values = torch.tensor([-1 for _ in patterns])
            t_values = torch.tensor([pat[3] for pat in patterns])
            return x_values, y_values, z_values, t_values

        def my_collate_fn2(patterns):
            x_values = torch.stack([pat[0] for pat in patterns], 0)
            y_values = torch.tensor([pat[1] for pat in patterns]) + 2
            z_values = torch.tensor([-2 for _ in patterns])
            t_values = torch.tensor([pat[3] for pat in patterns])
            return x_values, y_values, z_values, t_values

        whole_dataset = TensorDataset(tensor_x, tensor_y, tensor_z)
        dataset = AvalancheDataset(whole_dataset, collate_fn=my_collate_fn)
        inherited = AvalancheDataset(dataset, collate_fn=my_collate_fn2)  # Ok

        x, y, z, t = inherited[0:5]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x[0:5], x))
        self.assertTrue(torch.equal(tensor_y[0:5] + 2, y))
        self.assertTrue(torch.equal(torch.full((5,), -2, dtype=torch.long), z))
        self.assertTrue(torch.equal(torch.zeros(5, dtype=torch.long), t))

        classification_dataset = AvalancheDataset(
            whole_dataset, dataset_type=AvalancheDatasetType.CLASSIFICATION)

        with self.assertRaises(ValueError):
            bad_inherited = AvalancheDataset(
                classification_dataset, collate_fn=my_collate_fn)
        ok_inherited_classification = AvalancheDataset(classification_dataset)
        self.assertEqual(AvalancheDatasetType.CLASSIFICATION,
                         ok_inherited_classification.dataset_type)

    def test_avalanche_concat_dataset_collate_fn_inheritance(self):
        tensor_x = torch.rand(200, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (200,))
        tensor_z = torch.randint(0, 100, (200,))

        tensor_x2 = torch.rand(200, 3, 28, 28)
        tensor_y2 = torch.randint(0, 100, (200,))
        tensor_z2 = torch.randint(0, 100, (200,))

        def my_collate_fn(patterns):
            x_values = torch.stack([pat[0] for pat in patterns], 0)
            y_values = torch.tensor([pat[1] for pat in patterns]) + 1
            z_values = torch.tensor([-1 for _ in patterns])
            t_values = torch.tensor([pat[3] for pat in patterns])
            return x_values, y_values, z_values, t_values

        def my_collate_fn2(patterns):
            x_values = torch.stack([pat[0] for pat in patterns], 0)
            y_values = torch.tensor([pat[1] for pat in patterns]) + 2
            z_values = torch.tensor([-2 for _ in patterns])
            t_values = torch.tensor([pat[3] for pat in patterns])
            return x_values, y_values, z_values, t_values

        dataset1 = TensorDataset(tensor_x, tensor_y, tensor_z)
        dataset2 = AvalancheTensorDataset(tensor_x2, tensor_y2, tensor_z2,
                                          collate_fn=my_collate_fn)
        concat = AvalancheConcatDataset([dataset1, dataset2],
                                        collate_fn=my_collate_fn2)  # Ok

        x, y, z, t = dataset2[0:5]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x2[0:5], x))
        self.assertTrue(torch.equal(tensor_y2[0:5] + 1, y))
        self.assertTrue(torch.equal(torch.full((5,), -1, dtype=torch.long), z))
        self.assertTrue(torch.equal(torch.zeros(5, dtype=torch.long), t))

        x2, y2, z2, t2 = concat[0:5]
        self.assertIsInstance(x2, Tensor)
        self.assertTrue(torch.equal(tensor_x[0:5], x2))
        self.assertTrue(torch.equal(tensor_y[0:5] + 2, y2))
        self.assertTrue(torch.equal(torch.full((5,), -2, dtype=torch.long), z2))
        self.assertTrue(torch.equal(torch.zeros(5, dtype=torch.long), t2))

        dataset1_classification = AvalancheTensorDataset(
            tensor_x, tensor_y, tensor_z,
            dataset_type=AvalancheDatasetType.CLASSIFICATION)

        dataset2_segmentation = AvalancheDataset(
            dataset2, dataset_type=AvalancheDatasetType.SEGMENTATION)

        with self.assertRaises(ValueError):
            bad_concat_types = dataset1_classification + dataset2_segmentation

        with self.assertRaises(ValueError):
            bad_concat_collate = AvalancheConcatDataset(
                [dataset1, dataset2_segmentation], collate_fn=my_collate_fn)

        ok_concat_classification = dataset1_classification + dataset2
        self.assertEqual(AvalancheDatasetType.CLASSIFICATION,
                         ok_concat_classification.dataset_type)

        ok_concat_classification2 = dataset2 + dataset1_classification
        self.assertEqual(AvalancheDatasetType.CLASSIFICATION,
                         ok_concat_classification2.dataset_type)

    def test_avalanche_concat_dataset_recursion(self):

        def gen_random_tensors(n):
            return torch.rand(n, 3, 28, 28),\
                   torch.randint(0, 100, (n,)),\
                   torch.randint(0, 100, (n,))

        tensor_x, tensor_y, tensor_z = \
            gen_random_tensors(200)

        tensor_x2, tensor_y2, tensor_z2 = \
            gen_random_tensors(200)

        tensor_x3, tensor_y3, tensor_z3 = \
            gen_random_tensors(200)

        tensor_x4, tensor_y4, tensor_z4 = \
            gen_random_tensors(200)

        tensor_x5, tensor_y5, tensor_z5 = \
            gen_random_tensors(200)

        tensor_x6, tensor_y6, tensor_z6 = \
            gen_random_tensors(200)

        tensor_x7, tensor_y7, tensor_z7 = \
            gen_random_tensors(200)

        dataset1 = TensorDataset(tensor_x, tensor_y, tensor_z)
        dataset2 = AvalancheTensorDataset(tensor_x2, tensor_y2, tensor_z2,
                                          task_labels=1)
        dataset3 = AvalancheTensorDataset(tensor_x3, tensor_y3, tensor_z3,
                                          task_labels=2)

        dataset4 = AvalancheTensorDataset(tensor_x4, tensor_y4, tensor_z4,
                                          task_labels=3)
        dataset5 = AvalancheTensorDataset(tensor_x5, tensor_y5, tensor_z5,
                                          task_labels=4)
        dataset6 = AvalancheTensorDataset(tensor_x6, tensor_y6, tensor_z6)
        dataset7 = AvalancheTensorDataset(tensor_x7, tensor_y7, tensor_z7)

        # This will test recursion on both PyTorch ConcatDataset and
        # AvalancheConcatDataset
        concat = ConcatDataset([dataset1, dataset2])

        # Beware of the explicit task_labels=5 that *must* override the
        # task labels set in dataset4 and dataset5

        def transform_target_to_constant(ignored_target_value):
            return 101

        def transform_target_to_constant2(ignored_target_value):
            return 102

        concat2 = AvalancheConcatDataset(
            [dataset4, dataset5], task_labels=5,
            target_transform=transform_target_to_constant)

        concat3 = AvalancheConcatDataset(
            [dataset6, dataset7],
            target_transform=transform_target_to_constant2).freeze_transforms()
        concat_uut = AvalancheConcatDataset(
            [concat, dataset3, concat2, concat3])

        self.assertEqual(400, len(concat))
        self.assertEqual(400, len(concat2))
        self.assertEqual(400, len(concat3))
        self.assertEqual(1400, len(concat_uut))

        x, y, z, t = concat_uut[0]
        x2, y2, z2, t2 = concat_uut[200]
        x3, y3, z3, t3 = concat_uut[400]
        x4, y4, z4, t4 = concat_uut[600]
        x5, y5, z5, t5 = concat_uut[800]
        x6, y6, z6, t6 = concat_uut[1000]
        x7, y7, z7, t7 = concat_uut[1200]

        self.assertTrue(torch.equal(x, tensor_x[0]))
        self.assertTrue(torch.equal(y, tensor_y[0]))
        self.assertTrue(torch.equal(z, tensor_z[0]))
        self.assertEqual(0, t)

        self.assertTrue(torch.equal(x2, tensor_x2[0]))
        self.assertTrue(torch.equal(y2, tensor_y2[0]))
        self.assertTrue(torch.equal(z2, tensor_z2[0]))
        self.assertEqual(1, t2)

        self.assertTrue(torch.equal(x3, tensor_x3[0]))
        self.assertTrue(torch.equal(y3, tensor_y3[0]))
        self.assertTrue(torch.equal(z3, tensor_z3[0]))
        self.assertEqual(2, t3)

        self.assertTrue(torch.equal(x4, tensor_x4[0]))
        self.assertEqual(101, y4)
        self.assertTrue(torch.equal(z4, tensor_z4[0]))
        self.assertEqual(5, t4)

        self.assertTrue(torch.equal(x5, tensor_x5[0]))
        self.assertEqual(101, y5)
        self.assertTrue(torch.equal(z5, tensor_z5[0]))
        self.assertEqual(5, t5)

        self.assertTrue(torch.equal(x6, tensor_x6[0]))
        self.assertEqual(102, y6)
        self.assertTrue(torch.equal(z6, tensor_z6[0]))
        self.assertEqual(0, t6)

        self.assertTrue(torch.equal(x7, tensor_x7[0]))
        self.assertEqual(102, y7)
        self.assertTrue(torch.equal(z7, tensor_z7[0]))
        self.assertEqual(0, t7)

    def test_avalanche_pytorch_subset_recursion(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[1010]

        subset = Subset(dataset_mnist, indices=[3000, 8, 4, 1010, 12])

        dataset = AvalancheSubset(
            subset, indices=[0, 3])

        self.assertEqual(5, len(subset))
        self.assertEqual(2, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertTrue(pil_images_equal(x, x3))
        self.assertEqual(y, y3)
        self.assertEqual(0, t3)
        self.assertTrue(pil_images_equal(x2, x4))
        self.assertEqual(y2, y4)
        self.assertEqual(0, t4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

        def transform_target_to_constant(ignored_target_value):
            return 101

        subset = Subset(dataset_mnist, indices=[3000, 8, 4, 1010, 12])

        dataset = AvalancheSubset(
            subset, indices=[0, 3],
            target_transform=transform_target_to_constant,
            task_labels=5)

        self.assertEqual(5, len(subset))
        self.assertEqual(2, len(dataset))

        x5, y5, t5 = dataset[0]
        x6, y6, t6 = dataset[1]
        self.assertTrue(pil_images_equal(x, x5))
        self.assertEqual(101, y5)
        self.assertEqual(5, t5)
        self.assertTrue(pil_images_equal(x2, x6))
        self.assertEqual(101, y6)
        self.assertEqual(5, t6)
        self.assertFalse(pil_images_equal(x, x6))
        self.assertFalse(pil_images_equal(x2, x5))

    def test_avalanche_pytorch_subset_recursion_no_indices(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[8]

        subset = Subset(dataset_mnist, indices=[3000, 8, 4, 1010, 12])

        dataset = AvalancheSubset(subset)

        self.assertEqual(5, len(subset))
        self.assertEqual(5, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertTrue(pil_images_equal(x, x3))
        self.assertEqual(y, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        self.assertEqual(y2, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

    def test_avalanche_avalanche_subset_recursion_no_indices_transform(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[8]

        def transform_target_to_constant(ignored_target_value):
            return 101

        def transform_target_plus_one(target_value):
            return target_value+1

        subset = AvalancheSubset(dataset_mnist,
                                 indices=[3000, 8, 4, 1010, 12],
                                 transform=ToTensor(),
                                 target_transform=transform_target_to_constant)

        dataset = AvalancheSubset(subset,
                                  target_transform=transform_target_plus_one)

        self.assertEqual(5, len(subset))
        self.assertEqual(5, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertIsInstance(x3, Tensor)
        self.assertIsInstance(x4, Tensor)
        self.assertTrue(torch.equal(ToTensor()(x), x3))
        self.assertEqual(102, y3)
        self.assertTrue(torch.equal(ToTensor()(x2), x4))
        self.assertEqual(102, y4)
        self.assertFalse(torch.equal(ToTensor()(x), x4))
        self.assertFalse(torch.equal(ToTensor()(x2), x3))

    def test_avalanche_avalanche_subset_recursion_transform(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[1010]

        def transform_target_to_constant(ignored_target_value):
            return 101

        def transform_target_plus_one(target_value):
            return target_value+2

        subset = AvalancheSubset(dataset_mnist,
                                 indices=[3000, 8, 4, 1010, 12],
                                 target_transform=transform_target_to_constant)

        dataset = AvalancheSubset(subset,
                                  indices=[0, 3, 1],
                                  target_transform=transform_target_plus_one)

        self.assertEqual(5, len(subset))
        self.assertEqual(3, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]

        self.assertTrue(pil_images_equal(x, x3))
        self.assertEqual(103, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        self.assertEqual(103, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

    def test_avalanche_avalanche_subset_recursion_frozen_transform(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[1010]

        def transform_target_to_constant(ignored_target_value):
            return 101

        def transform_target_plus_two(target_value):
            return target_value+2

        subset = AvalancheSubset(dataset_mnist,
                                 indices=[3000, 8, 4, 1010, 12],
                                 target_transform=transform_target_to_constant)
        subset = subset.freeze_transforms()

        dataset = AvalancheSubset(subset,
                                  indices=[0, 3, 1],
                                  target_transform=transform_target_plus_two)

        self.assertEqual(5, len(subset))
        self.assertEqual(3, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]

        self.assertTrue(pil_images_equal(x, x3))
        self.assertEqual(103, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        self.assertEqual(103, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

        dataset = AvalancheSubset(subset,
                                  indices=[0, 3, 1],
                                  target_transform=transform_target_plus_two)
        dataset = dataset.replace_transforms(None, None)

        x5, y5, t5 = dataset[0]
        x6, y6, t6 = dataset[1]

        self.assertTrue(pil_images_equal(x, x5))
        self.assertEqual(101, y5)
        self.assertTrue(pil_images_equal(x2, x6))
        self.assertEqual(101, y6)
        self.assertFalse(pil_images_equal(x, x6))
        self.assertFalse(pil_images_equal(x2, x5))

    def test_avalanche_avalanche_subset_recursion_modified_transforms(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[1010]

        def transform_target_to_constant(ignored_target_value):
            return 101

        def transform_target_to_constant2(ignored_target_value):
            return 102

        def transform_target_plus_two(target_value):
            return target_value+2

        subset = AvalancheSubset(dataset_mnist,
                                 indices=[3000, 8, 4, 1010, 12],
                                 target_transform=transform_target_to_constant)
        subset.target_transform = transform_target_to_constant2

        dataset = AvalancheSubset(subset,
                                  indices=[0, 3, 1],
                                  target_transform=transform_target_plus_two)

        self.assertEqual(5, len(subset))
        self.assertEqual(3, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]

        self.assertTrue(pil_images_equal(x, x3))
        self.assertEqual(104, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        self.assertEqual(104, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

    def test_avalanche_avalanche_subset_recursion_sub_class_mapping(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[1010]

        class_mapping = list(range(10))
        random.shuffle(class_mapping)

        subset = AvalancheSubset(dataset_mnist,
                                 indices=[3000, 8, 4, 1010, 12],
                                 class_mapping=class_mapping)

        dataset = AvalancheSubset(subset,
                                  indices=[0, 3, 1])

        self.assertEqual(5, len(subset))
        self.assertEqual(3, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]

        self.assertTrue(pil_images_equal(x, x3))
        expected_y3 = class_mapping[y]
        self.assertEqual(expected_y3, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        expected_y4 = class_mapping[y2]
        self.assertEqual(expected_y4, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

    def test_avalanche_avalanche_subset_recursion_up_class_mapping(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[1010]

        class_mapping = list(range(10))
        random.shuffle(class_mapping)

        subset = AvalancheSubset(dataset_mnist,
                                 indices=[3000, 8, 4, 1010, 12])

        dataset = AvalancheSubset(subset,
                                  indices=[0, 3, 1],
                                  class_mapping=class_mapping)

        self.assertEqual(5, len(subset))
        self.assertEqual(3, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]

        self.assertTrue(pil_images_equal(x, x3))
        expected_y3 = class_mapping[y]
        self.assertEqual(expected_y3, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        expected_y4 = class_mapping[y2]
        self.assertEqual(expected_y4, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

    def test_avalanche_avalanche_subset_recursion_mix_class_mapping(self):
        dataset_mnist = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/",
            download=True)
        x, y = dataset_mnist[3000]
        x2, y2 = dataset_mnist[1010]

        class_mapping = list(range(10))
        class_mapping2 = list(range(10))
        random.shuffle(class_mapping)
        random.shuffle(class_mapping2)

        subset = AvalancheSubset(dataset_mnist,
                                 indices=[3000, 8, 4, 1010, 12],
                                 class_mapping=class_mapping)

        dataset = AvalancheSubset(subset,
                                  indices=[0, 3, 1],
                                  class_mapping=class_mapping2)

        self.assertEqual(5, len(subset))
        self.assertEqual(3, len(dataset))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]

        self.assertTrue(pil_images_equal(x, x3))
        expected_y3 = class_mapping2[class_mapping[y]]
        self.assertEqual(expected_y3, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        expected_y4 = class_mapping2[class_mapping[y2]]
        self.assertEqual(expected_y4, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

    def test_avalanche_avalanche_subset_concat_stack_overflow(self):
        d_sz = 25
        tensor_x = torch.rand(d_sz, 3, 28, 28)
        tensor_y = torch.randint(0, 10, (d_sz,))
        tensor_t = torch.randint(0, 10, (d_sz,))
        dataset = AvalancheTensorDataset(
            tensor_x, tensor_y, task_labels=tensor_t)
        dataset_hierarchy_depth = 500

        rolling_indices: List[List[int]] = []
        expect_indices: List[List[int]] = []

        for _ in range(dataset_hierarchy_depth):
            idx_permuted = list(range(d_sz))
            random.shuffle(idx_permuted)
            rolling_indices.append(idx_permuted)

        forward_indices = range(d_sz)
        expect_indices.append(list(forward_indices))
        for idx in range(dataset_hierarchy_depth):
            forward_indices = [forward_indices[x] for x in rolling_indices[idx]]
            expect_indices.append(forward_indices)

        expect_indices = list(reversed(expect_indices))

        leaf = dataset

        for idx in range(dataset_hierarchy_depth):
            intermediate_idx_test = (dataset_hierarchy_depth - 1) - idx
            subset = AvalancheSubset(leaf, indices=rolling_indices[idx])
            leaf = AvalancheConcatDataset((subset, leaf))

            # Regression test for #616 (second bug)
            # https://github.com/ContinualAI/avalanche/issues/616#issuecomment-848852287
            all_targets = []
            for c_dataset in leaf._dataset_list:
                all_targets += c_dataset.targets

            all_targets = torch.tensor(all_targets)

            for idx_internal in range(idx+1):
                leaf_range = range(idx_internal * d_sz,
                                   (idx_internal + 1) * d_sz)
                permuted = expect_indices[idx_internal+intermediate_idx_test]
                self.assertTrue(torch.equal(tensor_y[permuted],
                                            all_targets[leaf_range]))

            self.assertTrue(torch.equal(tensor_y,
                                        all_targets[-d_sz:]))

        self.assertEqual(d_sz * dataset_hierarchy_depth + d_sz, len(leaf))

        for idx in range(dataset_hierarchy_depth):
            leaf_range = range(idx*d_sz, (idx+1) * d_sz)
            permuted = expect_indices[idx]
            self.assertTrue(torch.equal(tensor_x[permuted],
                                        leaf[leaf_range][0]))
            self.assertTrue(torch.equal(tensor_y[permuted],
                                        leaf[leaf_range][1]))
            self.assertTrue(torch.equal(tensor_y[permuted],
                                        torch.tensor(leaf.targets)[leaf_range]))
            self.assertTrue(torch.equal(tensor_t[permuted],
                                        leaf[leaf_range][2]))

        self.assertTrue(torch.equal(tensor_x,
                                    leaf[d_sz*dataset_hierarchy_depth:][0]))
        self.assertTrue(torch.equal(tensor_y,
                                    leaf[d_sz*dataset_hierarchy_depth:][1]))
        self.assertTrue(torch.equal(
            tensor_y,
            torch.tensor(leaf.targets)[d_sz * dataset_hierarchy_depth:]))
        self.assertTrue(torch.equal(tensor_t,
                                    leaf[d_sz*dataset_hierarchy_depth:][2]))

class TransformationSubsetTests(unittest.TestCase):
    def test_avalanche_subset_transform(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[0]
        dataset = AvalancheSubset(dataset_mnist, transform=ToTensor())
        x2, y2, t2 = dataset[0]
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, int)
        self.assertIsInstance(t2, int)
        self.assertTrue(torch.equal(ToTensor()(x), x2))
        self.assertEqual(y, y2)
        self.assertEqual(0, t2)

    def test_avalanche_subset_composition(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True, transform=RandomCrop(16))
        x, y = dataset_mnist[0]
        self.assertIsInstance(x, Image)
        self.assertEqual([x.width, x.height], [16, 16])
        self.assertIsInstance(y, int)

        dataset = AvalancheSubset(
            dataset_mnist, transform=ToTensor(),
            target_transform=lambda target: -1)

        x2, y2, t2 = dataset[0]
        self.assertIsInstance(x2, Tensor)
        self.assertEqual(x2.shape, (1, 16, 16))
        self.assertIsInstance(y2, int)
        self.assertIsInstance(t2, int)
        self.assertEqual(y2, -1)
        self.assertEqual(0, t2)

    def test_avalanche_subset_indices(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[1000]
        x2, y2 = dataset_mnist[1007]

        dataset = AvalancheSubset(
            dataset_mnist, indices=[1000, 1007])

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertTrue(pil_images_equal(x, x3))
        self.assertEqual(y, y3)
        self.assertTrue(pil_images_equal(x2, x4))
        self.assertEqual(y2, y4)
        self.assertFalse(pil_images_equal(x, x4))
        self.assertFalse(pil_images_equal(x2, x3))

    def test_avalanche_subset_mapping(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        _, y = dataset_mnist[1000]

        mapping = list(range(10))
        other_classes = list(mapping)
        other_classes.remove(y)

        swap_y = random.choice(other_classes)

        mapping[y] = swap_y
        mapping[swap_y] = y

        dataset = AvalancheSubset(dataset_mnist, class_mapping=mapping)

        _, y2, _ = dataset[1000]
        self.assertEqual(y2, swap_y)

    def test_avalanche_subset_uniform_task_labels(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[1000]
        x2, y2 = dataset_mnist[1007]

        # First, test by passing len(task_labels) == len(dataset_mnist)
        dataset = AvalancheSubset(
            dataset_mnist, indices=[1000, 1007],
            task_labels=[1] * len(dataset_mnist))

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertEqual(y, y3)
        self.assertEqual(1, t3)
        self.assertEqual(y2, y4)
        self.assertEqual(1, t4)

        # Secondly, test by passing len(task_labels) == len(indices)
        dataset = AvalancheSubset(
            dataset_mnist, indices=[1000, 1007],
            task_labels=[1, 1])

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertEqual(y, y3)
        self.assertEqual(1, t3)
        self.assertEqual(y2, y4)
        self.assertEqual(1, t4)

    def test_avalanche_subset_mixed_task_labels(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        x, y = dataset_mnist[1000]
        x2, y2 = dataset_mnist[1007]

        full_task_labels = [1] * len(dataset_mnist)
        full_task_labels[1000] = 2
        # First, test by passing len(task_labels) == len(dataset_mnist)
        dataset = AvalancheSubset(
            dataset_mnist, indices=[1000, 1007],
            task_labels=full_task_labels)

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertEqual(y, y3)
        self.assertEqual(2, t3)
        self.assertEqual(y2, y4)
        self.assertEqual(1, t4)

        # Secondly, test by passing len(task_labels) == len(indices)
        dataset = AvalancheSubset(
            dataset_mnist, indices=[1000, 1007],
            task_labels=[3, 5])

        x3, y3, t3 = dataset[0]
        x4, y4, t4 = dataset[1]
        self.assertEqual(y, y3)
        self.assertEqual(3, t3)
        self.assertEqual(y2, y4)
        self.assertEqual(5, t4)

    def test_avalanche_subset_task_labels_inheritance(self):
        dataset_mnist = MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                              download=True)
        random_task_labels = [random.randint(0, 10)
                              for _ in range(len(dataset_mnist))]
        dataset_orig = AvalancheDataset(dataset_mnist, transform=ToTensor(),
                                        task_labels=random_task_labels)

        dataset_child = AvalancheSubset(dataset_orig,
                                        indices=[1000, 1007])
        _, _, t2 = dataset_orig[1000]
        _, _, t5 = dataset_orig[1007]
        _, _, t3 = dataset_child[0]
        _, _, t6 = dataset_child[1]

        self.assertEqual(random_task_labels[1000], t2)
        self.assertEqual(random_task_labels[1007], t5)
        self.assertEqual(random_task_labels[1000], t3)
        self.assertEqual(random_task_labels[1007], t6)

        self.assertListEqual(random_task_labels,
                             list(dataset_orig.targets_task_labels))

        self.assertListEqual([random_task_labels[1000],
                              random_task_labels[1007]],
                             list(dataset_child.targets_task_labels))

    def test_avalanche_subset_collate_fn_inheritance(self):
        tensor_x = torch.rand(200, 3, 28, 28)
        tensor_y = torch.randint(0, 100, (200,))
        tensor_z = torch.randint(0, 100, (200,))

        def my_collate_fn(patterns):
            x_values = torch.stack([pat[0] for pat in patterns], 0)
            y_values = torch.tensor([pat[1] for pat in patterns]) + 1
            z_values = torch.tensor([-1 for _ in patterns])
            t_values = torch.tensor([pat[3] for pat in patterns])
            return x_values, y_values, z_values, t_values

        def my_collate_fn2(patterns):
            x_values = torch.stack([pat[0] for pat in patterns], 0)
            y_values = torch.tensor([pat[1] for pat in patterns]) + 2
            z_values = torch.tensor([-2 for _ in patterns])
            t_values = torch.tensor([pat[3] for pat in patterns])
            return x_values, y_values, z_values, t_values

        whole_dataset = TensorDataset(tensor_x, tensor_y, tensor_z)
        dataset = AvalancheDataset(whole_dataset, collate_fn=my_collate_fn)
        inherited = AvalancheSubset(dataset, indices=list(range(5, 150)),
                                    collate_fn=my_collate_fn2)  # Ok

        x, y, z, t = inherited[0:5]
        self.assertIsInstance(x, Tensor)
        self.assertTrue(torch.equal(tensor_x[5:10], x))
        self.assertTrue(torch.equal(tensor_y[5:10] + 2, y))
        self.assertTrue(torch.equal(torch.full((5,), -2, dtype=torch.long), z))
        self.assertTrue(torch.equal(torch.zeros(5, dtype=torch.long), t))

        classification_dataset = AvalancheDataset(
            whole_dataset, dataset_type=AvalancheDatasetType.CLASSIFICATION)

        with self.assertRaises(ValueError):
            bad_inherited = AvalancheSubset(
                classification_dataset, indices=list(range(5, 150)),
                collate_fn=my_collate_fn)
        ok_inherited_classification = AvalancheSubset(
            classification_dataset, indices=list(range(5, 150)))
        self.assertEqual(AvalancheDatasetType.CLASSIFICATION,
                         ok_inherited_classification.dataset_type)

class TransformationTensorDatasetTests(unittest.TestCase):
    def test_tensor_dataset_helper_tensor_y(self):

        train_exps = [[torch.rand(50, 32, 32), torch.randint(0, 100, (50,))]
                      for _ in range(5)]
        test_exps = [[torch.rand(23, 32, 32), torch.randint(0, 100, (23,))]
                     for _ in range(5)]

        cl_benchmark = create_generic_benchmark_from_tensor_lists(
            train_tensors=train_exps, test_tensors=test_exps,
            task_labels=[0] * 5)

        self.assertEqual(5, len(cl_benchmark.train_stream))
        self.assertEqual(5, len(cl_benchmark.test_stream))
        self.assertEqual(5, cl_benchmark.n_experiences)

        for exp_id in range(cl_benchmark.n_experiences):
            scenario_train_x, scenario_train_y, _ = \
                load_all_dataset(cl_benchmark.train_stream[exp_id].dataset)
            scenario_test_x, scenario_test_y, _ = \
                load_all_dataset(cl_benchmark.test_stream[exp_id].dataset)

            self.assertTrue(torch.all(torch.eq(
                train_exps[exp_id][0],
                scenario_train_x)))
            self.assertTrue(torch.all(torch.eq(
                train_exps[exp_id][1],
                scenario_train_y)))
            self.assertSequenceEqual(
                train_exps[exp_id][1].tolist(),
                cl_benchmark.train_stream[exp_id].dataset.targets)
            self.assertEqual(0, cl_benchmark.train_stream[exp_id].task_label)

            self.assertTrue(torch.all(torch.eq(
                test_exps[exp_id][0],
                scenario_test_x)))
            self.assertTrue(torch.all(torch.eq(
                test_exps[exp_id][1],
                scenario_test_y)))
            self.assertSequenceEqual(
                test_exps[exp_id][1].tolist(),
                cl_benchmark.test_stream[exp_id].dataset.targets)
            self.assertEqual(0, cl_benchmark.test_stream[exp_id].task_label)

    def test_tensor_dataset_helper_list_y(self):
        train_exps = [(torch.rand(50, 32, 32),
                       torch.randint(0, 100, (50,)).tolist()) for _ in range(5)]
        test_exps = [(torch.rand(23, 32, 32),
                      torch.randint(0, 100, (23,)).tolist()) for _ in range(5)]

        cl_benchmark = create_generic_benchmark_from_tensor_lists(
            train_tensors=train_exps, test_tensors=test_exps,
            task_labels=[0] * 5)

        self.assertEqual(5, len(cl_benchmark.train_stream))
        self.assertEqual(5, len(cl_benchmark.test_stream))
        self.assertEqual(5, cl_benchmark.n_experiences)

        for exp_id in range(cl_benchmark.n_experiences):
            scenario_train_x, scenario_train_y, _ = \
                load_all_dataset(cl_benchmark.train_stream[exp_id].dataset)
            scenario_test_x, scenario_test_y, _ = \
                load_all_dataset(cl_benchmark.test_stream[exp_id].dataset)

            self.assertTrue(torch.all(torch.eq(
                train_exps[exp_id][0],
                scenario_train_x)))
            self.assertSequenceEqual(
                train_exps[exp_id][1],
                scenario_train_y.tolist())
            self.assertSequenceEqual(
                train_exps[exp_id][1],
                cl_benchmark.train_stream[exp_id].dataset.targets)
            self.assertEqual(0, cl_benchmark.train_stream[exp_id].task_label)

            self.assertTrue(torch.all(torch.eq(
                test_exps[exp_id][0],
                scenario_test_x)))
            self.assertSequenceEqual(
                test_exps[exp_id][1],
                scenario_test_y.tolist())
            self.assertSequenceEqual(
                test_exps[exp_id][1],
                cl_benchmark.test_stream[exp_id].dataset.targets)
            self.assertEqual(0, cl_benchmark.test_stream[exp_id].task_label)

class AvalancheDatasetTransformOpsTests(unittest.TestCase):
    def test_avalanche_inherit_groups(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True
        )

        def plus_one_target(target):
            return target+1

        transform_groups = dict(
            train=(ToTensor(), None),
            eval=(None, plus_one_target)
        )
        x, y = original_dataset[0]
        dataset = AvalancheDataset(original_dataset,
                                   transform_groups=transform_groups)

        x2, y2, _ = dataset[0]
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, int)
        self.assertTrue(torch.equal(ToTensor()(x), x2))
        self.assertEqual(y, y2)

        dataset_eval = dataset.eval()

        x3, y3, _ = dataset_eval[0]
        self.assertIsInstance(x3, PIL.Image.Image)
        self.assertIsInstance(y3, int)
        self.assertEqual(y+1, y3)

        # Regression test for #565
        dataset_inherit = AvalancheDataset(dataset_eval)

        x4, y4, _ = dataset_inherit[0]
        self.assertIsInstance(x4, PIL.Image.Image)
        self.assertIsInstance(y4, int)
        self.assertEqual(y + 1, y4)

        # Regression test for #566
        dataset_sub_train = AvalancheSubset(dataset)
        dataset_sub_eval = dataset_sub_train.eval()
        dataset_sub = AvalancheSubset(dataset_sub_eval, indices=[0])

        x5, y5, _ = dataset_sub[0]
        self.assertIsInstance(x5, PIL.Image.Image)
        self.assertIsInstance(y5, int)
        self.assertEqual(y + 1, y5)
        # End regression tests

        concat_dataset = AvalancheConcatDataset([dataset_sub_eval, dataset_sub])

        x6, y6, _ = concat_dataset[0]
        self.assertIsInstance(x6, PIL.Image.Image)
        self.assertIsInstance(y6, int)
        self.assertEqual(y + 1, y6)

        concat_dataset_no_inherit_initial = \
            AvalancheConcatDataset([dataset_sub_eval, dataset])

        x7, y7, _ = concat_dataset_no_inherit_initial[0]
        self.assertIsInstance(x7, Tensor)
        self.assertIsInstance(y7, int)
        self.assertEqual(y, y7)

    def test_freeze_transforms(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True
        )
        x, y = original_dataset[0]
        dataset = AvalancheDataset(original_dataset, transform=ToTensor())
        dataset_frozen = dataset.freeze_transforms()
        dataset_frozen.transform = None

        x2, y2, _ = dataset_frozen[0]
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(y2, int)
        self.assertTrue(torch.equal(ToTensor()(x), x2))
        self.assertEqual(y, y2)

        dataset.transform = None
        x2, y2, _ = dataset[0]
        self.assertIsInstance(x2, Image)

        x2, y2, _ = dataset_frozen[0]
        self.assertIsInstance(x2, Tensor)

    def test_freeze_transforms_chain(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True,
            transform=ToTensor()
        )
        x, *_ = original_dataset[0]
        self.assertIsInstance(x, Tensor)

        dataset_transform = AvalancheDataset(original_dataset,
                                             transform=ToPILImage())
        x, *_ = dataset_transform[0]
        self.assertIsInstance(x, Image)

        dataset_frozen = dataset_transform.freeze_transforms()

        x2, *_ = dataset_frozen[0]
        self.assertIsInstance(x2, Image)

        dataset_transform.transform = None

        x2, *_ = dataset_transform[0]
        self.assertIsInstance(x2, Tensor)

        dataset_frozen.transform = ToTensor()

        x2, *_ = dataset_frozen[0]
        self.assertIsInstance(x2, Tensor)

        dataset_frozen2 = dataset_frozen.freeze_transforms()

        x2, *_ = dataset_frozen2[0]
        self.assertIsInstance(x2, Tensor)

        dataset_frozen.transform = None

        x2, *_ = dataset_frozen2[0]
        self.assertIsInstance(x2, Tensor)
        x2, *_ = dataset_frozen[0]
        self.assertIsInstance(x2, Image)

    def test_add_transforms(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True
        )
        x, _ = original_dataset[0]
        dataset = AvalancheDataset(original_dataset, transform=ToTensor())
        dataset_added = dataset.add_transforms(ToPILImage())
        x2, *_ = dataset[0]
        x3, *_ = dataset_added[0]
        self.assertIsInstance(x, Image)
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(x3, Image)

    def test_add_transforms_chain(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True
        )
        x, _ = original_dataset[0]
        dataset = AvalancheDataset(original_dataset, transform=ToTensor())
        dataset_added = AvalancheDataset(dataset, transform=ToPILImage())
        x2, *_ = dataset[0]
        x3, *_ = dataset_added[0]
        self.assertIsInstance(x, Image)
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(x3, Image)

    def test_transforms_freeze_add_mix(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        x, _ = original_dataset[0]
        dataset = AvalancheDataset(original_dataset, transform=ToTensor())
        dataset_frozen = dataset.freeze_transforms()
        dataset_added = dataset_frozen.add_transforms(ToPILImage())

        self.assertEqual(None, dataset_frozen.transform)

        x2, *_ = dataset[0]
        x3, *_ = dataset_added[0]
        self.assertIsInstance(x, Image)
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(x3, Image)

        dataset_frozen = dataset_added.freeze_transforms()
        dataset_added.transform = None

        x4, *_ = dataset_frozen[0]
        x5, *_ = dataset_added[0]
        self.assertIsInstance(x4, Image)
        self.assertIsInstance(x5, Tensor)

    def test_replace_transforms(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        x, y = original_dataset[0]
        dataset = AvalancheDataset(original_dataset, transform=ToTensor())
        x2, *_ = dataset[0]
        dataset_reset = dataset.replace_transforms(None, None)
        x3, *_ = dataset_reset[0]

        self.assertIsInstance(x, Image)
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(x3, Image)

        dataset_reset.transform = ToTensor()

        x4, *_ = dataset_reset[0]
        self.assertIsInstance(x4, Tensor)

        dataset_reset.replace_transforms(None, None)

        x5, *_ = dataset_reset[0]
        self.assertIsInstance(x5, Tensor)

        dataset_other = AvalancheDataset(dataset_reset)
        dataset_other = dataset_other.replace_transforms(None, lambda l: l + 1)

        _, y6, _ = dataset_other[0]
        self.assertEqual(y+1, y6)

    def test_transforms_replace_freeze_mix(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        x, _ = original_dataset[0]
        dataset = AvalancheDataset(original_dataset, transform=ToTensor())
        x2, *_ = dataset[0]
        dataset_reset = dataset.replace_transforms(None, None)
        x3, *_ = dataset_reset[0]

        self.assertIsInstance(x, Image)
        self.assertIsInstance(x2, Tensor)
        self.assertIsInstance(x3, Image)

        dataset_frozen = dataset.freeze_transforms()

        x4, *_ = dataset_frozen[0]
        self.assertIsInstance(x4, Tensor)

        dataset_frozen_reset = dataset_frozen.replace_transforms(None, None)

        x5, *_ = dataset_frozen_reset[0]
        self.assertIsInstance(x5, Tensor)

    def test_transforms_groups_base_usage(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        dataset = AvalancheDataset(
            original_dataset,
            transform_groups=dict(train=(ToTensor(), None),
                                  eval=(None, Lambda(lambda t: float(t)))))

        x, y, _ = dataset[0]
        self.assertIsInstance(x, Tensor)
        self.assertIsInstance(y, int)

        dataset_test = dataset.eval()

        x2, y2, _ = dataset_test[0]
        x3, y3, _ = dataset[0]
        self.assertIsInstance(x2, Image)
        self.assertIsInstance(y2, float)
        self.assertIsInstance(x3, Tensor)
        self.assertIsInstance(y3, int)

        dataset_train = dataset.train()
        dataset.transform = None

        x4, y4, _ = dataset_train[0]
        x5, y5, _ = dataset[0]
        self.assertIsInstance(x4, Tensor)
        self.assertIsInstance(y4, int)
        self.assertIsInstance(x5, Image)
        self.assertIsInstance(y5, int)

    def test_transforms_groups_constructor_error(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        with self.assertRaises(Exception):
            # Test tuple has only one element
            dataset = AvalancheDataset(
                original_dataset,
                transform_groups=dict(train=(ToTensor(), None),
                                      eval=(Lambda(lambda t: float(t)))))

        with self.assertRaises(Exception):
            # Test is not a tuple has only one element
            dataset = AvalancheDataset(
                original_dataset,
                transform_groups=dict(train=(ToTensor(), None),
                                      eval=[None, Lambda(lambda t: float(t))]))

        with self.assertRaises(Exception):
            # Train is None
            dataset = AvalancheDataset(
                original_dataset,
                transform_groups=dict(train=None,
                                      eval=(None, Lambda(lambda t: float(t)))))

        with self.assertRaises(Exception):
            # transform_groups is not a dictionary
            dataset = AvalancheDataset(
                original_dataset,
                transform_groups='Hello world!')

    def test_transforms_groups_alternative_default_group(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        dataset = AvalancheDataset(
            original_dataset,
            transform_groups=dict(train=(ToTensor(), None), eval=(None, None)),
            initial_transform_group='eval')

        x, *_ = dataset[0]
        self.assertIsInstance(x, Image)

        dataset_test = dataset.eval()

        x2, *_ = dataset_test[0]
        x3, *_ = dataset[0]
        self.assertIsInstance(x2, Image)
        self.assertIsInstance(x3, Image)

    def test_transforms_groups_partial_constructor(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        dataset = AvalancheDataset(
            original_dataset, transform_groups=dict(train=(ToTensor(), None)))

        x, *_ = dataset[0]
        self.assertIsInstance(x, Tensor)

        dataset = dataset.eval()
        x2, *_ = dataset[0]
        self.assertIsInstance(x2, Tensor)

    def test_transforms_groups_multiple_groups(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        dataset = AvalancheDataset(
            original_dataset,
            transform_groups=dict(
                train=(ToTensor(), None),
                eval=(None, None),
                other=(Compose([ToTensor(),
                               Lambda(lambda tensor: tensor.numpy())]), None)))

        x, *_ = dataset[0]
        self.assertIsInstance(x, Tensor)

        dataset = dataset.eval()
        x2, *_ = dataset[0]
        self.assertIsInstance(x2, Image)

        dataset = dataset.with_transforms('other')
        x3, *_ = dataset[0]
        self.assertIsInstance(x3, np.ndarray)

    def test_transforms_add_group(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        dataset = AvalancheDataset(original_dataset)

        with self.assertRaises(Exception):
            # Can't add existing groups
            dataset = dataset.add_transforms_group('train', ToTensor(), None)

        with self.assertRaises(Exception):
            # Can't add group with bad names (must be str)
            dataset = dataset.add_transforms_group(123, ToTensor(), None)

        # Can't add group with bad names (must be str)
        dataset = dataset.add_transforms_group('other', ToTensor(), None)
        dataset_other = dataset.with_transforms('other')

        x, *_ = dataset[0]
        x2, *_ = dataset_other[0]
        self.assertIsInstance(x, Image)
        self.assertIsInstance(x2, Tensor)

        dataset_other2 = AvalancheDataset(dataset_other)

        # Checks that the other group is used on dataset_other2
        x3, *_ = dataset_other2[0]
        self.assertIsInstance(x3, Tensor)

        with self.assertRaises(Exception):
            # Can't add group if it already exists
            dataset_other2 = dataset_other2.add_transforms_group(
                'other', ToTensor(), None)

        # Check that the above failed method didn't change the 'other' group
        x4, *_ = dataset_other2[0]
        self.assertIsInstance(x4, Tensor)

    def test_transformation_concat_dataset(self):
        original_dataset = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)
        original_dataset2 = MNIST(
            root=expanduser("~") + "/.avalanche/data/mnist/", download=True)

        dataset = AvalancheConcatDataset([original_dataset,
                                          original_dataset2])

        self.assertEqual(len(original_dataset) + len(original_dataset2),
                         len(dataset))

    def test_transformation_concat_dataset_groups(self):
        original_dataset = AvalancheDataset(
            MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                  download=True),
            transform_groups=dict(eval=(None, None), train=(ToTensor(), None)))
        original_dataset2 = AvalancheDataset(
            MNIST(root=expanduser("~") + "/.avalanche/data/mnist/",
                  download=True),
            transform_groups=dict(train=(None, None), eval=(ToTensor(), None)))

        dataset = AvalancheConcatDataset([original_dataset,
                                          original_dataset2])

        self.assertEqual(len(original_dataset) + len(original_dataset2),
                         len(dataset))

        x, *_ = dataset[0]
        x2, *_ = dataset[len(original_dataset)]
        self.assertIsInstance(x, Tensor)
        self.assertIsInstance(x2, Image)

        dataset = dataset.eval()

        x3, *_ = dataset[0]
        x4, *_ = dataset[len(original_dataset)]
        self.assertIsInstance(x3, Image)
        self.assertIsInstance(x4, Tensor)

if __name__ == '__main__':
    unittest.main()


text:
#!/usr/bin/env python

# libraries
from __future__ import print_function, division
import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import os
import h5py
import re
import shutil
import copy
import time
import random
import warnings
import operator
from datetime import datetime
from tqdm import tqdm, tqdm_notebook
from matplotlib.pyplot import specgram
import torch
import torchvision
from torchvision import datasets, models, transforms
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from torchtext import data
from sklearn.metrics import precision_score, f1_score, recall_score, accuracy_score, confusion_matrix
from sklearn.metrics.cluster import contingency_matrix
import joblib
import csv
import warnings
warnings.filterwarnings('ignore')
#%matplotlib inline

### 1. LOAD DATA:

### Manual inputs (denoted by #-markers throughout):
subject = sys.argv[1] ###### CHANGE: subject name (make sure name does NOT have "_" in it; used for subject col and exportname)
eegfileloading = 'sub-'+subject+'_eegdata.csv' ############################################################### CHANGE: filename here

### LOAD EEG DATA --- rows = channels, cols = timepoints
eegdir = sys.argv[2] ################################################################################ CHANGE: dir here
# load file, checking for header
input_csv_file = eegdir+'/'+eegfileloading
with open(input_csv_file, 'rb') as csvfile:
    csv_test_bytes = csvfile.read(10)  # grab sample of .csv for format detection
    headertest = csv_test_bytes.decode("utf-8")
    if any(c.isalpha() for c in headertest) == True:
        data = pd.read_csv(input_csv_file, header=0)
        channels = data.columns
    else:
        data = pd.read_csv(input_csv_file, header=None)

### quick check: transpose if not in proper format (rows = chans, cols = timepoints) - build on this later.
if len(data) > len(data.columns):
    data = data.T
    if type(data[0][0]) == str:
        data = data.drop(data.columns[0], axis=1)
    data = data.astype(float)
    print('CHECK: Number of channels ~ %d' % len(data))
else:
    data = data.astype(float)

### AUTO DUMP IED IMAGES: clears dir containing spectrograms if produced in previous iteration
spectdir = eegdir+'/SPECTS/IEDS/' ################################################################## CHANGE: dir here
os.makedirs(spectdir, exist_ok = True)
for filename in os.listdir(spectdir):
    file_path = os.path.join(spectdir, filename)
    try:
        if os.path.isfile(file_path) or os.path.islink(file_path):
            os.unlink(file_path)
        elif os.path.isdir(file_path):
            shutil.rmtree(file_path)
    except Exception as e:
        print('Failed to delete %s. Reason: %s' % (file_path, e))

### 2. LOAD TEMPLATE-MATCHING DETECTOR FUNCTIONS:

def detect_peaks(x, mph=None, mpd=1, threshold=0, edge='rising',
                kpsh=False, valley=False, show=False, ax=None):
    """Detect peaks in data based on their amplitude and other features."""
    x = np.atleast_1d(x).astype('float64')
    if x.size < 3:
        return np.array([], dtype=int)
    if valley:
        x = -x
    # find indices of all peaks
    dx = x[1:] - x[:-1]
    indnan = np.where(np.isnan(x))[0]
    if indnan.size:
        x[indnan] = np.inf
        dx[np.where(np.isnan(dx))[0]] = np.inf
    ine, ire, ife = np.array([[], [], []], dtype=int)
    if not edge:
        ine = np.where((np.hstack((dx, 0)) < 0) & (np.hstack((0, dx)) > 0))[0]
    else:
        if edge.lower() in ['rising', 'both']:
            ire = np.where((np.hstack((dx, 0)) <= 0) & (np.hstack((0, dx)) > 0))[0]
        if edge.lower() in ['falling', 'both']:
            ife = np.where((np.hstack((dx, 0)) < 0) & (np.hstack((0, dx)) >= 0))[0]
    ind = np.unique(np.hstack((ine, ire, ife)))
    if ind.size and indnan.size:
        # NaN's and values close to NaN's cannot be peaks
        ind = ind[np.in1d(ind, np.unique(np.hstack((indnan, indnan-1, indnan+1))), invert=True)]
    # first and last values of x cannot be peaks
    if ind.size and ind[0] == 0:
        ind = ind[1:]
    if ind.size and ind[-1] == x.size-1:
        ind = ind[:-1]
    # remove peaks < minimum peak height
    if ind.size and mph is not None:
        ind = ind[x[ind] >= mph]
    # remove peaks - neighbors < threshold
    if ind.size and threshold > 0:
        dx = np.min(np.vstack([x[ind]-x[ind-1], x[ind]-x[ind+1]]), axis=0)
        ind = np.delete(ind, np.where(dx < threshold)[0])
    # detect small peaks closer than minimum peak distance
    if ind.size and mpd > 1:
        ind = ind[np.argsort(x[ind])][::-1]  # sort ind by peak height
        idel = np.zeros(ind.size, dtype=bool)
        for i in range(ind.size):
            if not idel[i]:
                # keep peaks with the same height if kpsh is True
                idel = idel | (ind >= ind[i] - mpd) & (ind <= ind[i] + mpd) \
                    & (x[ind[i]] > x[ind] if kpsh else True)
                idel[i] = 0  # Keep current peak
        # remove the small peaks and sort back the indices by their occurrence
        ind = np.sort(ind[~idel])

    if show:
        if indnan.size:
            x[indnan] = np.nan
        if valley:
            x = -x
        _plot(x, mph, mpd, threshold, edge, valley, ax, ind)
    return ind

def _plot(x, mph, mpd, threshold, edge, valley, ax, ind):
    """Plot results of the detect_peaks function, see its help."""
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print('matplotlib is not available.')
    else:
        if ax is None:
            _, ax = plt.subplots(1, 1, figsize=(8, 4))

        ax.plot(x, 'b', lw=1)
        if ind.size:
            label = 'valley' if valley else 'peak'
            label = label + 's' if ind.size > 1 else label
            ax.plot(ind, x[ind], '+', mfc=None, mec='r', mew=2, ms=8,
                    label='%d %s' % (ind.size, label))
            ax.legend(loc='best', framealpha=.5, numpoints=1)
        ax.set_xlim(-.02*x.size, x.size*1.02-1)
        ymin, ymax = x[np.isfinite(x)].min(), x[np.isfinite(x)].max()
        yrange = ymax - ymin if ymax > ymin else 1
        ax.set_ylim(ymin - 0.1*yrange, ymax + 0.1*yrange)
        ax.set_xlabel('Data #', fontsize=14)
        ax.set_ylabel('Amplitude', fontsize=14)
        mode = 'Valley detection' if valley else 'Peak detection'
        ax.set_title("%s (mph=%s, mpd=%d, threshold=%s, edge='%s')"
                    % (mode, str(mph), mpd, str(threshold), edge))
        # plt.grid()
        plt.show()

def locate_downsample_freq(sample_freq, min_freq=200, max_freq=340):
    min_up_factor = np.inf
    best_candidate_freq = None
    for candidate in range(min_freq, max_freq+1):
        down_samp_rate = sample_freq / float(candidate)
        down_factor, up_factor = down_samp_rate.as_integer_ratio()
        if up_factor <= min_up_factor:
            min_up_factor = up_factor
            best_candidate_freq = candidate
    return best_candidate_freq

def butter_bandpass(low_limit, high_limit, samp_freq, order=5):
    nyquist_limit = samp_freq / 2
    low_prop = low_limit / nyquist_limit
    high_prop = high_limit / nyquist_limit
    b, a = signal.butter(order, [low_prop, high_prop], btype='band')
    def bb_filter(data):
        return signal.filtfilt(b, a, data)
    return bb_filter

def detect(channel, samp_freq, return_eeg=False, temp_func=None, signal_func=None):
    # assume that eeg is [channels x samples]
    # Round samp_freq to the nearest integer if it is large
    if samp_freq > 100:
        samp_freq = int(np.round(samp_freq))
    down_samp_freq = locate_downsample_freq(samp_freq)
    template = signal.triang(np.round(down_samp_freq * 0.06))
    kernel = np.array([-2, -1, 1, 2]) / float(8)
    template = np.convolve(kernel, np.convolve(template, kernel, 'valid') ,'full')
    if temp_func:
        template = temp_func(template, samp_freq)
    if signal_func:
        channel = signal_func(channel, samp_freq)

    down_samp_rate = samp_freq / float(down_samp_freq)
    down_samp_factor, up_samp_factor = down_samp_rate.as_integer_ratio()
    channel = signal.detrend(channel, type='constant')
    results = template_match(channel, template, down_samp_freq)
    up_samp_results = [np.round(spikes * down_samp_factor / float(up_samp_factor)).astype(int) for spikes in results]
    if return_eeg:
        return up_samp_results, [channel[start:end] for start, end in results]
    else:
        return up_samp_results

def template_match(channel, template, down_samp_freq, thresh=7, min_spacing=0): #######@@@############################## CHANGE: d:7,0
    template_len = len(template)
    cross_corr = np.convolve(channel, template, 'valid')
    cross_corr_std = med_std(cross_corr, down_samp_freq)
    detections = []
    # catch empty channels
    if cross_corr_std > 0:
        # normalize the cross-correlation
        cross_corr_norm = ((cross_corr - np.mean(cross_corr)) / cross_corr_std)
        cross_corr_norm[1] = 0
        cross_corr_norm[-1] = 0
        # find regions with high cross-corr
        if np.any(abs(cross_corr_norm > thresh)):
            peaks = detect_peaks(abs(cross_corr_norm), mph=thresh, mpd=template_len)
            peaks += int(np.ceil(template_len / 2.)) # center detection on template
            peaks = [peak for peak in peaks if peak > template_len and peak <= len(channel)-template_len]
            if peaks:
                # find peaks that are at least (min_spacing) secs away
                distant_peaks = np.diff(peaks) > min_spacing * down_samp_freq
                # always keep the first peak
                to_keep = np.insert(distant_peaks, 0, True)
                peaks = [peaks[x] for x in range(len(peaks)) if to_keep[x] == True]
                detections = [(peak-template_len, peak+template_len) for peak in peaks]
    return np.array(detections)

def med_std(signal, window_len):
    window = np.zeros(window_len) + (1 / float(window_len))
    std = np.sqrt(np.median(np.convolve(np.square(signal), window, 'valid') - np.square(np.convolve(signal, window, 'valid'))))
    return std

### 3. RUN TEMPLATE-MATCHING DETECTOR:

def autoDetect(eegdata, samp_freq = 200, subject = subject):
    """
    AUTODETECT: DETECTS ALL SPIKES IN EACH CHANNEL
         INPUT: raw eeg file (preprocessed signal)
        OUTPUT: all_detections (list containing a list of arrays for all detections),
                channel_names (eeg channel names corresponding to each detection list)
    """
    ### DETECT SPIKES:
    all_detections = []
    channel_names = []
    for i in range(eegdata.shape[0]):
        channel = eegdata.iloc[i,:].astype(float) # run on each row (chan)
        detections = detect(channel, samp_freq, return_eeg=False, temp_func=None, signal_func=None)
        all_detections.append(detections)
        channel_names.append(int(float((eegdata.columns[i]))))

    ### REFORMAT SPIKES:
    detections = pd.DataFrame(all_detections)
    channels = pd.DataFrame(channel_names)
    spikes = pd.concat([channels,detections], axis = 1)
    newspikes = spikes.transpose()
    newspikes.columns = newspikes.iloc[0]
    newspikes = newspikes.iloc[1:] # remove duplicate channel_name row
    ### AUTO LONG-FORMATTING OF SPIKES
    spikeDf = pd.DataFrame() # empty df to store final spikes and spikeTimes
    for idx, col in enumerate(newspikes.columns):
        # extract spikes for each column
        tempSpikes = newspikes.iloc[:,idx].dropna() # column corresponding to channel with all spikes
        tempSpikes2 = tempSpikes.tolist() # convert series to list
        # extract channel name for each spike (duplicate based on the number of spikes)
        tempName = tempSpikes.name # channel name
        tempName2 = [tempName] * len(tempSpikes) # repeat col name by the number of spikes in this channel
        tempDf = pd.DataFrame({'channel': tempName2, 'spikeTime': tempSpikes2})
        # save and append to final df
        spikeDf = spikeDf.append(tempDf)
        spikeDf['fs'] = samp_freq
        spikeDf['subject'] = subject
    return(spikeDf)

spikes = autoDetect(data) ### eegfile, Fs, sessionname; kleen_fs=200, preprocess_fs=200

print("SPIKES DETECTED (TEMP MATCH) = ", len(spikes))
print("")
print(spikes[:3])

### 4. GENERATE INPUT IMAGES FOR CNN:

def spectimgs(eegdata, spikedf):
    """
    SPECTS: GENERATE SPECTS FOR CNN
        INPUT: 1) eegdata, 2) spikedf (df from automated template-matching spike detector)
        OUTPUT: spects within ./SPECTS/IEDS
    """
    for i in tqdm(range(0,len(spikedf))):
        samp_freq = int(float(spikedf.fs.values[0]))
        #######################################
        pad = 1 # d:1 number of seconds for window
        dpi_setting = 300 # d:300
        Nfft = 128*(samp_freq/500) # d: 128
        h = 3
        w = 3
        #######################################
        try:
            subject = spikedf.subject.values[0]
            chan_name = int(spikedf.channel.values[i]) # zero idxed -1
            spikestart = spikedf.spikeTime.values[i][0] # start spike
            ### select eeg data row
            ecogclip = eegdata.iloc[chan_name]
            ### filter out line noise
            b_notch, a_notch = signal.iirnotch(60.0, 30.0, samp_freq)
            ecogclip = pd.Series(signal.filtfilt(b_notch, a_notch, ecogclip))

            ### trim eeg clip based on cushion
            ### mean imputation if missing indices
            end = int(float((spikestart+int(float(pad*samp_freq)))))
            start = int(float((spikestart-int(float(pad*samp_freq)))))
            if end > max(ecogclip.index):
                temp = list(ecogclip[list(range(spikestart-int(float(pad*samp_freq)), max(ecogclip.index)))])
                cushend = [np.mean(ecogclip)]*(end - max(ecogclip.index))
                temp = np.array(temp + cushend)
            elif start < min(ecogclip.index):
                temp = list(ecogclip[list(range(min(ecogclip.index), spikestart+pad*samp_freq))])
                cushstart = [np.mean(ecogclip)]*(min(ecogclip.index)-start)
                temp = np.array(cushstart, temp)
            else:
                temp = np.array(ecogclip[list(range(spikestart-int(float(pad*samp_freq)),
                                         spikestart+int(float(pad*samp_freq))))])

            ### PLOT AND EXPORT:
            plt.figure(figsize=(h,w))
            specgram(temp, NFFT = int(Nfft), Fs = samp_freq, noverlap=int(Nfft/2), detrend = "linear", cmap = "YlOrRd")
            plt.axis("off")
            plt.xlim(0, pad*2)
            plt.ylim(0,100)
            plt.savefig(spectdir+subject+"_"+str(spikestart)+"_"+str(chan_name)+".png", dpi = dpi_setting)
            plt.close()
        except Exception as e:
            print(e)
            print("ERROR with IED portion:", i)
            plt.close()
            continue

spectimgs(data, spikes)

### 5. ResNet-18 CNN DETECTOR:

### A: LOAD ALL DATA --- extract clip_id from path
model_dir = os.path.dirname(os.path.abspath(__file__))+"/../" # dir with trained model ################################################## CHANGE: dir here
proj_dir = eegdir+"/" # dir with main project script ############################################# CHANGE: dir here
imgs = 'SPECTS' # dir with IED / NONIED image dirs (name of subdir)

data_transforms = {
    imgs: transforms.Compose([
        transforms.Resize(224),
        transforms.Pad(1, fill=0, padding_mode='constant'),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])}

class ImageFolderWithPaths(datasets.ImageFolder):
    """Custom dataset that includes image file paths.
    Extends torchvision.datasets.ImageFolder
    """
    # override the __getitem__ method. this is the method that dataloader calls
    def __getitem__(self, index):
        # this is what ImageFolder normally returns
        original_tuple = super(ImageFolderWithPaths, self).__getitem__(index)
        # the image file path
        path = self.imgs[index][0]
        # make a new tuple that includes original and the path
        tuple_with_path = (original_tuple + (path,))
        return (tuple_with_path)

image_datasets = {x: ImageFolderWithPaths(os.path.join(proj_dir, x),
                                          data_transforms[x]) for x in [imgs]}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=1, # use batch=1, shuffle=F
                                             shuffle=False, num_workers=0) for x in [imgs]}
class_names = image_datasets[imgs].classes
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# extract image paths
path_names = []
for images,labels,paths in dataloaders[imgs]:
    path_names.append(paths)

#in windows, have back-slashes and so need to convert them to forward-slashes:
if sys.platform=='win32':
    for it in range(len(path_names)):
        temp=list(path_names[it])
        temp[0]=re.sub('\\\\','/',temp[0])
        path_names[it]=tuple(temp)

# convert list of paths to dataframe col
df = pd.DataFrame(path_names)
df.columns = ['clip_ids']
df[['clip_ids','clip']] = df['clip_ids'].str.split('IEDS/',expand=True)
df['clip'] = df['clip'].str.rstrip('.png')
df[['subject','start','chan']] = df['clip'].str.split('_',expand=True)

##############################################

### B: LOAD PRETRAINED MODEL
try:
    model = torch.load(model_dir+'model_aied.pt')
    # model.eval() # model architecture
except ImportError:
    print('TRAINED MODEL NOT FOUND: Check that trained model is in eegdir and name matches: model_aied.pt')

###############################################

### C: RUN MODEL
y_pred = []
with torch.no_grad():
    for inputs,labels,paths in dataloaders[imgs]:
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = model.forward(inputs)
        _,predicted = torch.max(outputs, 1)
        pred = predicted.numpy()
        lab = labels.numpy()
        y_pred.append(pred)

# reformat outputs:
y_pred_flat = np.concatenate((y_pred),axis=0)
# classes = ['class 0', 'class 1']
df['predicted_class'] = y_pred_flat
# # export as .csv
# df.to_csv(proj_dir+'all_predictions.csv', encoding='utf-8', index=False)
# print(df[:3]) # here, 1 = nonied, 0 = ied (df)

# 6. CLEAN SPIKE DF FOR EXPORT:

def dataCleaner(df, samp_freq = spikes.fs.values[0], win = 0.5):
    """
    CLEANS SPIKE DATA FOR EXPORT
    INPUT: df from resnet model, win = number of seconds allowed for spike overlap
                                (i.e., spikes within 3s of ea.other = single event))
    OUTPUT: clean df with subjectid, spikeStart,
    channels (where spikes detected) - if contactName present, change to rownames,
    numChannels (# channels spike detected)
    """
    ### only keep spikes: predicted_class = 0
    df = df[df.predicted_class == 0]
    df['start'] = df['start'].astype(int) # convert from str to int
    ### sort start times in df:
    df = df.sort_values(by = 'start', ascending = True)
    ### dedupe spikes by col and time
    bins =  np.arange(min(df.start.values), max(df.start.values), samp_freq*win)
    spikebins = np.digitize(df['start'], bins)
    cleandf = df.groupby(spikebins)['start'].describe()
    chanlist = df.groupby(spikebins)['chan'].apply(lambda x: x.values.tolist())
    chanlist = [list(set(x)) for x in chanlist]
    chancounts = [len(l) for l in chanlist]
    meanspikestart = (cleandf['min']).astype(int)
    subjectid = [subject]*len(meanspikestart)
    ### reformat into new df
    finaldf = pd.DataFrame({'subject': subjectid, 'spikeStart': meanspikestart,
                            'channels': chanlist, 'numChannels': chancounts})
    ### reject spikes detected in >= 12 channels within time window
    finaldf = finaldf[finaldf.numChannels < 12]
    return (finaldf)

finaldf = dataCleaner(df)

#also save raw data about spikes:
df['start'] = df['start'].astype(int) # convert from str to int
### sort start times in df:
df = df.sort_values(by = 'start', ascending = True)
df.to_csv(proj_dir+'sub-'+subject+'_allspikes.csv', encoding='utf-8', index=False)

### export as .csv
finaldf.to_csv(proj_dir+'sub-'+subject+'_finalspikes.csv', encoding='utf-8', index=False) ############################ CHANGE: export name
print(finaldf[:3])
print("")
print("FINAL SPIKES DETECTED = ", len(finaldf))

ext: py

max_stars_repo_name: Willie-Lab/aied

max_issues_repo_name: Willie-Lab/aied

max_forks_repo_name: Willie-Lab/aied

reason: import numpy,from scipy