import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import matplotlib.pyplot as plt



plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['axes.labelweight'] = 'bold'
plt.rcParams['axes.titleweight'] = 'bold'
plt.rcParams['font.weight'] = 'bold'


def main():
    # Load data
    file_path = r""
    data = pd.read_excel(file_path, header=None)
    data.columns = ['Year_2008', 'Year_2016', 'Year_2023']

    # Create a mask to identify non-5 data points
    mask_2008_2016 = (data['Year_2008'] != 5) & (data['Year_2016'] != 5)
    mask_2016_2023 = (data['Year_2016'] != 5) & (data['Year_2023'] != 5)
    combined_mask = mask_2008_2016 & mask_2016_2023

    # Create change features, only using non-5 data
    data['Change_2008_2016'] = np.where(mask_2008_2016, data['Year_2016'] - data['Year_2008'], np.nan)
    data['Change_2016_2023'] = np.where(mask_2016_2023, data['Year_2023'] - data['Year_2016'], np.nan)

    # Use only non-5 data for training and prediction
    train_data = data.loc[combined_mask, ['Change_2008_2016', 'Change_2016_2023']]
    train_labels = data.loc[combined_mask, 'Year_2023'].map({1: 0, 2: 1, 8: 2})

    # Data preprocessing
    scaler = MinMaxScaler(feature_range=(0, 1))
    train_scaled = scaler.fit_transform(train_data)

    # Prepare LSTM input data
    X_train = train_scaled.reshape(-1, 2, 1)
    y_train = train_labels.values

    # Use the changes of all data as the test set
    test_data = data[['Change_2008_2016', 'Change_2016_2023']].fillna(0)  # Fill NaN with 0 to ensure consistency
    test_scaled = scaler.transform(test_data)
    X_test = test_scaled.reshape(-1, 2, 1)

    # Build LSTM model
    model = Sequential()
    model.add(LSTM(50, input_shape=(2, 1)))
    model.add(Dense(3, activation='softmax'))  # Three categories
    model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

    # Train the model
    model.fit(X_train, y_train, epochs=10, batch_size=1, verbose=2)

    # Make predictions
    predictions = model.predict(X_test)
    predictions_classes = np.argmax(predictions, axis=1)

    # Restore 5s to their original positions
    final_predictions = np.where(data['Year_2023'] == 5, 5, predictions_classes)

    # Reverse map the prediction results
    reverse_mapping = {0: 1, 1: 2, 2: 8}
    final_predictions = np.vectorize(reverse_mapping.get)(final_predictions)

    # Calculate metrics
    mask_combined_test = combined_mask.values
    y_test = data.loc[mask_combined_test, 'Year_2023'].map({1: 0, 2: 1, 8: 2}).values
    accuracy = accuracy_score(y_test, predictions_classes[mask_combined_test])
    precision = precision_score(y_test, predictions_classes[mask_combined_test], average='weighted')
    recall = recall_score(y_test, predictions_classes[mask_combined_test], average='weighted')
    f1 = f1_score(y_test, predictions_classes[mask_combined_test], average='weighted')

    # Save prediction results to Excel
    predicted_df = pd.DataFrame(final_predictions, columns=['Predicted'])
    output_path = r""
    predicted_df.to_excel(output_path, index=False)
    print(f"File saved to: {output_path}")

    return y_test, final_predictions, accuracy, precision, recall, f1

def plot_results(y_test, final_predictions):
    plt.figure(figsize=(12, 6))
    plt.plot(y_test, label='Actual Values', marker='o')
    plt.plot(final_predictions, label='Predicted Values', marker='x')
    plt.title('Comparison of Actual and Predicted Values')
    plt.xlabel('Index')
    plt.ylabel('Values')
    plt.legend()
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    y_test, final_predictions, accuracy, precision, recall, f1 = main()
    plot_results(y_test, final_predictions)
    print(f"Accuracy: {accuracy}")
    print(f"Precision: {precision}")
    print(f"Recall: {recall}")
    print(f"F1 Score: {f1}")
