from re import X

import scipy.io

import numpy as np

import torch

from torch.utils.data import Dataset, DataLoader

from torch.utils.data.dataloader import default_collate

from sklearn.preprocessing import OrdinalEncoder

class dataset(Dataset):

    def __init__(self, X, y, train=True):

        self.X = X

        self.y = y

        self.train=train

    def __len__(self):

        return len(self.y)

    # def __getitem__(self, idx):

    #     rng = np.random.randint(0, high=200)

    #     if self.train:

    #         x = self.X[idx][:, rng:rng + 600]

    #     else:

    #         x = self.X[idx][:, 200: 800]

    #     return x, self.y[idx]

    def __getitem__(self, idx):

        x = self.X[idx]

        if self.train:

#             rn = np.random.randint(0, high=500)

#             x = x[:, rn:rn+4000]

            x = x[:, 0:4000]

        else:

            x = x[:, 0:4000]

        return x, self.y[idx]

def generate_x_train(mat):

    # out: num_data_points * chl * trial_length

    data = []

    last_label = False

    for i in range(0, len(mat['mrk'][0][0][0][0])-1):

        start_idx = mat['mrk'][0][0][0][0][i]

        end_idx = mat['mrk'][0][0][0][0][i+1]

        # to resolve shape issues, we use a shifted window 

        # (possible overlapping but acceptable given it's trivial)

        end_idx += (8000 + start_idx - end_idx) 

        data.append(mat['cnt'][start_idx: end_idx,].T)

    # add the last datapoint

    if len(mat['cnt']) - mat['mrk'][0][0][0][0][-1] >= 8000:

        last_label = True

        start_idx = mat['mrk'][0][0][0][0][-1]

        end_idx = start_idx + 8000

        data.append(mat['cnt'][start_idx: end_idx,].T)

    return np.array(data), last_label

def generate_y_train(mat, last_label):

    # out: 1 * num_labels

    class1, class2 = mat['nfo']['classes'][0][0][0][0][0], mat['nfo']['classes'][0][0][0][1][0]

    mapping = {-1: class1, 1: class2}

    labels = np.vectorize(mapping.get)(mat['mrk'][0][0][1])[0]

    if not last_label:

        labels = labels[:-1]

    return labels

def generate_data(files):

    X, y = [], []

    for file in files:

        print(file)

        mat = scipy.io.loadmat(file)

        X_batch, last_label = generate_x_train(mat)

        X.append(X_batch)

        y.append(generate_y_train(mat, last_label))

    X, y = np.concatenate(X, axis=0), np.concatenate(y)

    y = OrdinalEncoder().fit_transform(y.reshape(-1, 1))

    return X, y

def split_data(X, y):

    def get_idx():

        np.random.seed(seed=42)

        rng = np.random.choice(len(y), len(y), replace=False)

        return rng

    train_size, val_size, test_size = 1000, 197, 200

    indices = get_idx()

    train_idx, val_idx, test_idx = indices[0: train_size], \

                indices[train_size: train_size + val_size], indices[train_size + val_size:]

    train_X, train_y, val_X, val_y, test_X, test_y = \

        X[train_idx], y[train_idx], X[val_idx], y[val_idx], X[test_idx], y[test_idx]

    return train_X, train_y, val_X, val_y, test_X, test_y

def get_loaders(train_X, train_y, val_X, val_y, test_X, test_y):

    train_set, val_set, test_set = dataset(train_X, train_y, True), dataset(val_X, val_y, False), dataset(test_X, test_y, False)

    data_loader_train = torch.utils.data.DataLoader(

        train_set, 

        batch_size=1, 

        num_workers=1,

        pin_memory=True, 

        drop_last=False,

    )

    data_loader_val = torch.utils.data.DataLoader(

            val_set, 

            batch_size=1, 

            num_workers=1,

            pin_memory=True, 

            drop_last=False,

    )

    data_loader_test = torch.utils.data.DataLoader(

            test_set, 

            batch_size=1, 

            num_workers=1,

            pin_memory=True, 

            drop_last=False,

    )

    dataloaders = {

        'train': data_loader_train,

        'val': data_loader_val,

        'test': data_loader_test

    }

    return dataloaders