
import pandas as pd
import numpy as np

from keras import optimizers, losses, Model
from keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from keras.engine.saving import load_model
from keras.layers import Dense
from sklearn.metrics import f1_score, accuracy_score
from sklearn.model_selection import train_test_split


def get_transferred_model_frozen():
    # load pretrained MIT model:
    old_model = load_model("own_lstm_mitbih.h5")
    # freeze the complete old model:
    for layer in old_model.layers:
        layer.trainable = False
    # create new model with the pretrained frozen LSTM layers from the old model:
    # CAUTION: the layer name lstm_2 is hardcoded and autogenerated (not a good combination...)
    new_dense = Dense(50, activation="relu")(old_model.get_layer("lstm_2").output)
    new_dense = Dense(50, activation="relu")(new_dense)
    # new_dense = Dense(50, activation="relu")(new_dense)
    new_out = Dense(1, activation="sigmoid")(new_dense)

    model2 = Model(inputs=old_model.input, outputs=new_out)
    model2.compile(loss=losses.binary_crossentropy, optimizer=optimizers.Adam(lr=0.0008), metrics=['acc'])
    model2.summary()

    return model2


def get_transferred_model_trainable():
    # load pretrained MIT model:
    old_model = load_model("own_lstm_mitbih.h5")
    # create new model with the pretrained frozen LSTM layers from the old model:
    # CAUTION: the layer name lstm_2 is hardcoded and autogenerated (not a good combination...)
    new_dense = Dense(50, activation="relu")(old_model.get_layer("lstm_2").output)
    new_dense = Dense(50, activation="relu")(new_dense)
    # new_dense = Dense(50, activation="relu")(new_dense)
    new_out = Dense(1, activation="sigmoid")(new_dense)

    model2 = Model(inputs=old_model.input, outputs=new_out)
    model2.compile(loss=losses.binary_crossentropy, optimizer=optimizers.Adam(lr=0.0008), metrics=['acc'])
    model2.summary()

    return model2


if __name__ == "__main__":
    # define dataset-specific constants:
    timepoints_per_sample = 187

    data_path_normal = "./input/ptbdb_normal.csv"
    data_path_abnormal = "./input/ptbdb_abnormal.csv"

    df_1 = pd.read_csv(data_path_normal, header=None)
    df_2 = pd.read_csv(data_path_abnormal, header=None)
    df = pd.concat([df_1, df_2])

    df_train, df_test = train_test_split(df, test_size=0.2, random_state=1337, stratify=df[187])

    Y = np.array(df_train[187].values).astype(np.int8)
    X = np.array(df_train[list(range(187))].values)[..., np.newaxis]

    Y_test = np.array(df_test[187].values).astype(np.int8)
    X_test = np.array(df_test[list(range(187))].values)[..., np.newaxis]

    # normalize without taking the zero-padding at the end into account:
    for sample_idx in range(X.shape[0]):
        first_zero_sample = timepoints_per_sample
        while X[sample_idx, first_zero_sample - 1, 0] == 0:
            first_zero_sample -= 1
        X[sample_idx, 0: first_zero_sample, 0] -= np.mean(X[sample_idx, 0: first_zero_sample, 0])
        X[sample_idx, 0: first_zero_sample, 0] /= np.std(X[sample_idx, 0: first_zero_sample, 0])

    for sample_idx in range(X_test.shape[0]):
        first_zero_sample = timepoints_per_sample
        while X_test[sample_idx, first_zero_sample - 1, 0] == 0:
            first_zero_sample -= 1
        X_test[sample_idx, 0: first_zero_sample, 0] -= np.mean(X_test[sample_idx, 0: first_zero_sample, 0])
        X_test[sample_idx, 0: first_zero_sample, 0] /= np.std(X_test[sample_idx, 0: first_zero_sample, 0])

    freeze = False
    if freeze:
        model = get_transferred_model_frozen()
        file_path = "own_lstm_ptb_transfer_frozen.h5"
    else:
        model = get_transferred_model_trainable()
        file_path = "own_lstm_ptb_transfer_trainable.h5"

    checkpoint = ModelCheckpoint(file_path, monitor='val_acc', verbose=1, save_best_only=True, mode='max')
    early = EarlyStopping(monitor="val_acc", mode="max", patience=8, verbose=1)
    redonplat = ReduceLROnPlateau(monitor="val_acc", mode="max", patience=3, factor=0.3, verbose=2)
    callbacks_list = [checkpoint, early, redonplat]

    model.fit(X, Y, epochs=150, batch_size=160, verbose=2, callbacks=callbacks_list, validation_split=0.1)

    model.load_weights(file_path)
    pred_test = model.predict(X_test)
    pred_test[pred_test >= 0.5] = 1
    pred_test[pred_test < 0.5] = 0

    f1 = f1_score(Y_test, pred_test, average="macro")
    print("Test f1 score : ", f1)
    acc = accuracy_score(Y_test, pred_test)
    print("Test accuracy score : ", acc)

    '''
    First run frozen:
    Test f1 score :  0.7234967216932953
    Test accuracy score :  0.7993816557883888
    
    Second run frozen:
    Test f1 score :  0.73243422763509
    Test accuracy score :  0.8028169014084507
    
    Third run frozen:
    Test f1 score :  0.7299105071821994
    Test accuracy score :  0.7997251803503951
    
    Average: 
    Test f1 score :  0.72861
    Test accuracy score :  0.80064
    
    First run trainable:
    Test f1 score :  0.956425546420248
    Test accuracy score :  0.9653040192373755
    
    Second run trainable:
    Test f1 score :  0.9632134579854976
    Test accuracy score :  0.9708004122294744
    
    Third run trainable:
    Test f1 score :  0.9704821271667332
    Test accuracy score :  0.9762968052215734
    
    Average:
    Test f1 score :  0.96337
    Test accuracy score :  0.97080
    '''
    '''
    Note: performed in transfer learning notebook
    To repeat the experiment please re-run
    "Frozen and then trainable LSTM layers" section
    with perform_training=True
    
    First run frozen then trainable:
    Test accuracy score :  0.9783579525936105
    Test f1 score :  0.9730078352384828
    
    Second run frozen then trainable:
    Test accuracy score :  0.9780144280316042
    Test f1 score :  0.9725689192945755
    
    Third run frozen then trainable:
    Test accuracy score :  0.9776709034695981
    Test f1 score :  0.9721296598659295
    
    Average:
    Test accuracy score :  0.9780
    Test f1 score :  0.9726
    '''
