import time

import argparse

import tensorflow as tf

import numpy as np

from tensorflow.keras.layers import Dense, Flatten, Conv1D, BatchNormalization, MaxPool1D, Dropout

from tensorflow.keras.metrics import CategoricalAccuracy

from sklearn.model_selection import train_test_split

from sklearn.metrics import precision_score, recall_score, confusion_matrix

from utils import get_labels, get_datasets, check_processed_dir_existance

par = argparse.ArgumentParser(description="ECG Convolutional " +

                                           "Neural Network implementation with Tensorflow 2.0")

par.add_argument("-lr", dest="learning_rate",

                 type=float, default=0.001,

                 help="Learning rate used by the model")

par.add_argument("-e", dest="epochs",

                 type=int, default=50,

                 help="The number of epochs the model will train for")

par.add_argument("-bs", dest="batch_size",

                 type=int, default=32,

                 help="The batch size of the model")

par.add_argument("--display-step", dest="display_step",

                 type=int, default=10,

                 help="The display step")

par.add_argument("--dropout", type=float, default=0.5,

                 help="Dropout probability")

par.add_argument("--restore", dest="restore_model",

                 action="store_true", default=False,

                 help="Restore the model previously saved")

par.add_argument("--freeze", dest="freeze",

                 action="store_true", default=False,

                 help="Freezes the model")

par.add_argument("--heart-diseases", nargs="+",

                 dest="heart_diseases",

                 default=["apnea-ecg", "svdb", "afdb"],

                 choices=["apnea-ecg", "mitdb", "nsrdb", "svdb", "afdb"],

                 help="Select the ECG diseases for the model")

par.add_argument("--verbose", dest="verbose",

                 action="store_true", default=False,

                 help="Display information about minibatches")

args = par.parse_args()

# Parameters

learning_rate = args.learning_rate

epochs = args.epochs

batch_size = args.batch_size

display_step = args.display_step

dropout = args.dropout

restore_model = args.restore_model

freeze = args.freeze

heart_diseases = args.heart_diseases

verbose = args.verbose

# Network Parameters

n_inputs = 350

n_classes = len(heart_diseases)

check_processed_dir_existance()

class CNN:

    def __init__(self):

        self.datasets = get_datasets(heart_diseases, n_inputs)

        self.label_data = get_labels(self.datasets)

        self.callbacks = []

        # Initialize callbacks

        tensorboard_logs_path = "tensorboard_data/cnn/"

        tb_callback = tf.keras.callbacks.TensorBoard(log_dir=tensorboard_logs_path,

                                                     histogram_freq=1, write_graph=True,

                                                     embeddings_freq=1)

        # load_weights_on_restart will read the filepath of the weights if it exists and it will

        # load the weights into the model

        cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath="saved_models/cnn/model.hdf5",

                                                         save_best_only=True,

                                                         save_weights_only=True,

                                                         load_weights_on_restart=restore_model)

        self.callbacks.extend([tb_callback, cp_callback])

        self.set_data()

        self.define_model()

    def set_data(self):

        dataset_len = []

        for dataset in self.datasets:

            dataset_len.append(len(dataset))

        # validation on 10% of the training data

        validation_size = 0.1

        print("Validation percentage: {}%".format(validation_size*100))

        print("Total samples: {}".format(sum(dataset_len)))

        print("Heart diseases: {}".format(', '.join(heart_diseases)))

        concat_dataset = np.concatenate(self.datasets)

        self.split_data(concat_dataset, validation_size)

        # Reshape input so that we can feed it to the conv layer

        self.X_train = tf.reshape(self.X_train, shape=[-1, n_inputs, 1])

        self.X_test = tf.reshape(self.X_test, shape=[-1, n_inputs, 1])

        self.X_val = tf.reshape(self.X_val, shape=[-1, n_inputs, 1])

        if verbose:

            print("X_train shape: {}".format(self.X_train.shape))

            print("Y_train shape: {}".format(self.Y_train.shape))

            print("X_test shape: {}".format(self.X_test.shape))

            print("Y_test shape: {}".format(self.Y_test.shape))

            print("X_val shape: {}".format(self.X_val.shape))

            print("Y_val shape: {}".format(self.Y_val.shape))

    def define_model(self):

         inputs = tf.keras.Input(shape=(n_inputs, 1), name='input')

         # 64 filters, 10 kernel size

         x = Conv1D(64, 10, activation='relu')(inputs)

         x = MaxPool1D()(x)

         x = BatchNormalization()(x)

         x = Conv1D(128, 10, activation='relu')(x)

         x = MaxPool1D()(x)

         x = BatchNormalization()(x)

         x = Conv1D(128, 10, activation='relu')(x)

         x = MaxPool1D()(x)

         x = BatchNormalization()(x)

         x = Conv1D(256, 10, activation='relu')(x)

         x = MaxPool1D()(x)

         x = BatchNormalization()(x)

         x = Flatten()(x)

         x = Dense(1024, activation='relu', name='dense_1')(x)

         x = BatchNormalization()(x)

         x = Dropout(dropout)(x)

         x = Dense(2048, activation='relu', name='dense_2')(x)

         x = BatchNormalization()(x)

         x = Dropout(dropout)(x)

         outputs = Dense(n_classes, activation='softmax', name='predictions')(x)

         self.cnn_model = tf.keras.Model(inputs=inputs, outputs=outputs)

         optimizer = tf.keras.optimizers.Adam(lr=learning_rate)

         accuracy = CategoricalAccuracy()

         self.cnn_model.compile(optimizer=optimizer, loss='categorical_crossentropy',

                                metrics=[accuracy])

    def split_data(self, dataset, validation_size):

        """

        Suffle then split training, testing and validation sets

        """

        # In order to use statify in train_test_split we can't use one hot encodings,

        # so we convert to array of labels

        label_data = np.argmax(self.label_data, axis=1)

        # Splitting the dataset into train and test datasets

        res = train_test_split(dataset, label_data,

                               test_size=validation_size, shuffle=True,

                               stratify=label_data)

        self.X_train, self.X_test, self.Y_train, self.Y_test = res

        # From the training dataset we further split it to obtain the validation dataset

        res = train_test_split(self.X_train, self.Y_train,

                               test_size=validation_size, stratify=self.Y_train)

        self.X_train, self.X_val, self.Y_train, self.Y_val = res

        # Convert the array of labels back into one hot encodings to be able to do training

        self.Y_train = tf.keras.utils.to_categorical(self.Y_train)

        self.Y_test = tf.keras.utils.to_categorical(self.Y_test)

        self.Y_val = tf.keras.utils.to_categorical(self.Y_val)

    def get_data(self):

        return (self.X_train, self.X_test, self.X_val,

                self.Y_train, self.Y_test, self.Y_val)

def main():

    # Construct model

    model = CNN()

    X_train, X_test, X_val, Y_train, Y_test, Y_val = model.get_data()

    # Set start time

    total_time = time.time()

    print("-"*50)

    if restore_model:

        print("Restoring model: {}".format('saved_models/cnn/model.hdf5'))

    # Train

    model.cnn_model.fit(X_train, Y_train, batch_size=batch_size,

                        epochs=epochs, validation_data=(X_val, Y_val),

                        callbacks=model.callbacks)

    print("-"*50)

    # Total training time

    print("Total training time: {0:.2f}s".format(time.time() - total_time))

    # Test

    model.cnn_model.evaluate(X_test, Y_test, batch_size=batch_size)

    print("-"*50)

    print("Testing results:")

    y_pred = model.cnn_model.predict(X_test, batch_size=batch_size)

    # The following scikit-learn methods only accept array of labels, not one hot encodings

    y_pred = np.argmax(y_pred, axis=1)

    y_true = np.argmax(Y_test, axis=1)

    # Precision and recall could also be done as callbacks in the evaluate or fit function

    print("Precision: {}".format(precision_score(y_true, y_pred, average='micro')))

    print("Recall: {}".format(recall_score(y_true, y_pred, average='micro')))

    print("Confusion matrix: \n{}".format(confusion_matrix(y_true, y_pred, labels=[0,1,2])))

    disease_indexes = list(range(len(heart_diseases)))

    print("Indexes {} correspond to labels {}".format(disease_indexes, [x for x in heart_diseases]))

    print("-"*50)

if __name__ == "__main__":

    main()