#
# (c) All rights reserved. ECOLE POLYTECHNIQUE FÉDÉRALE DE LAUSANNE,
# Switzerland, Laboratory of Prof. Mackenzie W. Mathis (UPMWMATHIS) and
# original authors: Steffen Schneider, Jin H Lee, Mackenzie W Mathis. 2023.
#
# Source code:
# https://github.com/AdaptiveMotorControlLab/CEBRA
#
# Please see LICENSE.md for the full license document:
# https://github.com/AdaptiveMotorControlLab/CEBRA/LICENSE.md
#
"""Allen single mouse dataset.

References:
    *Deitch, Daniel, Alon Rubin, and Yaniv Ziv. "Representational drift in the mouse visual cortex." Current biology 31.19 (2021): 4327-4339.
    *de Vries, Saskia EJ, et al. "A large-scale standardized physiological survey reveals functional organization of the mouse visual cortex." Nature neuroscience 23.1 (2020): 138-151.
    *https://github.com/zivlab/visual_drift
    *http://observatory.brain-map.org/visualcoding

"""
import glob
import hashlib
import os

import h5py
import joblib
import numpy as np
import pandas as pd
import scipy.io
import torch
from numpy.random import Generator
from numpy.random import PCG64
from sklearn.decomposition import PCA

import cebra.data
from cebra.datasets import get_datapath
from cebra.datasets import init
from cebra.datasets import parametrize
from cebra.datasets import register

_SINGLE_SESSION_CA = (
    get_datapath(
        "allen/visual_drift/data/calcium_excitatory/VISp/680156909.mat"),
    get_datapath(
        "allen/visual_drift/data/calcium_excitatory/VISp/511510779.mat"),
    get_datapath(
        "allen/visual_drift/data/calcium_excitatory/VISp/679702882.mat"),
    get_datapath(
        "allen/visual_drift/data/calcium_excitatory/VISp/688678764.mat"),
)


@parametrize(
    "allen-movie1-ca-single-session-{session_id}",
    session_id=range(len(_SINGLE_SESSION_CA)),
)
class SingleSessionAllenCa(cebra.data.SingleSessionDataset):
    """A single mouse 30Hz calcium events dataset during the allen MOVIE1 stimulus.

    A dataset of a single mouse 30Hz calcium events from the excitatory neurons in the primary visual cortex
    during the 10 repeats of the MOVIE1 stimulus in session type A. The preprocessed data from *Deitch et al. (2021) are used.
    The continuous labels corresponding to a DINO embedding of each stimulus frame.

    Args:
        session_id: The integer value to pick a session among 4 sessions with the largest number of recorded neruons. Choose between 0-3.
        frame_feature_path: The path of the movie frame features.

    """

    def __init__(
        self,
        session_id: int,
        frame_feature_path: str = get_datapath(
            "allen/features/allen_movies/vit_base/8/movie_one_image_stack.npz/testfeat.pth"
        ),
        pca: bool = False,
    ):
        self.path = _SINGLE_SESSION_CA[session_id]
        traces = scipy.io.loadmat(self.path)
        if pca:
            pca_ = PCA()
            neural = pca_.fit_transform(
                traces["filtered_traces_days_events"][0,
                                                      0].transpose(1,
                                                                   0))[:, :32]
        else:
            neural = traces["filtered_traces_days_events"][0, 0].transpose(1, 0)
        self.neural = torch.from_numpy(neural).float()
        frame_feature = torch.load(frame_feature_path)
        self.index = frame_feature.repeat(10, 1)

    def __len__(self):
        return self.neural.size(0)

    @property
    def continuous_index(self):
        return self.index

    @property
    def input_dimension(self):
        return self.neural.size(1)

    def __getitem__(self, index):
        index = self.expand_index(index)
        return self.neural[index].transpose(2, 1)


@parametrize(
    "allen-movie1-ca-single-session-corrupt-{session_id}",
    session_id=range(len(_SINGLE_SESSION_CA)),
)
class SingleSessionAllenCa(cebra.data.SingleSessionDataset):
    """A corrupted single mouse 30Hz calcium events dataset during the allen MOVIE1 stimulus.

    A dataset of a single mouse 30Hz calcium events from the excitatory neurons in the primary visual cortex
    during the 10 repeats of the MOVIE1 stimulus in session type A. The preprocessed data from *Deitch et al. (2021) are used.
    The continuous labels corresponding to a DINO embedding of each stimulus frame, but in randomly shuffled order.

    Args:
        session_id: The integer value to pick a session among 4 sessions with the largest number of recorded neruons. Choose between 0-3.
        frame_feature_path: The path of the movie frame features.

    """

    def __init__(
        self,
        session_id: int,
        frame_feature_path: str = get_datapath(
            "allen/features/allen_movies/vit_base/8/movie_one_image_stack.npz/testfeat.pth"
        ),
        pca: bool = False,
    ):
        self.path = _SINGLE_SESSION_CA[session_id]
        traces = scipy.io.loadmat(self.path)
        if pca:
            pca_ = PCA()
            neural = pca_.fit_transform(
                traces["filtered_traces_days_events"][0,
                                                      0].transpose(1,
                                                                   0))[:, :32]
        else:
            neural = traces["filtered_traces_days_events"][0, 0].transpose(1, 0)

        self.neural = torch.from_numpy(neural).float()
        frame_feature = torch.load(frame_feature_path)
        self.frame_index = np.arange(900)
        rng = np.random.Generator(np.random.PCG64(111))
        rng.shuffle(self.frame_index)
        self.index = frame_feature.repeat(10, 1)[np.tile(self.frame_index, 10)]

    def __len__(self):
        return self.neural.size(0)

    @property
    def continuous_index(self):
        return self.index

    @property
    def input_dimension(self):
        return self.neural.size(1)

    def __getitem__(self, index):
        index = self.expand_index(index)
        return self.neural[index].transpose(2, 1)


@parametrize(
    "allen-movie1-ca-single-session-time-{session_id}",
    session_id=range(len(_SINGLE_SESSION_CA)),
)
class SingleSessionAllenCaTime(SingleSessionAllenCa):
    """A single mouse 30Hz calcium events dataset during the allen MOVIE1 stimulus.

    A dataset of a single mouse 30Hz calcium events from the excitatory neurons in the primary visual cortex
    during the 10 repeats of the MOVIE1 stimulus in session type A. The preprocessed data from *Deitch et al. (2021) are used.
    There is no behavioral variable used as the continuous label.

    Args:
        session_id: The integer value to pick a session among 4 sessions with the largest number of recorded neruons. Choose between 0-3.
        frame_feature_path: The path of the movie frame features.

    """

    @property
    def continuous_index(self):
        return torch.arange(len(self))


@parametrize(
    "allen-movie1-ca-single-session-decoding-{session_id}-repeat-{repeat_no}-{split_flag}",
    session_id=range(len(_SINGLE_SESSION_CA)),
    repeat_no=np.arange(10),
    split_flag=["train", "test"],
)
class SingleSessionAllenCaDecoding(cebra.data.SingleSessionDataset):
    """A single mouse 30Hz calcium events dataset during the allen MOVIE1 stimulus with train/test splits.

    A dataset of a single mouse 30Hz calcium events from the excitatory neurons in the primary visual cortex
    during the 10 repeats of the MOVIE1 stimulus in session type A. The preprocessed data from *Deitch et al. (2021) are used.
    The continuous labels corresponding to a DINO embedding of each stimulus frame.
    A neural recording during the chosen repeat is used as a test set and the remaining 9 repeats are used as a train set.

    Args:
        session_id: The integer value to pick a session among 4 sessions with the largest number of recorded neruons. Choose between 0-3.
        repeat_no: The nth repeat to use as the test set. Choose between 0-9.
        split_flag: The `train`/`test` split to load.
        frame_feature_path: The path of the movie frame features.
        pca: If true, 32 principal components of the PCA transformed calcium data are used as neural input. Default value is `False`.

    """

    def __init__(
        self,
        session_id: int,
        repeat_no: int,
        split_flag: str,
        frame_feature_path: str = get_datapath(
            "allen/features/allen_movies/vit_base/8/movie_one_image_stack.npz/testfeat.pth"
        ),
        pca: bool = False,
    ):
        self.path = _SINGLE_SESSION_CA[session_id]
        traces = scipy.io.loadmat(self.path)
        if pca:
            pca_ = PCA()
            neural = pca_.fit_transform(
                traces["filtered_traces_days_events"][0,
                                                      0].transpose(1,
                                                                   0))[:, :32]
        else:
            neural = traces["filtered_traces_days_events"][0, 0].transpose(1, 0)

        test_idx = np.arange(900 * repeat_no, 900 * (repeat_no + 1))
        train_idx = np.delete(np.arange(9000), test_idx)
        frame_feature = torch.load(frame_feature_path)
        if split_flag == "train":
            self.neural = torch.from_numpy(neural[train_idx]).float()
            self.index = frame_feature.repeat(9, 1)
        elif split_flag == "test":
            self.neural = torch.from_numpy(neural[test_idx]).float()
            self.index = frame_feature

    def __len__(self):
        return self.neural.size(0)

    @property
    def continuous_index(self):
        return self.index

    @property
    def input_dimension(self):
        return self.neural.size(1)

    def __getitem__(self, index):
        index = self.expand_index(index)
        return self.neural[index].transpose(2, 1)


@parametrize(
    "allen-movie1-ca-single-session-leave2out-{session_id}-repeat-{repeat_no}-{split_flag}",
    session_id=range(len(_SINGLE_SESSION_CA)),
    repeat_no=[0, 2, 4, 6, 8],
    split_flag=["train", "valid", "test"],
)
class SingleSessionAllenCaDecodingLeave2Out(cebra.data.SingleSessionDataset):

    def __init__(
        self,
        session_id,
        repeat_no,
        split_flag,
        frame_feature_path=get_datapath(
            "allen/features/allen_movies/vit_base/8/movie_one_image_stack.npz/testfeat.pth"
        ),
    ):
        self.path = _SINGLE_SESSION_CA[session_id]
        traces = scipy.io.loadmat(self.path)
        neural = traces["filtered_traces_days_events"][0, 0].transpose(1, 0)

        valid_idx = np.arange(900 * repeat_no, 900 * (repeat_no + 1))
        test_idx = np.arange(900 * (repeat_no + 1), 900 * (repeat_no + 2))
        train_idx = np.delete(np.arange(9000),
                              np.concatenate([valid_idx, test_idx]))
        frame_feature = torch.load(frame_feature_path)
        if split_flag == "train":
            self.neural = torch.from_numpy(neural[train_idx]).float()
            self.index = frame_feature.repeat(9, 1)
        elif split_flag == "valid":
            self.neural = torch.from_numpy(neural[valid_idx]).float()
            self.index = frame_feature
        elif split_flag == "test":
            self.neural = torch.from_numpy(neural[test_idx]).float()
            self.index = frame_feature
        else:
            raise ValueError(split_flag)

    def __len__(self):
        return self.neural.size(0)

    @property
    def continuous_index(self):
        return self.index

    @property
    def input_dimension(self):
        return self.neural.size(1)

    def __getitem__(self, index):
        index = self.expand_index(index)
        return self.neural[index].transpose(2, 1)


@parametrize(
    "allen-movie1-ca-multi-session-decoding-repeat-{repeat_no}-{split_flag}",
    repeat_no=np.arange(10),
    split_flag=["train", "test"],
)
class MultiSessionAllenCaDecoding(cebra.data.DatasetCollection):

    def __init__(self, repeat_no, split_flag):
        super().__init__(
            *[
                init(
                    f"allen-movie1-ca-single-session-decoding-{session_id}-repeat-{repeat_no}-{split_flag}"
                ) for session_id in range(4)
            ],)


@parametrize(
    "allen-movie1-ca-multi-session-leave2out-repeat-{repeat_no}-{split_flag}",
    repeat_no=[0, 2, 4, 6, 8],
    split_flag=["train", "valid", "test"],
)
class MultiSessionAllenCaLeave2Out(cebra.data.DatasetCollection):

    def __init__(self, repeat_no, split_flag):
        super().__init__(
            *[
                init(
                    f"allen-movie1-ca-single-session-leave2out-{session_id}-repeat-{repeat_no}-{split_flag}"
                ) for session_id in range(4)
            ],)


@parametrize(
    "allen-movie1-ca-single-session-decoding-corrupt-{session_id}-repeat-{repeat_no}-{split_flag}",
    session_id=range(len(_SINGLE_SESSION_CA)),
    repeat_no=[9],
    split_flag=["train", "test"],
)
class SingleSessionAllenCaDecoding(cebra.data.SingleSessionDataset):
    """A corrupted single mouse 30Hz calcium events dataset during the allen MOVIE1 stimulus with train/test splits.

    A dataset of a single mouse 30Hz calcium events from the excitatory neurons 
    in the primary visual cortex during the 10 repeats of the MOVIE1 stimulus 
    in session type A. The preprocessed data from *Deitch et al. (2021) are used.
    The continuous labels corresponding to a DINO embedding of each stimulus frame, 
    but in randomly shuffled order.
    A neural recording during the chosen repeat is used as a test set and the 
    remaining 9 repeats are used as a train set.

    Args:
        session_id: The integer value to pick a session among 4 sessions with the 
            largest number of recorded neruons. Choose between 0-3.
        repeat_no: The nth repeat to use as the test set. Choose between 0-9.
        split_flag: The `train`/`test` split to load.
        frame_feature_path: The path of the movie frame features.

    """

    def __init__(
        self,
        session_id: int,
        repeat_no: int,
        split_flag: str,
        frame_feature_path: str = get_datapath(
            "allen/features/allen_movies/vit_base/8/movie_one_image_stack.npz/testfeat.pth"
        ),
        pca: bool = False,
    ):
        self.path = _SINGLE_SESSION_CA[session_id]
        traces = scipy.io.loadmat(self.path)
        if pca:
            pca_ = PCA()
            neural = pca_.fit_transform(
                traces["filtered_traces_days_events"][0,
                                                      0].transpose(1,
                                                                   0))[:, :32]
        else:
            neural = traces["filtered_traces_days_events"][0, 0].transpose(1, 0)

        test_idx = np.arange(900 * repeat_no, 900 * (repeat_no + 1))
        train_idx = np.delete(np.arange(9000), test_idx)
        frame_feature = torch.load(frame_feature_path)
        rng = np.random.Generator(np.random.PCG64(111))
        if split_flag == "train":
            self.neural = torch.from_numpy(neural[train_idx]).float()
            self.frame_index = np.arange(900)
            rng.shuffle(self.frame_index)
            self.index = frame_feature.repeat(9,
                                              1)[np.tile(self.frame_index, 9)]
        elif split_flag == "test":
            self.neural = torch.from_numpy(neural[test_idx]).float()
            self.index = frame_feature
            self.frame_index = np.arange(900)
            rng.shuffle(self.frame_index)
            self.index = frame_feature[self.frame_index]

    def __len__(self):
        return self.neural.size(0)

    @property
    def continuous_index(self):
        return self.index

    @property
    def input_dimension(self):
        return self.neural.size(1)

    def __getitem__(self, index):
        index = self.expand_index(index)
        return self.neural[index].transpose(2, 1)
