import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.metrics import mean_absolute_error, mean_squared_error
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from sklearn.model_selection import GridSearchCV, KFold, TimeSeriesSplit, train_test_split
from sklearn.preprocessing import StandardScaler
import pickle
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LinearRegression, Lasso, LassoCV
import warnings
import seaborn as sns
from sklearn.neural_network import MLPRegressor
import tensorflow as tf
from statsmodels.regression.quantile_regression import QuantReg
warnings.filterwarnings('ignore')
season_colors = {'Winter': '#8bb3e6', 'Spring': '#90d17f', 'Summer': '#ff9a87', 'Fall': '#ffa04c'}
model_colors_pred = {'ANN': '#4FB075', 'LGBM': '#3070CF', 'LR': '#CF8F30'}
model_colors_real = {'ANN': '#060e00', 'LGBM': '#060e00', 'LR': '#060e00'}
binary_colours = {"blue": "#ff0000", "orange": "#9500ff"}
basic_colours = {"blue": "#3070CF", "orange": "#CF8F30"}
def calculate_and_plot_mse(X_train, y_train, model,target_col):
    kf = KFold(n_splits=5, shuffle=True, random_state=42)
    fold = 0
    for train_idx, test_idx in kf.split(X_train):
        fold += 1
        X_tr, X_te = X_train.iloc[train_idx], X_train.iloc[test_idx]
        y_tr, y_te = y_train.iloc[train_idx], y_train.iloc[test_idx]

        y_pred = model.predict(X_te)

        mse = mean_squared_error(y_te, y_pred)
        print(f'Fold {fold} - MSE: {mse:.2f}')

        # Plot predicted vs actual for each fold
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(y_te.index, y_te, label="Real Data")
        ax.plot(y_te.index, y_pred, label="Model Prediction")
        ax.set_title(f'{target_col} Fold {fold} Actual vs Predicted')
        ax.set_xlabel('Index')
        ax.set_ylabel(target_col)
        ax.legend()

def plot_feature_importance(feature_importance):
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(feature_importance['Feature'], feature_importance['Importance'])
    ax.set_title('Feature Importance')
    ax.set_xlabel('Features')
    ax.set_ylabel('Importance')
    ax.tick_params(axis='x', labelrotation=90)
    plt.show()


def plot_actual_vs_predicted(y_valid, y_pred, target_col):
    mae = mean_absolute_error(y_valid, y_pred)
    rel_error = np.mean(np.abs((y_valid - y_pred) / y_valid)) * 100
    
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(y_valid.index, y_valid, label="Real Data")
    ax.plot(y_valid.index, y_pred, label="Model Prediction", linewidth=2)
    ax.fill_between(y_valid.index, y_pred, y_valid, alpha=0.2)
    ax.set_title(f'{target_col} Actual vs Predicted\nMAE: {mae:.2f}  Relative Error: {rel_error:.2f}%')
    ax.set_xlabel('Index')
    ax.set_ylabel(target_col)
    ax.legend()
import os

def save_figure(fig, title):
    # Create the folder if it doesn't exist
    if not os.path.exists('DP_plots'):
        os.makedirs('DP_plots')

    # Check if the file already exists and add incremental suffix if it does
    suffix = 0
    filename = f'DP_plots/{title}.svg'
    while os.path.exists(filename):
        suffix += 1
        filename = f'DP_plots/{title}_{suffix}.svg'

    # Save the figure in SVG and PNG formats
    fig.savefig(filename, format='svg', dpi=300)

    # Reset suffix for PNG format
    suffix = 0
    filename = f'DP_plots/{title}.png'
    while os.path.exists(filename):
        suffix += 1
        filename = f'DP_plots/{title}_{suffix}.png'

    fig.savefig(filename, format='png', dpi=300)


def plot_actual_vs_predicted_colours(y_valid, y_pred, target_col, model_name):
    mae = mean_absolute_error(y_valid, y_pred)
    rel_error = np.mean(np.abs((y_valid - y_pred) / y_valid)) * 100

    fig, ax = plt.subplots(figsize=(10, 6))
    sns.lineplot(x=y_valid.index, y=y_valid, label="Real Data", ax=ax, color=model_colors_real[model_name])
    sns.lineplot(x=y_valid.index, y=y_pred, label="Model Prediction", linewidth=2, ax=ax, color=model_colors_pred[model_name])
    ax.fill_between(y_valid.index, y_pred, y_valid, alpha=0.2, color=model_colors_pred[model_name])
    ax.set_title(f'{target_col} Actual vs Predicted\nMAE: {mae:.2f}  Relative Error: {rel_error:.2f}%')
    ax.set_xlabel('Index')
    ax.set_ylabel(target_col)
    ax.legend()
    title = f'{target_col} Actual vs Predicted\nMAE: {mae:.2f}  Relative Error: {rel_error:.2f}%'
    save_figure(fig, f"Actual vs Predicted model{model_name}")
    plt.show()

def plot_fold_actual_vs_predicted(y_te, y_pred, target_col, fold):
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(y_te.index, y_te, label="Real Data")
    ax.plot(y_te.index, y_pred, label="Model Prediction")
    ax.set_title(f'{target_col} Fold {fold} Actual vs Predicted')
    ax.set_xlabel('Index')
    ax.set_ylabel(target_col)
    ax.legend()
from sklearn.preprocessing import StandardScaler

def train_lgbm(train_df, target_col, start=24, end=48):
    # Prepare data
    train_df = train_df.copy()
    split_index = int(0.8 * len(train_df))

    # Add lagged columns
    train_df = add_lagged_columns(train_df, target_col, start, end)

    # Split data into training and validation sets
    X_train = train_df.iloc[:split_index].drop(columns=[target_col]).reset_index(drop=True)
    cols_names = X_train.columns.tolist()
    y_train = train_df.iloc[:split_index][target_col]
    X_valid = train_df.iloc[split_index:].drop(columns=[target_col]).reset_index(drop=True)
    y_valid = train_df.iloc[split_index:][target_col]

    # Scale the X factors
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_valid = scaler.transform(X_valid)

    # Split train and validation sets using time-series cross-validation
    tscv = TimeSeriesSplit(n_splits=5)
    best_model = None
    best_score = float('-inf')
    for train_index, valid_index in tscv.split(X_train):
        X_train_cv, X_valid_cv = X_train[train_index], X_train[valid_index]
        y_train_cv, y_valid_cv = y_train.iloc[train_index], y_train.iloc[valid_index]

        # Create LightGBM dataset
        train_data = lgb.Dataset(X_train_cv, label=y_train_cv)
        valid_data = lgb.Dataset(X_valid_cv, label=y_valid_cv)

        # Set parameter grid for grid search
        param_grid = {
            'num_leaves': [5, 50, 100],
            'learning_rate': [0.01, 0.1],
            'feature_fraction': [0.9],
            'min_data_in_leaf': [5, 50],
            'reg_alpha': [0.1, 0.5],
            'reg_lambda': [0.1, 1.0],
            'max_depth': [5, 10, -1],
            'n_estimators': [100, 500, 1000],
            'min_child_samples': [5, 10, 20],
            'subsample': [0.5, 0.7, 0.9]
        }

        # Create LightGBM model
        lgbm = lgb.LGBMRegressor(objective='regression', metric='mse')

        # Perform grid search with cross-validation
        grid_search = GridSearchCV(lgbm, param_grid, cv=5, n_jobs=-1, scoring='mse')
        grid_search.fit(X_train_cv, y_train_cv)

        # Get best parameters from grid search
        best_params = grid_search.best_params_

        # Update parameters with best_params
        params = {
            'objective': 'regression',
            'metric': 'mse',
            **best_params
        }

        # Train model with best parameters on the full training data
        model = lgb.train(params, train_data, valid_sets=[valid_data], num_boost_round=1000, early_stopping_rounds=10, verbose_eval=100)

        # Evaluate model on the validation set
        y_pred = model.predict(X_valid_cv)
        score = mean_absolute_error(y_valid_cv, y_pred)

        # Save best model
        if score > best_score:
            best_score = score
            best_model = model

    # Get feature importance
    feature_importance = pd.DataFrame({'Feature': cols_names, 'Importance': best_model.feature_importance()}).sort_values('Importance', ascending=False)


    # Train final model on the full training data
    final_train_data = lgb.Dataset(X_train, label=y_train)
    final_model = lgb.train(params, final_train_data, num_boost_round=best_model.best_iteration)

    y_pred = final_model.predict(X_valid)
    plot_feature_importance(feature_importance)
    plot_actual_vs_predicted(y_valid, y_pred, target_col)
    plot_actual_vs_predicted_colours(y_valid, y_pred, target_col, "LGBM")
    return  scaler, final_model

def predict_LGBM(cls, real_data_df):
    # Get scaler and model
    scaler = cls.scaler_dict[cls.column_name + "_LGBM"]
    model = cls.model_dict[cls.column_name + "_LGBM"]
    # Scale the data
    X = real_data_df.reset_index(drop=True)
    X = scaler.transform(X)
    # Make predictions
    y_pred = model.predict(X)
    # Create a new DataFrame with the predicted values and the original index
    pred_df = pd.DataFrame({'y_pred': y_pred}, index=real_data_df.index)
    # Return the predicted values
    return pred_df
def plot_residuals(cls, real_data_df):
    # Load the scaler and model from the class dictionary
    real_data_df = real_data_df[cls.COI].copy()

    real_data_df = add_lagged_columns(real_data_df, cls.column_name)
    scaler = cls.scaler_dict['LGBM_scaler']
    model = cls.model_dict['LGBM_model']


    # Scale the data
    X = real_data_df.drop(columns=[cls.column_name]).reset_index(drop=True)
    X = scaler.transform(X)

    # Make predictions
    y_pred = model.predict(X)

    # Create a new DataFrame with the predicted values and the original index
    pred_df = pd.DataFrame({'y_pred': y_pred}, index=real_data_df.index)

    # Calculate the residuals
    residuals = (real_data_df[cls.column_name] - pred_df['y_pred']) / real_data_df[cls.column_name]

    # Plot the histogram of residuals
    plt.hist(residuals, bins=50)
    plt.xlabel('(y_real - y_pred) / y_real')
    plt.ylabel('Frequency')
    plt.title(f'Histogram of Residuals for {cls.column_name}')
    plt.show()
    
    return pred_df

from sklearn.model_selection import RandomizedSearchCV
from scipy.stats import randint as sp_randint

from sklearn.model_selection import RandomizedSearchCV
from scipy.stats import randint as sp_randint

def train_lgbm2(train_df, target_col, start=24, end=48):
    # Prepare data
    train_df = train_df[:"2015"].copy()
    split_index = int(0.8 * len(train_df))

    # Add lagged columns
    if target_col != "W_gen" and target_col != "W_trans":
        train_df = add_lagged_columns(train_df, target_col, start, end)

    # Split data into training and validation sets
    X_train = train_df.iloc[:split_index].drop(columns=[target_col]).reset_index(drop=True)
    cols_names = X_train.columns.tolist()
    y_train = train_df.iloc[:split_index][target_col]
    X_valid = train_df.iloc[split_index:].drop(columns=[target_col]).reset_index(drop=True)
    y_valid = train_df.iloc[split_index:][target_col]

    # Scale the X factors
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_valid = scaler.transform(X_valid)

    # Set parameter distributions for random search
    param_dist = {
        'num_leaves': sp_randint(5, 101),
        'learning_rate': [0.01, 0.1],
        'feature_fraction': [0.9],
        'min_data_in_leaf': sp_randint(5, 51),
        'reg_alpha': [0.1, 0.5],
        'reg_lambda': [0.1, 1.0],
        'max_depth': [5, 10, -1],
        'n_estimators': sp_randint(100, 1001),
        'min_child_samples': sp_randint(5, 21),
        'subsample': [0.5, 0.7, 0.9]
    }

    # Create LightGBM model
    lgbm = lgb.LGBMRegressor(objective='regression', metric='mae')

    # Perform random search with cross-validation
    random_search = RandomizedSearchCV(lgbm, param_distributions=param_dist, n_iter=75, cv=5, n_jobs=-1, scoring='neg_mean_absolute_error')
    random_search.fit(X_train, y_train)

    # Get best parameters from random search
    best_params = random_search.best_params_

    # Update parameters with best_params
    params = {
        'objective': 'regression',
        'metric': 'mae',
        **best_params
    }

    # Train model with best parameters on the full training data
    model = lgb.train(params, lgb.Dataset(X_train, label=y_train), num_boost_round=1000, early_stopping_rounds=10, verbose_eval=100, 
                      valid_sets=[lgb.Dataset(X_valid, label=y_valid)], valid_names=['validation'])

    # Evaluate model on the validation set
    y_pred = model.predict(X_valid)





    # Train final model on the full training data
    final_train_data = lgb.Dataset(X_train, label=y_train)
    final_model = lgb.train(params, final_train_data, num_boost_round=model.best_iteration)

    # Get feature importance
    feature_importance = pd.DataFrame({'Feature': cols_names, 'Importance': final_model.feature_importance()}).sort_values('Importance', ascending=False)

    y_pred = final_model.predict(X_valid)
    plot_feature_importance(feature_importance)
    plot_actual_vs_predicted(y_valid, y_pred, target_col)
    plot_actual_vs_predicted_colours(y_valid, y_pred, target_col, "LGBM")
    top_cols = feature_importance.head(8)['Feature'].tolist()
    return  scaler, final_model, top_cols

#---------------------------------------------------------------------------------------------------------------------------------------------------
def turbine_model_lasso(df):
    y_column = "W_gen"
    x_columns = ["mst_TG_calc", "t_avg", "mst_ex","h_avg"]


    X_train = df[:"2015"][x_columns].reset_index(drop=True)
    y_train = df[:"2015"][y_column]

    X_test = df["2016"][x_columns].reset_index(drop=True)
    y_test = df["2016"][y_column]


    # Scale the numerical columns in the data
    scaler = StandardScaler().fit(X_train)
    X_train_scaled = scaler.transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # Create the LassoCV regression model and fit the data
    model = LassoCV(cv=5, random_state=0, max_iter=1000, alphas=[0.0001, 0.001, 0.01, 0.1, 1])
    model.fit(X_train_scaled, y_train)

    lasso_best = Lasso(alpha=model.alpha_)
    lasso_best.fit(X_train_scaled, y_train)

    print('R squared training set', round(lasso_best.score(X_train_scaled, y_train)*100, 2))
    print('R squared test set', round(lasso_best.score(X_test_scaled, y_test)*100, 2))

    print(mean_squared_error(y_test, lasso_best.predict(X_test_scaled)))

    # Obtain the coefficients and print the regression equation
    coeffs = list(zip(lasso_best.coef_, x_columns))
    intercept = lasso_best.intercept_

    # Create the Lasso regression equation string
    eq_str = "y = {:.2f}".format(intercept)
    for coef, feat in coeffs:
        eq_str += " + {:.2f}*{}".format(coef, feat)

    # Plot the mean square error on each fold for different alpha values
    fig, ax = plt.subplots(figsize=(10, 6))

    ax.semilogx(model.alphas_, model.mse_path_, ":")
    ax.plot(
        model.alphas_ ,
        model.mse_path_.mean(axis=-1),
        "k",
        label="Average across the folds",
        linewidth=2,
    )
    ax.axvline(
        model.alpha_, linestyle="--", color="k", label="alpha: CV estimate"
    )

    ax.legend()
    ax.set_xlabel("alphas")
    ax.set_ylabel("Mean square error")
    ax.set_title("Mean square error on each fold\n{}".format(eq_str))  # Add the equation to the title
    ax.axis("tight")
    save_figure(fig, "Mean square error on each fold")
    plt.show()

    # Generate predictions for the test data
    y_pred = lasso_best.predict(X_test_scaled)

    # Create a scatter plot of the test data vs. predictions
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.scatter(y_test, y_pred, alpha=0.5)
    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=3)
    ax.set_xlabel('Actual W_gen')
    ax.set_ylabel('Predicted W_gen')
    ax.set_title('Test Data vs. Prediction')
    plt.show()

    # Compute MAE and relative error
    mae = mean_absolute_error(y_test, y_pred)
    relative_error = np.mean(np.abs((y_test - y_pred) / y_test))

    # Create a plot of the real data, predicted data, and MAE
    fig, ax = plt.subplots(figsize=(10, 6))

    ax.plot(y_test.index, y_test, label="Real Data")
    ax.plot(y_test.index, y_pred, label="Model Prediction")

    ax.fill_between(
        y_test.index,
        y_pred - mae,
        y_pred + mae,
        alpha=0.2,
        color="orange",
        label="MAE",
    )

    ax.set_xlabel("Index")
    ax.set_ylabel(y_column)
    ax.set_title(f"LassoCV Model (MAE={mae:.2f}, Relative Error={relative_error:.2%})\n{eq_str}")
    ax.legend()

    plt.show()
    plot_actual_vs_predicted_colours(y_test, y_pred, y_column, 'LR')
    # Plot the distribution of residuals
    residuals = y_test - y_pred
    fig, ax = plt.subplots(figsize=(8, 8))
    sns.histplot(residuals, ax=ax, kde=True)
    ax.set_xlabel('Residuals')
    ax.set_ylabel('Count')
    ax.set_title('Distribution of Residuals')
    plt.show()

    return scaler, lasso_best

def predict_turbine_power(cls, df, target_column = "W_gen"):
    # Split the data into real_X and real_y for the given target column
    y_column = "W_gen"
    x_columns = ["mst_TG_calc", "t_avg", "mst_ex"]

    scaler = cls.scaler_dict[cls.column_name + "_LGBM"]
    model = cls.model_dict[cls.column_name + "_LGBM"]
    real_X = df[x_columns]
    real_y = df[y_column]

    # Scale the real_X using the scaler
    real_X_scaled = scaler.transform(real_X)

    # Make predictions using the Lasso model
    predictions = model.predict(real_X_scaled)

    # Calculate the mean absolute error and relative error
    mae = mean_absolute_error(real_y, predictions)
    relative_error = np.mean(np.abs((real_y - predictions) / real_y))

    # Plot the real data, predicted data, and MAE
    fig, ax = plt.subplots(figsize=(10, 6))

    ax.plot(real_y.index, real_y, label="Real Data")
    ax.plot(real_y.index, predictions, label="Model Prediction")

    ax.fill_between(
        real_y.index,
        predictions - mae,
        predictions + mae,
        alpha=0.2,
        color="orange",
        label="MAE",
    )

    ax.set_xlabel("Index")
    ax.set_ylabel(target_column)
    ax.set_title(f"LassoCV Model (MAE={mae:.2f}, Relative Error={relative_error:.2%})")
    ax.legend()

    plt.show()

    return predictions

def add_lagged_columns(df, column_name, start=24, end=48):
    for i in range(start, end+1):
        col_name = f"{column_name}_lag{i}"
        df[col_name] = df[column_name].shift(periods=i)

    df = df.dropna()

    return df


from sklearn.linear_model import LassoCV, Lasso
from sklearn.metrics import mean_squared_error, mean_absolute_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
def model_lasso(df, target_col,main_cols, start=24, end=48):

    df = df.copy()
    # Add lagged columns
    if target_col != "W_gen" and target_col != "W_trans":
        df = add_lagged_columns(df, target_col, start, end)



    y_column = target_col
    #x_columns = ["mst_TG_calc", "t_avg", "mst_ex"]
    x_columns = main_cols
    # Split the data into train and test sets
    test_size = 0.2
    X_train, X_test, y_train, y_test = train_test_split(df[x_columns], df[y_column], test_size=test_size, random_state=42, shuffle=False)

    # Create the LassoCV regression model and fit the data
    model = LassoCV(cv=5, random_state=0, max_iter=1000, alphas=[0.0001, 0.001, 0.01, 0.1, 1])
    model.fit(X_train, y_train)

    lasso_best = Lasso(alpha=model.alpha_)
    lasso_best.fit(X_train, y_train)

    print('R squared training set', round(lasso_best.score(X_train, y_train)*100, 2))
    print('R squared test set', round(lasso_best.score(X_test, y_test)*100, 2))

    print(mean_squared_error(y_test, lasso_best.predict(X_test)))

    # Obtain the coefficients and print the regression equation
    coeffs = list(zip(lasso_best.coef_, x_columns))
    intercept = lasso_best.intercept_

    # Create the Lasso regression equation string
    eq_str = "y = {:.2f}".format(intercept)
    for coef, feat in coeffs:
        eq_str += " + {:.2f}*{}".format(coef, feat)

    # Plot the mean square error on each fold for different alpha values
    fig, ax = plt.subplots(figsize=(10, 6))

    ax.semilogx(model.alphas_, model.mse_path_, ":")
    ax.plot(
        model.alphas_ ,
        model.mse_path_.mean(axis=-1),
        "k",
        label="Average across the folds",
        linewidth=2,
    )
    ax.axvline(
        model.alpha_, linestyle="--", color="k", label="alpha: CV estimate"
    )

    ax.legend()
    ax.set_xlabel("alphas")
    ax.set_ylabel("Mean square error")
    ax.set_title("Mean square error on each fold\n{}".format(eq_str))  # Add the equation to the title
    ax.axis("tight")
    plt.show()

    # Generate predictions for the test data
    y_pred = lasso_best.predict(X_test)

    # Create a scatter plot of the test data vs. predictions
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.scatter(y_test, y_pred, alpha=0.5)
    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=3)
    ax.set_xlabel('Actual ' + str(y_column))
    ax.set_ylabel('Predicted ' + str(y_column))
    ax.set_title('Test Data vs. Prediction')
    plt.show()

    # Compute MAE and relative error
    mae = mean_absolute_error(y_test, y_pred)
    relative_error = np.mean(np.abs((y_test - y_pred) / y_test))

    # Create a plot of the real data, predicted data, and MAE
    fig, ax = plt.subplots(figsize=(10, 6))

    ax.plot(y_test.index, y_test, label="Real Data")
    ax.plot(y_test.index, y_pred, label="Model Prediction")

    ax.fill_between(
        y_test.index,
        y_pred - mae,
        y_pred + mae,
        alpha=0.2,
        color="orange",
        label="MAE",
    )

    ax.set_xlabel("Index")
    ax.set_ylabel(y_column)
    ax.set_title(f"LassoCV Model (MAE={mae:.2f}, Relative Error={relative_error:.2%})\n{eq_str}")
    ax.legend()

    plt.show()
    plot_actual_vs_predicted_colours(y_test, y_pred, y_column, 'LR')
    # Plot the residuals
    fig, ax = plt.subplots(figsize=(10, 6))
    residuals = y_test - y_pred
    sns.histplot(residuals, ax=ax, kde=True)
    ax.set_xlabel("Residuals")
    ax.set_ylabel("Count")
    ax.set_title("Histogram of Residuals")
    plt.show()

    # Plot the residuals vs. predicted values
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.scatter(y_pred, residuals, alpha=0.5)
    ax.axhline(y=0, color='r', linestyle='--', lw=3)
    ax.set_xlabel("Predicted Values")
    ax.set_ylabel("Residuals")
    ax.set_title("Residuals vs. Predicted Values")
    plt.show()
    return lasso_best

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPRegressor
from sklearn.model_selection import cross_val_score
from sklearn.metrics import mean_absolute_error

def ANN_model(df_train,feature_cols,target_col, start=24, end=48):
    df = df_train.copy()
    if target_col != "W_gen" and target_col != "W_trans":
        df = add_lagged_columns(df, target_col, start, end)
    df_train = df.dropna()

    df = df_train[feature_cols + [target_col]]
    split_index = int(0.8 * len(df_train))

    X_train = df[:"2015"].drop(columns=[target_col]).reset_index(drop=True)
    cols_names = X_train.columns.tolist()
    y_train = df[:"2015"][target_col]

    X_test = df["2016"].drop(columns=[target_col]).reset_index(drop=True)
    y_test = df["2016"][target_col]




    # Scale the features using StandardScaler
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train.reset_index(drop=True))
    X_test_scaled = scaler.transform(X_test.reset_index(drop=True))

    # Set parameter distributions for random search
    param_dist = {
        'hidden_layer_sizes': [(4,), (8,), (16,), (32,)],
        'activation': ['tanh', 'relu'],
        'solver': ['adam', 'sgd'],
        'learning_rate': ['constant', 'adaptive'],
        'alpha': [0.0001, 0.001, 0.01]
    }
    # Create the MLP regressor model
    model = MLPRegressor(max_iter=1000, random_state=42)

    # Perform random search with cross-validation
    random_search = RandomizedSearchCV(model, param_distributions=param_dist, n_iter=20, cv=5, n_jobs=-1, scoring='neg_mean_absolute_error')
    random_search.fit(X_train_scaled, y_train)

    # Get best parameters from random search
    best_params = random_search.best_params_

    # Create the final MLP regressor model with best parameters
    final_model = MLPRegressor(max_iter=1000, random_state=42, **best_params)

    # Evaluate the model with cross-validation
    scores = cross_val_score(final_model, X_train_scaled, y_train, cv=5, scoring='neg_mean_absolute_error')

    print("Cross-validation scores:", scores)
    print("Mean MAE:", np.mean(scores))

    # Fit the model to the training data
    final_model.fit(X_train_scaled, y_train)

    # Generate predictions for the train and test data
    y_train_pred = final_model.predict(X_train_scaled)
    y_test_pred = final_model.predict(X_test_scaled)

    # Calculate the MAE for the train and test sets
    train_mae = mean_absolute_error(y_train, y_train_pred)
    test_mae = mean_absolute_error(y_test, y_test_pred)



    print(f"Train MAE: {train_mae:.2f}")
    print(f"Test MAE: {test_mae:.2f}")
    # Create a scatter plot of the test data vs. predictions
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.scatter(y_test, y_test_pred, alpha=0.5)
    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=3)
    ax.set_xlabel('Actual '+target_col)
    ax.set_ylabel('Predicted '+ target_col)
    ax.set_title('Test Data vs. Prediction')
    plt.show()

    # Create a distribution plot of the residuals
    residuals = y_test - y_test_pred
    fig, ax = plt.subplots(figsize=(8, 8))
    sns.histplot(residuals, ax=ax, kde=True)
    ax.set_xlabel('Residuals')
    ax.set_ylabel('Frequency')
    ax.set_title('Distribution of Residuals')
    plt.show()

    mae = mean_absolute_error(y_test, y_test_pred)
    rel_error = np.mean(np.abs((y_test - y_test_pred) / y_test)) * 100

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(y_test.index, y_test, label="Real Data")
    ax.plot(y_test.index, y_test_pred, label="Model Prediction", linewidth=2)
    ax.fill_between(y_test.index, y_test_pred, y_test, alpha=0.2)
    ax.set_title(f'{target_col} Actual vs Predicted\nMAE: {mae:.2f}  Relative Error: {rel_error:.2f}%')
    ax.set_xlabel('Index')
    ax.set_ylabel(target_col)
    ax.legend()
    plot_actual_vs_predicted_colours(y_test, y_test_pred, target_col, 'ANN')
    return scaler,final_model

def quantile_LGBM(train_df, target_col, Q=0.90, start=24, end=48):
    # Prepare data
    train_df = train_df[:"2015"].copy()
    split_index = int(0.8 * len(train_df))

    # Add lagged columns
    if target_col != "W_gen":
        train_df = add_lagged_columns(train_df, target_col, start, end)

    # Split data into training and validation sets
    X_train = train_df.iloc[:split_index].drop(columns=[target_col]).reset_index(drop=True)
    cols_names = X_train.columns.tolist()
    y_train = train_df.iloc[:split_index][target_col]
    X_valid = train_df.iloc[split_index:].drop(columns=[target_col]).reset_index(drop=True)
    y_valid = train_df.iloc[split_index:][target_col]

    # Scale the X factors
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_valid = scaler.transform(X_valid)

    # Set parameter distributions for random search
    param_dist = {
        'num_leaves': sp_randint(5, 101),
        'learning_rate': [0.01, 0.1],
        'feature_fraction': [0.9],
        'min_data_in_leaf': sp_randint(5, 51),
        'reg_alpha': [0.1, 0.5],
        'reg_lambda': [0.1, 1.0],
        'max_depth': [5, 10, -1],
        'n_estimators': sp_randint(100, 1001),
        'min_child_samples': sp_randint(5, 21),
        'subsample': [0.5, 0.7, 0.9]
    }

    # Create LightGBM model
    lgbm = lgb.LGBMRegressor(objective='quantile', alpha=Q)

    # Perform random search with cross-validation
    random_search = RandomizedSearchCV(lgbm, param_distributions=param_dist, n_iter=75, cv=5, n_jobs=-1, scoring='neg_mean_absolute_error')
    random_search.fit(X_train, y_train)

    # Get best parameters from random search
    best_params = random_search.best_params_

    # Update parameters with best_params
    params = {
        'objective': 'quantile',
        'alpha': Q,
        **best_params
    }

    # Train model with best parameters on the full training data
    model = lgb.train(params, lgb.Dataset(X_train, label=y_train), num_boost_round=1000, early_stopping_rounds=10, verbose_eval=100, 
                      valid_sets=[lgb.Dataset(X_valid, label=y_valid)], valid_names=['validation'])

    # Evaluate model on the validation set
    y_pred = model.predict(X_valid)


    # Get feature importance
    feature_importance = pd.DataFrame({'Feature': cols_names, 'Importance': model.feature_importance()}).sort_values('Importance', ascending=False)


    # Train final model on the full training data
    final_train_data = lgb.Dataset(X_train, label=y_train)
    final_model = lgb.train(params, final_train_data, num_boost_round=model.best_iteration)

    y_pred = final_model.predict(X_valid)
    plot_feature_importance(feature_importance)
    plot_actual_vs_predicted(y_valid, y_pred, target_col)
    plot_actual_vs_predicted_colours(y_valid, y_pred, target_col, 'LGBM')
    return  scaler, final_model


def predict_ANN_W(cls, real_data_df):
    # Get scaler and model
    scaler = cls.scaler_dict[cls.column_name + "_ANN"]
    model = cls.model_dict[cls.column_name + "_ANN"]
    # Get feature columns

    # Scale the data
    X = real_data_df.reset_index(drop=True)
    X = scaler.transform(X)
    # Make predictions
    y_pred = model.predict(X)

    return y_pred


def predict_LR_W(cls, real_data_df):
    # Split the data into real_X and real_y for the given target column
    scaler = cls.scaler_dict[cls.column_name + "_LR"]
    model = cls.model_dict[cls.column_name + "_LR"]
    # Get feature columns

    # Scale the data
    X = real_data_df.reset_index(drop=True)
    X = scaler.transform(X)
    # Make predictions
    y_pred = model.predict(X)

    return y_pred


    plt.show()




def ANN_quantile_model2(df_train, feature_cols, target_col, quantile=0.95, start=24, end=48):
    df = df_train.copy()
    if target_col != "W_gen" and target_col != "W_trans":
        df = add_lagged_columns(df, target_col, start, end)
    df_train = df.dropna()

    df = df_train[feature_cols + [target_col]]
    split_index = int(0.8 * len(df_train))

    X_train = df[:"2015"].drop(columns=[target_col]).reset_index(drop=True)
    cols_names = X_train.columns.tolist()
    y_train = df[:"2015"][target_col]

    X_test = df["2016"].drop(columns=[target_col]).reset_index(drop=True)
    y_test = df["2016"][target_col]

    # Scale the features using StandardScaler
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train.reset_index(drop=True))
    X_test_scaled = scaler.transform(X_test.reset_index(drop=True))

    # Define custom quantile loss function
    def quantile_loss(q, y_true, y_pred):
        error = tf.subtract(y_true, y_pred)
        return tf.keras.backend.mean(tf.keras.backend.maximum(q * error, (q - 1) * error), axis=-1)

    # Build the model with quantile loss function
    inputs = tf.keras.Input(shape=(len(cols_names),))
    x = tf.keras.layers.Dense(8, activation='relu')(inputs)
    outputs = tf.keras.layers.Dense(1)(x)

    model = tf.keras.Model(inputs=inputs, outputs=outputs)
    model.compile(loss=lambda y_true, y_pred: quantile_loss(quantile, y_true, y_pred), optimizer='adam')

    # Fit the model to the training data
    model.fit(X_train_scaled, y_train, epochs=50, batch_size=64, verbose=0)

    # Generate predictions for the train and test data
    y_train_pred = model.predict(X_train_scaled).reshape(-1)
    y_test_pred = model.predict(X_test_scaled).reshape(-1)


    # Calculate the MAE for the train and test sets
    train_mae = mean_absolute_error(y_train, y_train_pred)
    test_mae = mean_absolute_error(y_test, y_test_pred)

    # Print the MAE for the
    print(f"Train MAE: {train_mae:.2f}")
    print(f"Test MAE: {test_mae:.2f}")

    return scaler,model
#%%
from keras_tuner import RandomSearch
from tensorflow.keras import regularizers
from tensorflow import keras
from tensorflow.keras import regularizers
import keras_tuner as kt
import tensorflow as tf



def ANN_quantile_model(df_train, feature_cols, target_col, quantile=0.95, start=24, end=48):
    df = df_train.copy()
    if target_col != "W_gen" and target_col != "W_trans":
        df = add_lagged_columns(df, target_col, start, end)
    df_train = df.dropna()

    df = df_train[feature_cols + [target_col]]
    split_index = int(0.8 * len(df_train))

    X_train = df[:"2015"].drop(columns=[target_col]).reset_index(drop=True)
    cols_names = X_train.columns.tolist()
    y_train = df[:"2015"][target_col]

    X_test = df["2016"].drop(columns=[target_col]).reset_index(drop=True)
    y_test = df["2016"][target_col]
    # Scale the features using StandardScaler
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train.reset_index(drop=True))
    X_test_scaled = scaler.transform(X_test.reset_index(drop=True))

    # Define custom quantile loss function
    def quantile_loss(q, y_true, y_pred):
        error = tf.subtract(y_true, y_pred)
        return tf.keras.backend.mean(tf.keras.backend.maximum(q * error, (q - 1) * error), axis=-1)

    # Define model builder function for Keras Tuner
    def build_model(hp):
        inputs = tf.keras.Input(shape=(len(cols_names),))
        x = tf.keras.layers.Dense(
            units=hp.Int('units', min_value=32, max_value=512, step=32),
            activation=hp.Choice('dense_activation', values=['relu', 'tanh', 'sigmoid']),
            kernel_regularizer=regularizers.l2(hp.Choice('l2_regularization', values=[0.01, 0.001, 0.0001]))
        )(inputs)
        outputs = tf.keras.layers.Dense(1)(x)
        model = tf.keras.Model(inputs=inputs, outputs=outputs)
        model.compile(
            loss=lambda y_true, y_pred: quantile_loss(quantile, y_true, y_pred),
            optimizer=tf.keras.optimizers.Adam(hp.Choice('learning_rate', values=[1e-2, 1e-3, 1e-4]))
        )
        return model

    tuner = RandomSearch(
        build_model,
        objective='val_loss',
        max_trials=5,
        executions_per_trial=3,
        directory='random_search',
        project_name='ANN_quantile_model'
    )

    tuner.search(X_train_scaled, y_train, epochs=50, validation_split=0.2, verbose=1)

    # Get the optimal hyperparameters
    best_hps = tuner.get_best_hyperparameters(num_trials=1)[0]

    # Build the model with the optimal hyperparameters
    model = tuner.hypermodel.build(best_hps)
    model.fit(X_train_scaled, y_train, epochs=50, batch_size=64, verbose=0)

    # Generate predictions for the train and test data
    y_train_pred = model.predict(X_train_scaled).reshape(-1)
    y_test_pred = model.predict(X_test_scaled).reshape(-1)

    # Calculate the MAE for the train and test sets
    train_mae = mean_absolute_error(y_train, y_train_pred)
    test_mae = mean_absolute_error(y_test, y_test_pred)

    # Print the MAE for the
    print(f"Train MAE: {train_mae:.2f}")
    print(f"Test MAE: {test_mae:.2f}")

    return scaler, model




def LR_model_quantile(df_train, feature_cols, target_col, quantile=0.5, start=24, end=48):
    df = df_train.copy()
    if target_col != "W_gen" and target_col != "W_trans":
        df = add_lagged_columns(df, target_col, start, end)
    df_train = df.dropna()

    df = df_train[feature_cols + [target_col]]
    split_index = int(0.8 * len(df_train))

    X_train = df[:"2015"].drop(columns=[target_col]).reset_index(drop=True)

    x_columns = X_train.columns.tolist()
    cols_names = X_train.columns.tolist()


    X_test = df["2016"].drop(columns=[target_col]).reset_index(drop=True)

    y_train = df[:"2015"][target_col].reset_index(drop=True)
    y_test =  df["2016"][target_col].reset_index(drop=True)
    # Scale the features using StandardScaler
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train.reset_index(drop=True))
    X_test_scaled = scaler.transform(X_test.reset_index(drop=True))

    # Add a constant term to the scaled data
    X_train_scaled = pd.DataFrame(X_train_scaled, columns=x_columns)
    X_test_scaled = pd.DataFrame(X_test_scaled, columns=x_columns)
    X_train_scaled["const"] = 1
    X_test_scaled["const"] = 1

    # Create the quantile regression model and fit the data
    model = QuantReg(y_train, X_train_scaled).fit(q=quantile)



    return scaler, model