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 keras.layers import Activation


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



#----- Model parameter tuning ----------
BATCH_SIZE = 64 # number of images per train
TOTAL_ITERATIONS = 80000 #Increase  if less bubbles were detected through the training, vice versa
dropout = 0.0
#Direct folder, file names, etc. 
extension = '.npy'
testlog = f'/path/to/excel/log/BL_log_paper.xlsx'
model_cond = f'IPF_condition'

#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 = []
X = []
ytemp = []
images = []
images_loop = [] 
error_est = []
y_plot_temp = []
error_est_ML = []

#Read samples
filtered_data = [] 
images_loop = []
basename = f'35C_10GPM_1HP_T'


image_dirs = [
    f'/path/to/audio/spectrum/{basename}',
]

# Gather all filenames and sample IDs across all dirs
image_xxx_values = set()
for d in image_dirs:
    if not os.path.isdir(d):
        print(f'directory not found :{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)

# Intersect with data log
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))],
}

# Load images for each sample id from ANY of the provided dirs
for xxx_value in common_xxx_values:
    for i in range(0, 202):
        image_filename = f"STFT_rmbg_ML_{xxx_value}_{i}.png"
        img_loaded = False
        for d in image_dirs:
            img_path = os.path.join(d, image_filename)
            if i == 0 and os.path.exists(img_path):
                print(img_path)
            if os.path.exists(img_path):
                img = image.load_img(img_path, target_size=(224, 224), color_mode='grayscale')
                img_array = image.img_to_array(img) / 255.0
                images_loop.append(img_array)
                img_loaded = True
                break  # stop at the first dir that has the file


images_loop = np.array(images_loop)
images.append(images_loop)


# Targets and repeats (unchanged)

y_loop_full = 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)
y_plot = y_loop_full
y_plot_temp.append(y_plot)
yt_loop = y_loop_full
ytemp.append(yt_loop)
y_loop_full = np.repeat(y_loop_full, 202, axis=0)

error_est_temp = 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)
error_est_ML.append(error_est_temp)
error_est_temp = np.repeat(error_est_temp, 202, axis=0)
error_est.append(error_est_temp)

if np.isnan(X).any() or np.isnan(y_loop_full).any():
    print("There are NaN values in the data!")
else:
    print("No NaN values found.")
y.append(y_loop_full)


error_est = np.concatenate(error_est, axis=0).astype(np.float32)


for i, im in enumerate(images):
    print(i, type(im), np.shape(im))
X = np.concatenate(images, axis=0).astype(np.float32)     
y = np.concatenate(y, axis=0).astype(np.float32)      
ytemp = np.concatenate(ytemp, axis=0).astype(np.float32)  
y_plot_temp = np.concatenate(y_plot_temp, axis=0)
NUM_SAMPLES = len(error_est[:, 0]) # The number of images that had been manually masked
print(error_est_ML)

# Sanity checks
assert X.shape[0] == y.shape[0], \
    f"Lens mismatch known: X={X.shape[0]} y={y.shape[0]}"
u_q  = error_est[:, 0]
u_SPT= error_est[:, 1]
u_d  = error_est[:, 2]
u_f  = error_est[:, 3]
u_tw = error_est[:, 4]
u_N  = error_est[:, 5]
print(u_N)
# Define a small value to replace zero uncertainty
epsilon = 1e-6
u_q  = np.where(u_q  == 0, epsilon, u_q)
u_SPT= np.where(u_SPT  == 0, epsilon, u_SPT)
u_d  = np.where(u_d  == 0, epsilon, u_d)
u_f  = np.where(u_f  == 0, epsilon, u_f)
u_tw = np.where(u_tw == 0, epsilon, u_tw)
u_N  = np.where(u_N  == 0, epsilon, u_N)

w_q  = 1.0 / (u_q  ** 2)
w_SPT  = 1.0 / (u_SPT  ** 2)
w_d  = 1.0 / (u_d  ** 2)
w_f  = 1.0 / (u_f  ** 2)
w_tw = 1.0 / (u_tw ** 2)
w_N  = 1.0 / (u_N  ** 2)
reg_weights_all = np.column_stack([w_q, w_SPT, w_d, w_f, w_tw, w_N]).astype(np.float32)



X_train, X_test, \
y_train, y_test, \
w_train, w_test = train_test_split(
    X, y[:, :], reg_weights_all[:,:],
    test_size=0.3, random_state=42, shuffle=True
)

print("Train lens:", len(X_train), len(y_train))
print(" Test lens:", len(X_test),  len(y_test))

HF = y_plot_temp[:,0]
SPT = y_plot_temp[:,1]
BD = y_plot_temp[:,2]
BFreq = y_plot_temp[:,3]
tw     = y_plot_temp[:,4] 
NSD = y_plot_temp[:,5]




def find_nearest_error(val, ytemp_array, error_array):
    idx = np.where(np.isclose(ytemp_array, val, atol=1e-5))[0]
    if len(idx) == 0:
        raise ValueError(f"No match found for value: {val}")
    return error_array[idx[0]]

ytemp = np.array(ytemp)
error_est_ML = np.concatenate(error_est_ML, axis=0)


HFerror = np.array([find_nearest_error(val, ytemp[:, 0], error_est_ML[:, 0]) for val in y_test[:, 0]])
HFerror = HFerror.flatten()

SPTerror = np.array([find_nearest_error(val, ytemp[:, 0], error_est_ML[:, 1]) for val in y_test[:, 0]])
SPTerror = SPTerror.flatten()

BDerror = np.array([find_nearest_error(val, ytemp[:, 0], error_est_ML[:, 2]) for val in y_test[:, 0]])
BDerror = BDerror.flatten()

BFerror = np.array([find_nearest_error(val, ytemp[:, 0], error_est_ML[:, 3]) for val in y_test[:, 0]])
BFerror = BFerror.flatten()

twerror = np.array([find_nearest_error(val, ytemp[:, 0], error_est_ML[:, 4]) for val in y_test[:, 0]])
twerror = twerror.flatten()

NSDerror = np.array([find_nearest_error(val, ytemp[:, 0], error_est_ML[:, 5]) for val in y_test[:, 0]])
NSDerror = NSDerror.flatten()





# Replace zero uncertainty values with the small epsilon value to avoid division by zero
HFerror = np.where(HFerror == 0, epsilon, HFerror)
SPTerror = np.where(SPTerror == 0, epsilon, SPTerror)
BDerror = np.where(BDerror == 0, epsilon, BDerror)
BFerror = np.where(BFerror == 0, epsilon, BFerror)
twerror = np.where(twerror == 0, epsilon, twerror)
NSDerror = np.where(NSDerror == 0, epsilon, NSDerror)


HFerror_W = 1 / HFerror
HFerror_WN = HFerror_W / np.sum(HFerror_W)
SPTerror_W = 1 / SPTerror
SPTerror_WN = SPTerror_W / np.sum(SPTerror_W)
BDerror_W = 1 / BDerror
BDerror_WN = BDerror_W / np.sum(BDerror_W)
BFerror_W = 1 / BFerror
BFerror_WN = BFerror_W / np.sum(BFerror_W)
twerror_W = 1 / twerror
twerror_WN = twerror_W / np.sum(twerror_W)
NSDerror_W = 1 / NSDerror
NSDerror_WN = NSDerror_W / np.sum(NSDerror_W)


y_error = np.column_stack([HFerror,SPTerror,BDerror,BFerror,twerror,NSDerror]).astype(np.float32)
weights = np.column_stack([HFerror_WN,SPTerror_WN,BDerror_WN,BFerror_WN,twerror_WN,NSDerror_WN]).astype(np.float32)

def se_block(x, ratio=8):
    filters = x.shape[-1]
    se = layers.GlobalAveragePooling2D()(x)
    se = layers.Dense(max(4, filters // ratio), activation='relu')(se)
    se = layers.Dense(filters, activation='sigmoid')(se)
    se = layers.Reshape((1, 1, filters))(se)
    return layers.Multiply()([x, se])



# --- Inputs ---
input_signal = tf.keras.Input(shape=(224, 224, 1), name="spectrogram_input")

# --- CNN Feature Extractor --- # These parameters and architecture can be changed as well
x = layers.Conv2D(32, (3, 3), strides=(1, 1), padding='same')(input_signal)  # was 12
x = se_block(x)
x = layers.BatchNormalization()(x)
x = tf.keras.activations.swish(x)
x = layers.MaxPooling2D((2, 2))(x)

x = layers.Conv2D(64, (3, 3), strides=(1, 1), padding='same')(x)  # was 24
x = se_block(x)
x = layers.BatchNormalization()(x)
x = tf.keras.activations.swish(x)
x = layers.MaxPooling2D((2, 2))(x)


x = layers.Conv2D(128, (3, 3), strides=(1, 1), padding='same')(x)  # was 48
x = se_block(x)
x = layers.BatchNormalization()(x)
x = tf.keras.activations.swish(x)
x = layers.MaxPooling2D((2, 2))(x)

x = layers.Conv2D(128, (3, 3), strides=(1, 1), padding='same')(x)  # keep same, add depth
x = se_block(x)
x = layers.BatchNormalization()(x)
x = tf.keras.activations.swish(x)
x = layers.MaxPooling2D((2, 2))(x)


x = layers.Conv2D(128, (3, 3), strides=(1, 1), padding='same')(x)  # keep same, add depth
x = se_block(x)
x = layers.BatchNormalization()(x)
x = tf.keras.activations.swish(x)
x = layers.MaxPooling2D((2, 2))(x)


x = layers.Conv2D(256, (3, 3), strides=(1, 1), padding='same')(x)  # widen the last layer
x = se_block(x)
x = layers.BatchNormalization()(x)
x = tf.keras.activations.swish(x)
x = layers.GlobalMaxPooling2D()(x)


# --- Combine CNN ---
steps_per_epoch = NUM_SAMPLES // BATCH_SIZE
EPOCHS = TOTAL_ITERATIONS // steps_per_epoch
features = x
shared = layers.Dense(512, use_bias=False, kernel_regularizer=regularizers.l2(0.005))(features)
shared = layers.LayerNormalization()(shared)
shared = tf.keras.activations.swish(shared)
shared = layers.Dropout(0.0)(shared)

# --- Regressor Head ---
output_regression = layers.Dense(6, name="regression_output")(shared)


# --- Model ---
# Final model
model = tf.keras.Model(
    inputs=[input_signal],
    outputs=[output_regression],
    name="hydrophone_CNN_model"
)

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
    loss={
        "regression_output": custom_loss  # your batch-safe loss
    },
    loss_weights={
        "regression_output": 1.0,  # main task
    }
)

model.summary()
# Build y with weights if using Option B
y_train_w = np.concatenate([y_train, w_train], axis=1).astype(np.float32)
y_test_w  = np.concatenate([y_test,  w_test ], axis=1).astype(np.float32)


model.fit(
    x=X_train,
    y=y_train_w,
    validation_data=(X_test, y_test_w),
    batch_size=BATCH_SIZE,
    epochs=EPOCHS,
    shuffle=True,
    verbose=1
)

#-------- Save model name -------------
model.save(f'{model_cond}.h5')

# Make predictions on the test set
y_pred  = model.predict([X_test])



def find_nearest_value(val, ref_array, target_array):
    idx = np.where(np.isclose(ref_array, val, atol=1e-5))[0]
    if len(idx) == 0:
        raise ValueError(f"No match found for value: {val}")
    return target_array[idx[0]]


HF = np.array(HF)
SPT = np.array(SPT)
y_pred_reg = y_pred  # regression output




data = {
    'q_pred': y_pred_reg[:,0],
    'q_data': y_test[:,0],
    'q_std':  y_error[:,0],
    'SPT_pred': y_pred_reg[:,1],
    'SPT_data': y_test[:,1],
    'SPT_std':y_error[:,1],
    'd_pred': y_pred_reg[:,2],
    'd_data': y_test[:,2],
    'd_std':  y_error[:,2],
    'f_pred': y_pred_reg[:,3],
    'f_data': y_test[:,3],
    'f_std':y_error[:,3],
    'tw_pred': y_pred_reg[:,4],
    'tw_data': y_test[:,4],
    'tw_std': y_error[:,4],
    'N_pred': y_pred_reg[:,5],
    'N_data': y_test[:,5],
    'N_std':y_error[:,5]
}



ss_total_q  = np.sum(HFerror_W)  # Weighted total sum of squares
ss_residual_q = np.sum(HFerror_W * (y_test[:,0] - y_pred_reg[:,0])**2)  # Weighted residual sum of squares

ss_total_SPT = np.sum(SPTerror_W)  # Weighted total sum of squares
ss_residual_SPT = np.sum(SPTerror_W * (y_test[:,1] - y_pred_reg[:,1])**2)  # Weighted residual sum of squares

ss_total_d = np.sum(BDerror_W)  # Weighted total sum of squares
ss_residual_d = np.sum(BDerror_W * (y_test[:,2] - y_pred_reg[:,2])**2)  # Weighted residual sum of squares

ss_total_f = np.sum(BFerror_W) # Weighted total sum of squares
ss_residual_f = np.sum(BFerror_W * (y_test[:,3] - y_pred_reg[:,3])**2)  # Weighted residual sum of squares

ss_total_tw = np.sum(twerror_W)  # Weighted total sum of squares
ss_residual_tw = np.sum(twerror_W * (y_test[:,4] - y_pred_reg[:,4])**2)  # Weighted residual sum of square

ss_total_N = np.sum(NSDerror_W)  # Weighted total sum of squares
ss_residual_N = np.sum(NSDerror_W * (y_test[:,5] - y_pred_reg[:,5])**2)  # Weighted residual sum of squares





metrics1 = {
    'q_pred': {
        'MAE': np.sum(HFerror_W * np.abs(y_test[:,0] - y_pred_reg[:,0])) / np.sum(HFerror_W),
        'MSE': np.sum(HFerror_W * (y_test[:,0] - y_pred_reg[:,0])**2) / np.sum(HFerror_W),
        'RMSE': np.sqrt(np.sum(HFerror_W * (y_test[:,0] - y_pred_reg[:,0])**2) / np.sum(HFerror_W)),
        'R2': 1 - ss_residual_q / ss_total_q,
    },
    'SPT_pred': {
        'MAE': np.sum(SPTerror_W * np.abs(y_test[:,1] - y_pred_reg[:,1])) / np.sum(SPTerror_W),
        'MSE': np.sum(SPTerror_W * (y_test[:,1] - y_pred_reg[:,1])**2) / np.sum(SPTerror_W),
        'RMSE': np.sqrt(np.sum(BDerror_W * (y_test[:,1] - y_pred_reg[:,1])**2) / np.sum(SPTerror_W)),
        'R2': 1 - ss_residual_SPT / ss_total_SPT,
    },
    'd_pred': {
        'MAE': np.sum(BDerror_W * np.abs(y_test[:,2] - y_pred_reg[:,2])) / np.sum(BDerror_W),
        'MSE': np.sum(BDerror_W * (y_test[:,2] - y_pred_reg[:,2])**2) / np.sum(BDerror_W),
        'RMSE': np.sqrt(np.sum(BDerror_W * (y_test[:,2] - y_pred_reg[:,2])**2) / np.sum(BDerror_W)),
        'R2': 1 - ss_residual_d / ss_total_d,
    },
    'f_pred': {
        'MAE': np.sum(BFerror_W * np.abs(y_test[:,3] - y_pred_reg[:,3])) / np.sum(BFerror_W),
        'MSE': np.sum(BFerror_W * (y_test[:,3] - y_pred_reg[:,3])**2) / np.sum(BFerror_W),
        'RMSE': np.sqrt(np.sum(BFerror_W * (y_test[:,3] - y_pred_reg[:,3])**2) / np.sum(BFerror_W)),
        'R2': 1 - ss_residual_f / ss_total_f,
    },
    'tw_pred': {
        'MAE': np.sum(twerror_W * np.abs(y_test[:,4] - y_pred_reg[:,4])) / np.sum(twerror_W),
        'MSE': np.sum(twerror_W * (y_test[:,4] - y_pred_reg[:,4])**2) / np.sum(twerror_W),
        'RMSE': np.sqrt(np.sum(twerror_W * (y_test[:,4] - y_pred_reg[:,4])**2) / np.sum(twerror_W)),
        'R2': 1 - ss_residual_tw / ss_total_tw,
    },
    'N_pred': {
        'MAE': np.sum(NSDerror_W * np.abs(y_test[:,5] - y_pred_reg[:,5])) / np.sum(NSDerror_W),
        'MSE': np.sum(NSDerror_W * (y_test[:,5] - y_pred_reg[:,5])**2) / np.sum(NSDerror_W),
        'RMSE': np.sqrt(np.sum(NSDerror_W * (y_test[:,5] - y_pred_reg[:,5])**2) / np.sum(NSDerror_W)),
        'R2': 1 - ss_residual_N / ss_total_N,
    }
}
metrics2 = {
    'q_pred': {
        'MAE': mean_absolute_error(y_test[:,0], y_pred_reg[:,0]),
        'MSE': mean_squared_error(y_test[:,0], y_pred_reg[:,0]),
        'RMSE': np.sqrt(mean_squared_error(y_test[:,0], y_pred_reg[:,0])),
        'R2': r2_score(y_test[:,0], y_pred_reg[:,0]),
    },
    'SPT_pred': {
        'MAE': mean_absolute_error(y_test[:,1], y_pred_reg[:,1]),
        'MSE': mean_squared_error(y_test[:,1], y_pred_reg[:,1]),
        'RMSE': np.sqrt(mean_squared_error(y_test[:,1], y_pred_reg[:,1])),
        'R2': r2_score(y_test[:,1], y_pred_reg[:,1]),
    },
    'd_pred': {
        'MAE': mean_absolute_error(y_test[:,2], y_pred_reg[:,2]),
        'MSE': mean_squared_error(y_test[:,2], y_pred_reg[:,2]),
        'RMSE': np.sqrt(mean_squared_error(y_test[:,2], y_pred_reg[:,2])),
        'R2': r2_score(y_test[:,2], y_pred_reg[:,2]),
    },
    'f_pred': {
        'MAE': mean_absolute_error(y_test[:,3], y_pred_reg[:,3]),
        'MSE': mean_squared_error(y_test[:,3], y_pred_reg[:,3]),
        'RMSE': np.sqrt(mean_squared_error(y_test[:,3], y_pred_reg[:,3])),
        'R2': r2_score(y_test[:,3], y_pred_reg[:,3]),
    },
    'tw_pred': {
        'MAE': mean_absolute_error(y_test[:,4], y_pred_reg[:,4]),
        'MSE': mean_squared_error(y_test[:,4], y_pred_reg[:,4]),
        'RMSE': np.sqrt(mean_squared_error(y_test[:,4], y_pred_reg[:,4])),
        'R2': r2_score(y_test[:,4], y_pred_reg[:,4]),
    },
    'N_pred': {
        'MAE': mean_absolute_error(y_test[:,5], y_pred_reg[:,5]),
        'MSE': mean_squared_error(y_test[:,5], y_pred_reg[:,5]),
        'RMSE': np.sqrt(mean_squared_error(y_test[:,5], y_pred_reg[:,5])),
        'R2': r2_score(y_test[:,5], y_pred_reg[:,5]),
    }
}


df = pd.DataFrame(data)
dew = pd.DataFrame(metrics1)
de = pd.DataFrame(metrics2)
dew['Metric'] = dew.index
dew = dew[['Metric'] + [col for col in dew.columns if col != 'Metric']]
de['Metric'] = de.index
de = de[['Metric'] + [col for col in de.columns if col != 'Metric']]

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

# Print out the metrics for each prediction
for pred in metrics2:
    print(f"Metrics for {pred}:")
    for metric, value in metrics2[pred].items():
        print(f"{metrics2}: {value:.4f}")
    print()



del common_xxx_values 
del image_xxx_values
