import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from tensorflow import keras
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]
    X_train, X_test, y_train, y_test = train_test_split(data1, target, test_size=0.2, random_state=345)
    return X_train, y_train, X_test, y_test


def normalize(x_train, x_test):
    """
    Normalize data
    """
    scaler = StandardScaler()
    x_train = scaler.fit_transform(x_train)
    x_test = scaler.transform(x_test)
    return x_train, x_test


def one_hot_encode(labels, num_classes):
    encoded = np.zeros((len(labels), num_classes))
    for i, label in enumerate(labels):
        if 1 <= label <= num_classes:
            encoded[i, label - 1] = 1
        else:
            encoded[i, 0] = 1
    return encoded


def BPModel():
    """
    Build BP Neural Network model
    """
    model = keras.models.Sequential([
        keras.layers.Dense(128, activation='relu', input_shape=(21,)),
        keras.layers.Dropout(0.4),
        keras.layers.Dense(64, activation='relu'),
        keras.layers.Dense(32, activation='relu'),
        keras.layers.Dense(7, activation='softmax')
    ])
    model.summary()
    return model


if __name__ == '__main__':
    file = "C:/Users/asus/Desktop/fdy_dataset_prediction7.xlsx"
    dataset = read_data(file)
    train_X, train_Y, test_X, test_Y = split_train_test(dataset)
    x_train, x_test = normalize(train_X, test_X)

    num_classes = 7
    y_train = one_hot_encode(train_Y, num_classes)
    y_test = one_hot_encode(test_Y, num_classes)

    model = BPModel()
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss='categorical_crossentropy',
        metrics=['accuracy'])

    history = model.fit(x_train, y_train, batch_size=32, epochs=80,
                        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("Confusion Matrix:")
    print(cm)
    label = ["Label {}".format(i) for i in range(1, cm.shape[0] + 1)]
    df = pd.DataFrame(cm, index=label, columns=label)

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

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

    # ROC Curve
    y_true_classes = np.argmax(y_test, axis=1)
    df = pd.DataFrame()
    pre_score = y_pred
    df['y_test'] = y_true_classes.tolist()
    for i in range(num_classes):
        df[f'pre_score{i + 1}'] = pre_score[:, i]

    pre_list = [np.array(df[f'pre_score{i + 1}']) for i in range(num_classes)]
    y_list = df['y_test'].to_list()
    label_names = [str(i) for i in range(1, num_classes + 1)]
    colors1 = ["r", "b", "g", "m", "c", "y", "k"]
    linestyles = ["-", "--", ":", "-.", "-", "--", ":"]

    plt.figure(figsize=(10, 5), facecolor='w')
    for i in range(num_classes):
        fpr, tpr, threshold = roc_curve(y_list, pre_list[i], pos_label=i)
        roc_auc = auc(fpr, tpr)
        plt.text(0.3, 0.01 + i * 0.1,
                 "class " + label_names[i] + ' :ROC curve (area = %0.2f)' % roc_auc)
        plt.plot(fpr, tpr, color=colors1[i], linestyle=linestyles[i], linewidth=3,
                 label="class:" + label_names[i])
        plt.stackplot(fpr, tpr, colors="red", 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")
    plt.show()
