
from sklearn.neural_network import MLPRegressor
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.metrics import mean_absolute_error, mean_squared_error, silhouette_score,r2_score
import numpy as np
import pandas as pd


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
}

def mean_mae(y_val, y_pred, y_train,scaler,cont_features):
    """
    Calculate normalized mean absolute error
    
    Parameters:
    -----------
    y_true: DataFrame/array - Actual target values
    y_pred: array - Predicted values
    y_train: DataFrame/array - Training data (for std calculation)
    
    Returns:
    --------
    float: Negative mean normalized MAE
    """

    normalized_maes = []
    y_val_orig = y_val.copy()
    y_pred_df = pd.DataFrame(
        y_pred, 
        columns=[col for col in y_val_orig.columns if col != 'THICKNESS_TULUS [mm]']
    ).reset_index(drop=True)

    thickness = y_val_orig['THICKNESS_TULUS [mm]'].copy()
    y_pred_df['THICKNESS_TULUS [mm]'] = thickness

    # Inverse transform
    y_val_inv = y_val_orig.copy()
    y_pred_inv = y_pred_df.copy()
    y_val_inv[cont_features] = scaler.inverse_transform(y_val_inv[cont_features])
    y_pred_inv[cont_features] = scaler.inverse_transform(y_pred_inv[cont_features])

    # Drop target thickness
    y_val_eval = y_val_inv.drop(columns=['THICKNESS_TULUS [mm]'])
    y_pred_eval = y_pred_inv.drop(columns=['THICKNESS_TULUS [mm]'])
    print("y_val_eval columns:", y_val_eval.columns)
    print("y_pred_eval columns:", y_pred_eval.columns)
    # Calcola mismatch % per NOZZLE_TYPE
    nozzles = [col for col in y_val_eval.columns if col.startswith('NOZZLE_TYPE_')]
    if nozzles:
        # Indici nozzle veri e predetti
        true_nozzle_idx = y_val_eval[nozzles].values.argmax(axis=1)
        pred_nozzle_idx = y_pred_eval[nozzles].values.argmax(axis=1)

        # Trova gli indici dei mismatch
        mismatch_mask = true_nozzle_idx != pred_nozzle_idx
        print(f"Mismatch mask: {mismatch_mask}")
        mismatch_indices = np.where(mismatch_mask)[0]

        # Stampa righi mismatchati
        print("Righe con mismatch tra nozzle predetto e vero:")
        for idx in mismatch_indices:
            true_class = nozzles[true_nozzle_idx[idx]]
            pred_class = nozzles[pred_nozzle_idx[idx]]
            print(f"Riga {idx}: VERO = {true_class}, PRED = {pred_class}")
       
        true_nozzle_idx = y_val_eval[nozzles].values.argmax(axis=1)
        pred_nozzle_idx = y_pred_eval[nozzles].values.argmax(axis=1)
        mismatch_pct = np.mean(true_nozzle_idx != pred_nozzle_idx) * 100
        print(f"NOZZLE_TYPE mismatch: {mismatch_pct:.2f}%")
    y_pred_eval=round_values(y_pred_eval.copy(), rounding_factors)
    # Calcolo MAE normalizzato
    for i in range(y_pred_eval.shape[1]):
       
        mae_i = mean_absolute_error(y_val_eval.iloc[:, i], y_pred_eval.iloc[:, i])
        std_i = np.std(y_train[y_val_eval.columns].values[:, i])
        norm_mae_i = mae_i / std_i if std_i != 0 else mae_i
        print(f"MAE for {y_val_eval.columns[i]}: {mae_i:.4f}")
        normalized_maes.append(norm_mae_i)

    mean_normalized_mae = np.mean(normalized_maes)
    return -mean_normalized_mae  # negativo per ottimizzazione

class MultiOutputMLPWrapper(BaseEstimator, RegressorMixin):
    def __init__(self, hidden_layer_sizes=(100,), activation='relu', 
                 solver='adam', alpha=0.0001, batch_size='auto',
                 learning_rate='constant', learning_rate_init=0.001,
                 max_iter=500, shuffle=True, random_state=42, 
                 early_stopping=False, validation_fraction=0.2,
                 n_iter_no_change=10,momentum=0.9,nesterovs_momentum=True):
        self.hidden_layer_sizes = hidden_layer_sizes
        self.activation = activation
        self.solver = solver
        self.alpha = alpha
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.learning_rate_init = learning_rate_init
        self.max_iter = max_iter
        self.shuffle = shuffle
        self.random_state = random_state
        self.early_stopping = early_stopping
        self.validation_fraction = validation_fraction
        self.n_iter_no_change = n_iter_no_change
        self.momentum = momentum
        self.nesterovs_momentum = nesterovs_momentum
        self.model = None
       
        
    def fit(self, X, y):
        # Create a multi-output MLP regressor
        self.model = MLPRegressor(
            hidden_layer_sizes=self.hidden_layer_sizes,
            activation=self.activation,
            solver=self.solver,
            alpha=self.alpha,
            batch_size=self.batch_size,
            learning_rate=self.learning_rate,
            learning_rate_init=self.learning_rate_init,
            max_iter=self.max_iter,
            shuffle=self.shuffle,
            random_state=self.random_state,
            early_stopping=self.early_stopping,
            validation_fraction=self.validation_fraction,
            n_iter_no_change=self.n_iter_no_change
        )
        
        # Fit the model to our data
        self.model.fit(X, y)
        return self
    
    def predict(self, X):
        return self.model.predict(X)
    
    def score(self, X_val, y_val, y_train,scaler,cont_features):
        y_pred = self.predict(X_val)
        return mean_mae(y_val, y_pred, y_train,scaler,cont_features)