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 create_sequences(data, time_steps=5):
    X = []
    for i in range(len(data) - time_steps + 1):
        X.append(data[i:i + time_steps])
    return np.array(X)


def GRUModel(input_shape, num_classes):
    """
    Build GRU model
    """
    model = keras.models.Sequential([
        keras.layers.GRU(128, return_sequences=True, input_shape=input_shape),
        keras.layers.GRU(32),
        keras.layers.Dense(num_classes, activation='softmax')
    ])
    model.summary()
    return model


if __name__ == '__main__':
    file = "C:/Users/asus/Desktop/fdy_dataset-0-predict66.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)

    time_steps = 5
    train_X_sequences = create_sequences(x_train, time_steps)
    test_X_sequences = create_sequences(x_test, time_steps)
    y_train_sequences = y_train[time_steps - 1:]
    y_test_sequences = y_test[time_steps - 1:]

    input_shape = (time_steps, x_train.shape[1])
    model = GRUModel(input_shape=input_shape, num_classes=num_classes)
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss='categorical_crossentropy',
        metrics=['accuracy'])

    history = model.fit(train_X_sequences, y_train_sequences, batch_size=1, epochs=2,
                        validation_data=(test_X_sequences, y_test_sequences), verbose=1)

    model.save('EEG_4class_trained_model_GRU.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(test_X_sequences)
    y_pred_classes = np.argmax(y_pred, axis=1)
    y_true_classes = np.argmax(y_test_sequences, 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()

    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_sequences, 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]

    pre_list = [df['pre_score1'].to_numpy(),
                df['pre_score2'].to_numpy(),
                df['pre_score3'].to_numpy(),
                df['pre_score4'].to_numpy()]

    y_list = df['y_test'].to_list()
    print(y_list, len(pre_list))
    class_labels = ['1', '2', '3', '4']
    colors1 = ["r", "b", "g", "b"]
    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 " + class_labels[i] + ' :ROC curve (area = %0.2f)' % roc_auc)
        plt.plot(fpr, tpr, color=colors1[i], linestyle=linestyles[i], linewidth=3,
                 label="class:" + class_labels[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()
