import pandas as pd
from iapws import IAPWS95
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 general_file_preprocessing(df):
    for col in ['Q_PARA150_OKAM', 'Q_VYST_PARA_OKAM',  '10QHS01CF001/AIN', 'FI0224/AIN', 'FI0216/AIN',"Q_VODA_OKAM","R2_6_11_PP2","FI0226"]:
        df[col] = pd.to_numeric(df[col], errors='coerce')
        
    name_mapping = [
        ["začátek hod.","DateTime"],
        ['F1015_COMM/F_1015', 'Boiler_1'],
        ['F2015_COMM/F_2015', 'Boiler_2'],
        ['F3015_COMM/F_3015', 'Boiler_3'],
        ['F4015_COMM/F_4015', 'Boiler_4'],
        ['20LBA53CF001', 'mst_TG'],
        ['DIF', 'mst_TG_dif'],
        ['20MAA10CP006', 'p_ex'],
        ['20LBS10CF001', 'mst_ex'],
        ['20LBS10CT001', 'T_ex'],
        ['Q_PARA150_OKAM', 'mst_external_1'],
        ['Q_VYST_PARA_OKAM', 'mst_external_2'],
        ['FI0224/AIN', 'mst_da'],
        ['FI0216/AIN', 'mst_da_tg'],
        ['FI0226', 'mst_sf_demi'],
        ['20MKA10CE001', 'W_gen'],
        ['R2_6_11_PP2', 'W_trans'],
        ['20LCA10CF001', 'm_cond'],
        ['20LCA17CT001', 'T_condz_pumps'],
        ['20MAG10CP001', 'P_turb'],
        ['Q_VODA_OKAM', 'Q_DHS'],
        ['20MAG10CT010', 'T_amb'],
        ['20NDA14CT002', 'T_HE1'],
        ['20NDB10CT001', 'T_in_ZEVO'],
        ['20NDA14CF001', 'm_w_in_DHS'],
        ['20NDB15CP001', 'p_from_DHS'],
        ['20LCA32CT001', 'T_condz'],
        ['20NDA12CT001', 'T_HE2'],
        ['20NDA14CP002', 'p_HE12'],
        ['20LBF54CT001', 'T_BP'],
        ['20LBF54CP001', 'P_BP'],
        ["10QHS01CF001/AIN","m_st_bo"],
        ["20LBA54AA201/AUT","bp_vent_proc"]
    ]

    old_col_names = [old for old, new in name_mapping]

    # Select only the columns with the old column names using the loc method
    df = df.loc[:, old_col_names]

    # Rename the columns using the rename() method and the name_dict dictionary
    name_dict = {old: new for old, new in name_mapping}
    df = df.rename(columns=name_dict)


    cols = ["Boiler_1", "Boiler_2","Boiler_3","Boiler_4"]
    df['total_steam'] = df[cols].sum(axis=1)
    df['boilers_active'] = df[cols].apply(lambda x: (x > 1).sum(), axis=1)
    df['maintenance_period'] = df['boilers_active'].apply(lambda x: 1 if x == 0 else 0)
    df['year'] = df['DateTime'].dt.year
    df['month'] = df['DateTime'].dt.month
    df['day'] = df['DateTime'].dt.day
    df['hour'] = df['DateTime'].dt.hour
    df['DD-MM-YYYY'] = df['day'].astype(str) + '-' + df['month'].astype(str) + '-' + df['year'].astype(str)
    df_weather = pd.read_excel('weather_Day.xlsx')
    df = pd.merge(df, df_weather, on='DD-MM-YYYY', how='left')

    seasons = ['Winter', 'Spring', 'Summer', 'Fall']
    month_days = [(12, 21), (3, 20), (6, 21), (9, 22)]
    month_day_map = dict(zip(range(1, 13), month_days))

    # create a boolean mask indicating whether the day is greater than or equal to the season start date
    # calculate the season using vectorized operations
    df = df.set_index('DateTime',inplace=False)
    mask = df.index.day >= pd.Series(df.index.month).map(month_day_map).apply(pd.Series)[1]
    season_indices = (df.index.month + mask.astype(int)) // 3 % 4
    season_map = dict(zip(range(4), seasons))
    df['season'] = season_indices.map(season_map)
    def calculate_T_min(row):
        if row['season'] == 'Winter':
            return 125
        elif row['season'] == 'Summer':
            return 115
        elif row['season'] == 'Fall':
            return 120
        elif row['season'] == 'Spring':
            return 120
        else:
            return None
    df['T_min'] = df.apply(calculate_T_min, axis=1)
    df = pd.concat([df, pd.get_dummies(df['season'])], axis=1)
    df = df.drop(['rok', 'Měsíc', 'Den',"season","DD-MM-YYYY"], axis=1)

    df[seasons] = df[seasons].astype('float')

    df["external_cons"] =df["mst_external_1"] + df["mst_external_2"]
    df["de_air"] = df["mst_da"] + df["mst_da_tg"] 


    external_consumton = ["de_air","external_cons","m_st_bo","mst_sf_demi"]
    df["sum_externals"]= df[external_consumton].sum(axis=1)

    df["lin_bp"] = (df["total_steam"]-df["sum_externals"]) * df["bp_vent_proc"]/100
    return df

import numpy as np

def clean_data(df,known_cols,column_name):
    df_lgbm = df[known_cols].copy()
    zscore_threshold=2
    quantile_threshold=0.01
    zscore = np.abs((df_lgbm[column_name] - df_lgbm[column_name].mean()) / df_lgbm[column_name].std())
    # Remove rows with z-score greater than threshold
    df_lgbm = df_lgbm[zscore < zscore_threshold]
    # Remove rows with value less than quantile threshold
    quantile_value = df_lgbm[column_name].quantile(quantile_threshold)
    df_lgbm = df_lgbm[df_lgbm[column_name] > quantile_value]
    quantile_value = df_lgbm[column_name].quantile(1-quantile_threshold)
    df_lgbm = df_lgbm[df_lgbm[column_name] < quantile_value]

    return df_lgbm

import matplotlib.pyplot as plt
import numpy as np
import Models as models
import seaborn as sns
def plot_column_distribution(df, column_name, bins=10, zscore_threshold=3, lower_quantile=0.01, upper_quantile=0.99):
    # Set seaborn style
    sns.set(style="whitegrid")
    dusty_teal = basic_colours["blue"]
    soft_coral= basic_colours['orange']
    # Create a histogram of the column values
    fig, ax = plt.subplots()
    hist, bin_edges, _ = ax.hist(df[column_name], bins=bins, color=dusty_teal, edgecolor="black")

    # Add overflow bins at lower and upper quantiles
    lower_val = df[column_name].quantile(lower_quantile)
    upper_val = df[column_name].quantile(upper_quantile)
    if df[column_name].min() < lower_val:
        bin_edges = np.insert(bin_edges, 0, df[column_name].min())
        hist = np.insert(hist, 0, np.sum(df[column_name] < lower_val))
    if df[column_name].max() > upper_val:
        bin_edges = np.append(bin_edges, df[column_name].max())
        hist = np.append(hist, np.sum(df[column_name] > upper_val))

    # Add horizontal lines at z-score threshold
    mean = df[column_name].mean()
    std = df[column_name].std()
    ax.axvline(mean - zscore_threshold * std, color=soft_coral, linestyle='--', label=f'Z-score = {zscore_threshold}')
    ax.axvline(mean + zscore_threshold * std, color=soft_coral, linestyle='--')

    # Set axis labels and title
    ax.set_xlabel(column_name)
    ax.set_ylabel('Frequency')
    ax.set_title(f'Distribution of {column_name}')
    ax.legend()
    models.save_figure(fig, f'Distribution of {column_name}')
    plt.show()



def clean_dataframe(df, column_name, zscore_threshold=3):
    # Create a copy of the DataFrame for training
    train_data = df.copy()

    # Remove outliers for the column
    zscore = np.abs((train_data[column_name] - train_data[column_name].mean()) / train_data[column_name].std())
    # Remove rows with z-score greater than threshold
    removed_rows = train_data[zscore >= zscore_threshold]
    train_data = train_data[zscore < zscore_threshold]

    # Get the number of removed points and the min/max of the removed values
    num_removed = len(removed_rows)
    removed_min = removed_rows[column_name].min()
    removed_max = removed_rows[column_name].max()

    # Create a plot of the original data with the removed data highlighted
    fig, ax = plt.subplots(figsize=(8, 6))
    fig.suptitle(f'Original {column_name} Data with {num_removed} Outliers Removed (Min={removed_min:.2f}, Max={removed_max:.2f})')
    
    # Set colors
    dusty_teal = basic_colours["blue"]
    soft_coral= basic_colours['orange']
    
    ax.plot(df.index, df[column_name], label=column_name, color=dusty_teal)
    ax.set_ylabel(column_name)
    if not removed_rows.empty:
        ax.scatter(removed_rows.index, removed_rows[column_name], c=soft_coral, s=20, alpha=0.5)

    plt.show()
    models.save_figure(fig, f'Original {column_name} Data with {num_removed}')
    return train_data



from iapws import IAPWS95

def calculate_mst_ex2(Q_req, mst_TG, mst_ex_sc,T_min = 125):
    m_st_bp = 0 
    mst_ex = 10
    T_st_T = 70 + 273.15  # convert temperature to Kelvin (K)
    water_T = IAPWS95(T=T_st_T, x=0)
    steam_T = IAPWS95(T=235+273.15, P=1.22)
    hst_T = water_T.h  # specific enthalpy of 70C degree water (kJ/kg)
    hst_BP =steam_T.h
    m_st_ex_aux = -1  # initialize to a value different from mst_ex
    bp_valve = 0 #proc
    mst_TG_og = mst_TG
    mst_ex_o = 0
    try:
        count = 0
        while abs(mst_ex - m_st_ex_aux) >= 0.01:
            m_st_ex_aux = mst_ex
            pst_ex = (-90.5+5.6*mst_TG-3*mst_ex+0.01*mst_TG**2+0.07*mst_ex**2-0.08*mst_ex*mst_TG)
            Tst_ex = (142.2+0.4*mst_TG-1.1*mst_ex)  
            h_st_ex = IAPWS95(P=pst_ex/1000, T=Tst_ex + 273.15 ).h
            m_dhs = ((Q_req * 1000) / (h_st_ex - hst_T)) * 3600 * (1/1000) 
            mst_ex = m_dhs + mst_ex_sc
            Temp_mix = IAPWS95(P=pst_ex/1000, h=h_st_ex).T -273.15
            mst_ex_o = mst_ex
            count += 1
            if count > 20:
                raise RuntimeError("Exceeded maximum iteration count in first loop")
    except RuntimeError as e:
        print(f"Error occurred: {e}")
        mst_ex = mst_ex_o
        m_st_bp = 0
        Temp_mix = T_min+1
        mst_ex = 0
    try:
        count = 0
        if Temp_mix > T_min:
            test = 1
        else:
            for m_st_bp in range(0,int(0.5*mst_TG*10)):
                m_st_bp /= 10
                mst_TG = mst_TG_og-m_st_bp
                mst_ex = 1
                m_st_ex_aux = -1 
                while abs(mst_ex - m_st_ex_aux) >= 0.01:
                    m_st_ex_aux = mst_ex
                    pst_ex = (-90.5+5.6*mst_TG-3*mst_ex+0.01*mst_TG**2+0.07*mst_ex**2-0.08*mst_ex*mst_TG)
                    Tst_ex = (142.2+0.4*mst_TG-1.1*mst_ex)  
                    h_st_ex = IAPWS95(P=pst_ex/1000, T=Tst_ex + 273.15 ).h
                    h_total = (mst_ex*h_st_ex+hst_BP*m_st_bp)/(m_st_bp+mst_ex)
                    Temp_mix = IAPWS95(P=pst_ex/1000, h=h_total).T -273.15
                    m_dhs = ((Q_req * 1000) / (h_total - hst_T)) * 3600 * (1/1000) 
                    mst_ex = m_dhs + mst_ex_sc
                    if count > 20:
                        raise RuntimeError("Exceeded maximum iteration count in second loop")
                if Temp_mix > T_min:
                    break
    except RuntimeError as e:
        print(f"Error occurred: {e}")
        mst_ex = mst_ex_o
        m_st_bp = 0
    return mst_ex, m_st_bp

import numpy as np
from scipy.interpolate import interp1d
from scipy.optimize import brentq

def calculate_mst_ex3(Q_req, mst_TG, mst_ex_sc, T_min=125):
    # Convert temperature to Kelvin (K)
    T_st_T = 70 + 273.15
    T_st_BP = 235 + 273.15
    
    # Calculate specific enthalpy of water and steam
    water_T = IAPWS95(T=T_st_T, x=0)
    steam_T = IAPWS95(T=T_st_BP, P=1.22)
    h_st_T = water_T.h
    h_st_BP = steam_T.h
    
    # Initialize variables
    m_st_bp = 0
    mst_ex = 1
    m_st_ex_aux = -1
    bp_valve = 0
    mst_TG_og = mst_TG
    mst_ex_o = 0
    
    # Cache intermediate values
    count = 0
    pst_ex_arr = np.zeros((21,))
    Tst_ex_arr = np.zeros((21,))
    h_st_ex_arr = np.zeros((21,))
    
    try:
        while abs(mst_ex - m_st_ex_aux) >= 0.01:
            m_st_ex_aux = mst_ex
            pst_ex = (-90.5 + 5.6*mst_TG - 3*mst_ex + 0.01*mst_TG**2
                      + 0.07*mst_ex**2 - 0.08*mst_ex*mst_TG)
            Tst_ex = (142.2 + 0.4*mst_TG - 1.1*mst_ex)
            h_st_ex = IAPWS95(P=pst_ex/1000, T=Tst_ex + 273.15).h
            m_dhs = ((Q_req * 1000) / (h_st_ex - h_st_T)) * 3600 * (1/1000)
            mst_ex = m_dhs + mst_ex_sc
            pst_ex_arr[count] = pst_ex
            Tst_ex_arr[count] = Tst_ex
            h_st_ex_arr[count] = h_st_ex
            mst_ex_o = mst_ex
            count += 1
            if count > 20:
                raise RuntimeError("Exceeded maximum iteration count in first loop")
           
        # Use vectorization for faster computation
        m_st_bp_arr = np.arange(0, int(0.5*mst_TG_og*10))/10
        mst_TG_arr = mst_TG_og - m_st_bp_arr[:, np.newaxis]
        mst_ex_arr = np.zeros_like(mst_TG_arr)
        T_mix_arr = np.zeros_like(mst_TG_arr)

        for i in range(21):
            interp_func = interp1d(mst_TG_arr[:, i], mst_ex_arr[:, i], kind='cubic')
            roots = brentq(lambda x: interp_func(x) - mst_ex_o, 0, mst_TG_og)
            m_st_bp = m_st_bp_arr[np.where(np.isclose(mst_TG_arr[:, i], roots, rtol=0.01))[0][0]]
            mst_TG = mst_TG_og - m_st_bp
            mst_ex_arr[:, i] = mst_ex_o
            m_st_ex_aux = -1
            count = 0
            while abs(mst_ex_arr[-1, i] - m_st_ex_aux) >= 0.01:
                m_st_ex_aux = mst_ex_arr[-1, i]
                pst_ex = (-90.5 + 5.6*mst_TG - 3*mst_ex_arr[-1, i] + 0.01*mst_TG**2
                          + 0.07*mst_ex_arr[-1, i]**2 - 0.08*mst_ex_arr[-1, i]*mst_TG)
                Tst_ex = (142.2 + 0.4*mst_TG - 1.1*mst_ex_arr[-1, i])
                h_st_ex = h_st_ex_arr[count, i]
                h_total = (mst_ex_arr[-1, i]*h_st_ex + h_st_BP*m_st_bp)/(m_st_bp+mst_ex_arr[-1, i])
                T_mix = IAPWS95(P=pst_ex/1000, h=h_total).T - 273.15
                m_dhs = ((Q_req * 1000) / (h_total - h_st_T)) * 3600 * (1/1000)
                mst_ex_arr[-1, i] = m_dhs + mst_ex_sc
                T_mix_arr[-1, i] = T_mix
                count += 1
                if count > 20:
                    raise RuntimeError("Exceeded maximum iteration count in second loop")
            if T_mix_arr[-1, i] > T_min:
                break

        if T_mix_arr[-1, i] <= T_min:
            m_st_bp = 0
            mst_ex_arr[-1, i] = mst_ex_o
            T_mix_arr[-1, i] = T_min+1

    except RuntimeError as e:
        print(f"Error occurred: {e}")
        m_st_bp = 0
        mst_ex_arr[-1, :] = mst_ex_o
        T_mix_arr[-1, :] = T_min+1

    # Return result
    mst_ex = mst_ex_arr[-1, i]
    return mst_ex, m_st_bp

# %%
def calculate_mst_ex(Q_req, mst_TG, mst_ex_sc,m_st_bp = 0,T_min = 125):
    m_st_bp = 0 
    mst_ex = 1
    T_st_T = 70 + 273.15  # convert temperature to Kelvin (K)
    water_T = IAPWS95(T=T_st_T, x=0)
    steam_T = IAPWS95(T=235+273.15, P=1.22)
    hst_T = water_T.h  # specific enthalpy of 70C degree water (kJ/kg)
    hst_BP =steam_T.h
    m_st_ex_aux = -1  # initialize to a value different from mst_ex
    bp_valve = 0 #proc
    mst_TG_og = mst_TG
    mst_ex_o = 0
    try:
        count = 0
        while abs(mst_ex - m_st_ex_aux) >= 0.01:
            m_st_ex_aux = mst_ex
            pst_ex = (-90.5+5.6*mst_TG-3*mst_ex+0.01*mst_TG**2+0.07*mst_ex**2-0.08*mst_ex*mst_TG)
            Tst_ex = (142.2+0.4*mst_TG-1.1*mst_ex)  
            h_st_ex = IAPWS95(P=pst_ex/1000, T=Tst_ex + 273.15 ).h
            m_dhs = ((Q_req * 1000) / (h_st_ex - hst_T)) * 3600 * (1/1000) 
            mst_ex = m_dhs + mst_ex_sc
            Temp_mix = IAPWS95(P=pst_ex/1000, h=h_st_ex).T -273.15
            mst_ex_o = mst_ex
            count += 1
            if count > 20:
                raise RuntimeError("Exceeded maximum iteration count in first loop")
    except RuntimeError as e:
        print(f"Error occurred: {e}")
        mst_ex = mst_ex_o
        m_st_bp = 0
        Temp_mix = T_min+1
        mst_ex = 0
    try:
        count = 0

        if Temp_mix > T_min:
            test = 1
                # Set initial lower and upper bounds for m_st_bp
        else:
            
            lower_bound = 0
            upper_bound = int(0.5 * mst_TG)

            # Set a small increment for updating the lower bound
            increment = 0.5
            # Perform binary search to find the minimum m_st_bp that satisfies Temp_mix > T_min
            while upper_bound >= (lower_bound+0.5):
                
                # Calculate midpoint between lower and upper bounds
                midpoint = (lower_bound + upper_bound) / 2

                # Update m_st_bp and other variables based on new value
                m_st_bp = midpoint 
                mst_TG = mst_TG_og - m_st_bp
                mst_ex = 1
                m_st_ex_aux = -1
                count = 0
                
                # Perform iterative calculation to find new value of Temp_mix
                while abs(mst_ex - m_st_ex_aux) >= 0.01:
                    m_st_ex_aux = mst_ex
                    pst_ex = (-90.5 + 5.6 * mst_TG - 3 * mst_ex + 0.01 * mst_TG**2 + 0.07 * mst_ex**2 - 0.08 * mst_ex * mst_TG)
                    Tst_ex = (142.2 + 0.4 * mst_TG - 1.1 * mst_ex)
                    h_st_ex = IAPWS95(P=pst_ex / 1000, T=Tst_ex + 273.15).h
                    h_total = (mst_ex * h_st_ex + hst_BP * m_st_bp) / (m_st_bp + mst_ex)
                    Temp_mix = IAPWS95(P=pst_ex / 1000, h=h_total).T - 273.15
                    m_dhs = ((Q_req * 1000) / (h_total - hst_T)) * 3600 * (1 / 1000)
                    mst_ex = m_dhs + mst_ex_sc
                    count += 1
                    if count > 20:
                        raise RuntimeError("Exceeded maximum iteration count in second loop")
                    
                # Update bounds based on value of Temp_mix
                if Temp_mix > T_min:
                    upper_bound = midpoint - increment

                else:
                    lower_bound = midpoint + increment


            # The final value of m_st_bp is the upper bound
            m_st_bp = upper_bound 
    except RuntimeError as e:
        print(f"Error occurred: {e}")
        mst_ex = mst_ex_o
        m_st_bp = 0
    return mst_ex, m_st_bp

