# Neural Network Model (Dec 8th)

from sklearn.model_selection import KFold


def examplar_epoch_data(epochs_in, N, navg):
    epo = epochs_in.copy()[range(N)]
    fold_num = N // navg
    if N % navg > 0: fold_num += 1

    X1 = epo.get_data()  # signals: n_epochs, n_channels, n_times
    y1 = epo.events[:, 2]

    n_trl, n_ch, n_times = epo.get_data().shape

    kf = KFold(n_splits=fold_num, shuffle= True)
    kf_idx = [x for x in kf.split(X1, y1)]
    train_y = [y1[x[0]] for x in kf_idx] # get y labels of each training-fold
    test_y = [x[1] for x in kf_idx] # get y labels of each testing-fold

    exampler_data = np.zeros((fold_num, n_ch, n_times))
    for i in range(fold_num):
        temp = epo[test_y[i]].average().data
        exampler_data[i,:,:] = temp
        pass
    y = np.ones(fold_num)
    return exampler_data, y

features_ = epochs.get_data()
labels = epochs.events[:, 2]
print(features_.shape)

## get subsets of epochs

epoch_subset = epochs[['Famous', 'Scrambled']]
features_ = epoch_subset.get_data()
labels = epoch_subset.events[:, 2]
print(features_.shape)

c1 = np.where(labels<10) # Famous face
c2 = np.where(labels>10) # Scrambled face
#c2 = [i for i, y_ in zip(range(len(labels)),labels) if (12<y_ and y_ < 16)] #Scrambled face
#c3 = np.where(labels>16) # Unfamiliar

labels[c1] = 0
labels[c2] = 1
labels

from sklearn.metrics import roc_curve, auc
from sklearn.model_selection import StratifiedKFold

N = 10

trl_numbers = features_.shape[0] # serial numbers from 0 to N, N = len(trlinfo) - 1
X = [i for i in range(trl_numbers)]
y_ = labels
skf = StratifiedKFold(n_splits=N, shuffle=True)
kf_idx = [x for x in skf.split(X, y_)]

train_y = [y_[x[0]] for x in kf_idx] # get y labels of each training-fold
test_y = [y_[x[1]] for x in kf_idx] # get y labels of each testing-fold

train_x = [features_[x[0],:,:] for x in kf_idx] # get y labels of each training-fold
test_x = [features_[x[1],:,:] for x in kf_idx] # get y labels of each testing-fold