#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from tensorflow import keras
from keras.layers import LeakyReLU
import tensorflow as tf
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.metrics import confusion_matrix
import seaborn as sns
from sklearn.metrics import roc_curve, auc
import matplotlib

matplotlib.rc("font", family='Arial')


def read_data(file):
    """
    Read data from Excel
    """
    data = pd.read_excel(file, sheet_name='Sheet1')
    return data


def split_train_test(data):
    """
    Split dataset into training and testing sets
    """
    target = data['Distraction']
    data1 = data.iloc[:, :-1]
    print("Before split, data1 shape:", data1.shape)
    X_train, X_test, y_train, y_test = train_test_split(data1, target, test_size=0.2, random_state=345)
    print("After split, X_train shape:", X_train.shape)
    print("After split, X_test shape:", X_test.shape)
    return X_train, y_train, X_test, y_test


def normalize(x_train, x_test):
    """
    Perform data normalization
    """
    scaler = StandardScaler()
    print("Before normalize, x_train shape:", x_train.shape)
    x_train = scaler.fit_transform(x_train)
    print("After normalize, x_train shape:", x_train.shape)
    print("Before normalize, x_test shape:", x_test.shape)
    x_test = scaler.fit_transform(x_test)
    print("After normalize, x_test shape:", x_test.shape)
    return x_train, x_test


def DnnModel():
    """
    Define model using Sequential, total 4 layers:
    first layer 32, second layer 32, third layer 32, last layer 5-class output.
    """
    model = keras.models.Sequential([
        keras.layers.Dense(32, activation='relu', input_shape=(9,)),
        keras.layers.Dropout(0.3),
        keras.layers.Dense(32, activation=LeakyReLU()),
        keras.layers.Dense(32, activation='relu'),
        keras.layers.Dense(5, activation='softmax')
    ])
    model.summary()
    return model


if __name__ == '__main__':
    file = "data_train_88888.xlsx"
    dataset = read_data(file)
    train_X, train_Y, test_X, test_Y = split_train_test(dataset)  # read dataset
    x_train, x_test = normalize(train_X, test_X)  # normalization
    y_train = tf.keras.utils.to_categorical(train_Y - 1, 5)
    y_test = tf.keras.utils.to_categorical(test_Y - 1, 5)
    model = DnnModel()
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss='categorical_crossentropy',
        metrics=['accuracy'])

    history = model.fit(x_train, y_train, batch_size=32, epochs=100,
                        validation_data=(x_test, y_test), verbose=1)

    plt.plot(history.history['loss'], label='train_loss')
    plt.plot(history.history['val_loss'], label='val_loss')
    plt.plot(history.history['accuracy'], label='train_accuracy')
    plt.plot(history.history['val_accuracy'], label='val_accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Loss/Accuracy')
    plt.title('Training and Validation Loss and Accuracy')
    plt.legend()
    plt.show()

    y_pred = model.predict(x_test)
    y_pred_classes = np.argmax(y_pred, axis=1)
    y_true_classes = np.argmax(y_test, axis=1)
    cm = confusion_matrix(y_true_classes, y_pred_classes)

    print(cm)
    label = ["Label {}".format(i) for i in range(1, cm.shape[0] + 1)]
    df = pd.DataFrame(cm, index=label, columns=label)

    for i in range(cm.shape[0]):
        accuracy = cm[i][i] / np.sum(cm[:, i])
        print(f"Accuracy of class {i + 1}: {accuracy:.2%}")

    plt.figure(figsize=(7.5, 6.3))
    ax = sns.heatmap(df, xticklabels=df.corr().columns,
                     yticklabels=df.corr().columns, cmap='magma_r',
                     linewidths=6, annot=True)

    plt.xticks(fontsize=10, family='Times New Roman')
    plt.yticks(fontsize=10, family='Times New Roman')

    plt.tight_layout()
    plt.show()

    y_true_classes = np.argmax(y_test, axis=1)
    df = pd.DataFrame()
    pre_score = y_pred
    df['y_test'] = y_true_classes.tolist()

    df['pre_score1'] = pre_score[:, 0]
    df['pre_score2'] = pre_score[:, 1]
    df['pre_score3'] = pre_score[:, 2]
    df['pre_score4'] = pre_score[:, 3]
    df['pre_score5'] = pre_score[:, 4]

    pre1 = np.array(df['pre_score1'])
    pre2 = np.array(df['pre_score2'])
    pre3 = np.array(df['pre_score3'])
    pre4 = np.array(df['pre_score4'])
    pre5 = np.array(df['pre_score5'])

    y_list = df['y_test'].to_list()
    pre_list = [pre1, pre2, pre3, pre4, pre5]
    print(y_list, len(pre_list))
    label_names = ['1', '2', '3', '4', '5']
    colors1 = ["r", "b", "g", 'gold', 'pink']
    colors2 = 'red'
    my_list = []
    linestyles = ["-", "--", ":", "-", "--"]

    plt.figure(figsize=(10, 5), facecolor='w')
    for i in range(5):
        if i == 0:
            fpr, tpr, threshold = roc_curve(y_list, pre_list[i], pos_label=0)
            roc_auc = auc(fpr, tpr)
            plt.text(0.3, 0.01, "class " + label_names[i] + ' :ROC curve (area = %0.2f)' % roc_auc)
        elif i == 1:
            fpr, tpr, threshold = roc_curve(y_list, pre_list[i], pos_label=1)
            roc_auc = auc(fpr, tpr)
            plt.text(0.3, 0.11, "class " + label_names[i] + ' :ROC curve (area = %0.2f)' % roc_auc)
        elif i == 2:
            fpr, tpr, threshold = roc_curve(y_list, pre_list[i], pos_label=2)
            roc_auc = auc(fpr, tpr)
            plt.text(0.3, 0.21, "class " + label_names[i] + ' :ROC curve (area = %0.2f)' % roc_auc)
        elif i == 3:
            fpr, tpr, threshold = roc_curve(y_list, pre_list[i], pos_label=3)
            roc_auc = auc(fpr, tpr)
            plt.text(0.3, 0.31, "class " + label_names[i] + ' :ROC curve (area = %0.2f)' % roc_auc)
        elif i == 4:
            fpr, tpr, threshold = roc_curve(y_list, pre_list[i], pos_label=4)
            roc_auc = auc(fpr, tpr)
            plt.text(0.3, 0.41, "class " + label_names[i] + ' :ROC curve (area = %0.2f)' % roc_auc)

        my_list.append(roc_auc)
        plt.plot(fpr, tpr, color=colors1[i], linestyle=linestyles[i], linewidth=3,
                 label="class:" + label_names[i])
        plt.stackplot(fpr, tpr, colors=colors2, alpha=0.5, edgecolor=colors1[i])

    plt.plot([0, 1], [0, 1], color='black', linestyle='--', linewidth=3)
    plt.xlabel('1 - Specificity')
    plt.ylabel('Sensitivity')
    plt.grid()
    plt.legend()
    plt.title("ROC Curve and AUC Values")
    plt.show()
