
from re import X
from catboost import CatBoostRegressor, CatBoostClassifier
from matplotlib import pyplot as plt
from sklearn.model_selection import ParameterGrid, train_test_split
import random
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from tqdm import tqdm
from utils import flatten_categorical_columns
import pandas as pd
import numpy as np
from sklearn.cluster import KMeans,DBSCAN
from sklearn.metrics import mean_absolute_error, mean_squared_error, silhouette_score,r2_score
from kmodes.kprototypes import KPrototypes
from bad_good_mapper import pair_configs
from sklearn.multioutput import MultiOutputRegressor
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import KFold,StratifiedKFold
import json
random.seed(42)

#--------------------BURR---------------------------------------------------------------------
#Best classifier accuracy: 0.8041
#best classifier parameters: {'bagging_temperature': 0.5, 'depth': 15, 'iterations': 200, 'l2_leaf_reg': 1, 'learning_rate': 0.001, 'random_strength': 3}
#Best CV MAE score: 0.1004
#Best params: {'bagging_temperature': 0.5, 'depth': 8, 'early_stopping_rounds': 20, 'iterations': 800, 'l2_leaf_reg': 1, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 2}
#---------------------------------------------------------------------------------------------
#--------------------CUTTING LOSS-------------------------------------------------------------
#Best CV MAE score: 0.0914
#Best params: {'bagging_temperature': 0.5, 'depth': 4, 'early_stopping_rounds': 20, 'iterations': 1000, 'l2_leaf_reg': 5, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 1}
#----------------------------------------------------------------------------------------------
#--------------------CUTTING TORN-------------------------------------------------------------
#Best CV MAE score: 0.0483
#Best params: {'bagging_temperature': 0.5, 'depth': 6, 'early_stopping_rounds': 20, 'iterations': 800, 'l2_leaf_reg': 1, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 1}
#----------------------------------------------------------------------------------------------
#--------------------PLASMA-------------------------------------------------------------
#Best classifier accuracy: 0.9900
#Best classifier params: {'bagging_temperature': 2, 'depth': 8, 'iterations': 500, 'l2_leaf_reg': 5, 'learning_rate': 0.01, 'random_strength': 1}
#Best CV MAE score: 0.1206
#Best params: {'bagging_temperature': 0.5, 'depth': 6, 'early_stopping_rounds': 20, 'iterations': 1000, 'l2_leaf_reg': 3, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 2}
#--------------------------------------------------------------PLASMA V2----------------------------
#Best classifier params: {'bagging_temperature': 1, 'depth': 12, 'iterations': 200, 'l2_leaf_reg': 3, 'learning_rate': 0.0005, 'random_strength': 1}
#Best classifier accuracy: 0.9900
#Best CV MAE score: 0.1392
#Best params: {'bagging_temperature': 0.5, 'depth': 8, 'early_stopping_rounds': 20, 'iterations': 800, 'l2_leaf_reg': 1, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 1}
#----------------------------------------------------------------------------------------------

def round_values(df, rounding_factors):
    for col, factor in rounding_factors.items():
        if col in df.columns:
            df[col] = (df[col] / factor).round() * factor
    return df

rounding_factors = {
    "CONTOUR_SPEED [mm/min]": 50,     # Arrotonda a multipli di 50
    "LASER_POWER [W]": 1000,          # Arrotonda a multipli di 1000
    "CONTOUR_GAS_PRESSURE [bar]": 0.5, # Arrotonda a multipli di 0.5 (= 1/2)
    "CONTOUR_NOZZLE_DISTANCE [mm]": 0.1, # Arrotonda a multipli di 0.1
    "CONTOUR_FOCAL [mm]": 0.5 ,        # Arrotonda a multipli di 0.5
}
best_params_catboost={
    'Burr': {'reg':{'bagging_temperature': 0.5, 'depth': 8, 'early_stopping_rounds': 20, 'iterations': 800, 'l2_leaf_reg': 1, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 2,'eval_metric':'MultiRMSE'},'class':{"bagging_temperature": 0.5,"depth": 12,"early_stopping_rounds": 10,"iterations": 200,"l2_leaf_reg": 3,"learning_rate": 0.05,"random_strength": 1}},
    'Plasma': {'reg':{'bagging_temperature': 0.5, 'depth': 8, 'early_stopping_rounds': 20, 'iterations': 800, 'l2_leaf_reg': 1, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 1},'class': {'bagging_temperature': 1, 'depth': 12, 'iterations': 200, 'l2_leaf_reg': 3, 'learning_rate': 0.0005, 'random_strength': 1}},
    'Cutting loss': {'reg':{'bagging_temperature': 0.5, 'depth': 4, 'early_stopping_rounds': 20, 'iterations': 1000, 'l2_leaf_reg': 5, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 1,'eval_metric':'MultiRMSE'},'class':{}},
    'Cutting torn':  {'reg':{'bagging_temperature': 0.5, 'depth': 4, 'early_stopping_rounds': 20, 'iterations': 1000, 'l2_leaf_reg': 5, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 1,'eval_metric':'MultiRMSE'},'class':{}},
}
best_params_catboost_v2={
    'Burr': {'reg': {'bagging_temperature': 0, 'depth': 12, 'early_stopping_rounds': 20, 'eval_metric': 'MultiRMSE', 'iterations': 800, 'l2_leaf_reg': 1, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 3},'class':{"bagging_temperature": 0.5,"depth": 12,"early_stopping_rounds": 10,"iterations": 200,"l2_leaf_reg": 3,"learning_rate": 0.05,"random_strength": 1}},
    'Plasma': {'reg':{"bagging_temperature": 0,"depth": 12,"early_stopping_rounds": 20,"eval_metric": "MultiRMSE","iterations": 800,"l2_leaf_reg": 1,"learning_rate": 0.05,"loss_function": "MultiRMSE","random_strength": 3},'class':{"bagging_temperature": 0,"depth": 8,"early_stopping_rounds": 20,"iterations": 200,"l2_leaf_reg": 5,"learning_rate": 0.05,"random_strength": 2}},
    'Cutting loss': {'reg':{"bagging_temperature": 0,"depth": 12,"early_stopping_rounds": 10,"eval_metric": "MultiRMSE","iterations": 1000,"l2_leaf_reg": 1,"learning_rate": 0.05,"loss_function": "MultiRMSE","random_strength": 2},'class':{}},
    'Cutting torn': {'reg':{"bagging_temperature": 0,"depth": 12,"early_stopping_rounds": 20,"eval_metric": "MultiRMSE","iterations": 800,"l2_leaf_reg": 1,"learning_rate": 0.05,"loss_function": "MultiRMSE","random_strength": 3},'class':{}},
    'full':{'reg':{'bagging_temperature': 0, 'depth': 8, 'early_stopping_rounds': 20, 'eval_metric': 'MultiRMSE', 'iterations': 1000, 'l2_leaf_reg': 3, 'learning_rate': 0.05, 'loss_function': 'MultiRMSE', 'random_strength': 2},'class':{'bagging_temperature': 1, 'depth': 6, 'early_stopping_rounds': 10, 'iterations': 500, 'l2_leaf_reg': 1, 'learning_rate': 0.05, 'random_strength': 2}}
}
best_params_catboost_original={
    'full':{'reg':{'bagging_temperature': 0, 'depth': 8, 'early_stopping_rounds': 20, 'eval_metric': 'MultiRMSE', 'iterations': 1000, 'l2_leaf_reg': 1, 'learning_rate': 0.01, 'loss_function': 'MultiRMSE', 'random_strength': 1},'class':{"bagging_temperature": 2,"depth": 12,"early_stopping_rounds": 5,"iterations": 800,"l2_leaf_reg": 3,"learning_rate": 0.05,"random_strength": 1}}
}

cont_features = ['THICKNESS_TULUS [mm]', 'CONTOUR_SPEED [mm/min]', 'LASER_POWER [W]',
                        'CONTOUR_GAS_PRESSURE [bar]', 'CONTOUR_NOZZLE_DISTANCE [mm]', 'CONTOUR_FOCAL [mm]']
cat_features=['NOZZLE_TYPE','MATERIAL_NAME_TULUS','DEFECT_TYPE']


final_data2=pd.read_excel('outlier_detection_output/final_data_cleaned.xlsx').sample(frac=1, random_state=42).reset_index(drop=True)
#original_dataset=pd.read_excel("merged_files.xlsx")
#original_dataset=original_dataset[final_data2.columns]
#final_data2=original_dataset.copy()


print(final_data2.dtypes)
print(f"final_data2 length: {len(final_data2)}")
final_data2=final_data2.drop_duplicates(subset=[col for col in final_data2.columns if col not in ['DEFECT_TYPE', 'QUALITY_CUT']], keep='first', inplace=False).reset_index(drop=True)
print(f"final_data2 length after cleaning: {len(final_data2)}")
final_data2=final_data2[cont_features+cat_features]

defect='Cutting torn'

scaler=StandardScaler()
duplicates = final_data2.columns[final_data2.columns.duplicated()]
print("Colonne duplicate:", duplicates)
scaler.fit(final_data2[final_data2['DEFECT_TYPE'] == 'No Defects'][cont_features])
print(final_data2.columns)
materials=[material for material in final_data2['MATERIAL_NAME_TULUS'].unique()]
thicknesses=list(final_data2['THICKNESS_TULUS [mm]'].unique())

columns_order=final_data2.columns
mapping_bad_to_good_train = []  
mapping_bad_to_good_test = []  # Lista per memorizzare le coppie di configurazioni buone e cattive
for material in materials:
    data_mat = final_data2[(final_data2['MATERIAL_NAME_TULUS'] == material)]
    for thickness in thicknesses:
        # Separazione good e bad per questo materiale e spessore
        data_good = data_mat[(data_mat['THICKNESS_TULUS [mm]'] == thickness) & (data_mat['DEFECT_TYPE'] == 'No Defects')].copy().reset_index(drop=True)
        data_bad = data_mat[(data_mat['THICKNESS_TULUS [mm]'] == thickness) & (data_mat['DEFECT_TYPE'] != 'No Defects')].reset_index(drop=True)
        
        if len(data_good) != 0 and len(data_bad) != 0:
            # Assicura riproducibilità
            data_good = data_good.sample(frac=1, random_state=42)
            data_bad = data_bad.sample(frac=1, random_state=42)

            # Calcola quanti elementi per test (almeno 1 se possibile)
            test_size_bad = max(1, int(len(data_bad) * 0.1))
            
            if len(data_bad) <= 3:
                test_size_bad = 1

            if len(data_bad) == 1:
                test_size_bad = 0

            # Split train/test per good
            good_train = data_good.copy()
            good_test = data_good.copy()

            # Split train/test per bad
            bad_train = data_bad.iloc[test_size_bad:]
            bad_test = data_bad.iloc[:test_size_bad]

            # Aggiungi al set di training se ci sono dati
            if len(good_train) > 0 and len(bad_train) > 0:
                mapping_bad_to_good_train.append({
                    "material": material,
                    "thickness": thickness,
                    "bad_configs": bad_train.copy(),
                    "good_configs": good_train.copy()
                })

            # Aggiungi al set di test se ci sono dati
            if len(good_test) > 0 and len(bad_test) > 0:
                mapping_bad_to_good_test.append({
                    "material": material,
                    "thickness": thickness,
                    "bad_configs": bad_test.copy(),
                    "good_configs": good_test.copy()
                })

# Stampa statistiche
print(f"Training set: {len(mapping_bad_to_good_train)} combinazioni materiale/spessore")
print(f"Test set: {len(mapping_bad_to_good_test)} combinazioni materiale/spessore")

# Opzionale: stampa la distribuzione delle configurazioni
total_train_bad = sum(len(entry["bad_configs"]) for entry in mapping_bad_to_good_train)
total_train_good = sum(len(entry["good_configs"]) for entry in mapping_bad_to_good_train)
total_test_bad = sum(len(entry["bad_configs"]) for entry in mapping_bad_to_good_test)
total_test_good = sum(len(entry["good_configs"]) for entry in mapping_bad_to_good_test)

print(f"Training: {total_train_bad} bad configs, {total_train_good} good configs")
print(f"Test: {total_test_bad} bad configs, {total_test_good} good configs")


training_pairs = []  # Ogni elemento sarà una tupla (bad_config, good_config)
test_pairs=[]

for entry in mapping_bad_to_good_train:
    bad_configs = entry["bad_configs"]
    good_configs = entry["good_configs"]
    print(f'MATERIAL: {entry["material"]}, THICKNESS: {entry["thickness"]}')
    # usiamo il clustering per trovare la configurazione 'good' più adatta in termini di distanza di Gower e a garantire che a configurazioni bad vicine, corrispondano configurazioni good vicine
    pairs=pair_configs(bad_configs,good_configs,[f for f in cont_features if f not in ['THICKNESS_TULUS [mm]']],[f for f in cat_features if f not in ['DEFECT_TYPE','MATERIAL_NAME_TULUS']])
    training_pairs.extend(pairs)
    print(f"Training pairs: {len(pairs)}")

train_fake_bad_len=len(training_pairs)*0.05 # 5%
fake_bad=map(lambda x: (x[1],x[1]),list(random.sample(training_pairs, int(train_fake_bad_len))))
training_pairs.extend(fake_bad)

for entry in mapping_bad_to_good_test:
    bad_configs = entry["bad_configs"]
    good_configs = entry["good_configs"]
    print(f'MATERIAL: {entry["material"]}, THICKNESS: {entry["thickness"]}')
    # usiamo il clustering per trovare la configurazione 'good' più adatta in termini di cluster e a garantire che a configurazioni bad vicine, corrispondano configurazioni good vicine
    pairs=pair_configs(bad_configs,good_configs,[f for f in cont_features if f not in ['THICKNESS_TULUS [mm]']],[f for f in cat_features if f not in ['DEFECT_TYPE','MATERIAL_NAME_TULUS']])
    test_pairs.extend(pairs)
    print(f"test_pairs: {len(pairs)}")

test_fake_bad_len=len(test_pairs)*0.05 # 5%
fake_bad_test=map(lambda x: (x[1],x[1]),list(random.sample(test_pairs, int(test_fake_bad_len))))
test_pairs.extend(fake_bad_test)

random.shuffle(training_pairs)
random.shuffle(test_pairs)


print(f"training_pairs: {len(training_pairs)}")
X_train = pd.concat([pair[0] for pair in training_pairs], ignore_index=True)
y_train = pd.concat([pair[1] for pair in training_pairs], ignore_index=True)
#X_train=pd.DataFrame(X_train, columns=columns_order).drop(columns=['DEFECT_TYPE'])
X_train=pd.DataFrame(X_train, columns=columns_order)

X_train=X_train.copy().astype({'THICKNESS_TULUS [mm]': 'float64',
                                'CONTOUR_SPEED [mm/min]': 'float64',
                                'LASER_POWER [W]': 'float64',
                                'CONTOUR_GAS_PRESSURE [bar]': 'float64',
                                'CONTOUR_NOZZLE_DISTANCE [mm]': 'float64',
                                'CONTOUR_FOCAL [mm]': 'float64'})

y_train= y_train.copy().astype({'THICKNESS_TULUS [mm]': 'float64',
                                'CONTOUR_SPEED [mm/min]': 'float64',
                                'LASER_POWER [W]': 'float64',
                                'CONTOUR_GAS_PRESSURE [bar]': 'float64',
                                'CONTOUR_NOZZLE_DISTANCE [mm]': 'float64',
                                'CONTOUR_FOCAL [mm]': 'float64'})

y_train=pd.DataFrame(y_train, columns=columns_order).drop(columns=['DEFECT_TYPE'])
X_test = pd.concat([pair[0] for pair in test_pairs], ignore_index=True)
y_test = pd.concat([pair[1] for pair in test_pairs], ignore_index=True)
#X_test=pd.DataFrame(X_test, columns=columns_order).drop(columns=['DEFECT_TYPE'])
X_test=pd.DataFrame(X_test, columns=columns_order)
X_test=X_test.copy().astype({'THICKNESS_TULUS [mm]': 'float64',
                                'CONTOUR_SPEED [mm/min]': 'float64',
                                'LASER_POWER [W]': 'float64',
                                'CONTOUR_GAS_PRESSURE [bar]': 'float64',
                                'CONTOUR_NOZZLE_DISTANCE [mm]': 'float64',
                                'CONTOUR_FOCAL [mm]': 'float64'})

y_test= y_test.copy().astype({'THICKNESS_TULUS [mm]': 'float64',
                                'CONTOUR_SPEED [mm/min]': 'float64',
                                'LASER_POWER [W]': 'float64',
                                'CONTOUR_GAS_PRESSURE [bar]': 'float64',
                                'CONTOUR_NOZZLE_DISTANCE [mm]': 'float64',
                                'CONTOUR_FOCAL [mm]': 'float64'})
y_test=pd.DataFrame(y_test, columns=columns_order).drop(columns=['DEFECT_TYPE'])
y_test_materials=y_test['MATERIAL_NAME_TULUS'].copy()
y_test_thicknesses=y_test['THICKNESS_TULUS [mm]'].copy()
print(X_train["NOZZLE_TYPE"].unique())
print(X_test["NOZZLE_TYPE"].unique())
print(y_train.columns)
print(y_test.columns)


y_test = y_test.drop(columns=[ 'MATERIAL_NAME_TULUS'])
y_train = y_train.drop(columns=[ 'MATERIAL_NAME_TULUS'])

# Definisco un wrapper che costruisce ogni volta un MultiOutputRegressor con CatBoostRegressor


class MultiOutputCatBoostWrapper:
    def __init__(self, reg_params={}, clf_params={},scaler=None, scaled_cont_features=None,eval_set=None, eval_fraction=0.2):
        self.regressor = None
        self.classifier = None
        self.reg_params = reg_params
        self.valid_reg_columns = []
        self.clf_params = clf_params
        self.reg_columns = []
        self.cat_columns = []
        self.scaler = scaler
        self.scaled_cont_features = scaled_cont_features
        self.eval_set = eval_set
        if eval_set is not None:
            self.eval_fraction = 0.0
        else:
            self.eval_fraction = eval_fraction

    def fit(self, X, y, reg_columns, cat_columns, cat_features):
        self.reg_columns = reg_columns
        self.cat_columns = cat_columns
        # Fit regressor
        valid_reg_columns = [col for col in reg_columns if y[col].nunique() > 1]
        skipped_reg = [col for col in reg_columns if y[col].nunique() <= 1]
        print(f"Valid regression columns: {valid_reg_columns}")
        if valid_reg_columns:
            self.valid_reg_columns = valid_reg_columns
            self.eval_set = (self.eval_set[0], self.eval_set[1][valid_reg_columns]) if self.eval_set is not None else None
            self.regressor = CatBoostRegressor(**self.reg_params, verbose=0, random_state=42, eval_fraction=self.eval_fraction)
            if self.scaler:
                print(f"✅ Fitting regressor with columns: {valid_reg_columns}")
                X_scaled = X.copy()
                y_scaled = y.copy()
                #X_scaled[self.scaled_cont_features] = self.scaler.transform(X_scaled[self.scaled_cont_features])
                #y_scaled[self.scaled_cont_features] = self.scaler.transform(y_scaled[self.scaled_cont_features])
                y_scaled = y_scaled.drop(columns=['THICKNESS_TULUS [mm]'], errors='ignore')
                self.regressor.fit(X_scaled, y_scaled[valid_reg_columns], cat_features=cat_features,eval_set=self.eval_set)
            else:
                print(f"✅ Fitting regressor with columns: {valid_reg_columns}")
                self.regressor.fit(X, y[valid_reg_columns], cat_features=cat_features,eval_set=self.eval_set)

        if skipped_reg:
            print(f"⚠️ Regressor: skipped columns with only one unique value: {skipped_reg}")

        # Fit classifier
        if cat_columns:
            valid_cat_columns = [col for col in cat_columns if X[col].nunique() > 1 and y[col].nunique() > 1]
            dropped_cat_columns = [col for col in cat_columns if col not in valid_cat_columns]

            if dropped_cat_columns:
                print(f"⚠️  Colonne escluse dal fit del classificatore (solo un valore unico): {dropped_cat_columns}")

            if valid_cat_columns:
                self.cat_columns = valid_cat_columns
                self.classifier = CatBoostClassifier(**self.clf_params, verbose=0, random_state=42, eval_fraction=self.eval_fraction)
                print(f"✅ Fitting classifier con colonne: {valid_cat_columns}")
                self.eval_set=(self.eval_set[0], self.eval_set[1][valid_cat_columns]) if self.eval_set is not None else None
                self.classifier.fit(X, y[valid_cat_columns].astype(str), cat_features=cat_features,eval_set=self.eval_set)
            else:
                print("⛔ Nessuna colonna valida per il classificatore. Fit saltato.")
        return self

    def predict(self, X,testing=False):
       
        preds = []
        if self.regressor:
            reg_pred = self.regressor.predict(X)
            if len(self.valid_reg_columns) == 1:
                reg_pred = reg_pred.reshape(-1, 1)
            
            if reg_pred.shape[1] < len(self.reg_columns):
                reg_pred_df = pd.DataFrame(reg_pred, columns=self.valid_reg_columns)
                reg_pred_df[[col for col in self.reg_columns if col not in self.valid_reg_columns]] = X[[col for col in self.reg_columns if col not in self.valid_reg_columns]].copy()
                reg_pred_df = reg_pred_df[self.reg_columns]  # Reorder columns to match reg_columns
                reg_pred = reg_pred_df.values
            preds.append(reg_pred)
        if self.classifier:
            cat_pred = self.classifier.predict(X)
            if len(self.cat_columns) == 1:
                cat_pred = cat_pred.reshape(-1, 1)
            preds.append(cat_pred)
        # Concatenate predictions (columns)
        if not self.classifier and testing:
            cat_pred=X[self.cat_columns].values if len(self.cat_columns) > 1 else X[self.cat_columns].values.reshape(-1, 1)
            preds.append(cat_pred)
        return np.concatenate(preds, axis=1)

    def score(self, X, y_true):
        score = None

        if self.regressor:
            X_scaled = X.copy()
            """
            if self.scaler:
                X_scaled[self.scaled_cont_features] = self.scaler.transform(
                    X_scaled[self.scaled_cont_features]
                )
            """
            print("Scoring regressor...")
            reg_pred = self.regressor.predict(X_scaled)
        
            y_true_eval = y_true[self.valid_reg_columns].values
            y_pred_eval = reg_pred
            if len(self.clf_params)==0:
                mismatch_percent=0.0
            else:
                nozzle_preds = self.classifier.predict(X_scaled).ravel()
                mismatch_percent = ((y_true['NOZZLE_TYPE'].astype(str).values != nozzle_preds.astype(str)).mean() * 100)
                print(f"NOZZLE mismatch: {mismatch_percent:.2f}%")
            # Calcolo dei normalized MAE
            y_pred_eval=round_values(pd.DataFrame(y_pred_eval, columns=self.valid_reg_columns), rounding_factors).values
            normalized_maes = []
            for i in range(y_pred_eval.shape[1]):
                mae_i = mean_absolute_error(y_true_eval[:, i], y_pred_eval[:, i])
                print(f"MAE for {pd.DataFrame(y_pred_eval, columns=self.valid_reg_columns).columns[i]}: {mae_i:.4f}")
                std_i = np.std(y_train[self.valid_reg_columns].values[:, i])
                if std_i < 1e-8:
                    std_i = 1.0  # fallback: evita divisione per 0
                normalized_maes.append(mae_i / std_i)

            mean_normalized_mae = np.mean(normalized_maes)
            score = mean_normalized_mae

        return -score  # Per usarlo come punteggio da massimizzare in CV

"""
if len(y_train['NOZZLE_TYPE'].unique()) > 1:
    
    print("Tuning classifier parameters...")
    clf_param_grid = {
        "iterations": [ 200, 500, 800],
        "learning_rate": [0.0005,0.001,0.01, 0.05],
        "depth": [4, 6, 8, 12, 15],
        "l2_leaf_reg": [1, 3, 5],
        "bagging_temperature": [0, 0.5, 1,2],
        "random_strength": [1,2,3],
        "early_stopping_rounds": [5, 10, 20],
    }
    n_iter= 100
    best_clf_score = -np.inf
    best_clf_params = None
    all_params = list(ParameterGrid(clf_param_grid))
    sampled_params = random.sample(all_params, min(n_iter, len(all_params)))
    for i, params in enumerate(sampled_params):
        print(f"Random Search clf iteration {i+1}/{len(sampled_params)} - params: {params}")
        print(f"Current clf params: {params}")
        fold_scores = []
        kf = StratifiedKFold(n_splits=4, shuffle=True, random_state=42)

        for train_idx, val_idx in kf.split(X_train,y_train['NOZZLE_TYPE'].astype(str)):
            X_train_fold = X_train.iloc[train_idx]
            y_train_fold = y_train.iloc[train_idx]
            X_val_fold = X_train.iloc[val_idx]
            y_val_fold = y_train.iloc[val_idx]
           
            # Fit only classifier
            model = MultiOutputCatBoostWrapper(reg_params={}, clf_params=params,eval_set=(X_val_fold, y_val_fold), eval_fraction=0.0)
            model.fit(X_train_fold, y_train_fold,
                    reg_columns=[],  # no regression columns here
                    cat_columns=['NOZZLE_TYPE'],
                    cat_features=cat_features)

            # Classifier predict and score
            preds = model.classifier.predict(X_val_fold)
            preds=np.array(preds).ravel()
            print(f"preds len:{preds.shape} \n{preds}")
            print(f"real: \n{y_val_fold["NOZZLE_TYPE"].astype(str).values}")
            acc = np.mean(preds == y_val_fold["NOZZLE_TYPE"].astype(str).values)
            fold_scores.append(acc)

        avg_score = np.mean(fold_scores)
        print(f"  Avg classifier accuracy: {avg_score:.4f}")

        if avg_score > best_clf_score:
            best_clf_score = avg_score
            best_clf_params = params
            print("  New best classifier params found!")

    print(f"Best classifier params: {best_clf_params}")
    print(f"Best classifier accuracy: {best_clf_score:.4f}")
    best_params_catboost_original['full']['class'] = best_clf_params
    with open(f"catboost2_best_params/best_catboost_original_clf_RandomSearchCV.json", "w") as f:
        json.dump(best_clf_params, f, indent=4)
  
 

# --- tuning regressor second ---
print("Tuning regressor parameters...")
reg_param_grid = {
    'iterations': [200,500, 800, 1000],
    'learning_rate': [0.0005, 0.001, 0.01, 0.05,],
    'depth': [4, 6, 8, 12, 15],
    'l2_leaf_reg': [1, 3, 5],
    'bagging_temperature': [0, 0.5, 2],
    'random_strength': [1, 2, 3],
    'loss_function': ['MultiRMSE'],
    'eval_metric': ['MultiRMSE'],
    'early_stopping_rounds': [5,10,20],
}

best_reg_score = -np.inf  # MAE lower is better
best_reg_params = None
n_iter=200
all_params = list(ParameterGrid(reg_param_grid))
sampled_params = random.sample(all_params, min(n_iter, len(all_params)))
scaler= StandardScaler()
scaler.fit(X_train[cont_features].copy())

for params in tqdm(sampled_params,total=len(sampled_params), desc="Tuning regressor params"):
    print(f"Current reg params: {params}")
    fold_scores = []
    kf = StratifiedKFold(n_splits=4, shuffle=True, random_state=42)

    for train_idx, val_idx in kf.split(X_train,y_train['NOZZLE_TYPE'].astype(str)):
        X_train_fold = X_train.iloc[train_idx]
        y_train_fold = y_train.iloc[train_idx]
        X_val_fold = X_train.iloc[val_idx]
        y_val_fold = y_train.iloc[val_idx]

        fold_model = MultiOutputCatBoostWrapper(reg_params=params, clf_params={},scaler=scaler,scaled_cont_features=cont_features,eval_fraction=0.0,eval_set=(X_val_fold,y_val_fold))
        fold_model.fit(X_train_fold, y_train_fold,
                  reg_columns=[col for col in cont_features if col != 'THICKNESS_TULUS [mm]'],
                  cat_columns=[],
                  cat_features=cat_features)
        
        # Score on validation set
        fold_score = fold_model.score(X_val_fold, y_val_fold)
        fold_scores.append(fold_score)
    # Average score across all folds
    avg_score = np.mean(fold_scores)
    print(f"  Average CV MAE score: {-avg_score:.4f}")
    
    # Update best model if we found a better one
    if avg_score > best_reg_score:
        best_reg_score = avg_score
        best_reg_params = params
        print(f"  New best model found!")
        print(f"  Best params so far: {best_reg_params}")
        print(f"  Average CV MAE score: {-avg_score:.4f}")
        

print(f"Best CV MAE score: {-best_reg_score:.4f}")
print(f"Best params: {best_reg_params}")
best_params_catboost_original['full']['reg'] = best_reg_params
with open(f"catboost2_best_params/best_catboost_original_reg_RandomSearchCV.json", "w") as f:
    json.dump(best_reg_params, f, indent=4)
"""

# --- Train final model with best params ---

final_model = MultiOutputCatBoostWrapper(reg_params=best_params_catboost_original['full']['reg'], clf_params= ({} if len(y_train['NOZZLE_TYPE'].unique()) < 2 else best_params_catboost_original['full']['class']),)

final_model.fit(X_train, y_train,
                reg_columns=[col for col in cont_features if col != 'THICKNESS_TULUS [mm]'],
                cat_columns=['NOZZLE_TYPE'],
                cat_features=cat_features)
print("Training score....")
mae=final_model.score(X_train, y_train)
print("Test score....")
mae=final_model.score(X_test, y_test)
print("----------------------------------------------------------------------------------------------------------")
print(f"MAE on test set: {-mae:.4f}")

evals_result=final_model.regressor.get_evals_result()
plt.figure(figsize=(10, 6))
plt.plot(evals_result['learn']['MultiRMSE'], label='Train MultiRMSE')
plt.plot(evals_result['validation']['MultiRMSE'], label='Validation MultiRMSE')
plt.xlabel("Iterations")
plt.ylabel("MultiRMSE")
plt.title("CatBoost MultiRMSE over Iterations")
plt.legend()
plt.grid(True)
plt.savefig(f'catboost2_output/original_multi_rmse_loss.png')

print(f"X_test: {X_test}")
n_iter=10
# Predict on test
for i in range(n_iter):
    X_test = X_test.copy().astype({'THICKNESS_TULUS [mm]': 'float64',
                                    'CONTOUR_SPEED [mm/min]': 'float64',
                                    'LASER_POWER [W]': 'float64',
                                    'CONTOUR_GAS_PRESSURE [bar]': 'float64',
                                    'CONTOUR_NOZZLE_DISTANCE [mm]': 'float64',
                                    'CONTOUR_FOCAL [mm]': 'float64'})
    y_pred = pd.DataFrame(final_model.predict(X_test,testing=True),columns=(list(final_model.reg_columns) + list(final_model.cat_columns)))
    y_pred=y_pred.copy().astype({
                                        'CONTOUR_SPEED [mm/min]': 'float64',
                                        'LASER_POWER [W]': 'float64',
                                        'CONTOUR_GAS_PRESSURE [bar]': 'float64',
                                        'CONTOUR_NOZZLE_DISTANCE [mm]': 'float64',
                                        'CONTOUR_FOCAL [mm]': 'float64'})
    y_pred['THICKNESS_TULUS [mm]'] = X_test['THICKNESS_TULUS [mm]'].copy()
    y_pred['MATERIAL_NAME_TULUS'] = X_test['MATERIAL_NAME_TULUS'].copy()

    # Round the continuous parameters according to rounding factors
    for col, factor in rounding_factors.items():
        if col in y_pred.columns:
            y_pred[col] = (y_pred[col] / factor).round() * factor

    print(f"\n{n_iter+1} round...\n")
    y_pred["DEFECT_TYPE"] = X_test["DEFECT_TYPE"].copy()
    y_pred=y_pred[X_test.columns]
    mae=final_model.score(y_pred, y_test)
    print(f"MAE on test set {n_iter+1} round: {-mae:.4f}")
    X_test = y_pred.copy()

#y_pred['DEFECT_TYPE'] = defect
y_pred.to_excel(f'catboost2_output/original_predictions.xlsx', index=False)
#print(defect)
print(len(X_test))
print(len(X_train))
