"""This module for LSTM and all sections need updated"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping
import time

# Function to create dataset for LSTM
def create_dataset(data, look_back):
    X, Y = [], []
    for i in range(len(data) - look_back):
        X.append(data[i:(i + look_back), :])
        Y.append(data[i + look_back, 0])  # Assuming predicting Y
    return np.array(X), np.array(Y)

# Function to compute and print RMSE and other metrics
def compute_rmse(y_true, y_pred):
    mae = mean_absolute_error(y_true, y_pred)
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    r2 = r2_score(y_true, y_pred)

    print("mean y_test:", np.mean(y_true))
    print("MSE:", mse)
    print("MAE:", mae)
    print("RMSE:", rmse)
    print("R2 score:", r2)

    rmse_percent = rmse / np.mean(y_true) * 100

    print("MSE %:", mse / np.mean(y_true) * 100)
    print("MAE %:", mae / np.mean(y_true) * 100)
    print("RMSE %:", rmse_percent)

    return rmse_percent

# Function to build and compile the LSTM model
def build_lstm_model(input_shape):
    model = Sequential()
    model.add(LSTM(units=50, input_shape=input_shape))
    model.add(Dense(units=32, activation='relu'))
    model.add(Dense(units=1, activation='linear'))
    model.add(Dropout(0.1))
    model.add(Dense(units=1))
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

# Function to plot and save loss and accuracy
def plot_training_history(history, n_var, col, normx, area):
    loss = history.history['loss']
    plt.figure(figsize=(10, 6))
    plt.plot(range(len(loss)), loss)
    plt.xlabel('Epochs')
    plt.ylabel('Loss: Mean Squared Error')
    plt.title(f'Loss_{n_var}_INTPP_to_{col}_{normx}_{area}')
    plt.savefig(f'plots/LSTM/MULTI_926/LSTM_Loss_{n_var}_INTPP_to_{col}_{normx}_{area}.png')

    if 'accuracy' in history.history:
        acc = history.history['accuracy']
        plt.figure(figsize=(10, 6))
        plt.plot(range(len(acc)), acc)
        plt.xlabel('Epochs')
        plt.ylabel('Accuracy')
        plt.title(f'ACC_{n_var}_INTPP_to_{col}_{normx}_{area}')
        plt.savefig(f'plots/LSTM/MULTI_926/LSTM_Acc_{n_var}_INTPP_to_{col}_{normx}_{area}.png')

# Function to plot and save predictions
def plot_forecast(Y_test_transformed, y_pred_transformed, n_var, col, normx, area):
    plt.plot(Y_test_transformed, label='Test sequence', linestyle='dotted')
    plt.plot(y_pred_transformed, label='Predicted sequence')
    plt.title(f'ML LSTM Multvar. INTPP to {col.upper()} & {normx.upper()} @ {area} forecast. RMSE: {rmse:.2f} %')
    plt.xlabel('Time')
    plt.ylabel('Transformed values [ ]')
    plt.legend()
    plt.savefig(f'plots/LSTM/MULTI_926/LSTM_forecast_var_{n_var}_INTPP_to_{col}_{normx}_{area}.png')

# Main function to execute the LSTM process
def run_lstm(df_train, df_test, area_dict, normx_dict, columns, variables, n_forecast_steps):
    start_time_lstm = time.time()
    rmse_results = pd.DataFrame(columns=['NORM', 'COLUMNS', 'VAR', 'RMSE', 'AREA', 'MODEL', 'HYPERTUNE'])

    for area in area_dict:
        print('RUN FOR FOLLOWING LOCATION ', area)
        for normx in normx_dict:
            train_columns = {col: f'{variables[col]}_{normx}_train_{area}' for col in columns}
            test_columns = {col: f'{variables[col]}_{normx}_test_{area}' for col in columns}

            df_test_predictor = df_test['intpp_'+normx+'_test_'+area]
            df_train_predictor_shift = df_train['intpp_'+normx+'_train_'+area][1:]
            df_train_predictor = df_train['intpp_'+normx+'_train_'+area]

            df_new = pd.DataFrame({'Z': df_train_predictor_shift, 'Zt-1': df_train['intpp_lag1_'+normx+'_train_'+area][1:]})
            df_new_test = pd.DataFrame({'Z': df_test_predictor, 'Zt-1': df_test['intpp_lag1_'+normx+'_test_'+area]})

            for col in columns:
                df_new[col] = df_train[train_columns[col]][1:]
                df_new_test[col] = df_test[test_columns[col]]

                scaler = MinMaxScaler()
                scaled_data = scaler.fit_transform(df_new)
                scaled_test_data = scaler.transform(df_new_test)

                look_back = 10
                X, Y = create_dataset(scaled_data, look_back)
                X_test, Y_test = create_dataset(scaled_test_data, look_back)

                X_test = X_test[:n_forecast_steps]
                Y_test = Y_test[:n_forecast_steps]

                model = build_lstm_model(input_shape=(X.shape[1], X.shape[2]))
                history = model.fit(X, Y, epochs=100, batch_size=32, verbose=0)

                plot_training_history(history, str(columns.index(col)+1), variables[col], normx, area)

                y_pred = model.predict(X_test)
                min_val, max_val = scaler.data_min_[0], scaler.data_max_[0]
                y_pred_transformed = y_pred * (max_val - min_val) + min_val
                Y_test_transformed = Y_test * (max_val - min_val) + min_val

                rmse = compute_rmse(Y_test_transformed, y_pred_transformed)
                new_row = pd.DataFrame({'NORM': [normx], 'COLUMNS': [columns[:columns.index(col)+1]],
                                        'VAR': [variables[col]], 'RMSE': [rmse], 'AREA': [area], 'MODEL': 'LSTM', 'HYPERTUNE': '50,32,0.1,1,1'})
                rmse_results = pd.concat([rmse_results, new_row], ignore_index=True)

                plot_forecast(Y_test_transformed, y_pred_transformed, str(columns.index(col)+1), variables[col], normx, area)

    rmse_results.to_csv('plots/LSTM/MULTI_926/LSTM_RMSE.csv', index=False)

    end_time_lstm = time.time()
    time_used_lstm = end_time_lstm - start_time_lstm
    print("Time used for LSTM:", f'{time_used_lstm:.4f}', "Seconds")
    print("Time used for LSTM:", f'{time_used_lstm/60:.4f}', "Minutes")
    print("Time used for LSTM:", f'{time_used_lstm/3600:.4f}', "Hours")
