#%%
import numpy as np
import pandas as pd
import tensorflow as tf
import scipy.io
import random
import xgboost as xgb
import seaborn as sns
from tensorflow import keras
from tensorflow.keras import layers
from sklearn.model_selection import train_test_split
from sklearn import metrics, decomposition

import plotly.graph_objects as go
import matplotlib.pyplot as plt
from tensorflow.python.keras.backend import std
from tensorflow.python.keras.layers import embeddings
from tensorflow.python.keras.layers.core import RepeatVector
from tensorflow.python.keras.losses import mean_squared_error
from tensorflow.python.util.nest import _IF_SHALLOW_IS_SEQ_INPUT_MUST_BE_SEQ


from load_data import *


#%%
SAMPLE_LEN = 5
channels = 8
latent_dim = 3

data, label = split_session(SAMPLE_LEN)

x_train, x_test, y_train, y_test = train_test_split(session_data_1, session_data_1, train_size = 0.85)
x_val, x_test, y_val, y_test = train_test_split(x_test, y_test, train_size = 0.5)

#%%
new_batch = []
for i, batch in enumerate(session_data_1):
    print(batch.shape[1]//SAMPLE_LEN)
    new_batch.extend(np.array_split(batch, batch.shape[1]//SAMPLE_LEN, axis = 1)[1:batch.shape[1]//SAMPLE_LEN-2])


#%%
np.array(new_batch)

#%%

def create_model():

    input = keras.Input(shape = (sample_len,channels))  # Variable-length sequence of ints
    x = layers.LSTM(64, return_sequences=True)(input)
    x = layers.LSTM(64, return_sequences=True)(x)
    #x = layers.RepeatVector(sample_len)(x)
    #x = layers.LSTM(64, return_sequences=True)(x)
    x = layers.TimeDistributed(layers.Dense(16))(x)
    x = layers.Dense(latent_dim, name = 'embedding')(x)
    x = layers.Dense(16)(x)
    output = layers.Dense(channels, activation='linear')(x)

    model = keras.Model(inputs=[input],outputs=[output])

    return model

model = create_model()

keras.utils.plot_model(model, "keras_LSTM_autoencoder.png", show_shapes=True)


#%%
model.summary()

model.compile(optimizer=keras.optimizers.Adam(1e-3), loss='mse') #tf.keras.metrics.mean_squared_error

history = model.fit(x_train, y_train, 
                    validation_data = (x_val, y_val), 
                    epochs=100, batch_size=32)

plt.plot(history.history['loss'])
plt.yscale("log")
plt.show()


def get_embedding(sequences, model):
    
    encoder_model = keras.Model(inputs = model.input, 
                                outputs = model.get_layer('embedding').output)
    try:
        emb = encoder_model.predict(sequences)[:,-1,:]
    except:
        sequences = list(sequences)
        emb = encoder_model.predict(sequences)[:,-1,:]
    return emb

def plot_3D_embedding(sequences, model):
    emb_vec = get_embedding(sequences, model)

    try:
        emb_vec = emb_vec.reshape(3,len(list(emb_vec)))
    except:
        'Embedding dim != 3'

    fig = go.Figure(data=[go.Scatter3d(x=emb_vec[0], y=emb_vec[1], z=emb_vec[2],
                                    mode='markers')])
    fig.show()

def plot_curve_comparison(sequences, model, n_sample = 1, title = ''):

    for i in range(n_sample):
        pred = model.predict(sequences[i])

        plt.plot(sequences[i], label = 'Truth')
        plt.plot(pred.reshape(sample_len), label = 'Pred')
        plt.title(title)
        plt.legend()
        plt.show()

def plot_reconstruction_error(sequences, model, n_sample = 1, title = ''):

    res = model.predict(sequences)
    
    try:
        s_len = len(sequences)
    except:
        s_len = sequences.shape[0]

    for i in np.random.randint(0, s_len - 1, size = n_sample):
        plt.plot(res[i] - sequences[i])
        plt.title(title)
        plt.show()

#%%

plot_reconstruction_error(x_train, model, 2, title = 'Train: Rec - Truth')
plot_reconstruction_error(x_test, model, 2, title = 'Test: Rec - Truth')
plot_reconstruction_error(x_val, model, 2, title = 'Val: Rec - Truth')

if channels == 1:
    plot_curve_comparison(x_train, model, title = 'Train')
    plot_curve_comparison(x_test, model, title = 'Test')
    plot_curve_comparison(x_val, model, title = 'Validation')

if latent_dim == 3:
    plot_3D_embedding(x_train, model)
    plot_3D_embedding(x_test, model)
    plot_3D_embedding(x_val, model)

#%%

if channels == 1:
    def stress_seq(rand = 'normal'):
        ampl = np.random.uniform(0.1, 1.)
        freq = np.random.uniform(0.01, 0.1)
        high_noise = np.random.uniform(0.1, 0.5)
        low_noise  = np.random.uniform(0.1, 0.5)
        std = np.random.uniform(0.1,0.3)

        if rand == 'normal':
            recorded_data = ampl*np.random.normal(0, std, size = (channels, sample_len))
        else:
            recorded_data = ampl*np.random.uniform(-low_noise, high_noise, size=( channels, sample_len))
        
        recorded_data += ampl*np.cos(freq*np.arange(0,sample_len))

        #recorded_data = np.array([np.cumsum(j) for j in recorded_data])
        data = recorded_data.reshape(1, sample_len, channels)
        return recorded_data, data

    stress_test = [stress_seq()[1] for i in range(10)]

    for i in stress_test:
        preds = model.predict(i)

        for j in range(channels):
            plt.plot(i.reshape(sample_len))

        plt.plot(preds[0])
        plt.show()

# %%
if channels == 1:
    for patient,p in zip(data, range(len(data))):
        print('Patient number %d ' % (p + 1))
        for channel,i in zip(patient, range(len(patient))):
            
            pred = model.predict(channel)

            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1), pred.reshape(-1)))

            print('\tPred on channel %d has RMSE of %f' % (i, rmse))

# %%

if channels == 1:
    # SINGLE CHANNEL
    for patient,p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        for channel,i in zip(patient, range(len(patient))):
            
            data_serie = list([list(i) for i in channel[:,0]])
            label_serie = list([list(i) for i in channel[:,1]])

            emb_serie = get_embedding(data_serie, model)

            x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb_serie, label_serie, test_size = 0.2, stratify=label_serie)

            xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
            xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')

            xgb_preds = xgb_classifier.predict(x_b_test)

            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1), channel.reshape(-1)))
            
            print('\tChannel %d - Classifier ROC-AUC: %f' % (i,np.sqrt(metrics.roc_auc_score(y_b_test, xgb_preds))))
            
    ##################################
    # UNION OF CHANNELS

    for patient,p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        emb_serie = []
        for channel,i in zip(patient, range(len(patient))):
            
            data_serie = list([list(i) for i in channel[:,0]])
            label_serie = list([list(i) for i in channel[:,1]])

            emb = get_embedding(data_serie, model)
            emb = list(list(i) for i in emb)

            if i == 0:
                emb_serie = emb
            else:
                for j in range(len(channel)):
                    for k in emb[j]:
                        emb_serie[j].append(k)

        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(np.array(emb_serie), label_serie, test_size = 0.2, stratify=label_serie)

        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')

        xgb_preds = xgb_classifier.predict(x_b_test)
        
        print('\tClassifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(y_b_test, xgb_preds)))

    ####################################
    # ONE CHANNEL ALONE

    labeled_seq = np.array([[serie[i: i + sample_len],matrixS1['trig'][i]] for i in range(len(serie)) if matrixS1['trig'][i] != 0][:-20])

    data_serie = list([list(i) for i in labeled_seq[:,0]])
    label_serie = list([list(i) for i in labeled_seq[:,1]])

    emb_serie = get_embedding(data_serie, model)

    x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb_serie, label_serie, test_size = 0.2, stratify=label_serie)

    xgb_classifier = xgb.XGBClassifier(use_label_encoder=True)
    xgb_classifier.fit(x_b_train, y_b_train)

    xgb_preds = xgb_classifier.predict(x_b_test)
#%%
   
if channels != 1:
    def labeled_multichannel_data_factory():
        data = []
        labels = []
        for i in range(5):
            patient = []
            patient_lab = []
            matrixS = scipy.io.loadmat('./data/wetransfer-01a2dc/p300/S' + str(i+1) + '.mat')
            for serie in matrixS['y'].T:
                fs = matrixS['fs'][0,0]
                offset = int(0.3*fs)

                seq = np.array([serie[i + offset - int(sample_len/2): i + offset + int(sample_len/2)] for i in range(len(serie)) if matrixS['trig'][i] != 0][:-20])
                label = np.array([0.5*matrixS['trig'][i] + 0.5 for i in range(len(serie)) if matrixS['trig'][i] != 0][:-20])
                patient.append(seq)
                patient_lab.append(label)
            
            patient = np.array(patient).reshape(-1,sample_len,channels)
            label = np.array(patient_lab)
            data.append(patient)
            labels.append(label[0])
        return data, labels

    dat, lab = labeled_multichannel_data_factory()
    dat = dat/np.max(data[0])*10

    for patient,i in zip (dat, range(len(dat))):
        emb = get_embedding(patient, model)

        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb, lab[i], test_size = 0.2, stratify = lab[i])

        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')

        xgb_preds = xgb_classifier.predict(x_b_test)

        print('\nPatient number %d' % i)
        print('Classifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(y_b_test, xgb_preds)))

        confusion_matrix = metrics.confusion_matrix(y_b_test, xgb_preds)
        cm_plot = metrics.ConfusionMatrixDisplay(confusion_matrix)
        cm_plot.plot()

        plt.show()
        #xgb.plot_importance(xgb_classifier)


# %%

plot_pairplot = True

if channels != 1:
    for patient,i in zip (dat, range(len(dat))):
        emb = get_embedding(patient, model)
        pd_emb = pd.DataFrame(emb)
        pd_emb['signal'] = lab[i]

        if plot_pairplot:
            g = sns.pairplot(pd_emb, corner = True, hue = 'signal')
            g.map_lower(sns.kdeplot, levels=3, color=".2")
            plt.show()

        pca = decomposition.PCA(n_components=2)
        emb_pca = pca.fit(emb).transform(emb)

        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i][j] == 0], 
                    [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab[i][j] == 0],
                    label = 'No signal')
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i][j] == 1], 
                    [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab[i][j] == 1],
                    label = 'Signal')
        plt.legend()
        plt.title('PCA')
        plt.show()





# %%

87c917f4676f8ffc92f06eafa1f34211dffb5500
<mask token>


def get_embedding(sequences, model):
    encoder_model = keras.Model(inputs=model.input, outputs=model.get_layer
        ('embedding').output)
    try:
        emb = encoder_model.predict(sequences)[:, -1, :]
    except:
        sequences = list(sequences)
        emb = encoder_model.predict(sequences)[:, -1, :]
    return emb


def plot_3D_embedding(sequences, model):
    emb_vec = get_embedding(sequences, model)
    try:
        emb_vec = emb_vec.reshape(3, len(list(emb_vec)))
    except:
        """Embedding dim != 3"""
    fig = go.Figure(data=[go.Scatter3d(x=emb_vec[0], y=emb_vec[1], z=
        emb_vec[2], mode='markers')])
    fig.show()


def plot_curve_comparison(sequences, model, n_sample=1, title=''):
    for i in range(n_sample):
        pred = model.predict(sequences[i])
        plt.plot(sequences[i], label='Truth')
        plt.plot(pred.reshape(sample_len), label='Pred')
        plt.title(title)
        plt.legend()
        plt.show()


def plot_reconstruction_error(sequences, model, n_sample=1, title=''):
    res = model.predict(sequences)
    try:
        s_len = len(sequences)
    except:
        s_len = sequences.shape[0]
    for i in np.random.randint(0, s_len - 1, size=n_sample):
        plt.plot(res[i] - sequences[i])
        plt.title(title)
        plt.show()


<mask token>

<mask token>
for i, batch in enumerate(session_data_1):
    print(batch.shape[1] // SAMPLE_LEN)
    new_batch.extend(np.array_split(batch, batch.shape[1] // SAMPLE_LEN,
        axis=1)[1:batch.shape[1] // SAMPLE_LEN - 2])
np.array(new_batch)


def create_model():
    input = keras.Input(shape=(sample_len, channels))
    x = layers.LSTM(64, return_sequences=True)(input)
    x = layers.LSTM(64, return_sequences=True)(x)
    x = layers.TimeDistributed(layers.Dense(16))(x)
    x = layers.Dense(latent_dim, name='embedding')(x)
    x = layers.Dense(16)(x)
    output = layers.Dense(channels, activation='linear')(x)
    model = keras.Model(inputs=[input], outputs=[output])
    return model


<mask token>
keras.utils.plot_model(model, 'keras_LSTM_autoencoder.png', show_shapes=True)
model.summary()
model.compile(optimizer=keras.optimizers.Adam(0.001), loss='mse')
<mask token>
plt.plot(history.history['loss'])
plt.yscale('log')
plt.show()


def get_embedding(sequences, model):
    encoder_model = keras.Model(inputs=model.input, outputs=model.get_layer
        ('embedding').output)
    try:
        emb = encoder_model.predict(sequences)[:, -1, :]
    except:
        sequences = list(sequences)
        emb = encoder_model.predict(sequences)[:, -1, :]
    return emb


def plot_3D_embedding(sequences, model):
    emb_vec = get_embedding(sequences, model)
    try:
        emb_vec = emb_vec.reshape(3, len(list(emb_vec)))
    except:
        """Embedding dim != 3"""
    fig = go.Figure(data=[go.Scatter3d(x=emb_vec[0], y=emb_vec[1], z=
        emb_vec[2], mode='markers')])
    fig.show()


def plot_curve_comparison(sequences, model, n_sample=1, title=''):
    for i in range(n_sample):
        pred = model.predict(sequences[i])
        plt.plot(sequences[i], label='Truth')
        plt.plot(pred.reshape(sample_len), label='Pred')
        plt.title(title)
        plt.legend()
        plt.show()


def plot_reconstruction_error(sequences, model, n_sample=1, title=''):
    res = model.predict(sequences)
    try:
        s_len = len(sequences)
    except:
        s_len = sequences.shape[0]
    for i in np.random.randint(0, s_len - 1, size=n_sample):
        plt.plot(res[i] - sequences[i])
        plt.title(title)
        plt.show()


plot_reconstruction_error(x_train, model, 2, title='Train: Rec - Truth')
plot_reconstruction_error(x_test, model, 2, title='Test: Rec - Truth')
plot_reconstruction_error(x_val, model, 2, title='Val: Rec - Truth')
if channels == 1:
    plot_curve_comparison(x_train, model, title='Train')
    plot_curve_comparison(x_test, model, title='Test')
    plot_curve_comparison(x_val, model, title='Validation')
if latent_dim == 3:
    plot_3D_embedding(x_train, model)
    plot_3D_embedding(x_test, model)
    plot_3D_embedding(x_val, model)
if channels == 1:

    def stress_seq(rand='normal'):
        ampl = np.random.uniform(0.1, 1.0)
        freq = np.random.uniform(0.01, 0.1)
        high_noise = np.random.uniform(0.1, 0.5)
        low_noise = np.random.uniform(0.1, 0.5)
        std = np.random.uniform(0.1, 0.3)
        if rand == 'normal':
            recorded_data = ampl * np.random.normal(0, std, size=(channels,
                sample_len))
        else:
            recorded_data = ampl * np.random.uniform(-low_noise, high_noise,
                size=(channels, sample_len))
        recorded_data += ampl * np.cos(freq * np.arange(0, sample_len))
        data = recorded_data.reshape(1, sample_len, channels)
        return recorded_data, data
    stress_test = [stress_seq()[1] for i in range(10)]
    for i in stress_test:
        preds = model.predict(i)
        for j in range(channels):
            plt.plot(i.reshape(sample_len))
        plt.plot(preds[0])
        plt.show()
if channels == 1:
    for patient, p in zip(data, range(len(data))):
        print('Patient number %d ' % (p + 1))
        for channel, i in zip(patient, range(len(patient))):
            pred = model.predict(channel)
            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1),
                pred.reshape(-1)))
            print('\tPred on channel %d has RMSE of %f' % (i, rmse))
if channels == 1:
    for patient, p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        for channel, i in zip(patient, range(len(patient))):
            data_serie = list([list(i) for i in channel[:, 0]])
            label_serie = list([list(i) for i in channel[:, 1]])
            emb_serie = get_embedding(data_serie, model)
            x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(
                emb_serie, label_serie, test_size=0.2, stratify=label_serie)
            xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
            xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
            xgb_preds = xgb_classifier.predict(x_b_test)
            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1),
                channel.reshape(-1)))
            print('\tChannel %d - Classifier ROC-AUC: %f' % (i, np.sqrt(
                metrics.roc_auc_score(y_b_test, xgb_preds))))
    for patient, p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        emb_serie = []
        for channel, i in zip(patient, range(len(patient))):
            data_serie = list([list(i) for i in channel[:, 0]])
            label_serie = list([list(i) for i in channel[:, 1]])
            emb = get_embedding(data_serie, model)
            emb = list(list(i) for i in emb)
            if i == 0:
                emb_serie = emb
            else:
                for j in range(len(channel)):
                    for k in emb[j]:
                        emb_serie[j].append(k)
        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(np.
            array(emb_serie), label_serie, test_size=0.2, stratify=label_serie)
        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
        xgb_preds = xgb_classifier.predict(x_b_test)
        print('\tClassifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(
            y_b_test, xgb_preds)))
    labeled_seq = np.array([[serie[i:i + sample_len], matrixS1['trig'][i]] for
        i in range(len(serie)) if matrixS1['trig'][i] != 0][:-20])
    data_serie = list([list(i) for i in labeled_seq[:, 0]])
    label_serie = list([list(i) for i in labeled_seq[:, 1]])
    emb_serie = get_embedding(data_serie, model)
    x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb_serie,
        label_serie, test_size=0.2, stratify=label_serie)
    xgb_classifier = xgb.XGBClassifier(use_label_encoder=True)
    xgb_classifier.fit(x_b_train, y_b_train)
    xgb_preds = xgb_classifier.predict(x_b_test)
if channels != 1:

    def labeled_multichannel_data_factory():
        data = []
        labels = []
        for i in range(5):
            patient = []
            patient_lab = []
            matrixS = scipy.io.loadmat('./data/wetransfer-01a2dc/p300/S' +
                str(i + 1) + '.mat')
            for serie in matrixS['y'].T:
                fs = matrixS['fs'][0, 0]
                offset = int(0.3 * fs)
                seq = np.array([serie[i + offset - int(sample_len / 2):i +
                    offset + int(sample_len / 2)] for i in range(len(serie)
                    ) if matrixS['trig'][i] != 0][:-20])
                label = np.array([(0.5 * matrixS['trig'][i] + 0.5) for i in
                    range(len(serie)) if matrixS['trig'][i] != 0][:-20])
                patient.append(seq)
                patient_lab.append(label)
            patient = np.array(patient).reshape(-1, sample_len, channels)
            label = np.array(patient_lab)
            data.append(patient)
            labels.append(label[0])
        return data, labels
    dat, lab = labeled_multichannel_data_factory()
    dat = dat / np.max(data[0]) * 10
    for patient, i in zip(dat, range(len(dat))):
        emb = get_embedding(patient, model)
        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb,
            lab[i], test_size=0.2, stratify=lab[i])
        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
        xgb_preds = xgb_classifier.predict(x_b_test)
        print('\nPatient number %d' % i)
        print('Classifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(
            y_b_test, xgb_preds)))
        confusion_matrix = metrics.confusion_matrix(y_b_test, xgb_preds)
        cm_plot = metrics.ConfusionMatrixDisplay(confusion_matrix)
        cm_plot.plot()
        plt.show()
<mask token>
if channels != 1:
    for patient, i in zip(dat, range(len(dat))):
        emb = get_embedding(patient, model)
        pd_emb = pd.DataFrame(emb)
        pd_emb['signal'] = lab[i]
        if plot_pairplot:
            g = sns.pairplot(pd_emb, corner=True, hue='signal')
            g.map_lower(sns.kdeplot, levels=3, color='.2')
            plt.show()
        pca = decomposition.PCA(n_components=2)
        emb_pca = pca.fit(emb).transform(emb)
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i]
            [j] == 0], [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab
            [i][j] == 0], label='No signal')
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i]
            [j] == 1], [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab
            [i][j] == 1], label='Signal')
        plt.legend()
        plt.title('PCA')
        plt.show()

<mask token>
SAMPLE_LEN = 5
channels = 8
latent_dim = 3
data, label = split_session(SAMPLE_LEN)
x_train, x_test, y_train, y_test = train_test_split(session_data_1,
    session_data_1, train_size=0.85)
x_val, x_test, y_val, y_test = train_test_split(x_test, y_test, train_size=0.5)
new_batch = []
for i, batch in enumerate(session_data_1):
    print(batch.shape[1] // SAMPLE_LEN)
    new_batch.extend(np.array_split(batch, batch.shape[1] // SAMPLE_LEN,
        axis=1)[1:batch.shape[1] // SAMPLE_LEN - 2])
np.array(new_batch)


def create_model():
    input = keras.Input(shape=(sample_len, channels))
    x = layers.LSTM(64, return_sequences=True)(input)
    x = layers.LSTM(64, return_sequences=True)(x)
    x = layers.TimeDistributed(layers.Dense(16))(x)
    x = layers.Dense(latent_dim, name='embedding')(x)
    x = layers.Dense(16)(x)
    output = layers.Dense(channels, activation='linear')(x)
    model = keras.Model(inputs=[input], outputs=[output])
    return model


model = create_model()
keras.utils.plot_model(model, 'keras_LSTM_autoencoder.png', show_shapes=True)
model.summary()
model.compile(optimizer=keras.optimizers.Adam(0.001), loss='mse')
history = model.fit(x_train, y_train, validation_data=(x_val, y_val),
    epochs=100, batch_size=32)
plt.plot(history.history['loss'])
plt.yscale('log')
plt.show()


def get_embedding(sequences, model):
    encoder_model = keras.Model(inputs=model.input, outputs=model.get_layer
        ('embedding').output)
    try:
        emb = encoder_model.predict(sequences)[:, -1, :]
    except:
        sequences = list(sequences)
        emb = encoder_model.predict(sequences)[:, -1, :]
    return emb


def plot_3D_embedding(sequences, model):
    emb_vec = get_embedding(sequences, model)
    try:
        emb_vec = emb_vec.reshape(3, len(list(emb_vec)))
    except:
        """Embedding dim != 3"""
    fig = go.Figure(data=[go.Scatter3d(x=emb_vec[0], y=emb_vec[1], z=
        emb_vec[2], mode='markers')])
    fig.show()


def plot_curve_comparison(sequences, model, n_sample=1, title=''):
    for i in range(n_sample):
        pred = model.predict(sequences[i])
        plt.plot(sequences[i], label='Truth')
        plt.plot(pred.reshape(sample_len), label='Pred')
        plt.title(title)
        plt.legend()
        plt.show()


def plot_reconstruction_error(sequences, model, n_sample=1, title=''):
    res = model.predict(sequences)
    try:
        s_len = len(sequences)
    except:
        s_len = sequences.shape[0]
    for i in np.random.randint(0, s_len - 1, size=n_sample):
        plt.plot(res[i] - sequences[i])
        plt.title(title)
        plt.show()


plot_reconstruction_error(x_train, model, 2, title='Train: Rec - Truth')
plot_reconstruction_error(x_test, model, 2, title='Test: Rec - Truth')
plot_reconstruction_error(x_val, model, 2, title='Val: Rec - Truth')
if channels == 1:
    plot_curve_comparison(x_train, model, title='Train')
    plot_curve_comparison(x_test, model, title='Test')
    plot_curve_comparison(x_val, model, title='Validation')
if latent_dim == 3:
    plot_3D_embedding(x_train, model)
    plot_3D_embedding(x_test, model)
    plot_3D_embedding(x_val, model)
if channels == 1:

    def stress_seq(rand='normal'):
        ampl = np.random.uniform(0.1, 1.0)
        freq = np.random.uniform(0.01, 0.1)
        high_noise = np.random.uniform(0.1, 0.5)
        low_noise = np.random.uniform(0.1, 0.5)
        std = np.random.uniform(0.1, 0.3)
        if rand == 'normal':
            recorded_data = ampl * np.random.normal(0, std, size=(channels,
                sample_len))
        else:
            recorded_data = ampl * np.random.uniform(-low_noise, high_noise,
                size=(channels, sample_len))
        recorded_data += ampl * np.cos(freq * np.arange(0, sample_len))
        data = recorded_data.reshape(1, sample_len, channels)
        return recorded_data, data
    stress_test = [stress_seq()[1] for i in range(10)]
    for i in stress_test:
        preds = model.predict(i)
        for j in range(channels):
            plt.plot(i.reshape(sample_len))
        plt.plot(preds[0])
        plt.show()
if channels == 1:
    for patient, p in zip(data, range(len(data))):
        print('Patient number %d ' % (p + 1))
        for channel, i in zip(patient, range(len(patient))):
            pred = model.predict(channel)
            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1),
                pred.reshape(-1)))
            print('\tPred on channel %d has RMSE of %f' % (i, rmse))
if channels == 1:
    for patient, p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        for channel, i in zip(patient, range(len(patient))):
            data_serie = list([list(i) for i in channel[:, 0]])
            label_serie = list([list(i) for i in channel[:, 1]])
            emb_serie = get_embedding(data_serie, model)
            x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(
                emb_serie, label_serie, test_size=0.2, stratify=label_serie)
            xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
            xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
            xgb_preds = xgb_classifier.predict(x_b_test)
            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1),
                channel.reshape(-1)))
            print('\tChannel %d - Classifier ROC-AUC: %f' % (i, np.sqrt(
                metrics.roc_auc_score(y_b_test, xgb_preds))))
    for patient, p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        emb_serie = []
        for channel, i in zip(patient, range(len(patient))):
            data_serie = list([list(i) for i in channel[:, 0]])
            label_serie = list([list(i) for i in channel[:, 1]])
            emb = get_embedding(data_serie, model)
            emb = list(list(i) for i in emb)
            if i == 0:
                emb_serie = emb
            else:
                for j in range(len(channel)):
                    for k in emb[j]:
                        emb_serie[j].append(k)
        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(np.
            array(emb_serie), label_serie, test_size=0.2, stratify=label_serie)
        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
        xgb_preds = xgb_classifier.predict(x_b_test)
        print('\tClassifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(
            y_b_test, xgb_preds)))
    labeled_seq = np.array([[serie[i:i + sample_len], matrixS1['trig'][i]] for
        i in range(len(serie)) if matrixS1['trig'][i] != 0][:-20])
    data_serie = list([list(i) for i in labeled_seq[:, 0]])
    label_serie = list([list(i) for i in labeled_seq[:, 1]])
    emb_serie = get_embedding(data_serie, model)
    x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb_serie,
        label_serie, test_size=0.2, stratify=label_serie)
    xgb_classifier = xgb.XGBClassifier(use_label_encoder=True)
    xgb_classifier.fit(x_b_train, y_b_train)
    xgb_preds = xgb_classifier.predict(x_b_test)
if channels != 1:

    def labeled_multichannel_data_factory():
        data = []
        labels = []
        for i in range(5):
            patient = []
            patient_lab = []
            matrixS = scipy.io.loadmat('./data/wetransfer-01a2dc/p300/S' +
                str(i + 1) + '.mat')
            for serie in matrixS['y'].T:
                fs = matrixS['fs'][0, 0]
                offset = int(0.3 * fs)
                seq = np.array([serie[i + offset - int(sample_len / 2):i +
                    offset + int(sample_len / 2)] for i in range(len(serie)
                    ) if matrixS['trig'][i] != 0][:-20])
                label = np.array([(0.5 * matrixS['trig'][i] + 0.5) for i in
                    range(len(serie)) if matrixS['trig'][i] != 0][:-20])
                patient.append(seq)
                patient_lab.append(label)
            patient = np.array(patient).reshape(-1, sample_len, channels)
            label = np.array(patient_lab)
            data.append(patient)
            labels.append(label[0])
        return data, labels
    dat, lab = labeled_multichannel_data_factory()
    dat = dat / np.max(data[0]) * 10
    for patient, i in zip(dat, range(len(dat))):
        emb = get_embedding(patient, model)
        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb,
            lab[i], test_size=0.2, stratify=lab[i])
        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
        xgb_preds = xgb_classifier.predict(x_b_test)
        print('\nPatient number %d' % i)
        print('Classifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(
            y_b_test, xgb_preds)))
        confusion_matrix = metrics.confusion_matrix(y_b_test, xgb_preds)
        cm_plot = metrics.ConfusionMatrixDisplay(confusion_matrix)
        cm_plot.plot()
        plt.show()
plot_pairplot = True
if channels != 1:
    for patient, i in zip(dat, range(len(dat))):
        emb = get_embedding(patient, model)
        pd_emb = pd.DataFrame(emb)
        pd_emb['signal'] = lab[i]
        if plot_pairplot:
            g = sns.pairplot(pd_emb, corner=True, hue='signal')
            g.map_lower(sns.kdeplot, levels=3, color='.2')
            plt.show()
        pca = decomposition.PCA(n_components=2)
        emb_pca = pca.fit(emb).transform(emb)
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i]
            [j] == 0], [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab
            [i][j] == 0], label='No signal')
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i]
            [j] == 1], [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab
            [i][j] == 1], label='Signal')
        plt.legend()
        plt.title('PCA')
        plt.show()

import numpy as np
import pandas as pd
import tensorflow as tf
import scipy.io
import random
import xgboost as xgb
import seaborn as sns
from tensorflow import keras
from tensorflow.keras import layers
from sklearn.model_selection import train_test_split
from sklearn import metrics, decomposition
import plotly.graph_objects as go
import matplotlib.pyplot as plt
from tensorflow.python.keras.backend import std
from tensorflow.python.keras.layers import embeddings
from tensorflow.python.keras.layers.core import RepeatVector
from tensorflow.python.keras.losses import mean_squared_error
from tensorflow.python.util.nest import _IF_SHALLOW_IS_SEQ_INPUT_MUST_BE_SEQ
from load_data import *
SAMPLE_LEN = 5
channels = 8
latent_dim = 3
data, label = split_session(SAMPLE_LEN)
x_train, x_test, y_train, y_test = train_test_split(session_data_1,
    session_data_1, train_size=0.85)
x_val, x_test, y_val, y_test = train_test_split(x_test, y_test, train_size=0.5)
new_batch = []
for i, batch in enumerate(session_data_1):
    print(batch.shape[1] // SAMPLE_LEN)
    new_batch.extend(np.array_split(batch, batch.shape[1] // SAMPLE_LEN,
        axis=1)[1:batch.shape[1] // SAMPLE_LEN - 2])
np.array(new_batch)


def create_model():
    input = keras.Input(shape=(sample_len, channels))
    x = layers.LSTM(64, return_sequences=True)(input)
    x = layers.LSTM(64, return_sequences=True)(x)
    x = layers.TimeDistributed(layers.Dense(16))(x)
    x = layers.Dense(latent_dim, name='embedding')(x)
    x = layers.Dense(16)(x)
    output = layers.Dense(channels, activation='linear')(x)
    model = keras.Model(inputs=[input], outputs=[output])
    return model


model = create_model()
keras.utils.plot_model(model, 'keras_LSTM_autoencoder.png', show_shapes=True)
model.summary()
model.compile(optimizer=keras.optimizers.Adam(0.001), loss='mse')
history = model.fit(x_train, y_train, validation_data=(x_val, y_val),
    epochs=100, batch_size=32)
plt.plot(history.history['loss'])
plt.yscale('log')
plt.show()


def get_embedding(sequences, model):
    encoder_model = keras.Model(inputs=model.input, outputs=model.get_layer
        ('embedding').output)
    try:
        emb = encoder_model.predict(sequences)[:, -1, :]
    except:
        sequences = list(sequences)
        emb = encoder_model.predict(sequences)[:, -1, :]
    return emb


def plot_3D_embedding(sequences, model):
    emb_vec = get_embedding(sequences, model)
    try:
        emb_vec = emb_vec.reshape(3, len(list(emb_vec)))
    except:
        """Embedding dim != 3"""
    fig = go.Figure(data=[go.Scatter3d(x=emb_vec[0], y=emb_vec[1], z=
        emb_vec[2], mode='markers')])
    fig.show()


def plot_curve_comparison(sequences, model, n_sample=1, title=''):
    for i in range(n_sample):
        pred = model.predict(sequences[i])
        plt.plot(sequences[i], label='Truth')
        plt.plot(pred.reshape(sample_len), label='Pred')
        plt.title(title)
        plt.legend()
        plt.show()


def plot_reconstruction_error(sequences, model, n_sample=1, title=''):
    res = model.predict(sequences)
    try:
        s_len = len(sequences)
    except:
        s_len = sequences.shape[0]
    for i in np.random.randint(0, s_len - 1, size=n_sample):
        plt.plot(res[i] - sequences[i])
        plt.title(title)
        plt.show()


plot_reconstruction_error(x_train, model, 2, title='Train: Rec - Truth')
plot_reconstruction_error(x_test, model, 2, title='Test: Rec - Truth')
plot_reconstruction_error(x_val, model, 2, title='Val: Rec - Truth')
if channels == 1:
    plot_curve_comparison(x_train, model, title='Train')
    plot_curve_comparison(x_test, model, title='Test')
    plot_curve_comparison(x_val, model, title='Validation')
if latent_dim == 3:
    plot_3D_embedding(x_train, model)
    plot_3D_embedding(x_test, model)
    plot_3D_embedding(x_val, model)
if channels == 1:

    def stress_seq(rand='normal'):
        ampl = np.random.uniform(0.1, 1.0)
        freq = np.random.uniform(0.01, 0.1)
        high_noise = np.random.uniform(0.1, 0.5)
        low_noise = np.random.uniform(0.1, 0.5)
        std = np.random.uniform(0.1, 0.3)
        if rand == 'normal':
            recorded_data = ampl * np.random.normal(0, std, size=(channels,
                sample_len))
        else:
            recorded_data = ampl * np.random.uniform(-low_noise, high_noise,
                size=(channels, sample_len))
        recorded_data += ampl * np.cos(freq * np.arange(0, sample_len))
        data = recorded_data.reshape(1, sample_len, channels)
        return recorded_data, data
    stress_test = [stress_seq()[1] for i in range(10)]
    for i in stress_test:
        preds = model.predict(i)
        for j in range(channels):
            plt.plot(i.reshape(sample_len))
        plt.plot(preds[0])
        plt.show()
if channels == 1:
    for patient, p in zip(data, range(len(data))):
        print('Patient number %d ' % (p + 1))
        for channel, i in zip(patient, range(len(patient))):
            pred = model.predict(channel)
            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1),
                pred.reshape(-1)))
            print('\tPred on channel %d has RMSE of %f' % (i, rmse))
if channels == 1:
    for patient, p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        for channel, i in zip(patient, range(len(patient))):
            data_serie = list([list(i) for i in channel[:, 0]])
            label_serie = list([list(i) for i in channel[:, 1]])
            emb_serie = get_embedding(data_serie, model)
            x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(
                emb_serie, label_serie, test_size=0.2, stratify=label_serie)
            xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
            xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
            xgb_preds = xgb_classifier.predict(x_b_test)
            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1),
                channel.reshape(-1)))
            print('\tChannel %d - Classifier ROC-AUC: %f' % (i, np.sqrt(
                metrics.roc_auc_score(y_b_test, xgb_preds))))
    for patient, p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        emb_serie = []
        for channel, i in zip(patient, range(len(patient))):
            data_serie = list([list(i) for i in channel[:, 0]])
            label_serie = list([list(i) for i in channel[:, 1]])
            emb = get_embedding(data_serie, model)
            emb = list(list(i) for i in emb)
            if i == 0:
                emb_serie = emb
            else:
                for j in range(len(channel)):
                    for k in emb[j]:
                        emb_serie[j].append(k)
        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(np.
            array(emb_serie), label_serie, test_size=0.2, stratify=label_serie)
        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
        xgb_preds = xgb_classifier.predict(x_b_test)
        print('\tClassifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(
            y_b_test, xgb_preds)))
    labeled_seq = np.array([[serie[i:i + sample_len], matrixS1['trig'][i]] for
        i in range(len(serie)) if matrixS1['trig'][i] != 0][:-20])
    data_serie = list([list(i) for i in labeled_seq[:, 0]])
    label_serie = list([list(i) for i in labeled_seq[:, 1]])
    emb_serie = get_embedding(data_serie, model)
    x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb_serie,
        label_serie, test_size=0.2, stratify=label_serie)
    xgb_classifier = xgb.XGBClassifier(use_label_encoder=True)
    xgb_classifier.fit(x_b_train, y_b_train)
    xgb_preds = xgb_classifier.predict(x_b_test)
if channels != 1:

    def labeled_multichannel_data_factory():
        data = []
        labels = []
        for i in range(5):
            patient = []
            patient_lab = []
            matrixS = scipy.io.loadmat('./data/wetransfer-01a2dc/p300/S' +
                str(i + 1) + '.mat')
            for serie in matrixS['y'].T:
                fs = matrixS['fs'][0, 0]
                offset = int(0.3 * fs)
                seq = np.array([serie[i + offset - int(sample_len / 2):i +
                    offset + int(sample_len / 2)] for i in range(len(serie)
                    ) if matrixS['trig'][i] != 0][:-20])
                label = np.array([(0.5 * matrixS['trig'][i] + 0.5) for i in
                    range(len(serie)) if matrixS['trig'][i] != 0][:-20])
                patient.append(seq)
                patient_lab.append(label)
            patient = np.array(patient).reshape(-1, sample_len, channels)
            label = np.array(patient_lab)
            data.append(patient)
            labels.append(label[0])
        return data, labels
    dat, lab = labeled_multichannel_data_factory()
    dat = dat / np.max(data[0]) * 10
    for patient, i in zip(dat, range(len(dat))):
        emb = get_embedding(patient, model)
        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb,
            lab[i], test_size=0.2, stratify=lab[i])
        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')
        xgb_preds = xgb_classifier.predict(x_b_test)
        print('\nPatient number %d' % i)
        print('Classifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(
            y_b_test, xgb_preds)))
        confusion_matrix = metrics.confusion_matrix(y_b_test, xgb_preds)
        cm_plot = metrics.ConfusionMatrixDisplay(confusion_matrix)
        cm_plot.plot()
        plt.show()
plot_pairplot = True
if channels != 1:
    for patient, i in zip(dat, range(len(dat))):
        emb = get_embedding(patient, model)
        pd_emb = pd.DataFrame(emb)
        pd_emb['signal'] = lab[i]
        if plot_pairplot:
            g = sns.pairplot(pd_emb, corner=True, hue='signal')
            g.map_lower(sns.kdeplot, levels=3, color='.2')
            plt.show()
        pca = decomposition.PCA(n_components=2)
        emb_pca = pca.fit(emb).transform(emb)
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i]
            [j] == 0], [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab
            [i][j] == 0], label='No signal')
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i]
            [j] == 1], [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab
            [i][j] == 1], label='Signal')
        plt.legend()
        plt.title('PCA')
        plt.show()

#%%
import numpy as np
import pandas as pd
import tensorflow as tf
import scipy.io
import random
import xgboost as xgb
import seaborn as sns
from tensorflow import keras
from tensorflow.keras import layers
from sklearn.model_selection import train_test_split
from sklearn import metrics, decomposition

import plotly.graph_objects as go
import matplotlib.pyplot as plt
from tensorflow.python.keras.backend import std
from tensorflow.python.keras.layers import embeddings
from tensorflow.python.keras.layers.core import RepeatVector
from tensorflow.python.keras.losses import mean_squared_error
from tensorflow.python.util.nest import _IF_SHALLOW_IS_SEQ_INPUT_MUST_BE_SEQ


from load_data import *


#%%
SAMPLE_LEN = 5
channels = 8
latent_dim = 3

data, label = split_session(SAMPLE_LEN)

x_train, x_test, y_train, y_test = train_test_split(session_data_1, session_data_1, train_size = 0.85)
x_val, x_test, y_val, y_test = train_test_split(x_test, y_test, train_size = 0.5)

#%%
new_batch = []
for i, batch in enumerate(session_data_1):
    print(batch.shape[1]//SAMPLE_LEN)
    new_batch.extend(np.array_split(batch, batch.shape[1]//SAMPLE_LEN, axis = 1)[1:batch.shape[1]//SAMPLE_LEN-2])


#%%
np.array(new_batch)

#%%

def create_model():

    input = keras.Input(shape = (sample_len,channels))  # Variable-length sequence of ints
    x = layers.LSTM(64, return_sequences=True)(input)
    x = layers.LSTM(64, return_sequences=True)(x)
    #x = layers.RepeatVector(sample_len)(x)
    #x = layers.LSTM(64, return_sequences=True)(x)
    x = layers.TimeDistributed(layers.Dense(16))(x)
    x = layers.Dense(latent_dim, name = 'embedding')(x)
    x = layers.Dense(16)(x)
    output = layers.Dense(channels, activation='linear')(x)

    model = keras.Model(inputs=[input],outputs=[output])

    return model

model = create_model()

keras.utils.plot_model(model, "keras_LSTM_autoencoder.png", show_shapes=True)


#%%
model.summary()

model.compile(optimizer=keras.optimizers.Adam(1e-3), loss='mse') #tf.keras.metrics.mean_squared_error

history = model.fit(x_train, y_train, 
                    validation_data = (x_val, y_val), 
                    epochs=100, batch_size=32)

plt.plot(history.history['loss'])
plt.yscale("log")
plt.show()


def get_embedding(sequences, model):
    
    encoder_model = keras.Model(inputs = model.input, 
                                outputs = model.get_layer('embedding').output)
    try:
        emb = encoder_model.predict(sequences)[:,-1,:]
    except:
        sequences = list(sequences)
        emb = encoder_model.predict(sequences)[:,-1,:]
    return emb

def plot_3D_embedding(sequences, model):
    emb_vec = get_embedding(sequences, model)

    try:
        emb_vec = emb_vec.reshape(3,len(list(emb_vec)))
    except:
        'Embedding dim != 3'

    fig = go.Figure(data=[go.Scatter3d(x=emb_vec[0], y=emb_vec[1], z=emb_vec[2],
                                    mode='markers')])
    fig.show()

def plot_curve_comparison(sequences, model, n_sample = 1, title = ''):

    for i in range(n_sample):
        pred = model.predict(sequences[i])

        plt.plot(sequences[i], label = 'Truth')
        plt.plot(pred.reshape(sample_len), label = 'Pred')
        plt.title(title)
        plt.legend()
        plt.show()

def plot_reconstruction_error(sequences, model, n_sample = 1, title = ''):

    res = model.predict(sequences)
    
    try:
        s_len = len(sequences)
    except:
        s_len = sequences.shape[0]

    for i in np.random.randint(0, s_len - 1, size = n_sample):
        plt.plot(res[i] - sequences[i])
        plt.title(title)
        plt.show()

#%%

plot_reconstruction_error(x_train, model, 2, title = 'Train: Rec - Truth')
plot_reconstruction_error(x_test, model, 2, title = 'Test: Rec - Truth')
plot_reconstruction_error(x_val, model, 2, title = 'Val: Rec - Truth')

if channels == 1:
    plot_curve_comparison(x_train, model, title = 'Train')
    plot_curve_comparison(x_test, model, title = 'Test')
    plot_curve_comparison(x_val, model, title = 'Validation')

if latent_dim == 3:
    plot_3D_embedding(x_train, model)
    plot_3D_embedding(x_test, model)
    plot_3D_embedding(x_val, model)

#%%

if channels == 1:
    def stress_seq(rand = 'normal'):
        ampl = np.random.uniform(0.1, 1.)
        freq = np.random.uniform(0.01, 0.1)
        high_noise = np.random.uniform(0.1, 0.5)
        low_noise  = np.random.uniform(0.1, 0.5)
        std = np.random.uniform(0.1,0.3)

        if rand == 'normal':
            recorded_data = ampl*np.random.normal(0, std, size = (channels, sample_len))
        else:
            recorded_data = ampl*np.random.uniform(-low_noise, high_noise, size=( channels, sample_len))
        
        recorded_data += ampl*np.cos(freq*np.arange(0,sample_len))

        #recorded_data = np.array([np.cumsum(j) for j in recorded_data])
        data = recorded_data.reshape(1, sample_len, channels)
        return recorded_data, data

    stress_test = [stress_seq()[1] for i in range(10)]

    for i in stress_test:
        preds = model.predict(i)

        for j in range(channels):
            plt.plot(i.reshape(sample_len))

        plt.plot(preds[0])
        plt.show()

# %%
if channels == 1:
    for patient,p in zip(data, range(len(data))):
        print('Patient number %d ' % (p + 1))
        for channel,i in zip(patient, range(len(patient))):
            
            pred = model.predict(channel)

            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1), pred.reshape(-1)))

            print('\tPred on channel %d has RMSE of %f' % (i, rmse))

# %%

if channels == 1:
    # SINGLE CHANNEL
    for patient,p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        for channel,i in zip(patient, range(len(patient))):
            
            data_serie = list([list(i) for i in channel[:,0]])
            label_serie = list([list(i) for i in channel[:,1]])

            emb_serie = get_embedding(data_serie, model)

            x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb_serie, label_serie, test_size = 0.2, stratify=label_serie)

            xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
            xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')

            xgb_preds = xgb_classifier.predict(x_b_test)

            rmse = np.sqrt(metrics.mean_squared_error(channel.reshape(-1), channel.reshape(-1)))
            
            print('\tChannel %d - Classifier ROC-AUC: %f' % (i,np.sqrt(metrics.roc_auc_score(y_b_test, xgb_preds))))
            
    ##################################
    # UNION OF CHANNELS

    for patient,p in zip(data_labeled, range(len(data_labeled))):
        print('\n\nPatient number %d ' % (p + 1))
        emb_serie = []
        for channel,i in zip(patient, range(len(patient))):
            
            data_serie = list([list(i) for i in channel[:,0]])
            label_serie = list([list(i) for i in channel[:,1]])

            emb = get_embedding(data_serie, model)
            emb = list(list(i) for i in emb)

            if i == 0:
                emb_serie = emb
            else:
                for j in range(len(channel)):
                    for k in emb[j]:
                        emb_serie[j].append(k)

        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(np.array(emb_serie), label_serie, test_size = 0.2, stratify=label_serie)

        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')

        xgb_preds = xgb_classifier.predict(x_b_test)
        
        print('\tClassifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(y_b_test, xgb_preds)))

    ####################################
    # ONE CHANNEL ALONE

    labeled_seq = np.array([[serie[i: i + sample_len],matrixS1['trig'][i]] for i in range(len(serie)) if matrixS1['trig'][i] != 0][:-20])

    data_serie = list([list(i) for i in labeled_seq[:,0]])
    label_serie = list([list(i) for i in labeled_seq[:,1]])

    emb_serie = get_embedding(data_serie, model)

    x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb_serie, label_serie, test_size = 0.2, stratify=label_serie)

    xgb_classifier = xgb.XGBClassifier(use_label_encoder=True)
    xgb_classifier.fit(x_b_train, y_b_train)

    xgb_preds = xgb_classifier.predict(x_b_test)
#%%
   
if channels != 1:
    def labeled_multichannel_data_factory():
        data = []
        labels = []
        for i in range(5):
            patient = []
            patient_lab = []
            matrixS = scipy.io.loadmat('./data/wetransfer-01a2dc/p300/S' + str(i+1) + '.mat')
            for serie in matrixS['y'].T:
                fs = matrixS['fs'][0,0]
                offset = int(0.3*fs)

                seq = np.array([serie[i + offset - int(sample_len/2): i + offset + int(sample_len/2)] for i in range(len(serie)) if matrixS['trig'][i] != 0][:-20])
                label = np.array([0.5*matrixS['trig'][i] + 0.5 for i in range(len(serie)) if matrixS['trig'][i] != 0][:-20])
                patient.append(seq)
                patient_lab.append(label)
            
            patient = np.array(patient).reshape(-1,sample_len,channels)
            label = np.array(patient_lab)
            data.append(patient)
            labels.append(label[0])
        return data, labels

    dat, lab = labeled_multichannel_data_factory()
    dat = dat/np.max(data[0])*10

    for patient,i in zip (dat, range(len(dat))):
        emb = get_embedding(patient, model)

        x_b_train, x_b_test, y_b_train, y_b_test = train_test_split(emb, lab[i], test_size = 0.2, stratify = lab[i])

        xgb_classifier = xgb.XGBClassifier(use_label_encoder=False)
        xgb_classifier.fit(x_b_train, y_b_train, eval_metric='error')

        xgb_preds = xgb_classifier.predict(x_b_test)

        print('\nPatient number %d' % i)
        print('Classifier ROC-AUC: %f' % np.sqrt(metrics.roc_auc_score(y_b_test, xgb_preds)))

        confusion_matrix = metrics.confusion_matrix(y_b_test, xgb_preds)
        cm_plot = metrics.ConfusionMatrixDisplay(confusion_matrix)
        cm_plot.plot()

        plt.show()
        #xgb.plot_importance(xgb_classifier)


# %%

plot_pairplot = True

if channels != 1:
    for patient,i in zip (dat, range(len(dat))):
        emb = get_embedding(patient, model)
        pd_emb = pd.DataFrame(emb)
        pd_emb['signal'] = lab[i]

        if plot_pairplot:
            g = sns.pairplot(pd_emb, corner = True, hue = 'signal')
            g.map_lower(sns.kdeplot, levels=3, color=".2")
            plt.show()

        pca = decomposition.PCA(n_components=2)
        emb_pca = pca.fit(emb).transform(emb)

        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i][j] == 0], 
                    [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab[i][j] == 0],
                    label = 'No signal')
        plt.scatter([emb_pca.T[0][j] for j in range(len(emb_pca)) if lab[i][j] == 1], 
                    [emb_pca.T[1][j] for j in range(len(emb_pca)) if lab[i][j] == 1],
                    label = 'Signal')
        plt.legend()
        plt.title('PCA')
        plt.show()





# %%
