import tensorflow as tf
from keras import layers, models, regularizers, optimizers, losses, callbacks
from keras._tf_keras.keras.preprocessing import image
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import scipy.io
from scipy.signal import butter, lfilter
import pandas as pd
import math
import os, re
from scipy.interpolate import interp1d
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score 
from sklearn.preprocessing import StandardScaler
from sklearn.utils.class_weight import compute_class_weight
from tensorflow.keras.losses import MeanSquaredError



def custom_loss(y_true, y_pred):
    y_lbl = y_true[:, :6]
    w     = y_true[:, 6:]
    mse   = tf.square(y_lbl - y_pred)          # (batch, 4)
    wmse  = tf.reduce_mean(w * mse, axis=-1)   # (batch,)
    return wmse



class Silu(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(Silu, self).__init__(**kwargs)

    def call(self, inputs):
        return tf.keras.activations.swish(inputs)

    def get_config(self):
        config = super(Silu, self).get_config()
        return config

def silu_with_name(x, **kwargs):
    return tf.keras.activations.swish(x)


#----- Model parameter tuning ----------
#Direct folder, file names, etc. 
save_folder_path = f'/path/to/save/figure' 
testlog = f'/path/for/excel/log/BL_log_paper.xlsx'
model_cond = f'IPF_condition' #model name
log = []
data = []
#Read data log
log = pd.read_excel(testlog, sheet_name='Data_ML')
data = {
'data_no': log.iloc[:, 0].to_numpy(),   
'fr': log.iloc[:, 1].to_numpy(),   
'Tb': log.iloc[:, 3].to_numpy(),
'q': log.iloc[:, 4].to_numpy(),  
'q_std':  log.iloc[:, 5].to_numpy(), 
'SPT': log.iloc[:, 6].to_numpy(),  
'SPT_std':  log.iloc[:, 7].to_numpy(),  
'd': log.iloc[:, 8].to_numpy(),  
'd_std': log.iloc[:, 9].to_numpy(),  
'f': log.iloc[:, 10].to_numpy(),    
'f_std': log.iloc[:, 11].to_numpy(),  
'tw': log.iloc[:, 12].to_numpy(),  
'tw_std': log.iloc[:, 13].to_numpy(),  
'N': log.iloc[:, 14].to_numpy(),  
'N_std': log.iloc[:, 15].to_numpy(),   
}
df = pd.DataFrame(data)

#Initialize matrix

y_uk = []
images_uk = []
error_est_ML_uk = []
error_est_uk = []
ytemp_uk = [] 
filtered_data = [] 
images_test = []
basename = f'35C_10GPM_1HP_T'

image_dirs = [
    f'/path/to/audio spectrum/{basename}',
]
image_xxx_values = set()
for d in image_dirs:
    if not os.path.isdir(d):
            continue
    for filename in os.listdir(d):
        if filename.endswith('.png'):
            xxx_value = filename.split('STFT_rmbg_ML_')[1].split('_')[0]
            image_xxx_values.add(xxx_value)

common_xxx_values = sorted(set(data['data_no'].astype(str)).intersection(image_xxx_values))

filtered_data = {
    'data_no': data['data_no'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'fr': data['fr'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'Tb': data['Tb'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'q': data['q'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'q_std': data['q_std'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'SPT': data['SPT'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'SPT_std': data['SPT_std'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'd': data['d'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'd_std': data['d_std'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'f': data['f'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'f_std': data['f_std'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'tw': data['tw'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'tw_std': data['tw_std'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'N': data['N'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
    'N_std': data['N_std'][np.isin(data['data_no'].astype(str), list(common_xxx_values))],
}
for xxx_value in common_xxx_values:
    for i in range(0,202):
        image_filename = f"STFT_rmbg_ML_{xxx_value}_{i}.png"
        for d in image_dirs:
            img_path = os.path.join(d, image_filename)
            if i == 0:
                print(img_path)        
            if os.path.exists(img_path):
                    # Load the image
                img_test = image.load_img(img_path, target_size=(224, 224), color_mode='grayscale')  # Resize image to 224x224
                img_test_array = image.img_to_array(img_test) / 255.0  # Convert to array and normalize pixel values
                images_test.append(img_test_array)
                break  # first dir wins

images_test = np.array(images_test)                       
images_uk.append(images_test) 


# Prepare the target data (parameters) for regression (outputs)
error_est_temp_uk = np.column_stack([
    filtered_data['q']*0.25, (filtered_data['SPT']+120)*0.25,
    filtered_data['d_std'], filtered_data['f_std'],
    filtered_data['tw_std'], filtered_data['N_std']
]).astype(np.float32)  # (3636, 6)
y_loop_uk = np.column_stack([
    filtered_data['q'], (filtered_data['SPT']+120),
    filtered_data['d'], filtered_data['f'],
    filtered_data['tw'], filtered_data['N']
]).astype(np.float32)  # Ensure that y is of type float32

print(y_loop_uk)
error_est_ML_uk.append(error_est_temp_uk)
error_est_temp_uk = np.repeat(error_est_temp_uk, 202, axis=0)  
error_est_uk.append(error_est_temp_uk)
yt_loop_uk = y_loop_uk
ytemp_uk.append(yt_loop_uk)
y_loop_uk = np.repeat(y_loop_uk,202,axis=0)
y_uk.append(y_loop_uk)
    
X_uk_test = np.concatenate(images_uk, axis=0).astype(np.float32)  
error_est_ML_uk = np.concatenate(error_est_ML_uk, axis=0).astype(np.float32)
y_uk = np.concatenate(y_uk, axis=0).astype(np.float32)         
ytemp_uk = np.concatenate(ytemp_uk, axis=0).astype(np.float32)  
error_est_uk = np.concatenate(error_est_uk, axis=0).astype(np.float32)
HF_uk = ytemp_uk[:,0]
SPT_uk = ytemp_uk[:,1]
BD_uk = ytemp_uk[:,2]
BFreq_uk = ytemp_uk[:,3]
tw_uk     = ytemp_uk[:,4] 
NSD_uk = ytemp_uk[:,5]


# Load model
model_nm = f'{model_cond}.h5'

model = tf.keras.models.load_model(
    model_nm,
    custom_objects={
        "custom_loss": custom_loss,
        "Silu": Silu,  # only if you actually used this activation in the model
        "mse": MeanSquaredError()
    }
)


# Make predictions
y_pred_uk = model.predict([X_uk_test])

y_pred_uk_reg = y_pred_uk  # same unpacking
HF = y_pred_uk_reg[:,0]
SPT = y_pred_uk_reg[:,1]
d = y_pred_uk_reg[:,2]
f = y_pred_uk_reg[:,3]
tw = y_pred_uk_reg[:,4]
NSD = y_pred_uk_reg[:,5]
data_uk = {
    'HF_pred': HF,
    'HF_data': y_uk[:,0],
    'SPT_pred': SPT,
    'SPT_data': y_uk[:,1],
    'd_pred': d,
    'd_data': y_uk[:,2],
    'f_pred': f,
    'f_data': y_uk[:,3],
    'tw_pred': tw,
    'tw_data': y_uk[:,4],
    'N_pred': NSD,
    'N_data': y_uk[:,5],
    'd_std':  error_est_uk[:,2],
    'f_std':  error_est_uk[:,3],
    'tw_std':  error_est_uk[:,4],
    'N_std':  error_est_uk[:,5],
}

metrics = {}


# List of predictions with corresponding columns
pred_sets = {
    'HF_pred': (0, HF),
    'SPT_pred': (1, SPT),
    'd_pred': (2, d),
    'f_pred': (3, f),
    'tw_pred': (4, tw),
    'N_pred': (5, NSD)
}

for key, (col_idx, pred_data) in pred_sets.items():
    y_true = y_uk[:, col_idx]
    y_pred = pred_data
    model_std = error_est_uk[:, col_idx]

    mae = mean_absolute_error(y_true, y_pred)
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    r2 = r2_score(y_true, y_pred)
    mean_model_std = np.mean(model_std)
    mean_truth = np.mean(y_true)

    metrics[key] = {
        'STD': mean_model_std,
        'MAE': mae,
        'MSE': mse,
        'RMSE': rmse,
        'R2': r2,
        'NRMSE_by_model_std': rmse / mean_model_std if mean_model_std != 0 else np.nan,
        'NRMSE_by_truth_mean': rmse / mean_truth if mean_truth != 0 else np.nan
    }

df_uk = pd.DataFrame(data_uk)
de = pd.DataFrame(metrics)
de['Metric'] = de.index
de = de[['Metric'] + [col for col in de.columns if col != 'Metric']]


with pd.ExcelWriter(f'{model_cond}_{basename9}.xlsx', engine='openpyxl') as writer:  
    # Save the data to an Excel file
    df_uk.to_excel(writer,sheet_name=f'data_uk', index=False)
    de.to_excel(writer,sheet_name=f'error', index=False)


fullfilename = f'{model_cond}_{basename9}.png'  # or use a different identifier for your filename
file_path = os.path.join(save_folder_path, fullfilename)
fig, axes = plt.subplots(2, 2, figsize=(12, 12))
axes[0, 0].scatter(y_pred_uk_reg[:,0],y_pred_uk_reg[:,2],color = 'blue')
axes[0, 0].scatter(HF_uk, BD_uk, color='black', label='Line Plot with Error Bars')
axes[0, 0].errorbar(HF_uk,BD_uk, xerr=error_est_ML_uk[:,0], yerr=error_est_ML_uk[:,2], fmt='o', color='black', ecolor='black')
axes[0, 0].set_xlabel(r'HEAT FLUX (MW/m$^2$)', fontsize=20)
axes[0, 0].set_ylabel(r'BUBBLE DEPARTURE DIAMETER (mm)', fontsize=20)
axes[0, 0].tick_params(axis='both', labelsize=18)
axes[0, 0].grid()

axes[0, 1].scatter(y_pred_uk_reg[:,0],y_pred_uk_reg[:,3],color = 'blue')
axes[0, 1].scatter(HF_uk, BFreq_uk, color='black', label='Line Plot with Error Bars')
axes[0, 1].errorbar(HF_uk,BFreq_uk, xerr=error_est_ML_uk[:,0], yerr=error_est_ML_uk[:,3], fmt='o', color='black', ecolor='black')
axes[0, 1].set_xlabel(r'HEAT FLUX (MW/m$^2$)', fontsize=20)
axes[0, 1].set_ylabel(r'BUBBLE DEPARTURE FREQUENCY (Hz)', fontsize=20)
axes[0, 1].tick_params(axis='both', labelsize=18)
axes[0, 1].grid()

axes[1, 0].scatter(y_pred_uk_reg[:,0],y_pred_uk_reg[:,4],color = 'blue')
axes[1, 0].scatter(HF_uk, tw_uk, color='black', label='Line Plot with Error Bars')
axes[1, 0].errorbar(HF_uk,tw_uk, xerr=error_est_ML_uk[:,0], yerr=error_est_ML_uk[:,4], fmt='o', color='black', ecolor='black')
axes[1, 0].set_xlabel(r'HEAT FLUX (MW/m$^2$)', fontsize=20)
axes[1, 0].set_ylabel(r'BUBBLE WAIT TIME (ms)', fontsize=20)
axes[1, 0].tick_params(axis='both', labelsize=18)
axes[1, 0].grid()

axes[1, 1].scatter(y_pred_uk_reg[:,1],y_pred_uk_reg[:,5],color = 'blue')
axes[1, 1].scatter(SPT_uk, NSD_uk, color='black', label='Line Plot with Error Bars')
axes[1, 1].errorbar(SPT_uk,NSD_uk, xerr=error_est_ML_uk[:,1], yerr=error_est_ML_uk[:,5], fmt='o', color='black', ecolor='black')
axes[1, 1].set_xlabel(r'SUPERHEAT TEMPERATURE (K)', fontsize=20)
axes[1, 1].set_ylabel(r'NUCLEATION SITE DENSITY (#/cm$^2$)', fontsize=20)
axes[1, 1].tick_params(axis='both', labelsize=18)
axes[1, 1].grid()

plt.tight_layout()
plt.savefig(file_path)

