import os

from torch.utils.data import DataLoader
from continuum.datasets import CIFAR10, InMemoryDataset
from continuum.datasets import MNIST
import torchvision
from continuum.scenarios import TransformationIncremental
import pytest
import numpy as np

from continuum.transforms.bg_swap import BackgroundSwap

DATA_PATH = os.environ.get("CONTINUUM_DATA_PATH")

# Uncomment for debugging via image output
# import matplotlib.pyplot as plt


def test_bg_swap_fast():
    """
    Fast test for background swap.
    """
    bg_x = np.ones(shape=[2, 5, 5, 3]) * -1
    bg_y = np.random.rand(2)

    fg = np.random.normal(loc=.5, scale=.1, size=[5, 5])
    bg = InMemoryDataset(bg_x, bg_y)

    bg_swap = BackgroundSwap(bg, input_dim=(5, 5), normalize_bg=None)

    spliced_1_channel = bg_swap(fg)[:, :, 0]

    assert np.array_equal((spliced_1_channel <= -1), (fg <= .5))


@pytest.mark.slow
def test_background_swap_numpy():
    """
    Test background swap on a single ndarray input.
    """
    mnist = MNIST(DATA_PATH, download=True, train=True)
    cifar = CIFAR10(DATA_PATH, download=True, train=True)

    bg_swap = BackgroundSwap(cifar, input_dim=(28, 28))

    im = mnist.get_data()[0][0]
    im = bg_swap(im)

    # Uncomment for debugging
    # plt.imshow(im, interpolation='nearest')
    # plt.show()


@pytest.mark.slow
def test_background_swap_torch():
    """
    Test background swap on a single tensor input.
    """
    cifar = CIFAR10(DATA_PATH, download=True, train=True)

    mnist = torchvision.datasets.MNIST(DATA_PATH, train=True, download=True,
                                       transform=torchvision.transforms.Compose([
                                           torchvision.transforms.ToTensor()
                                       ]))

    bg_swap = BackgroundSwap(cifar, input_dim=(28, 28))
    im = mnist[0][0]

    im = bg_swap(im)

    # Uncomment for debugging
    # plt.imshow(im.permute(1, 2, 0), interpolation='nearest')
    # plt.show()


@pytest.mark.slow
def test_background_tranformation():
    """
    Example code using TransformationIncremental to create a setting with 3 tasks.
    """
    cifar = CIFAR10(DATA_PATH, train=True)
    mnist = MNIST(DATA_PATH, download=False, train=True)
    nb_task = 3
    list_trsf = []
    for i in range(nb_task):
        list_trsf.append([torchvision.transforms.ToTensor(), BackgroundSwap(cifar, bg_label=i, input_dim=(28, 28)),
                          torchvision.transforms.ToPILImage()])
    scenario = TransformationIncremental(mnist, base_transformations=[torchvision.transforms.ToTensor()],
                                         incremental_transformations=list_trsf)
    folder = "tests/samples/background_trsf/"
    if not os.path.exists(folder):
        os.makedirs(folder)
    for task_id, task_data in enumerate(scenario):
        task_data.plot(path=folder, title=f"background_{task_id}.jpg", nb_samples=100, shape=[28, 28, 3])
        loader = DataLoader(task_data)
        _, _, _ = next(iter(loader))

[(13, 'os.environ.get', 'os.environ.get', (['"""CONTINUUM_DATA_PATH"""'], {}), False, 'import os\n'), (24, 'numpy.random.rand', 'np.random.rand', (['(2)'], {}), True, 'import numpy as np\n'), (26, 'numpy.random.normal', 'np.random.normal', ([], {'loc': '(0.5)', 'scale': '(0.1)', 'size': '[5, 5]'}), True, 'import numpy as np\n'), (27, 'continuum.datasets.InMemoryDataset', 'InMemoryDataset', (['bg_x', 'bg_y'], {}), False, 'from continuum.datasets import CIFAR10, InMemoryDataset\n'), (29, 'continuum.transforms.bg_swap.BackgroundSwap', 'BackgroundSwap', (['bg'], {'input_dim': '(5, 5)', 'normalize_bg': 'None'}), False, 'from continuum.transforms.bg_swap import BackgroundSwap\n'), (33, 'numpy.array_equal', 'np.array_equal', (['(spliced_1_channel <= -1)', '(fg <= 0.5)'], {}), True, 'import numpy as np\n'), (41, 'continuum.datasets.MNIST', 'MNIST', (['DATA_PATH'], {'download': '(True)', 'train': '(True)'}), False, 'from continuum.datasets import MNIST\n'), (42, 'continuum.datasets.CIFAR10', 'CIFAR10', (['DATA_PATH'], {'download': '(True)', 'train': '(True)'}), False, 'from continuum.datasets import CIFAR10, InMemoryDataset\n'), (44, 'continuum.transforms.bg_swap.BackgroundSwap', 'BackgroundSwap', (['cifar'], {'input_dim': '(28, 28)'}), False, 'from continuum.transforms.bg_swap import BackgroundSwap\n'), (59, 'continuum.datasets.CIFAR10', 'CIFAR10', (['DATA_PATH'], {'download': '(True)', 'train': '(True)'}), False, 'from continuum.datasets import CIFAR10, InMemoryDataset\n'), (66, 'continuum.transforms.bg_swap.BackgroundSwap', 'BackgroundSwap', (['cifar'], {'input_dim': '(28, 28)'}), False, 'from continuum.transforms.bg_swap import BackgroundSwap\n'), (81, 'continuum.datasets.CIFAR10', 'CIFAR10', (['DATA_PATH'], {'train': '(True)'}), False, 'from continuum.datasets import CIFAR10, InMemoryDataset\n'), (82, 'continuum.datasets.MNIST', 'MNIST', (['DATA_PATH'], {'download': '(False)', 'train': '(True)'}), False, 'from continuum.datasets import MNIST\n'), (23, 'numpy.ones', 'np.ones', ([], {'shape': '[2, 5, 5, 3]'}), True, 'import numpy as np\n'), (91, 'os.path.exists', 'os.path.exists', (['folder'], {}), False, 'import os\n'), (92, 'os.makedirs', 'os.makedirs', (['folder'], {}), False, 'import os\n'), (95, 'torch.utils.data.DataLoader', 'DataLoader', (['task_data'], {}), False, 'from torch.utils.data import DataLoader\n'), (86, 'torchvision.transforms.ToTensor', 'torchvision.transforms.ToTensor', ([], {}), False, 'import torchvision\n'), (86, 'continuum.transforms.bg_swap.BackgroundSwap', 'BackgroundSwap', (['cifar'], {'bg_label': 'i', 'input_dim': '(28, 28)'}), False, 'from continuum.transforms.bg_swap import BackgroundSwap\n'), (87, 'torchvision.transforms.ToPILImage', 'torchvision.transforms.ToPILImage', ([], {}), False, 'import torchvision\n'), (88, 'torchvision.transforms.ToTensor', 'torchvision.transforms.ToTensor', ([], {}), False, 'import torchvision\n'), (63, 'torchvision.transforms.ToTensor', 'torchvision.transforms.ToTensor', ([], {}), False, 'import torchvision\n')]