import pandas as pd
import warnings
import importlib
import Models as models
import pandas as pd
import matplotlib.pyplot as plt
import pre_processing as pp
import Classes as CS
import pandas as pd
import numpy as np
from importlib import reload
models = reload(models)
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 fill_missing_Q_req(df, Q_req_file):
    """
    Fill missing Q_req values in a dataframe using the average values for each hour and month.

    Parameters:
    -----------
    df : pandas.DataFrame
        Dataframe containing Q_req values to be filled.
    Q_req_file : str
        File path for the excel file containing Q_req data.

    Returns:
    --------
    df : pandas.DataFrame
        Dataframe with missing Q_req values filled.
    """

    # read Q_req data
    df_qs = pd.read_excel(Q_req_file)
    df_qs = df_qs.set_index('DateTime')
    df_qs['m-d-h'] = df_qs.index.strftime('%m-%d-%H')

    df_orig = df.copy()

    # get unique m-d-h values and create a new dataframe with them
    unique_mdh = df_qs['m-d-h'].unique()
    new_data = {'m-d-h': unique_mdh}
    new_df = pd.DataFrame(data=new_data)

    # calculate average Q_req for each unique m-d-h value
    for mdh in unique_mdh:
        avg_qreq = df_qs.loc[df_qs['m-d-h'] == mdh, 'Qreq'].mean()
        new_df.loc[new_df['m-d-h'] == mdh, 'Q_req'] = avg_qreq

    # ignore warnings
    warnings.filterwarnings("ignore")

    df = df_orig.copy()

    # map Q_req values to m-d-h values in the dataframe
    df['m-d-h'] = df.index.strftime('%m-%d-%H')
    df['DateTime'] = df.index
    price_dict = new_df.set_index('m-d-h')['Q_req'].to_dict()
    df['Q_req'] = df['m-d-h'].map(price_dict)

    # calculate the mean for each hour and month
    mean_by_hour_month = df.groupby([df.index.month, df.index.hour]).transform('mean')
    # replace missing values with mean for same hour and month
    df['Q_req'] = df['Q_req'].fillna(mean_by_hour_month['Q_req'])

    return df



def create_steam_stream_classes(df, classes_names, main_columns, general_factors, zscore_threshold=3):
    """
    Create steam stream classes using the given dataframe and column names.

    Parameters:
    -----------
    df : pandas.DataFrame
        The input dataframe.
    classes_names : list of str
        List of names for the steam stream classes to be created.
    main_columns : list of str
        List of column names for the main columns in the steam stream classes.
    general_factors : list of str
        List of column names for the general factors to be included in the steam stream classes.
    zscore_threshold : float, optional
        The z-score threshold for outlier removal in the dataframe. Default is 3.

    Returns:
    --------
    classes : list of type CS.SteamStream
        List of steam stream classes created.
    df : pandas.DataFrame
        The pre-processed dataframe.
    """
    classes = []
    for class_name, column_name in zip(classes_names, main_columns):
        general_factors.append(column_name)
        new_class = type(class_name, (CS.SteamStream,), {})
        setattr(new_class, 'column_name', column_name)
        setattr(new_class, '__init__', lambda self, x_factors: CS.SteamStream.__init__(self, x_factors=general_factors, column_name=column_name))
        setattr(new_class, 'COI', general_factors.copy())
        setattr(new_class, 'model_dict', {})
        setattr(new_class, 'scaler_dict', {})
        globals()[class_name] = new_class
        print(f"Class SteamStream created: {class_name}")
        classes.append(new_class)

    # clean the dataframe


    df = pp.clean_dataframe(df, "total_steam", zscore_threshold=zscore_threshold)

    q_001 = df['mst_sf_demi'].quantile(0.01)
    to_remove = len(df[df['mst_sf_demi'] <= q_001])
    df = df[df['mst_sf_demi'] > q_001]

    
    q_001 = df['external_cons'].quantile(0.01)
    to_remove = len(df[df['external_cons'] <= q_001])
    df = df[df['external_cons'] > q_001]
    df = pp.clean_dataframe(df, "external_cons", zscore_threshold=4)

    for col in main_columns:
        df = pp.clean_dataframe(df, col, zscore_threshold=zscore_threshold)
        

    df["mst_TG_calc"] = df["total_steam"] - (df["m_st_bo"]+df["de_air"]+df["external_cons"]+df["lin_bp"])

    return classes, df

def preprocess_W_gen_and_W_trans(df, classes, general_factors, general_factors_TDW, main_columns):
    q_001 = df['W_gen'].quantile(0.01)
    df = pp.clean_dataframe(df, "W_gen", zscore_threshold=3)
    to_remove = len(df[df['W_gen'] <= q_001])
    df = df[df['W_gen'] > q_001]

    W_cols = ["mst_TG_calc",'t_avg', 'mst_ex',"h_avg"]
    W_cols.append("W_gen")
    general_factors.append("W_gen")
    column_name = "W_gen"
    W_cols = W_cols
    class_name = "W_tg_gen"
    new_class = type(class_name, (CS.SteamStream,), {})
    setattr(new_class, 'column_name', column_name)
    setattr(new_class, '__init__', lambda self, x_factors: CS.SteamStream.__init__(self, x_factors=W_cols, column_name=column_name))
    setattr(new_class, 'COI', W_cols.copy())
    setattr(new_class, 'model_dict', {})
    setattr(new_class, 'scaler_dict', {})
    globals()[class_name] = new_class
    print(f"Class SteamStream created: {class_name}")
    classes.append(new_class)

    q_001 = df['W_trans'].quantile(0.01)
    df = pp.clean_dataframe(df, "W_trans", zscore_threshold=3)
    to_remove = len(df[df['W_trans'] <= q_001])
    df = df[df['W_trans'] > q_001]

    general_factors.append("W_trans")
    general_factors_TDW.append("W_gen")
    general_factors_TDW.append("W_trans")
    W_cols =  general_factors_TDW
    class_name = "W_trans"
    column_name = "W_trans"
    new_class = type(class_name, (CS.SteamStream,), {})
    setattr(new_class, 'column_name', column_name)
    setattr(new_class, '__init__', lambda self, x_factors: CS.SteamStream.__init__(self, x_factors=W_cols, column_name=column_name))
    setattr(new_class, 'COI', W_cols.copy())
    setattr(new_class, 'model_dict', {})
    setattr(new_class, 'scaler_dict', {})
    globals()[class_name] = new_class
    print(f"Class SteamStream created: {class_name}")
    classes.append(new_class)

    main_columns.append("W_gen")
    main_columns.append("W_trans")
    return df, classes, general_factors, general_factors_TDW, main_columns




def model_creation(df, classes):
    train_df = df.copy()
    cls_steam = classes[0] 
    print(" models.quantile_LGBM(train_df[cls_steam.COI], cls_steam.column_name, 0.95)")
    print("------------------------------------------------------------------------------------------------------------------------------")
    scaler,model  = models.quantile_LGBM(train_df[cls_steam.COI], cls_steam.column_name, 0.95, start=24, end=48)
    cls_steam.model_dict[cls_steam.column_name + '_LGBM_max'] = model
    cls_steam.scaler_dict[cls_steam.column_name + '_LGBM_max'] = scaler

    print(" models.quantile_LGBM(train_df[cls_steam.COI], cls_steam.column_name, 0.05)")
    print("------------------------------------------------------------------------------------------------------------------------------")
    scaler,model  = models.quantile_LGBM(train_df[cls_steam.COI], cls_steam.column_name, 0.05, start=24, end=48)
    cls_steam.model_dict[cls_steam.column_name + '_LGBM_min'] = model
    cls_steam.scaler_dict[cls_steam.column_name + '_LGBM_min'] = scaler
    
    TEST = train_df[cls_steam.COI].iloc[:100].reset_index(drop=True).copy()
    TEST_list = models.add_lagged_columns(TEST, cls_steam.column_name,start=24, end=48).drop(columns=[cls_steam.column_name])
    cols_names = TEST_list.columns.tolist()
    feature_importance = pd.DataFrame({'Feature': cols_names, 'Importance': classes[0].model_dict[classes[0].column_name + "_LGBM_min"].feature_importance()}).sort_values('Importance', ascending=False)
    feature_importance_list =feature_importance.head(8)['Feature'].tolist()

    print(" models.LR_model_quantile(train_df[cls_steam.COI],cols_atm_aviable[:5], cls_steam.column_name,0.05)")
    print("------------------------------------------------------------------------------------------------------------------------------")
    scaler,model  = models.LR_model_quantile(train_df[cls_steam.COI],feature_importance_list[:], cls_steam.column_name,0.05, start=24, end=48)
    cls_steam.model_dict[cls_steam.column_name + '_LR_max'] = model
    cls_steam.scaler_dict[cls_steam.column_name + '_LR_max'] = scaler


    print(" models.ANN_quantile_model(train_df[cls_steam.COI],cols_atm_aviable[:5], cls_steam.column_name, 0.05)")
    print("------------------------------------------------------------------------------------------------------------------------------")
    scaler,model  = models.ANN_quantile_model(train_df[cls_steam.COI],feature_importance_list[:], cls_steam.column_name, 0.05, start=24, end=48)
    cls_steam.model_dict[cls_steam.column_name + '_ANN_max'] = model
    cls_steam.scaler_dict[cls_steam.column_name + '_ANN_max'] = scaler

    for cls in classes[:-1]:
        print("scaler,model,top_cols  = models.train_lgbm2(train_df[cls.COI], cls.column_name)")
        print("------------------------------------------------------------------------------------------------------------------------------")
        print(cls.column_name)
        print(cls.COI)
        # Train the model and scaler
        scaler,model,top_cols  = models.train_lgbm2(train_df[cls.COI], cls.column_name, start=24, end=48)
        # Save the model and scaler to the class dictionary
        cls.model_dict[cls.column_name + '_LGBM'] = model
        cls.scaler_dict[cls.column_name + '_LGBM'] = scaler
        print("Top Cols:")
        print(top_cols)
        print(" scaler,model  = models.ANN_model(train_df,top_cols,cls.column_name)")
        print("------------------------------------------------------------------------------------------------------------------------------")
        scaler,model  = models.ANN_model(train_df,top_cols,cls.column_name, start=24, end=48)
        cls.model_dict[cls.column_name + '_ANN'] = model
        cls.scaler_dict[cls.column_name + '_ANN'] = scaler
        print(" scaler,model  = models.LR_model_quantile(train_df,top_cols,cls.column_name,quantile = 0.5)")
        print("------------------------------------------------------------------------------------------------------------------------------")
        scaler,model  = models.LR_model_quantile(train_df,top_cols,cls.column_name,quantile = 0.5, start=24, end=48)
        cls.model_dict[cls.column_name + '_LR'] = model
        cls.scaler_dict[cls.column_name + '_LR'] = scaler

        #remove columns that are not good from the top cols (aka year and offsets)

    cls = classes[-2] 
    print(cls.column_name)
    print(top_cols)
    scaler,model  = models.ANN_model(train_df,top_cols,cls.column_name, start=24, end=48)
    cls.model_dict[cls.column_name + '_ANN'] = model
    cls.scaler_dict[cls.column_name + '_ANN'] = scaler
    scaler,model  = models.turbine_model_lasso(train_df.dropna())
    cls.model_dict[cls.column_name + '_LR'] = model
    cls.scaler_dict[cls.column_name + '_LR'] = scaler
    scaler,model  = models.LR_model_quantile(train_df,top_cols,cls.column_name,quantile = 0.95, start=24, end=48)
    cls.model_dict[cls.column_name + '_LR_Q95'] = model
    cls.scaler_dict[cls.column_name + '_LR_Q95'] = scaler
    

    cls = classes[-1] 
    scaler,model,top_cols  = models.train_lgbm2(train_df[cls.COI], cls.column_name, start=24, end=48)
    cls.model_dict[cls.column_name + '_LGBM'] = model
    cls.scaler_dict[cls.column_name + '_LGBM'] = scaler
    print(cls.column_name)
    print(top_cols)
    scaler,model  = models.ANN_model(train_df,top_cols,cls.column_name, start=24, end=48)
    cls.model_dict[cls.column_name + '_ANN'] = model
    cls.scaler_dict[cls.column_name + '_ANN'] = scaler

    scaler,model  = models.LR_model_quantile(train_df,top_cols,cls.column_name,0.5, start=24, end=48)
    cls.model_dict[cls.column_name + '_LR'] = model
    cls.scaler_dict[cls.column_name + '_LR'] = scaler

#scaler,model  = models.train_lgbm(train_df[W_col], "W_gen")
def calculate_df(pred_df, df, cls, model_name):
    real_total_steam = pred_df['total_steam'].values
    real_externals = pred_df['external_cons'].values
    real_da = pred_df['de_air'].values
    real_bo = pred_df['m_st_bo'].values
    real_w = pred_df['W_gen'].values
    plan_W = pred_df['W_plan'].values
    T_min = df['T_min'].loc[pred_df.index].values
    Q_real = df['Q_req'].loc[pred_df.index].values
    t_real = df['t_avg'].loc[pred_df.index].values
    h_real = df['h_avg'].loc[pred_df.index].values

    mst_ex_sc = pred_df['mst_sf_demi'].values

    results = []
    bo_reg = np.zeros_like(real_w)
    mst_ex_cal_reg = np.zeros_like(real_w)
    mst_TG = np.zeros_like(real_w)
    steam_T = IAPWS95(T=235+273.15, P=1.22)
    hst_BP =steam_T.h# specific enthalpy of 235 degree steam (MJ/ton)
    for i in range(len(real_w)):
        progress_pct = (i + 1) / len(real_w) * 100
        print(f"\rCalculating MST Ex correction... {progress_pct:.2f}% complete", end='')
        if plan_W[i] < (real_w[i] - 0.5):
            plan_adjustment = real_w[i]
            bo_reg_i = 0
            mst_TG[i] = real_total_steam[i] \
            - real_externals[i] \
            - real_da[i] \
            - real_bo[i]
            print(f"mstg 1 = {mst_TG[i]}")
            upper_bound = mst_TG[i] * 0.25  # set initial upper bound
            lower_bound = 0  # set initial lower bound
            while True:
                bo_reg_i = (upper_bound + lower_bound) / 2  # calculate midpoint
                print(f"bo_reg_i = {bo_reg_i}")

                mst_TG[i] = real_total_steam[i] \
                            - real_externals[i] \
                            - real_da[i] \
                            - real_bo[i] \
                            - bo_reg_i
                print(f"mstg 2 = {mst_TG[i]}")
                mst_ex_cal_reg[i], mst_bp_temp = pp.calculate_mst_ex(Q_real[i]-hst_BP*bo_reg_i/3600, mst_TG[i], mst_ex_sc[i],T_min = T_min[i])
                mst_TG[i] -= mst_bp_temp
                print(f"mstg 3 = {mst_TG[i]}")
                cols_atm_aviable = cls.scaler_dict[cls.column_name + model_name].get_feature_names_out().tolist()
                W_reg_data = [mst_ex_cal_reg[i],mst_TG[i]  ,t_real[i], h_real[i]]
                df_reg = pd.DataFrame([W_reg_data], columns=cols_atm_aviable)


                if model_name == "_LGBM":
                    plan_adjustment = models.predict_ANN_W(cls,df_reg[cols_atm_aviable]).values[0][0]
                elif model_name == "_ANN":
                    print("\n")
                    print(df_reg)
                    plan_adjustment = models.predict_ANN_W(cls,df_reg[cols_atm_aviable])[0]
                    print(plan_adjustment)
                elif model_name == "_LR":
                    plan_adjustment = models.predict_ANN_W(cls,df_reg[cols_atm_aviable]).values[0][0]


                d_w = plan_W[i] - plan_adjustment
                d_w_plan_real = plan_W[i] - real_w[i]
                if abs(d_w) <= 0.5 or upper_bound <= lower_bound:
                    break
                elif d_w > 0:
                    upper_bound = bo_reg_i - 0.5  # update upper bound
                else:
                    lower_bound = bo_reg_i + 0.5  # update lower bound

            bo_reg[i] = bo_reg_i
            results.append([plan_adjustment, bo_reg_i])
        else:
            results.append([plan_W[i], 0])
        


    columns = ['plan_adjustment', "by pass plan regulation"]
    result_df = pd.DataFrame(results, columns=columns)
    return result_df



def create_lagged_data(real_df, main_columns,general_factors, start, end):
    list_known = list(general_factors) + ['mst_TG_calc', 'mst_ex']
    real_data_X = real_df[list_known]
    
    real_data_Y = pd.DataFrame()
    for main_col in main_columns:
        for i in range(start, end+1):
            col_name = f"{main_col}_lag{i}"
            real_data_X[col_name] = real_data_X[main_col].shift(periods=i)
        real_data_Y[main_col] = real_data_X.pop(main_col)
    
    real_data_X = real_data_X.dropna()
    real_data_Y = real_data_Y.loc[real_data_X.index]
    pred_df = pd.DataFrame()
    
    return real_data_X, real_data_Y, pred_df
#%%
def predict_general2(model, scaler, real_data_df):
    # Get the columns from the dataframe that were used in training
    cols_used_in_training = scaler.get_feature_names_out().tolist()
    X = real_data_df[cols_used_in_training].reset_index(drop=True)
    X_scaled = scaler.transform(X)
    # Convert the scaled data back into a DataFrame
    X_scaled = pd.DataFrame(X_scaled, columns=cols_used_in_training)
    X_scaled["const"] = 1
    
    # Make predictions
    y_pred = model.predict(X_scaled)


    return y_pred

def predict_general(model,scaler, real_data_df):
    cols_atm_aviable = scaler.get_feature_names_out().tolist()
    X = real_data_df[cols_atm_aviable].reset_index(drop=True)
    X = scaler.transform(X)
    # Make predictions
    y_pred = model.predict(X)
    return y_pred

def predict_classes(real_data_X, real_data_Y, pred_df, classes,models_used):
    for cls,model_name in zip(classes[:],models_used):
        print(cls.column_name)
        print(cls.COI)
        col_name = cls.column_name + model_name

        scaler = cls.scaler_dict[col_name]
        model = cls.model_dict[col_name]
        
        cols_atm_aviable = cls.scaler_dict[cls.column_name + model_name].get_feature_names_out().tolist()
        real_df = real_data_X[cols_atm_aviable].copy()

        if '_ANN' in model_name:
            y_pred = predict_general(model,scaler, real_df)
            pred_df[cls.column_name+"_pred"]  = y_pred
        if '_LR' in model_name:
            if col_name == "W_gen_LR":
                y_pred = predict_general(model,scaler, real_df)
                pred_df[cls.column_name+"_pred"]  = y_pred.tolist()
            else:
                try:
                    y_pred = predict_general2(model, scaler, real_df)
                except Exception as e:
                    print("An error occurred:", e)
                    y_pred = predict_general(model, scaler, real_df)
                pred_df[cls.column_name + "_pred"] = y_pred.tolist()
        if '_LGBM' in model_name:
            y_pred = predict_general(model,scaler, real_df)
            pred_df[cls.column_name+"_pred"] = y_pred.tolist()


        pred_df = pred_df.merge(real_data_Y[cls.column_name], left_index=True, right_index=True)
        real_data_X[cls.column_name] = pred_df[cls.column_name+"_pred"]
    return real_data_X, real_data_Y, pred_df

def calculate_mst_TG(df):
    return (df['total_steam']
            - df['external_cons']
            - df['de_air']
            - df['m_st_bo']
            - df['mst_bp'])


def predict_general(model,scaler, real_data_df):
    cols_atm_aviable = scaler.get_feature_names_out().tolist()
    X = real_data_df[cols_atm_aviable].reset_index(drop=True)
    X = scaler.transform(X)
    # Make predictions
    y_pred = model.predict(X)
    return y_pred


def process_last_class(cls, df, real_data_X, real_data_Y, pred_df,model_name):
    w_gen_cols = ["total_steam","external_cons","de_air","m_st_bo",'t_avg',"h_avg"]
    real_df_W_max = df[w_gen_cols].loc[real_data_X.index].copy().dropna()
    Q_req = df['Q_req'].loc[real_df_W_max.index].values
    real_df_W_max["mst_bp"] = 0

    real_df_W_max['mst_TG_calc'] = calculate_mst_TG(real_df_W_max)
    #hell
    mst_TG =  real_df_W_max['mst_TG_calc'].values
    mst_ex_sc = df['mst_sf_demi'].loc[real_df_W_max.index].values
    mst_ex_cal = np.zeros_like(Q_req)
    mst_bp = np.zeros_like(Q_req)
    T_min = df['T_min'].loc[pred_df.index].values
    for i in range(len(Q_req)):
        progress_pct = (i+1)/len(Q_req)*100
        print(f"\rCalculating MST Ex original... {progress_pct:.2f}% complete", end="")
        mst_ex_cal[i],mst_bp[i] = pp.calculate_mst_ex(Q_req[i], mst_TG[i],  mst_ex_sc[i],T_min =T_min[i])
        print("\n")


    real_df_W_max['mst_ex']  = mst_ex_cal
    real_df_W_max['mst_bp'] = mst_bp
    real_df_W_max['mst_TG_calc'] = calculate_mst_TG(real_df_W_max)
    
    cols_atm_aviable = cls.scaler_dict[cls.column_name + model_name].get_feature_names_out().tolist()
    

    model_dict = classes[-2].model_dict
    scaler_dict = classes[-2].scaler_dict

    A_model = model_dict['W_gen_LGBM']
    A_scaler = scaler_dict['W_gen_LGBM']
    A_pred = predict_general(A_model,A_scaler, real_df_W_max)
    B_model = model_dict['W_gen_ANN']
    B_scaler = scaler_dict['W_gen_ANN']
    B_pred = predict_general(B_model,B_scaler, real_df_W_max)
    C_model = model_dict['W_gen_LR']
    C_scaler = scaler_dict['W_gen_LR']
    C_pred = predict_general(C_model,C_scaler, real_df_W_max)

    
    if model_name == "_LGBM":
                    pred_df[cls.column_name] = models.predict_LGBM(cls,real_df_W_max[cols_atm_aviable].loc[real_data_X.index])
    elif model_name == "_ANN":
                    pred_df[cls.column_name] = models.predict_ANN_W(cls,real_df_W_max[cols_atm_aviable].loc[real_data_X.index])
    elif model_name == "_LR":
                    pred_df[cls.column_name] = models.predict_LR_W(cls,real_df_W_max[cols_atm_aviable].loc[real_data_X.index])                



    pred_df['mst_TG_calc'] = pred_df['total_steam_pred'] \
                                     - pred_df['external_cons_pred'] \
                                     - pred_df['de_air_pred'] \
                                     - pred_df['m_st_bo_pred']
    pred_df['mst_TG'] = real_df_W_max['mst_TG_calc']

    #hell
    Q_req = df['Q_req'].loc[real_data_X.index].values
    mst_TG =  pred_df['mst_TG_calc'].values
    mst_ex_sc = pred_df['mst_sf_demi_pred'].values
    mst_ex_cal = np.zeros_like(Q_req)
    mst_bp = np.zeros_like(Q_req)
    for i in range(len(Q_req)):
        progress_pct = (i+1)/len(Q_req)*100
        print(f"\rCalculating MST Ex prediction... {progress_pct:.2f}% complete", end="")
        mst_ex_cal[i],mst_bp[i] = pp.calculate_mst_ex(Q_req[i], mst_TG[i],  mst_ex_sc[i],T_min =T_min[i])
        print("\n")

    pred_df['t_avg'] = df["t_avg"].loc[real_data_X.index]
    pred_df['mst_ex'] = mst_ex_cal
    pred_df["mst_bp_pred"]= mst_bp
    pred_df['h_avg'] = df["h_avg"].loc[real_data_X.index]
    
    pred_df['mst_TG_calc'] = pred_df['total_steam_pred'] \
                                    - pred_df['external_cons_pred'] \
                                    - pred_df['de_air_pred'] \
                                    - pred_df['m_st_bo_pred'] - pred_df["mst_bp_pred"]
    if model_name == "_LGBM":
                        pred_df[cls.column_name+"_pred"] = models.predict_LGBM(cls,pred_df[cols_atm_aviable])
    elif model_name == "_ANN":
                        pred_df[cls.column_name+"_pred"] = models.predict_ANN_W(cls,pred_df[cols_atm_aviable])
    elif model_name == "_LR":
                        pred_df[cls.column_name+"_pred"] = models.predict_LR_W(cls,pred_df[cols_atm_aviable])

    pred_df = pred_df.rename(columns={'mst_TG_calc': 'mst_TG_pred'})
    
    pred_df['mst_TG_calc'] = pred_df['total_steam_pred_max'] \
                                     - pred_df['external_cons_pred'] \
                                     - pred_df['de_air_pred'] \
                                     - pred_df['m_st_bo_pred'] - pred_df["mst_bp_pred"]
    

    if model_name == "_LGBM":
                            pred_df[cls.column_name+"_pred_max"] = models.predict_LGBM(cls,pred_df[cols_atm_aviable])
    elif model_name == "_ANN":
                            pred_df[cls.column_name+"_pred_max"] = models.predict_ANN_W(cls,pred_df[cols_atm_aviable])
    elif model_name == "_LR":
                            pred_df[cls.column_name+"_pred_max"] = models.predict_LR_W(cls,pred_df[cols_atm_aviable])


    pred_df = pred_df.rename(columns={'mst_TG_calc': 'mst_TG_calc_max'})

    pred_df['mst_TG_calc'] = pred_df['total_steam_pred_min'] \
                                     - pred_df['external_cons_pred'] \
                                     - pred_df['de_air_pred'] \
                                     - pred_df['m_st_bo_pred'] - pred_df["mst_bp_pred"]



    if model_name == "_LGBM":
        pred_df[cls.column_name+"_pred_min"] = models.predict_LGBM(cls,pred_df[cols_atm_aviable])
    elif model_name == "_ANN":
        pred_df[cls.column_name+"_pred_min"] = models.predict_ANN_W(cls,pred_df[cols_atm_aviable])
    elif model_name == "_LR":
        pred_df[cls.column_name+"_pred_min"] = models.predict_LR_W(cls,pred_df[cols_atm_aviable])

    pred_df = pred_df.rename(columns={'mst_TG_calc': 'mst_TG_calc_min'})

    
    pred_df['mst_ex_pred'] = pred_df['mst_ex']  
    pred_df["mst_ex"] = real_df_W_max['mst_ex'] 
    pred_df["mst_bp"] = real_df_W_max['mst_bp'] 

    pred_df["W_plan"] = 0
    # Define a lambda function to apply the logic to each row
    def calculate_total_steam_plan(row,cls):
        if abs(row[cls.column_name+"_pred"] - row[cls.column_name+"_pred_min"]) > 0.5:
            return (row[cls.column_name+"_pred"] + row[cls.column_name+"_pred_min"]) / 2
        else:
            return row[cls.column_name+"_pred"]

    pred_df['W_plan'] = pred_df.apply(lambda row: calculate_total_steam_plan(row,cls), axis=1)
    return pred_df




def plot_predictions(df, col_names, num_days):
    # Get the index of the last num_days days
    end_date = df.index[-1]
    start_date = end_date - pd.Timedelta(days=num_days)

    # Create a line plot of the last num_days days of the three columns
    plt.plot(df.loc[start_date:end_date, col_names[0]], label='Minimum Prediction')
    plt.plot(df.loc[start_date:end_date, col_names[1]], label='Maximum Prediction')
    plt.plot(df.loc[start_date:end_date, col_names[2]], label='Average Prediction')
    if len(col_names) > 3:
        reality = df.loc[start_date:end_date, col_names[3]]
        plt.plot(reality, label='Reality')
        # Add a dot for each time Average Prediction is higher than Reality + 0.5
        avg_prediction = df.loc[start_date:end_date, col_names[2]]
        diff = avg_prediction - reality - 0.5
        for i, val in enumerate(diff):
            if val > 0:
                plt.plot(reality.index[i], reality[i], 'ro')

    # Add a title and labels for the x and y axes
    plt.title('Total Steam Predictions (Last ' + str(num_days) + ' Days)')
    plt.xlabel('Date')
    plt.ylabel('Total Steam Prediction')

    # Add a legend
    plt.legend()

    # Show the plot
    plt.show()
def plot_predictions_W(df, col_names, num_days):
    # Get the index of the last num_days days
    end_date = df.index[-1]
    start_date = end_date - pd.Timedelta(days=num_days)

    # Create a line plot of the last num_days days of the three columns
    plt.plot(df.loc[start_date:end_date, col_names[0]], label='Minimum Prediction')
    plt.plot(df.loc[start_date:end_date, col_names[1]], label='Real tomorow adjusted for plan')
    plt.plot(df.loc[start_date:end_date, col_names[2]], label='Average Prediction')
    if len(col_names) > 3:
        reality = df.loc[start_date:end_date, col_names[3]]
        plt.plot(reality, label='Real W Prediction')
        # Add a dot for each time Average Prediction is higher than Reality + 0.5
        avg_prediction = df.loc[start_date:end_date, col_names[2]]
        diff = avg_prediction - reality - 0.5
        num_overestimates = sum(1 for val in diff if val > 0)
        pct_overestimates = round(num_overestimates / len(diff) * 100, 2)
        plt.title('Total W Predictions (Last ' + str(num_days) + ' Days), ' + str(pct_overestimates) + '% overestimates')
        for i, val in enumerate(diff):
            if val > 0:
                plt.plot(reality.index[i], reality[i], 'ro')
    else:
        plt.title('Total W Predictions (Last ' + str(num_days) + ' Days)')

    # Add labels for the x and y axes
    plt.xlabel('Date')
    plt.ylabel('Total W Prediction')

    # Add a legend
    plt.legend()

    # Show the plot
    plt.show()


def add_pred_to_df(cls_steam, real_data_X, pred_df):
    # Get the list of columns available in the models
    cols_available = cls_steam.scaler_dict[cls_steam.column_name + "_LGBM_min"].get_feature_names_out().tolist()

    # Select the columns from the real data that are available in the models
    real_data_sel = real_data_X[cols_available].copy()

    # Scale the data using the minimum and maximum scalers
    scaler_min = cls_steam.scaler_dict[cls_steam.column_name + "_LGBM_min"]
    X_min = scaler_min.transform(real_data_sel)

    scaler_max = cls_steam.scaler_dict[cls_steam.column_name + "_LGBM_max"]
    X_max = scaler_max.transform(real_data_sel)

    # Make predictions using the minimum and maximum models
    model_min = cls_steam.model_dict[cls_steam.column_name + "_LGBM_min"]
    y_pred_min = model_min.predict(X_min)

    model_max = cls_steam.model_dict[cls_steam.column_name + "_LGBM_max"]
    y_pred_max = model_max.predict(X_max)

    # Add the predicted minimum and maximum values to the prediction DataFrame
    pred_df[cls_steam.column_name + "_pred_min"] = pd.DataFrame({'y_pred': y_pred_min}, index=real_data_sel.index)
    pred_df[cls_steam.column_name + "_pred_max"] = pd.DataFrame({'y_pred': y_pred_max}, index=real_data_sel.index)

    # Add the predicted minimum and maximum values to the real data DataFrame
    real_data_X[cls_steam.column_name + "_min"] = pred_df[cls_steam.column_name + "_pred_min"]
    real_data_X[cls_steam.column_name + "_max"] = pred_df[cls_steam.column_name + "_pred_max"]

    return real_data_X,pred_df




def plot_heatmap(pred_df):
    # Plot W_gen_pred vs. W_gen as a heatmap
    fig, ax = plt.subplots(figsize=(8, 6))
    hb = ax.hexbin(pred_df['W_trans'], pred_df['W_trans_pred'], gridsize=50, cmap='binary')
    ax.set_xlabel('W_trans')
    ax.set_ylabel('W_trans_pred')
    ax.set_title('W_trans_pred vs. W_trans')
    cb = fig.colorbar(hb, ax=ax)
    cb.set_label('count')

    # Add boundary lines at W_gen + 0.5 and W_gen - 0.5
    ax.plot([min(pred_df['W_trans']), max(pred_df['W_trans'])], [min(pred_df['W_trans'])+0.5, max(pred_df['W_trans'])+0.5], 'r--')
    ax.plot([min(pred_df['W_trans']), max(pred_df['W_trans'])], [min(pred_df['W_trans'])-0.5, max(pred_df['W_trans'])-0.5], 'r--')
    # Calculate percentage of points within boundaries
    count_within = (pred_df['W_trans_pred'] - pred_df['W_trans']) > 0.5
    count_within = len(count_within[count_within == True])
    percent_within = (1 - (count_within / len(pred_df) ))* 100
    ax.set_title(f'W_trans_pred vs. W_trans\n{percent_within:.1f}% within boundaries')

    plt.show()

def compute_predictions(df, classes, main_columns,general_factors, start, end):
    cls_steam = classes[0]

    real_df = df.copy()
    real_data_X, real_data_Y, pred_df = create_lagged_data(real_df, main_columns,general_factors, start=24, end=48)
    real_data_X, pred_df = add_pred_to_df(cls_steam, real_data_X, pred_df)
    real_data_X, real_data_Y, pred_df = predict_classes(real_data_X, real_data_Y, pred_df, classes[:-2])

    # Iterate over each column pair and compute the error and residual
    cls = classes[-2]

    model_name = "_ANN"

    pred_df_ANN = process_last_class(cls, df, real_data_X, real_data_Y, pred_df, model_name)
    cls = classes[-1]
    pred_df_ANN = predict_and_plot(cls, model_name, real_data_X, pred_df_ANN,df)

    return pred_df_ANN



def predict_and_plot(cls, model_name, real_data_X, pred_df,df):
    cols_atm_aviable = cls.scaler_dict[cls.column_name + model_name].get_feature_names_out().tolist()
    real_data_X["W_gen"] = pred_df["W_plan"]
    pred_df["W_trans_pred"] = models.predict_ANN_W(cls, real_data_X[cols_atm_aviable])

    df_trans_plc = pd.DataFrame()
    df_trans_plc["W_gen"] = pred_df["W_gen"]
    df_trans_plc[cols_atm_aviable[1:]] = real_data_X[cols_atm_aviable[1:]].loc[real_data_X.index]
    pred_df["W_trans"] = models.predict_ANN_W(cls, df_trans_plc[cols_atm_aviable].loc[real_data_X.index])
    pred_df["W_trans_classic"] = df["W_trans"].loc[real_data_X.index]
    
    return pred_df

#%%------------------------model applicatiion------------------------------------------------------

import joblib
import lightgbm as lgb

def save_lgbm_model(model_dict):
    for cls in classes:
        for key in cls.model_dict.keys():
            try:
                print(cls.column_name)
                filename = key + '.txt'
                model = cls.model_dict[key]
                if hasattr(model, 'save_model'):
                    model.save_model(filename)
                elif hasattr(model, 'save'):
                    joblib.dump(model, filename)
                else:
                    print(f"No save method available for {key}")
            except Exception as e:
                print(f"Error saving model for {key}: {str(e)}")
                
def load_lgbm_model(model_dict):
    for cls in classes:
        for key in cls.model_dict.keys():
            try:
                print(cls.column_name)
                filename = key + '.txt'
                model = lgb.Booster(model_file=filename)
                cls.model_dict[key] = model
            except:
                pass



def error_assesemnt(pred_df_ANN, df, classes):
    pred_W_trnas = pred_df_ANN[["W_trans_pred", "W_trans", "W_trans_classic"]]
    pred_W_trnas["steam_error"] = (pred_df_ANN["total_steam"] - pred_df_ANN["total_steam_pred"]) / pred_df_ANN["total_steam"]
    pred_W_trnas["W_gen_error"] = (pred_df_ANN["W_gen"] - pred_df_ANN["W_plan"]) / pred_df_ANN["W_gen"]
    pred_W_trnas["W_trans_error"] = (pred_df_ANN["W_trans"] - pred_df_ANN["W_trans_pred"]) / pred_df_ANN["W_trans"]
    pred_W_trnas["gain"] = pred_W_trnas["W_trans_pred"] - pred_W_trnas["W_trans_classic"]
    pred_W_trnas["W_gen_classic"] = df["W_gen"].loc[pred_df_ANN.index]

    col_names = ['gain', 'W_trans_classic', 'W_trans_pred', "W_trans"]
    MfD.plot_predictions_W(pred_W_trnas, col_names, num_days=200)

    import matplotlib.pyplot as plt
    from scipy.stats import pearsonr

    # get the last 5 days of data from the dataframe
    last_five_days = pred_W_trnas.tail(5*24)

    # create a figure and axis object
    fig, ax = plt.subplots()

    # plot each column as a line for the last 5 days
    ax.plot(last_five_days.index, last_five_days["steam_error"], label="steam_error")
    ax.plot(last_five_days.index, last_five_days["W_gen_error"], label="W_gen_error")
    ax.plot(last_five_days.index, last_five_days["W_trans_error"], label="W_trans_error")

    # calculate the Pearson correlation coefficient between each pair of error columns
    corr_AB, _ = pearsonr(pred_W_trnas["steam_error"], pred_W_trnas["W_gen_error"])
    corr_AC, _ = pearsonr(pred_W_trnas["steam_error"], pred_W_trnas["W_trans_error"])
    corr_BC, _ = pearsonr(pred_W_trnas["W_gen_error"], pred_W_trnas["W_trans_error"])

    # create a text string for the Pearson correlation coefficients
    corr_text = f"Pearson corr AB: {corr_AB:.2f}\nPearson corr AC: {corr_AC:.2f}\nPearson corr BC: {corr_BC:.2f}"

    # add the Pearson correlation coefficients to the plot in a single box
    ax.text(0.05, 0.05, corr_text, transform=ax.transAxes, fontsize=12,
            verticalalignment='bottom', bbox=dict(boxstyle='round', facecolor='white', alpha=0.5))

    # move the legend to below the X axis in horizontal alignment
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.2), ncol=3)

    # set axis labels
    ax.set_xlabel("Date")
    ax.set_ylabel("Error")

    # show the plot
    plt.show()

    return pred_W_trnas

import os
import joblib

def save_load_models(classes, mode):
    # Define the folder where the models are saved
    models_folder = "saved_models"

    # Loop through each class
    for class_obj in classes:
        if mode == "save":
            # Get the model and scaler dictionaries for this class
            model_dict = class_obj.model_dict
            scaler_dict = class_obj.scaler_dict
            column_name = class_obj.column_name
            
            # Create a folder for this class's models if it doesn't exist
            class_folder = os.path.join(models_folder, column_name)
            if not os.path.exists(class_folder):
                os.makedirs(class_folder)
            
            # Initialize a dictionary to map generated model names to original names
            name_map = {}
            
            # Loop through each model and scaler in the dictionaries
            for model_name, model_obj in model_dict.items():
                # Save the model to a file in the class's folder
                model_filename = f"{model_name}.joblib"
                model_path = os.path.join(class_folder, model_filename)
                joblib.dump(model_obj, model_path)
                
                # Check if a scaler exists for this model
                if model_name in scaler_dict:
                    scaler_obj = scaler_dict[model_name]
                    scaler_filename = f"{model_name}_scaler.joblib"
                    scaler_path = os.path.join(class_folder, scaler_filename)
                    joblib.dump(scaler_obj, scaler_path)
                else:
                    scaler_filename = None
                
                # Add an entry to the name_map dictionary
                name_map[model_filename] = (model_name, scaler_filename)
            
            # Save the name_map dictionary to a file in the class's folder
            name_map_path = os.path.join(class_folder, "name_map.txt")
            with open(name_map_path, "w") as f:
                for filename, (model_name, scaler_filename) in name_map.items():
                    f.write(f"{filename}={model_name},{scaler_filename}\n")
        
        elif mode == "load":
            # Get the model and scaler dictionaries for this class
            model_dict = {}
            scaler_dict = {}
            column_name = class_obj.column_name
            
            # Load the name_map for this class
            class_folder = os.path.join(models_folder, column_name)
            name_map_path = os.path.join(class_folder, "name_map.txt")
            with open(name_map_path, "r") as f:
                for line in f:
                    # Split the line into the filename and model name
                    filename, model_name = line.strip().split("=")
                    
                    # Check if the file is a model or scaler and load it accordingly
                    if "_scaler" in filename:
                        scaler_filename = filename
                        scaler_path = os.path.join(class_folder, scaler_filename)
                        scaler_obj = joblib.load(scaler_path)
                        scaler_dict[model_name] = scaler_obj
                    else:
                        model_path = os.path.join(class_folder, filename)
                        model_obj = joblib.load(model_path)
                        model_dict[model_name] = model_obj
                
            # Add the model and scaler dictionaries to the class's dictionaries
            class_obj.model_dict.update(model_dict)
            class_obj.scaler_dict.update(scaler_dict)
            
        else:
            raise ValueError("Invalid mode. Mode must be 'save' or 'load'.")
        
