import json

import matplotlib.pyplot as plt
import numpy as np
import tensorflow.keras as keras
from sklearn.model_selection import train_test_split
from numpy.linalg import inv

class MlpClassifier():
    
    def __init__(self, config):
        """[Constructor get config json so that is possible to run the pre-process]

        Args:
            config ([dict]): [parameter used to adjust the model train]
        """
        self.config = config

    def run(self):
        """[run the model train with the parameter given by the config and save it inside the directory]
        """        
        list_of_combinations = [
            ["mel_spec"],
        ]
        for combination in list_of_combinations:
            
            # OLD DATASET(Only Music and Speech)
            data = self.load_data(self.config["pre_process"]["output"])
            np.random.shuffle(data["data"])
            X, y = self.mount_input(data, combination)
            X_train, X_validation, X_test, y_train, y_validation, y_test = self.split_database(X, y, self.config["model"]["test_size"], self.config["model"]["validation_size"])
            input_shape = (X_train.shape[1], X_train.shape[2], 1)
            model = self.build_model(input_shape)
            history = model.fit(X_train, y_train, validation_data=(X_validation, y_validation), batch_size=self.config["model"]["batch_size"], epochs=self.config["model"]["epochs"])
            suffix = "_".join(combination)
            self.plot_history(history,"old"+suffix)
            model.save('{}mlp_beat_box_{}'.format(self.config["model"]["output"], suffix+"_old"))
            del data, X, y, X_train, X_validation, X_test, y_train, y_validation, y_test, model
            keras.backend.clear_session()
            
            # YOUTUBE AUDIO SET (Every Type of Audio)
            data = self.load_data(self.config["youtube"]["output"])
            model = self.build_model(input_shape)
            np.random.shuffle(data["data"])
            X, y = self.mount_input(data, combination)
            X_train, X_validation, X_test, y_train, y_validation, y_test = self.split_database(X, y, self.config["model"]["test_size"], self.config["model"]["validation_size"])
            history = model.fit(X_train, y_train, validation_data=(X_validation, y_validation), batch_size=self.config["model"]["batch_size"], epochs=self.config["model"]["epochs"])
            self.plot_history(history,"yt"+suffix)


            # evaluate model on test set
            test_loss, test_acc = model.evaluate(X_test, y_test, verbose=2)
            print('\nTest accuracy: {}'.format(test_acc).center(60, "*"))
            print('Save Model'.center(60, "*"))
            model.save('{}mlp_beat_box_{}'.format(self.config["model"]["output"], suffix+"_yt"))
            keras.backend.clear_session() 
            del data, X, y, X_train, X_validation, X_test, y_train, y_validation, y_test, model
            
    def load_data(self,data_path):
        """[open the train data and split it over input(X) and output(Y)]

        Args:
            data_path ([string]): [path of the dataset generated by pre-process]

        Returns:
            [data]: [all data genereted on preprocess]
        """
        data = None      
        with open(data_path, "r") as fp:
            data = json.load(fp)
        print("Data was loaded!".center(60, "*"))
        return data

    def mount_input(self,data, combination):
        """[mount the input that will be used on the training stage]

        Args:
            data ([dict]): [features extracted from the audio on preprocess phase]
            combination ([list]): [list with the name of the features that will be used]
        """
        first_feature = combination[0]
        other_features =  list(filter(lambda x:first_feature not in x, combination))

        X = []
        y = []
        for row in data["data"]:
            c1 =  np.array(row[first_feature])
            for other_feature in other_features:
                c_other =  np.array(row[other_feature])
                c1 = np.concatenate([c1, c_other], axis = 0)
            X.append(c1.T)
            y.append(row["label"])
        X = np.array(X)
        y = np.array(y)
        print("Data was mounted!".center(60, "*"))
        return X, y


    def split_database(self,X, y, test_size, validation_size):
        """[split the input data between train, test and validation]

        Args:
            X ([np.array]): [input that will be given to the neural network]
            y ([np.array]): [output that will be given to the neural network]
            test_size ([float]): [the percentage of the samples that will be used to test the model]
            validation_size ([float]): [the percentage of the train samples that will be used to validate the model]

        Returns:
            [X_train]: [the samples that will be used to train the model]
            [X_validation]: [the samples that will be used for validation the model]
            [X_test]: [the samples that will be used to test the model]
            [y_train]: [the labels of the samples used to train the model]
            [y_validation]: [the labels of the samples used to validate the model]
            [y_test]: [the labels of the samples used to test the model]
        """        

        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size)
        X_train, X_validation, y_train, y_validation = train_test_split(X_train, y_train, test_size=validation_size)

        X_train = X_train[..., np.newaxis]
        X_validation = X_validation[..., np.newaxis]
        X_test = X_test[..., np.newaxis]

        return X_train, X_validation, X_test, y_train, y_validation, y_test

    def build_model(self, input_shape):
        """[mount neural network with the parameter informed on json file]

        Args:
            input_shape ([tuple]): [tuple with the shape of the entrance of the model]

        Returns:
            [model]: [keras model used build from parameters]
        """
        # build network topology
        model = keras.Sequential()

        # 1st conv layer
        model.add(keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=input_shape))
        model.add(keras.layers.MaxPooling2D((3, 3), strides=(2, 2), padding='same'))
        model.add(keras.layers.BatchNormalization())

        # 2nd conv layer
        model.add(keras.layers.Conv2D(32, (3, 3), activation='relu'))
        model.add(keras.layers.MaxPooling2D((3, 3), strides=(2, 2), padding='same'))
        model.add(keras.layers.BatchNormalization())

        # 3rd conv layer
        model.add(keras.layers.Conv2D(32, (2, 2), activation='relu'))
        model.add(keras.layers.MaxPooling2D((2, 2), strides=(2, 2), padding='same'))
        model.add(keras.layers.BatchNormalization())

        # flatten output and feed it into dense layer
        model.add(keras.layers.Flatten())
        model.add(keras.layers.Dense(64, activation='relu'))
        model.add(keras.layers.Dropout(0.2))

        # output layer
        model.add(keras.layers.Dense(2, activation='softmax'))

        optimizer = keras.optimizers.Adam(learning_rate=self.config["model"]["learning_rate"])
        model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
        model.summary()
        return model

    def plot_history(self,history, suffix):
        """[plot the evolution of the accuracy and the error of the model during the training]

        Args:
            history ([list]): [list with the error and accuracy of the model during training]
        """        
        fig, axs = plt.subplots(2,figsize=(15,15))
        axs[0].plot(history.history["accuracy"], label="train accuracy")
        axs[0].plot(history.history["val_accuracy"], label="test accuracy")
        axs[0].set_ylabel("Accuracy")
        axs[0].legend(loc="lower right")
        axs[0].set_title("Accuracy eval")

        axs[1].plot(history.history["loss"], label="train error")
        axs[1].plot(history.history["val_loss"], label="test error")
        axs[1].set_ylabel("Error")
        axs[1].set_xlabel("Epoch")
        axs[1].legend(loc="upper right")
        axs[1].set_title("Error eval")
        plt.savefig("history_{}.png".format(suffix))
        plt.clf()
