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):
    """
    Perform data normalization
    """
    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 DnnModel():
    """
    Define a 4-layer Sequential model:
    Dense(32) → Dropout(0.3) → Dense(32 LeakyReLU) → Dense(32 ReLU) → Dense(4 Softmax)
    """
    model = keras.models.Sequential([
        keras.layers.Dense(32, activation='relu', input_shape=(21,)),
        keras.layers.Dropout(0.3),
        keras.layers.Dense(32, activation=keras.layers.LeakyReLU()),
        keras.layers.Dense(32, activation='relu'),
        keras.layers.Dense(4, activation='softmax')
    ])
    model.summary()
    return model


if __name__ == '__main__':
    file = "C:/Users/fdy12/Desktop/fdy_dataset_0_pred6.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 = 4
    y_train = one_hot_encode(train_Y, num_classes)
    y_test = one_hot_encode(test_Y, num_classes)

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

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

    model.save('EEG_4class_trained_model_02.h5')

    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 = [f"Label {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))
    ax = sns.heatmap(df, xticklabels=df.corr().columns,
                     yticklabels=df.corr().columns, 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()

    label_names = {1: 'Decelerate', 2: 'LaneChangeLeft', 3: 'LaneChangeRight', 4: 'InvalidCommand'}
    for i in range(cm.shape[0]):
        label_acc = cm[i, i] / np.sum(cm[i, :]) if np.sum(cm[i, :]) != 0 else 0
        print(f"{label_names[i + 1]} accuracy: {label_acc * 100:.2f}%")

    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]

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

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

    plt.figure(figsize=(10, 5), facecolor='w')
    for i in range(4):
        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_seq[i] + f' :ROC curve (area = {roc_auc:0.2f})')
        plt.plot(fpr, tpr, color=colors1[i], linestyle=linestyles[i], linewidth=3,
                 label="class: " + label_names_seq[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")
    plt.show()
