# -*- coding: utf-8 -*-
##########################################################################
# NSAp - Copyright (C) CEA, 2019
# Distributed under the terms of the CeCILL-B license, as published by
# the CEA-CNRS-INRIA. Refer to the LICENSE file or to
# http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html
# for details.
##########################################################################

"""
Module that provides functions to prepare the Multimodal MNIST-SVHN dataset.
"""

# Imports
import os
import logging
import torchvision
import torch
import torchvision.transforms as transforms
from collections import namedtuple
import numpy as np
import pandas as pd
from pynet.datasets import Fetchers
from pynet.datasets.core import ArrayDataset, DataItem
import bisect


# Global parameters
Item = namedtuple("Item", ["input_path", "output_path", "metadata_path",
                           "labels"])
logger = logging.getLogger("pynet")


def get_matching_idx(labels1, labels2):
    """
    :param
    labels1: np.array, list of labels from dataset 1
    labels2: np.array, list of labels from dataset 2
    :return
    (idx1, idx2): 2 lists of indexes defining the unique correspondance between labels1 and labels2
    """
    assert isinstance(labels1, np.ndarray) and isinstance(labels2, np.ndarray), \
        "Wrong type: %s,%s"%(type(labels1), type(labels2))
    assert set(labels1) == set(labels2), "Missing labels"


    _idx1, _idx2 = [], []
    n1, n2 = len(labels1), len(labels2)
    for label in set(labels1):
        idx1 = np.arange(n1)[labels1 == label]
        idx2 = np.arange(n2)[labels2 == label]
        n = min(len(idx1), len(idx2))
        _idx1.extend(idx1[np.random.permutation(n)])
        _idx2.extend(idx2[np.random.permutation(n)])
    return np.stack(_idx1), np.stack(_idx2)

@Fetchers.register
def fetch_mnist_shvn(datasetdir):
    """ Fetch/prepare the MNIST-SHVN dataset for multi-modal dataset. Associate randomly one image from MNIST
    to one image from SVHN, sharing the same digit.

    Parameters
    ----------
    datasetdir: str
        the dataset destination folder.

    Returns
    -------
    item: namedtuple
        a named tuple containing 'input_path', 'output_path', and
        'metadata_path'.
    """
    logger.info("Loading MNIST/SHVN dataset.")
    classes = list(range(10))
    tf = transforms.ToTensor()

    if not os.path.isdir(datasetdir):
        os.mkdir(datasetdir)
    desc_path_mnist = os.path.join(datasetdir, "pynet_mnist.csv")
    input_path_mnist = os.path.join(datasetdir, "pynet_mnist.npy")
    desc_path_svhn = os.path.join(datasetdir, "pynet_svhn.csv")
    input_path_svhn = os.path.join(datasetdir, "pynet_svhn.npy")
    if not os.path.isfile(desc_path_mnist):
        mnist_trainset = torchvision.datasets.MNIST(
            root=datasetdir,
            train=True,
            download=True,
            transform=tf)
        mnist_testset = torchvision.datasets.MNIST(
            root=datasetdir,
            train=False,
            download=True,
            transform=tf)
        svhn_trainset = torchvision.datasets.SVHN(
            root=datasetdir,
            split='train',
            download=True,
            transform=tf)
        svhn_testset = torchvision.datasets.SVHN(
            root=datasetdir,
            split='test',
            download=True,
            transform=tf)


        mnist_train_idx, svhn_train_idx = get_matching_idx(mnist_trainset.targets.numpy(),
                                                           svhn_trainset.labels)
        mnist_test_idx, svhn_test_idx = get_matching_idx(mnist_testset.targets.numpy(),
                                                         svhn_testset.labels)

        data_mnist = torch.cat((torch.stack([mnist_trainset.data[i] for i in mnist_train_idx]),
                                torch.stack([mnist_testset.data[i] for i in mnist_test_idx])))
        labels_mnist = torch.cat((torch.stack([mnist_trainset.targets[i] for i in mnist_train_idx]),
                                  torch.stack([mnist_testset.targets[i] for i in mnist_test_idx])))
        data_svhn = np.concatenate((np.stack([svhn_trainset.data[i] for i in svhn_train_idx]),
                                    np.stack([svhn_testset.data[i] for i in svhn_test_idx])))
        labels_svhn = np.concatenate((np.stack([svhn_trainset.labels[i] for i in svhn_train_idx]),
                                        np.stack([svhn_testset.labels[i] for i in svhn_test_idx])))
        df_mnist = pd.DataFrame.from_dict(dict(digit=labels_mnist.numpy(), unique_key=np.arange(len(labels_mnist)),
                                               set=len(mnist_train_idx)*['train']+len(mnist_test_idx)*['test']))
        df_shvn = pd.DataFrame.from_dict(dict(digit=labels_svhn, unique_key=np.arange(len(labels_svhn)),
                                              set=len(svhn_train_idx)*['train']+len(svhn_test_idx)*['test']))

        np.save(input_path_mnist, data_mnist.detach().numpy())
        np.save(input_path_svhn, data_svhn)
        df_mnist.to_csv(desc_path_mnist, sep=",", index=False)
        df_shvn.to_csv(desc_path_svhn, sep=",", index=False)

        return True

def generate_mnist_svhn_examples():
    root = "/neurospin/psy_sbox/bd261576/vision_datasets/"
    data_mnist = np.load(root+"pynet_mnist.npy", mmap_mode='r')
    data_svhn = np.load(root+"pynet_svhn.npy", mmap_mode='r')
    df_mnist = pd.read_csv(root+"pynet_mnist.csv")
    df_svhn = pd.read_csv(root+"pynet_svhn.csv")

    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 5, figsize=(2*5, 2*2))
    idx = [100, 10000, 20000, 30000]
    axes[0,0].text(0.5, 0.5, "MNIST", horizontalalignment='center',
                   verticalalignment='center', transform=axes[0,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[1,0].text(0.5, 0.5, "SVHN", horizontalalignment='center',
                   verticalalignment='center', transform=axes[1,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[0,0].set_axis_off()
    axes[1,0].set_axis_off()
    for k, i in enumerate(idx, start=1):
        digit = df_mnist.iloc[i]['digit']
        idx_svhn = np.where(df_svhn['digit'].eq(digit))[0]
        axes[0,k].imshow(data_mnist[i])
        axes[1,k].imshow(data_svhn[idx_svhn[0]].transpose((1,2,0)))
        axes[0,k].set_axis_off()
        axes[1,k].set_axis_off()
    fig.subplots_adjust(wspace=0.05, hspace=0.05)
    fig.savefig("examples_mnist_svhn.png")
    plt.show()

def generate_colorfulmnist_svhn_examples():
    root = "/neurospin/psy_sbox/bd261576/vision_datasets/"
    data_mnist = np.load(root+"pynet_colorful_mnist.npy", mmap_mode='r')
    data_svhn = np.load(root+"pynet_svhn.npy", mmap_mode='r')
    df_mnist = pd.read_csv(root+"pynet_colorful_mnist.csv")
    df_svhn = pd.read_csv(root+"pynet_svhn.csv")

    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 5, figsize=(2*5, 2*2))
    idx = [0, 1, 2, 3]
    axes[0,0].text(0.5, 0.5, "Colorful MNIST", horizontalalignment='center',
                   verticalalignment='center', transform=axes[0,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[1,0].text(0.5, 0.5, "SVHN", horizontalalignment='center',
                   verticalalignment='center', transform=axes[1,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[0,0].set_axis_off()
    axes[1,0].set_axis_off()
    for k, i in enumerate(idx, start=1):
        digit = df_mnist.iloc[i]['mnist_digit']
        idx_svhn = np.where(df_svhn['digit'].eq(digit))[0]
        axes[0,k].imshow(data_mnist[i].transpose((1,2,0)))
        axes[1,k].imshow(data_svhn[idx_svhn[0]].transpose((1,2,0)))
        axes[0,k].set_axis_off()
        axes[1,k].set_axis_off()
    fig.subplots_adjust(wspace=0.05, hspace=0.05)
    fig.savefig("examples_colorfulmnist_svhn.png")
    plt.show()



@Fetchers.register
def fetch_colorful_mnist(datasetdir, svhn_df_pth=None):
    """ Fetch/prepare the Colorful MNIST dataset. To each MNIST image we add randomly a STL-10 image as a background to
    obtain an RGB image. Eventually, define a unique correspondance between SVHN image and ColorfulMNIST (MultiModal case)

    Parameters
    ----------
    datasetdir: str
        the dataset destination folder.
    svhn_df_pth: str
        path to SVHN pandas DataFrame, containing <unique_key>, <digit> and <set>
    Returns
    -------
    item: namedtuple
        a named tuple containing 'input_path', 'output_path', and
        'metadata_path'.
    """
    logger.info("Loading MNIST/STL-10 datasets.")
    classes = list(range(10))
    tf = None

    if not os.path.isdir(datasetdir):
        os.mkdir(datasetdir)

    desc_path_colorful_mnist = os.path.join(datasetdir, "pynet_colorful_mnist.csv")
    input_path_colorful_mnist = os.path.join(datasetdir, "pynet_colorful_mnist.npy")

    mnist_trainset = torchvision.datasets.MNIST(
        root=datasetdir,
        train=True,
        download=True,
        transform=tf)
    mnist_testset = torchvision.datasets.MNIST(
        root=datasetdir,
        train=False,
        download=True,
        transform=tf)
    stl10_trainset = torchvision.datasets.STL10(
        root=datasetdir,
        split='train',
        download=True,
        transform=tf)
    stl10_testset = torchvision.datasets.STL10(
        root=datasetdir,
        split='test',
        download=True,
        transform=tf)

    # Build the data-set by adding all labeled examples from STL-10 to MNIST (sample randomly from STL-10 for
    # each MNIST example)
    def get_color_imgs(big_d, small_d):
        from PIL import Image
        assert len(big_d) >= len(small_d)
        final_d = []
        label_d = {"mnist_digit": [], "stl10_class": []}
        tf = transforms.ToTensor()
        sampling_idx = np.arange(len(small_d))
        for i, (img, label) in enumerate(big_d):
            if i % len(small_d) == 0:
                np.random.shuffle(sampling_idx)
            (img2, label2) = small_d[sampling_idx[i % len(small_d)]]
            # Merge the 2 images
            mask = img.resize(img2.size)
            img = img.convert('RGB').resize(img2.size)  # MNIST
            final_img = tf(Image.composite(img, img2, mask))
            final_d.append(final_img)
            label_d["mnist_digit"].append(label)
            label_d["stl10_class"].append(label2)
        return (torch.stack(final_d), label_d)

    def set_unique_keys(this_df, other_df, this_data):
        """
        Map each <unique_key> in other df to a <unique_key> in this_df and make sure mnist_digit corresponds to digit.
        """
        this_df = this_df.copy()
        this_df["unique_key"] = np.nan
        for label in set(other_df.digit):
            for _set in ["train", "test"]:
                unique_keys = other_df[(other_df.digit == label) & other_df.set.eq(_set)].unique_key
                this_matching = this_df[(this_df.mnist_digit == label) & this_df.set.eq(_set)]
                assert len(this_matching) >= len(unique_keys)
                # take only the 1st len(unique_keys) matching samples
                this_mask = (this_df.mnist_digit == label) & this_df.set.eq(_set)
                this_mask[this_mask] = np.array(len(unique_keys) * [True] + (len(this_matching)-len(unique_keys)) * [False])
                assert np.all(this_mask & (this_df.mnist_digit == label) & this_df.set.eq(_set) == this_mask) \
                       and this_mask.sum() == len(unique_keys)
                this_df.loc[this_mask, "unique_key"] = unique_keys.values
        mask = ~this_df["unique_key"].isna()
        assert mask.sum() == len(other_df) and this_df[mask].unique_key.is_unique
        return this_df[mask], this_data[mask]

    colorful_mnist_train = get_color_imgs(mnist_trainset, stl10_trainset)
    colorful_mnist_test = get_color_imgs(mnist_testset, stl10_testset)
    data_colorful_mnist = torch.cat((colorful_mnist_train[0], colorful_mnist_test[0]))
    df_colorful_mnist = pd.DataFrame.from_dict(dict(mnist_digit=colorful_mnist_train[1]["mnist_digit"] +
                                                                colorful_mnist_test[1]["mnist_digit"],
                                                    stl10_class=colorful_mnist_train[1]["stl10_class"] +
                                                                colorful_mnist_test[1]["stl10_class"],
                                                    set=len(colorful_mnist_train[0]) * ['train'] +
                                                        len(colorful_mnist_test[0]) * ['test']))

    if svhn_df_pth is not None:
        df_svhn = pd.read_csv(svhn_df_pth)
        (df_colorful_mnist, data_colorful_mnist) = set_unique_keys(df_colorful_mnist, df_svhn, data_colorful_mnist)

    np.save(input_path_colorful_mnist, data_colorful_mnist.numpy())
    df_colorful_mnist.to_csv(desc_path_colorful_mnist, sep=",", index=False)

    return (data_colorful_mnist, df_colorful_mnist)


class MultiModalSimCLR_MNIST_SVHN_Dataset(ArrayDataset):
    """
        Generates pairs of images (v1^i, v2^j) from 2 modalities i \neq j such that digit(v1^i) == digit(v2^j). Each time an
        item k is required, always give the same v1^k with a matching v2^j randomly sampled among
        {v^j | digit(v^j) == digit(v^i)}
    """
    def __init__(self, inputs, *args, df=None, modalities=None, **kwargs):
        """
        :param inputs: a dict {mod: [np.arr]}
        :param *args, **kwargs: arguments to give to ArrayDataset
        :param df: a pandas DataFrame with {index_mnist, index_svhn, digit} in df.columns
        :param modalities: must be {'mnist', 'svhn'}
        """
        kwargs["concat_datasets"] = False
        super().__init__(None, *args, **kwargs)
        self.inputs = inputs
        self.modalities = modalities
        assert self.outputs is None, "Unknown output"
        assert set(self.modalities) == {'mnist', 'svhn'}, "Missing modalities: {}".format(self.modalities)
        if self.patch_size is not None or self.features_to_add is not None:
            raise NotImplementedError("Not yet implemented for multimodal dataset.")

        self.df = df[["index_%s" % m for m in modalities]+['digit']].values.copy()
        self.cumulative_sizes = {m: np.cumsum([len(inp) for inp in self.inputs[m]]) for m in modalities}

        # Transfos to sample from in the SimCLR framework
        s = 1
        color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)
        self.compose_transforms = {"svhn": transforms.Compose(
            [  # transforms.RandomApply([lambda x: add_blur(x, sigma=(0.2, 1))], p=0.5),
                lambda x: np.swapaxes(x, 0, 2),  # channel-last
                transforms.ToPILImage(mode='RGB'),
                transforms.RandomResizedCrop(size=32, scale=(0.5, 1)),
                transforms.RandomHorizontalFlip(),
                transforms.RandomApply([color_jitter], p=0.8),
                transforms.RandomGrayscale(p=0.2),
                transforms.ToTensor(),
                transforms.Normalize((.5, .5, .5), (.5, .5, .5))
            ]),

            "mnist": transforms.Compose(
                [  # transforms.RandomApply([lambda x: add_blur(x, sigma=(0.2, 1))], p=0.5),
                    transforms.ToPILImage(mode='L'),
                    transforms.RandomResizedCrop(size=28, scale=(0.5, 1)),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize((0.1307,), (0.3081,))  # mean, std of MNIST training set
                ])}

        self.minimum_tfs = {
            'svhn': transforms.Compose([lambda x: np.swapaxes(x, 0, 2),
                                        transforms.ToTensor(),
                                        transforms.Normalize((.5, .5, .5), (.5, .5, .5))]),
            'mnist': transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize((0.1307,), (0.3081,))])
        }


    def __getitem__(self, item):
        """
        :param item: int
        :return: DataItem(inputs, outputs=None, labels) where inputs and outputs are dict(mod: np.array) (if a modality is
        missing, then a None is put instead of np.array)
        """
        idx = self.indices[item]
        _inputs, _outputs, _labels = dict(), None, None

        if self.labels is not None:
            _labels = self.labels[idx]
            for tf in self.labels_transforms:
                _labels = tf(_labels)

        # 1) get the digit associated to the image
        digit = int(self.df[idx][-1])

        # 2) sample an image with the same digit
        idx_same_cls = np.random.choice([i for i in self.indices if self.df[i][-1]==digit])
        indices = [idx, idx_same_cls]

        for k, mod in enumerate(self.modalities):
            if np.isnan(self.df[indices[k]][k]):  # Missing modality
                _inputs[mod] = None
            else:
                mod_idx = int(self.df[indices[k]][k])
                dataset_idx = bisect.bisect_right(self.cumulative_sizes[mod], mod_idx)
                sample_idx = mod_idx - self.cumulative_sizes[mod][dataset_idx - 1] if dataset_idx > 0 else mod_idx
                _inputs[mod] = self.inputs[mod][dataset_idx][sample_idx]
                # Apply the transformations to the data
                for tf in self.input_transforms:
                    _inputs[mod] = tf(_inputs[mod])

        # Then generate pairs of data
        _inputs = self.generate_pairs(_inputs)

        return DataItem(inputs=_inputs, outputs=_outputs, labels=_labels)

    def generate_pairs(self, x):
        np.random.seed()
        out = dict()
        for mod in x:
            if x[mod] is None:
                out[mod] = None
            else:
                x1, x2, x3 = self.compose_transforms[mod](x[mod]), \
                             self.compose_transforms[mod](x[mod]), \
                             self.minimum_tfs[mod](x[mod])
                out[mod] = np.stack((x1, x2, x3), axis=0)
        return out

class MultiModalSimCLR_ColorfulMNIST_SVHN_Dataset(MultiModalSimCLR_MNIST_SVHN_Dataset):
    """
        Generates pairs of images (v1^i, v2^j) from 2 modalities i \neq j such that digit(v1^i) == digit(v2^j). Each time an
        item k is required, always give the same v1^k with a matching v2^j randomly sampled among
        {v^j | digit(v^j) == digit(v^i)}
    """
    def __init__(self, inputs, *args, df=None, modalities=None, self_supervision=None, **kwargs):
        """
        :param inputs: a dict {mod: [np.arr]}
        :param *args, **kwargs: arguments to give to ArrayDataset
        :param df: a pandas DataFrame with {index_mnist, index_svhn, digit} in df.columns
        :param modalities: must be {'mnist', 'svhn'}
        """
        kwargs["concat_datasets"] = False
        ArrayDataset.__init__(self, None, *args, **kwargs)
        self.inputs = inputs
        self.modalities = modalities
        assert self.outputs is None, "Unknown output"
        assert set(self.modalities) == {'colorful_mnist', 'svhn'}, "Missing modalities: {}".format(self.modalities)
        if self.patch_size is not None or self.features_to_add is not None:
            raise NotImplementedError("Not yet implemented for multimodal dataset.")

        self.df = df[["index_%s" % m for m in modalities] + ['mnist_digit']].values.copy()
        self.cumulative_sizes = {m: np.cumsum([len(inp) for inp in self.inputs[m]]) for m in modalities}

        # Transfos to sample from in the SimCLR framework
        s = 1
        color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)

        normalization = lambda x: (255.0 * (x -x.min())/(x.max()-x.min())).astype(np.uint8).swapaxes(0,2)

        self.compose_transforms = {"svhn": [ lambda x: np.swapaxes(x, 0, 2),  transforms.ToPILImage(mode='RGB')],
                                   "colorful_mnist": [normalization, transforms.ToPILImage(mode='RGB')]}

        self.minimum_tfs = {
            'svhn': transforms.Compose([lambda x: np.swapaxes(x, 0, 2),
                                        transforms.ToTensor(),
                                        transforms.Normalize((.5, .5, .5), (.5, .5, .5))]),
            'colorful_mnist': transforms.Compose([normalization,
                                                  transforms.ToTensor(),
                                                  transforms.Normalize((.5, .5, .5), (.5, .5, .5))])
        }

        self_supervision = set(self_supervision or [])
        if "crop" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomResizedCrop(size=32, scale=(0.5, 1)))
            self.compose_transforms["colorful_mnist"].append(transforms.RandomResizedCrop(size=96, scale=(0.2, 1)))
            self_supervision -= {"crop"}
        if "flip" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomHorizontalFlip())
            self.compose_transforms["colorful_mnist"].append(transforms.RandomHorizontalFlip())
            self_supervision -= {"flip"}
        if "jittering" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomApply([color_jitter], p=0.8))
            self.compose_transforms["colorful_mnist"].append(transforms.RandomApply([color_jitter], p=0.8))
            self_supervision -= {"jittering"}
        if "gray" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomGrayscale(p=0.2))
            self.compose_transforms["colorful_mnist"].append(transforms.RandomGrayscale(p=0.2))
            self_supervision -= {"gray"}

        self.compose_transforms["svhn"].extend([transforms.ToTensor(),
                                                transforms.Normalize((.5, .5, .5), (.5, .5, .5))])
        self.compose_transforms["colorful_mnist"].extend([transforms.ToTensor(),
                                                          transforms.Normalize((.5, .5, .5), (.5, .5, .5))])

        self.compose_transforms["svhn"] = transforms.Compose(self.compose_transforms["svhn"])
        self.compose_transforms["colorful_mnist"] = transforms.Compose(self.compose_transforms["colorful_mnist"])

        if len(self_supervision) > 0:
            logger.warning("Unknown transformations {}. Ignored.".format(self_supervision))
        logger.info("SimCLR Transformations:\n\t (svhn) {}\n\t (colorful MNIST) {}".
                    format(self.compose_transforms["svhn"], self.compose_transforms["colorful_mnist"]))

# -*- coding: utf-8 -*-
##########################################################################
# NSAp - Copyright (C) CEA, 2019
# Distributed under the terms of the CeCILL-B license, as published by
# the CEA-CNRS-INRIA. Refer to the LICENSE file or to
# http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html
# for details.
##########################################################################

"""
Module that provides functions to prepare the Multimodal MNIST-SVHN dataset.
"""

# Imports
import os
import logging
import torchvision
import torch
import torchvision.transforms as transforms
from collections import namedtuple
import numpy as np
import pandas as pd
from pynet.datasets import Fetchers
from pynet.datasets.core import ArrayDataset, DataItem
import bisect


# Global parameters
Item = namedtuple("Item", ["input_path", "output_path", "metadata_path",
                           "labels"])
logger = logging.getLogger("pynet")


def get_matching_idx(labels1, labels2):
    """
    :param
    labels1: np.array, list of labels from dataset 1
    labels2: np.array, list of labels from dataset 2
    :return
    (idx1, idx2): 2 lists of indexes defining the unique correspondance between labels1 and labels2
    """
    assert isinstance(labels1, np.ndarray) and isinstance(labels2, np.ndarray), \
        "Wrong type: %s,%s"%(type(labels1), type(labels2))
    assert set(labels1) == set(labels2), "Missing labels"


    _idx1, _idx2 = [], []
    n1, n2 = len(labels1), len(labels2)
    for label in set(labels1):
        idx1 = np.arange(n1)[labels1 == label]
        idx2 = np.arange(n2)[labels2 == label]
        n = min(len(idx1), len(idx2))
        _idx1.extend(idx1[np.random.permutation(n)])
        _idx2.extend(idx2[np.random.permutation(n)])
    return np.stack(_idx1), np.stack(_idx2)

@Fetchers.register
def fetch_mnist_shvn(datasetdir):
    """ Fetch/prepare the MNIST-SHVN dataset for multi-modal dataset. Associate randomly one image from MNIST
    to one image from SVHN, sharing the same digit.

    Parameters
    ----------
    datasetdir: str
        the dataset destination folder.

    Returns
    -------
    item: namedtuple
        a named tuple containing 'input_path', 'output_path', and
        'metadata_path'.
    """
    logger.info("Loading MNIST/SHVN dataset.")
    classes = list(range(10))
    tf = transforms.ToTensor()

    if not os.path.isdir(datasetdir):
        os.mkdir(datasetdir)
    desc_path_mnist = os.path.join(datasetdir, "pynet_mnist.csv")
    input_path_mnist = os.path.join(datasetdir, "pynet_mnist.npy")
    desc_path_svhn = os.path.join(datasetdir, "pynet_svhn.csv")
    input_path_svhn = os.path.join(datasetdir, "pynet_svhn.npy")
    if not os.path.isfile(desc_path_mnist):
        mnist_trainset = torchvision.datasets.MNIST(
            root=datasetdir,
            train=True,
            download=True,
            transform=tf)
        mnist_testset = torchvision.datasets.MNIST(
            root=datasetdir,
            train=False,
            download=True,
            transform=tf)
        svhn_trainset = torchvision.datasets.SVHN(
            root=datasetdir,
            split='train',
            download=True,
            transform=tf)
        svhn_testset = torchvision.datasets.SVHN(
            root=datasetdir,
            split='test',
            download=True,
            transform=tf)


        mnist_train_idx, svhn_train_idx = get_matching_idx(mnist_trainset.targets.numpy(),
                                                           svhn_trainset.labels)
        mnist_test_idx, svhn_test_idx = get_matching_idx(mnist_testset.targets.numpy(),
                                                         svhn_testset.labels)

        data_mnist = torch.cat((torch.stack([mnist_trainset.data[i] for i in mnist_train_idx]),
                                torch.stack([mnist_testset.data[i] for i in mnist_test_idx])))
        labels_mnist = torch.cat((torch.stack([mnist_trainset.targets[i] for i in mnist_train_idx]),
                                  torch.stack([mnist_testset.targets[i] for i in mnist_test_idx])))
        data_svhn = np.concatenate((np.stack([svhn_trainset.data[i] for i in svhn_train_idx]),
                                    np.stack([svhn_testset.data[i] for i in svhn_test_idx])))
        labels_svhn = np.concatenate((np.stack([svhn_trainset.labels[i] for i in svhn_train_idx]),
                                        np.stack([svhn_testset.labels[i] for i in svhn_test_idx])))
        df_mnist = pd.DataFrame.from_dict(dict(digit=labels_mnist.numpy(), unique_key=np.arange(len(labels_mnist)),
                                               set=len(mnist_train_idx)*['train']+len(mnist_test_idx)*['test']))
        df_shvn = pd.DataFrame.from_dict(dict(digit=labels_svhn, unique_key=np.arange(len(labels_svhn)),
                                              set=len(svhn_train_idx)*['train']+len(svhn_test_idx)*['test']))

        np.save(input_path_mnist, data_mnist.detach().numpy())
        np.save(input_path_svhn, data_svhn)
        df_mnist.to_csv(desc_path_mnist, sep=",", index=False)
        df_shvn.to_csv(desc_path_svhn, sep=",", index=False)

        return True

def generate_mnist_svhn_examples():
    root = "/neurospin/psy_sbox/bd261576/vision_datasets/"
    data_mnist = np.load(root+"pynet_mnist.npy", mmap_mode='r')
    data_svhn = np.load(root+"pynet_svhn.npy", mmap_mode='r')
    df_mnist = pd.read_csv(root+"pynet_mnist.csv")
    df_svhn = pd.read_csv(root+"pynet_svhn.csv")

    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 5, figsize=(2*5, 2*2))
    idx = [100, 10000, 20000, 30000]
    axes[0,0].text(0.5, 0.5, "MNIST", horizontalalignment='center',
                   verticalalignment='center', transform=axes[0,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[1,0].text(0.5, 0.5, "SVHN", horizontalalignment='center',
                   verticalalignment='center', transform=axes[1,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[0,0].set_axis_off()
    axes[1,0].set_axis_off()
    for k, i in enumerate(idx, start=1):
        digit = df_mnist.iloc[i]['digit']
        idx_svhn = np.where(df_svhn['digit'].eq(digit))[0]
        axes[0,k].imshow(data_mnist[i])
        axes[1,k].imshow(data_svhn[idx_svhn[0]].transpose((1,2,0)))
        axes[0,k].set_axis_off()
        axes[1,k].set_axis_off()
    fig.subplots_adjust(wspace=0.05, hspace=0.05)
    fig.savefig("examples_mnist_svhn.png")
    plt.show()

def generate_colorfulmnist_svhn_examples():
    root = "/neurospin/psy_sbox/bd261576/vision_datasets/"
    data_mnist = np.load(root+"pynet_colorful_mnist.npy", mmap_mode='r')
    data_svhn = np.load(root+"pynet_svhn.npy", mmap_mode='r')
    df_mnist = pd.read_csv(root+"pynet_colorful_mnist.csv")
    df_svhn = pd.read_csv(root+"pynet_svhn.csv")

    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 5, figsize=(2*5, 2*2))
    idx = [0, 1, 2, 3]
    axes[0,0].text(0.5, 0.5, "Colorful MNIST", horizontalalignment='center',
                   verticalalignment='center', transform=axes[0,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[1,0].text(0.5, 0.5, "SVHN", horizontalalignment='center',
                   verticalalignment='center', transform=axes[1,0].transAxes,
                   fontname='serif', fontweight='bold')
    axes[0,0].set_axis_off()
    axes[1,0].set_axis_off()
    for k, i in enumerate(idx, start=1):
        digit = df_mnist.iloc[i]['mnist_digit']
        idx_svhn = np.where(df_svhn['digit'].eq(digit))[0]
        axes[0,k].imshow(data_mnist[i].transpose((1,2,0)))
        axes[1,k].imshow(data_svhn[idx_svhn[0]].transpose((1,2,0)))
        axes[0,k].set_axis_off()
        axes[1,k].set_axis_off()
    fig.subplots_adjust(wspace=0.05, hspace=0.05)
    fig.savefig("examples_colorfulmnist_svhn.png")
    plt.show()



@Fetchers.register
def fetch_colorful_mnist(datasetdir, svhn_df_pth=None):
    """ Fetch/prepare the Colorful MNIST dataset. To each MNIST image we add randomly a STL-10 image as a background to
    obtain an RGB image. Eventually, define a unique correspondance between SVHN image and ColorfulMNIST (MultiModal case)

    Parameters
    ----------
    datasetdir: str
        the dataset destination folder.
    svhn_df_pth: str
        path to SVHN pandas DataFrame, containing <unique_key>, <digit> and <set>
    Returns
    -------
    item: namedtuple
        a named tuple containing 'input_path', 'output_path', and
        'metadata_path'.
    """
    logger.info("Loading MNIST/STL-10 datasets.")
    classes = list(range(10))
    tf = None

    if not os.path.isdir(datasetdir):
        os.mkdir(datasetdir)

    desc_path_colorful_mnist = os.path.join(datasetdir, "pynet_colorful_mnist.csv")
    input_path_colorful_mnist = os.path.join(datasetdir, "pynet_colorful_mnist.npy")

    mnist_trainset = torchvision.datasets.MNIST(
        root=datasetdir,
        train=True,
        download=True,
        transform=tf)
    mnist_testset = torchvision.datasets.MNIST(
        root=datasetdir,
        train=False,
        download=True,
        transform=tf)
    stl10_trainset = torchvision.datasets.STL10(
        root=datasetdir,
        split='train',
        download=True,
        transform=tf)
    stl10_testset = torchvision.datasets.STL10(
        root=datasetdir,
        split='test',
        download=True,
        transform=tf)

    # Build the data-set by adding all labeled examples from STL-10 to MNIST (sample randomly from STL-10 for
    # each MNIST example)
    def get_color_imgs(big_d, small_d):
        from PIL import Image
        assert len(big_d) >= len(small_d)
        final_d = []
        label_d = {"mnist_digit": [], "stl10_class": []}
        tf = transforms.ToTensor()
        sampling_idx = np.arange(len(small_d))
        for i, (img, label) in enumerate(big_d):
            if i % len(small_d) == 0:
                np.random.shuffle(sampling_idx)
            (img2, label2) = small_d[sampling_idx[i % len(small_d)]]
            # Merge the 2 images
            mask = img.resize(img2.size)
            img = img.convert('RGB').resize(img2.size)  # MNIST
            final_img = tf(Image.composite(img, img2, mask))
            final_d.append(final_img)
            label_d["mnist_digit"].append(label)
            label_d["stl10_class"].append(label2)
        return (torch.stack(final_d), label_d)

    def set_unique_keys(this_df, other_df, this_data):
        """
        Map each <unique_key> in other df to a <unique_key> in this_df and make sure mnist_digit corresponds to digit.
        """
        this_df = this_df.copy()
        this_df["unique_key"] = np.nan
        for label in set(other_df.digit):
            for _set in ["train", "test"]:
                unique_keys = other_df[(other_df.digit == label) & other_df.set.eq(_set)].unique_key
                this_matching = this_df[(this_df.mnist_digit == label) & this_df.set.eq(_set)]
                assert len(this_matching) >= len(unique_keys)
                # take only the 1st len(unique_keys) matching samples
                this_mask = (this_df.mnist_digit == label) & this_df.set.eq(_set)
                this_mask[this_mask] = np.array(len(unique_keys) * [True] + (len(this_matching)-len(unique_keys)) * [False])
                assert np.all(this_mask & (this_df.mnist_digit == label) & this_df.set.eq(_set) == this_mask) \
                       and this_mask.sum() == len(unique_keys)
                this_df.loc[this_mask, "unique_key"] = unique_keys.values
        mask = ~this_df["unique_key"].isna()
        assert mask.sum() == len(other_df) and this_df[mask].unique_key.is_unique
        return this_df[mask], this_data[mask]

    colorful_mnist_train = get_color_imgs(mnist_trainset, stl10_trainset)
    colorful_mnist_test = get_color_imgs(mnist_testset, stl10_testset)
    data_colorful_mnist = torch.cat((colorful_mnist_train[0], colorful_mnist_test[0]))
    df_colorful_mnist = pd.DataFrame.from_dict(dict(mnist_digit=colorful_mnist_train[1]["mnist_digit"] +
                                                                colorful_mnist_test[1]["mnist_digit"],
                                                    stl10_class=colorful_mnist_train[1]["stl10_class"] +
                                                                colorful_mnist_test[1]["stl10_class"],
                                                    set=len(colorful_mnist_train[0]) * ['train'] +
                                                        len(colorful_mnist_test[0]) * ['test']))

    if svhn_df_pth is not None:
        df_svhn = pd.read_csv(svhn_df_pth)
        (df_colorful_mnist, data_colorful_mnist) = set_unique_keys(df_colorful_mnist, df_svhn, data_colorful_mnist)

    np.save(input_path_colorful_mnist, data_colorful_mnist.numpy())
    df_colorful_mnist.to_csv(desc_path_colorful_mnist, sep=",", index=False)

    return (data_colorful_mnist, df_colorful_mnist)


class MultiModalSimCLR_MNIST_SVHN_Dataset(ArrayDataset):
    """
        Generates pairs of images (v1^i, v2^j) from 2 modalities i \neq j such that digit(v1^i) == digit(v2^j). Each time an
        item k is required, always give the same v1^k with a matching v2^j randomly sampled among
        {v^j | digit(v^j) == digit(v^i)}
    """
    def __init__(self, inputs, *args, df=None, modalities=None, **kwargs):
        """
        :param inputs: a dict {mod: [np.arr]}
        :param *args, **kwargs: arguments to give to ArrayDataset
        :param df: a pandas DataFrame with {index_mnist, index_svhn, digit} in df.columns
        :param modalities: must be {'mnist', 'svhn'}
        """
        kwargs["concat_datasets"] = False
        super().__init__(None, *args, **kwargs)
        self.inputs = inputs
        self.modalities = modalities
        assert self.outputs is None, "Unknown output"
        assert set(self.modalities) == {'mnist', 'svhn'}, "Missing modalities: {}".format(self.modalities)
        if self.patch_size is not None or self.features_to_add is not None:
            raise NotImplementedError("Not yet implemented for multimodal dataset.")

        self.df = df[["index_%s" % m for m in modalities]+['digit']].values.copy()
        self.cumulative_sizes = {m: np.cumsum([len(inp) for inp in self.inputs[m]]) for m in modalities}

        # Transfos to sample from in the SimCLR framework
        s = 1
        color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)
        self.compose_transforms = {"svhn": transforms.Compose(
            [  # transforms.RandomApply([lambda x: add_blur(x, sigma=(0.2, 1))], p=0.5),
                lambda x: np.swapaxes(x, 0, 2),  # channel-last
                transforms.ToPILImage(mode='RGB'),
                transforms.RandomResizedCrop(size=32, scale=(0.5, 1)),
                transforms.RandomHorizontalFlip(),
                transforms.RandomApply([color_jitter], p=0.8),
                transforms.RandomGrayscale(p=0.2),
                transforms.ToTensor(),
                transforms.Normalize((.5, .5, .5), (.5, .5, .5))
            ]),

            "mnist": transforms.Compose(
                [  # transforms.RandomApply([lambda x: add_blur(x, sigma=(0.2, 1))], p=0.5),
                    transforms.ToPILImage(mode='L'),
                    transforms.RandomResizedCrop(size=28, scale=(0.5, 1)),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize((0.1307,), (0.3081,))  # mean, std of MNIST training set
                ])}

        self.minimum_tfs = {
            'svhn': transforms.Compose([lambda x: np.swapaxes(x, 0, 2),
                                        transforms.ToTensor(),
                                        transforms.Normalize((.5, .5, .5), (.5, .5, .5))]),
            'mnist': transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize((0.1307,), (0.3081,))])
        }


    def __getitem__(self, item):
        """
        :param item: int
        :return: DataItem(inputs, outputs=None, labels) where inputs and outputs are dict(mod: np.array) (if a modality is
        missing, then a None is put instead of np.array)
        """
        idx = self.indices[item]
        _inputs, _outputs, _labels = dict(), None, None

        if self.labels is not None:
            _labels = self.labels[idx]
            for tf in self.labels_transforms:
                _labels = tf(_labels)

        # 1) get the digit associated to the image
        digit = int(self.df[idx][-1])

        # 2) sample an image with the same digit
        idx_same_cls = np.random.choice([i for i in self.indices if self.df[i][-1]==digit])
        indices = [idx, idx_same_cls]

        for k, mod in enumerate(self.modalities):
            if np.isnan(self.df[indices[k]][k]):  # Missing modality
                _inputs[mod] = None
            else:
                mod_idx = int(self.df[indices[k]][k])
                dataset_idx = bisect.bisect_right(self.cumulative_sizes[mod], mod_idx)
                sample_idx = mod_idx - self.cumulative_sizes[mod][dataset_idx - 1] if dataset_idx > 0 else mod_idx
                _inputs[mod] = self.inputs[mod][dataset_idx][sample_idx]
                # Apply the transformations to the data
                for tf in self.input_transforms:
                    _inputs[mod] = tf(_inputs[mod])

        # Then generate pairs of data
        _inputs = self.generate_pairs(_inputs)

        return DataItem(inputs=_inputs, outputs=_outputs, labels=_labels)

    def generate_pairs(self, x):
        np.random.seed()
        out = dict()
        for mod in x:
            if x[mod] is None:
                out[mod] = None
            else:
                x1, x2, x3 = self.compose_transforms[mod](x[mod]), \
                             self.compose_transforms[mod](x[mod]), \
                             self.minimum_tfs[mod](x[mod])
                out[mod] = np.stack((x1, x2, x3), axis=0)
        return out

class MultiModalSimCLR_ColorfulMNIST_SVHN_Dataset(MultiModalSimCLR_MNIST_SVHN_Dataset):
    """
        Generates pairs of images (v1^i, v2^j) from 2 modalities i \neq j such that digit(v1^i) == digit(v2^j). Each time an
        item k is required, always give the same v1^k with a matching v2^j randomly sampled among
        {v^j | digit(v^j) == digit(v^i)}
    """
    def __init__(self, inputs, *args, df=None, modalities=None, self_supervision=None, **kwargs):
        """
        :param inputs: a dict {mod: [np.arr]}
        :param *args, **kwargs: arguments to give to ArrayDataset
        :param df: a pandas DataFrame with {index_mnist, index_svhn, digit} in df.columns
        :param modalities: must be {'mnist', 'svhn'}
        """
        kwargs["concat_datasets"] = False
        ArrayDataset.__init__(self, None, *args, **kwargs)
        self.inputs = inputs
        self.modalities = modalities
        assert self.outputs is None, "Unknown output"
        assert set(self.modalities) == {'colorful_mnist', 'svhn'}, "Missing modalities: {}".format(self.modalities)
        if self.patch_size is not None or self.features_to_add is not None:
            raise NotImplementedError("Not yet implemented for multimodal dataset.")

        self.df = df[["index_%s" % m for m in modalities] + ['mnist_digit']].values.copy()
        self.cumulative_sizes = {m: np.cumsum([len(inp) for inp in self.inputs[m]]) for m in modalities}

        # Transfos to sample from in the SimCLR framework
        s = 1
        color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)

        normalization = lambda x: (255.0 * (x -x.min())/(x.max()-x.min())).astype(np.uint8).swapaxes(0,2)

        self.compose_transforms = {"svhn": [ lambda x: np.swapaxes(x, 0, 2),  transforms.ToPILImage(mode='RGB')],
                                   "colorful_mnist": [normalization, transforms.ToPILImage(mode='RGB')]}

        self.minimum_tfs = {
            'svhn': transforms.Compose([lambda x: np.swapaxes(x, 0, 2),
                                        transforms.ToTensor(),
                                        transforms.Normalize((.5, .5, .5), (.5, .5, .5))]),
            'colorful_mnist': transforms.Compose([normalization,
                                                  transforms.ToTensor(),
                                                  transforms.Normalize((.5, .5, .5), (.5, .5, .5))])
        }

        self_supervision = set(self_supervision or [])
        if "crop" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomResizedCrop(size=32, scale=(0.5, 1)))
            self.compose_transforms["colorful_mnist"].append(transforms.RandomResizedCrop(size=96, scale=(0.2, 1)))
            self_supervision -= {"crop"}
        if "flip" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomHorizontalFlip())
            self.compose_transforms["colorful_mnist"].append(transforms.RandomHorizontalFlip())
            self_supervision -= {"flip"}
        if "jittering" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomApply([color_jitter], p=0.8))
            self.compose_transforms["colorful_mnist"].append(transforms.RandomApply([color_jitter], p=0.8))
            self_supervision -= {"jittering"}
        if "gray" in self_supervision:
            self.compose_transforms["svhn"].append(transforms.RandomGrayscale(p=0.2))
            self.compose_transforms["colorful_mnist"].append(transforms.RandomGrayscale(p=0.2))
            self_supervision -= {"gray"}

        self.compose_transforms["svhn"].extend([transforms.ToTensor(),
                                                transforms.Normalize((.5, .5, .5), (.5, .5, .5))])
        self.compose_transforms["colorful_mnist"].extend([transforms.ToTensor(),
                                                          transforms.Normalize((.5, .5, .5), (.5, .5, .5))])

        self.compose_transforms["svhn"] = transforms.Compose(self.compose_transforms["svhn"])
        self.compose_transforms["colorful_mnist"] = transforms.Compose(self.compose_transforms["colorful_mnist"])

        if len(self_supervision) > 0:
            logger.warning("Unknown transformations {}. Ignored.".format(self_supervision))
        logger.info("SimCLR Transformations:\n\t (svhn) {}\n\t (colorful MNIST) {}".
                    format(self.compose_transforms["svhn"], self.compose_transforms["colorful_mnist"]))

# -*- coding: utf-8 -*- ########################################################################## # NSAp - Copyright (C) CEA, 2019 # Distributed under the terms of the CeCILL-B license, as published by # the CEA-CNRS-INRIA. Refer to the LICENSE file or to # http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html # for details. ########################################################################## Module that provides functions to prepare the Multimodal MNIST-SVHN dataset. # Imports # Global parameters :param     labels1: np.array, list of labels from dataset 1     labels2: np.array, list of labels from dataset 2     :return     (idx1, idx2): 2 lists of indexes defining the unique correspondance between labels1 and labels2 Fetch/prepare the MNIST-SHVN dataset for multi-modal dataset. Associate randomly one image from MNIST     to one image from SVHN, sharing the same digit.      Parameters     ----------     datasetdir: str         the dataset destination folder.      Returns     -------     item: namedtuple         a named tuple containing 'input_path', 'output_path', and         'metadata_path'. Fetch/prepare the Colorful MNIST dataset. To each MNIST image we add randomly a STL-10 image as a background to     obtain an RGB image. Eventually, define a unique correspondance between SVHN image and ColorfulMNIST (MultiModal case)      Parameters     ----------     datasetdir: str         the dataset destination folder.     svhn_df_pth: str         path to SVHN pandas DataFrame, containing <unique_key>, <digit> and <set>     Returns     -------     item: namedtuple         a named tuple containing 'input_path', 'output_path', and         'metadata_path'. # Build the data-set by adding all labeled examples from STL-10 to MNIST (sample randomly from STL-10 for # each MNIST example) # Merge the 2 images # MNIST Map each <unique_key> in other df to a <unique_key> in this_df and make sure mnist_digit corresponds to digit. # take only the 1st len(unique_keys) matching samples Generates pairs of images (v1^i, v2^j) from 2 modalities i \neq j such that digit(v1^i) == digit(v2^j). Each time an         item k is required, always give the same v1^k with a matching v2^j randomly sampled among         {v^j | digit(v^j) == digit(v^i)} :param inputs: a dict {mod: [np.arr]}         :param *args, **kwargs: arguments to give to ArrayDataset         :param df: a pandas DataFrame with {index_mnist, index_svhn, digit} in df.columns         :param modalities: must be {'mnist', 'svhn'} # Transfos to sample from in the SimCLR framework # transforms.RandomApply([lambda x: add_blur(x, sigma=(0.2, 1))], p=0.5), # channel-last # transforms.RandomApply([lambda x: add_blur(x, sigma=(0.2, 1))], p=0.5), # mean, std of MNIST training set :param item: int         :return: DataItem(inputs, outputs=None, labels) where inputs and outputs are dict(mod: np.array) (if a modality is         missing, then a None is put instead of np.array) # 1) get the digit associated to the image # 2) sample an image with the same digit # Missing modality # Apply the transformations to the data # Then generate pairs of data Generates pairs of images (v1^i, v2^j) from 2 modalities i \neq j such that digit(v1^i) == digit(v2^j). Each time an         item k is required, always give the same v1^k with a matching v2^j randomly sampled among         {v^j | digit(v^j) == digit(v^i)} :param inputs: a dict {mod: [np.arr]}         :param *args, **kwargs: arguments to give to ArrayDataset         :param df: a pandas DataFrame with {index_mnist, index_svhn, digit} in df.columns         :param modalities: must be {'mnist', 'svhn'} # Transfos to sample from in the SimCLR framework