import matplotlib.pyplot as plt

import numpy as np

import time

import tensorflow as tf

import pickle

import wfdb

from sklearn.utils import class_weight

from sklearn.model_selection import train_test_split

# Hyper-parameters

sequence_length = 240

epochs = 1000#int(input('Enter Number of Epochs (or enter default 1000): '))

FS = 100.0

def z_norm(result):

    result_mean = np.mean(result)

    result_std = np.std(result)

    result = (result - result_mean) / result_std

    return result

def split_data(X):

    X1 = []

    X2 = []

    for index in range(len(X)):

        X1.append([X[index][0], X[index][1]])

        X2.append([X[index][2], X[index][3]])

    return np.array(X1).astype('float64'), np.array(X2).astype('float64')

def get_data():

    with open('train_input.pickle','rb') as f: 

        X_train = np.asarray(pickle.load(f))

    with open('train_label.pickle','rb') as f: 

        y_train = np.asarray(pickle.load(f))

    with open('val_input.pickle','rb') as f: 

        X_val = np.asarray(pickle.load(f))

    with open('val_label.pickle','rb') as f: 

        y_val = np.asarray(pickle.load(f))

    with open('test_input.pickle','rb') as f: 

        X_test = np.asarray(pickle.load(f))

    with open('test_label.pickle','rb') as f: 

        y_test = np.asarray(pickle.load(f))

    #X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=42)

    '''

    X_train = X_train[:, 0, :]

    X_test = X_test[:, 0, :]

    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))

    X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))

    '''

    X_train1, X_train2 = split_data(X_train)

    X_val1, X_val2 = split_data(X_val)

    X_test1, X_test2 = split_data(X_test)

    X_train1 = np.transpose(X_train1, (0, 2, 1))

    #X_train2 = np.reshape(X_train2, (X_train2.shape[0], X_train2.shape[1], 1))

    X_test1 = np.transpose(X_test1, (0, 2, 1))

    #X_test2 = np.reshape(X_test2, (X_test2.shape[0], X_test2.shape[1], 1))

    X_val1 = np.transpose(X_val1, (0, 2, 1))

    return X_train1, X_train2, y_train, X_val1, X_val2, y_val, X_test1, X_test2, y_test

def build_model():

    layers = {'input': 2, 'hidden1': 256, 'hidden2': 256, 'hidden3': 256, 'output': 1}

    x1 = tf.keras.layers.Input(shape=(sequence_length, layers['input']))

    m1 = tf.keras.layers.LSTM(layers['hidden1'],                   

                    recurrent_dropout=0.5,

                   return_sequences=True)(x1)

    m1 = tf.keras.layers.LSTM(

            layers['hidden2'],

            recurrent_dropout=0.5,

            return_sequences=True)(m1)

    m1 = tf.keras.layers.LSTM(

            layers['hidden3'],

            recurrent_dropout=0.5,

            return_sequences=False)(m1)

    x2 = tf.keras.layers.Input(shape=(2,))

    m2 = tf.keras.layers.Dense(32)(x2)

    #merged = Merge([model1, model2], mode='concat')

    merged = tf.keras.layers.Concatenate(axis=1)([m1, m2])

    out = tf.keras.layers.Dense(8)(merged)

    out = tf.keras.layers.Dense(layers['output'], kernel_initializer='normal')(out)

    out = tf.keras.layers.Activation("sigmoid")(out)

    model = tf.keras.models.Model(inputs=[x1, x2], outputs=[out])

    start = time.time()

    model.compile(loss="binary_crossentropy", optimizer="adam",

                  metrics = ['accuracy'])

    print ("Compilation Time : ", time.time() - start)

    model.summary()

    return model

def run_network(model=None, data=None):

    global_start_time = time.time()

    print ('\nData Loaded. Compiling...\n')

    print('Loading data... ')

    X_train1, X_train2, y_train, X_val1, X_val2, y_val, X_test1, X_test2, y_test = get_data()

    class_w = class_weight.compute_class_weight(class_weight='balanced',

                                                     classes=np.unique(y_train),

                                                     y=y_train)

    print (class_w)

    if model is None:

        model = build_model()

    try:

        print("Training")

        class_w = {i : class_w[i] for i in range(2)}

        callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=3)

        history = model.fit([X_train1, X_train2], y_train, 

                            validation_data=([X_val1, X_val2], y_val),

                            callbacks=[callback],

                             epochs=epochs, batch_size=256, class_weight=class_w)

        import matplotlib.pyplot as plt

        '''

        plt.plot(history.losses)

        plt.ylabel('loss')

        plt.xlabel('epoch')

        plt.legend(['train'], loc='upper left')

        plt.show()

        '''

        # Evaluate Model

        y_pred = model.predict([X_test1, X_test2])

        scores = model.evaluate([X_test1, X_test2], y_test)

        print("%s: %.2f%%" % (model.metrics_names[1], scores[1] * 100))

    except KeyboardInterrupt:

        print("prediction exception")

        print ('Training duration (s) : ', time.time() - global_start_time)

        return model

    print ('Training duration (s) : ', time.time() - global_start_time)

    return model

run_network()