import mne
import numpy as np
import torch
import torch.nn as nn

class LoadData:
    def __init__(self,eeg_file_path: str):
        self.eeg_file_path = eeg_file_path

    def load_raw_data_gdf(self,file_to_load):
        self.raw_eeg_subject = mne.io.read_raw_gdf(self.eeg_file_path + '/' + file_to_load)
        return self

    def load_raw_data_mat(self,file_to_load):
        import scipy.io as sio
        self.raw_eeg_subject = sio.loadmat(self.eeg_file_path + '/' + file_to_load)

    def get_all_files(self,file_path_extension: str =''):
        if file_path_extension:
            return glob.glob(self.eeg_file_path+'/'+file_path_extension)
        return os.listdir(self.eeg_file_path)

class LoadBCIC(LoadData):
    '''Subclass of LoadData for loading BCI Competition IV Dataset 2a'''
    def __init__(self, file_to_load, *args):
        self.stimcodes=('769','770','771','772')
        # self.epoched_data={}
        self.file_to_load = file_to_load
        self.channels_to_remove = ['EOG-left', 'EOG-central', 'EOG-right']
        super(LoadBCIC,self).__init__(*args)

    def get_epochs(self, tmin=0,tmax=1,baseline=None):
        self.load_raw_data_gdf(self.file_to_load)
        raw_data = self.raw_eeg_subject
        # raw_downsampled = raw_data.copy().resample(sfreq=128)
        self.fs = raw_data.info.get('sfreq')
        events, event_ids = mne.events_from_annotations(raw_data)
        stims =[value for key, value in event_ids.items() if key in self.stimcodes]
        epochs = mne.Epochs(raw_data, events, event_id=stims, tmin=tmin, tmax=tmax, event_repeated='drop',
                            baseline=baseline, preload=True, proj=False, reject_by_annotation=False)
        epochs = epochs.drop_channels(self.channels_to_remove)
        self.y_labels = epochs.events[:, -1] - min(epochs.events[:, -1])
        self.x_data = epochs.get_data()*1e6
        eeg_data={'x_data':self.x_data,
                  'y_labels':self.y_labels,
                  'fs':self.fs}
        return eeg_data

data_path = "/home/pytorch/LiangXiaohan/MI_Dataverse/BCICIV_2a_gdf"
file_to_load = 'A01T.gdf'

'''for BCIC Dataset'''
bcic_data = LoadBCIC(file_to_load, data_path)
eeg_data = bcic_data.get_epochs() # {'x_data':, 'y_labels':, 'fs':}

X = eeg_data.get('x_data')
Y = eeg_data.get('y_labels')
Y.shape

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, random_state=0)
X_train.shape

from sklearn.model_selection import StratifiedKFold
train_idx = {}
eval_idx = {}

skf = StratifiedKFold(n_splits=4, shuffle=True)
i = 0
for train_indices, eval_indices in skf.split(X_train, y_train):
    train_idx.update({i: train_indices})
    eval_idx.update({i: eval_indices})
    i += 1
train_idx.get(1).shape

def split_xdata(eeg_data, train_idx, eval_idx):
    x_train=np.copy(eeg_data[train_idx,:,:])
    x_eval=np.copy(eeg_data[eval_idx,:,:])
    x_train = torch.from_numpy(x_train).to(torch.float32)
    x_eval = torch.from_numpy(x_eval).to(torch.float32)
    return x_train, x_eval

def split_ydata(y_true, train_idx, eval_idx):
    y_train = np.copy(y_true[train_idx])
    y_eval = np.copy(y_true[eval_idx])
    y_train = torch.from_numpy(y_train)
    y_eval = torch.from_numpy(y_eval)
    return y_train, y_eval

x_train, x_eval = split_xdata(X_train, train_idx.get(1), eval_idx.get(1))
y_train, y_eval = split_ydata(Y_train, train_idx.get(1), eval_idx.get(1))
y_train.shape

from torch.utils.data import Dataset, DataLoader, TensorDataset
from tqdm import tqdm
def BCICDataLoader(x_train, y_train, batch_size=64, num_workers=2, shuffle=True):
    
    data = TensorDataset(x_train, y_train)

    train_data = DataLoader(dataset=data, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

    return train_data

train_data = BCICDataLoader(x_train, y_train, batch_size=32)
for inputs, target in tqdm(train_data):
    print(target)