# used for SEED dataset

# train with 9 trials and test with 6 trials

import numpy as np

import scipy.io

root_path = '/Data/SEED/seed_syh/data/'

save_path = '/Data/SEED/seed_syh/data_1second/'

# i, j, k  subject, session, trial

for i in range(15):

    one_subject = []

    one_subject_label = []

    for j in range(3):

        one_session = []

        one_session_label = []

        # todo save training data 

        for k in range(9):

            one_trial = []

            trial_tmp = scipy.io.loadmat(root_path + 'S%d_%d_%d.mat' % (i+1, j+1, k+1))

            trial_data = trial_tmp['trial_data']

            trial_label = np.squeeze(trial_tmp['trial_label'])

            # change to using 2 seconds  200*2

            trial_number = np.int32(trial_data.shape[1]/200)

            for tmp_num in range(trial_number):

                one_trial.append(trial_data[:,tmp_num*200:(tmp_num+1)*200])

            one_trial_label = [trial_label]*trial_number

            one_session.append(one_trial)

            one_session_label.append(one_trial_label)

        one_session = np.concatenate(one_session)

        one_session_label = np.concatenate(one_session_label)

        data = np.transpose(one_session, [2, 1, 0])

        label = one_session_label

        np.save(save_path + 'S%d_session%dT.npy'%(i+1, j+1), data)

        np.save(save_path + 'S%d_session%dT_label'%(i+1, j+1), label)

        one_session = []

        one_session_label = []

        # todo save test data 

        for k in range(9, 15):

            one_trial = []

            trial_tmp = scipy.io.loadmat(root_path + 'S%d_%d_%d.mat' % (i+1, j+1, k+1))

            trial_data = trial_tmp['trial_data']

            trial_label = np.squeeze(trial_tmp['trial_label'])

            trial_number = np.int32(trial_data.shape[1]/200)

            for tmp_num in range(trial_number):

                one_trial.append(trial_data[:,tmp_num*200:(tmp_num+1)*200])

            one_trial_label = [trial_label]*trial_number

            one_session.append(one_trial)

            one_session_label.append(one_trial_label)

        one_session = np.concatenate(one_session)

        one_session_label = np.concatenate(one_session_label)

        data = np.transpose(one_session, [2, 1, 0])

        label = one_session_label

        np.save(save_path + 'S%d_session%dE.npy'%(i+1, j+1), data)

        np.save(save_path + 'S%d_session%dE_label'%(i+1, j+1), label)

        print('Finished Subject%d Session%d' % (i+1, j+1))