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()