"""Module de validation croisée pour les modèles de machine learning.
Ce module contient des fonctions pour évaluer les modèles avec validation croisée, en calculant des métriques telles que MAE, RMSE et R², et en retournant la moyenne et l'écart-type des scores sur les folds.
Il inclut également une fonction spécifique pour la validation croisée des modèles CatBoost, qui gère les features catégorielles de manière appropriée.
"""
from sklearn.model_selection import cross_validate
from sklearn.model_selection import cross_val_score
import numpy as np

from src.modeling.evaluate import evaluate_regression_model

def cross_validate_model(pipeline, X_train, y_train, cv=5):
    """
    Évalue un modèle avec validation croisée.
    Le modèle évalue les métriques sur l'échelle log :
    - MAE
    - RMSE
    - R²
    - Retourne la moyenne et l'écart-type des scores sur les folds.
    Args:
        pipeline : modèle avec préprocessing intégré (sklearn.pipeline.Pipeline)
        X_train : pd.DataFrame, features d'entraînement
        y_train : pd.Series ou np.array, target
        cv : int, nombre de folds

    Returns:
        dict: moyenne et std des scores
    """
    try:
        scoring = {
            "mae": "neg_mean_absolute_error",
            "rmse": "neg_root_mean_squared_error",
            "r2": "r2",
        }

        cv_results = cross_validate(pipeline, X_train, y_train, cv=cv, scoring=scoring)

        return {
            "mae_mean": -np.mean(cv_results["test_mae"]),
            "mae_std": np.std(cv_results["test_mae"]),
            "rmse_mean": -np.mean(cv_results["test_rmse"]),
            "rmse_std": np.std(cv_results["test_rmse"]),
            "r2_mean": np.mean(cv_results["test_r2"]),
            "r2_std": np.std(cv_results["test_r2"]),
        }

    except Exception as e:
        print(f"Erreur lors de la validation croisée : {e}")
        return None
 

def cross_validate_catboost(model, X_train, y_train, cv=5, random_state=42):
    try:
        from sklearn.model_selection import KFold

        categorical_features = X_train.select_dtypes(include=["object", "category"]).columns.tolist()

        kf = KFold(n_splits=cv, shuffle=True, random_state=random_state)

        mae_scores = []
        rmse_scores = []
        r2_scores = []

        for train_idx, val_idx in kf.split(X_train):
            X_tr = X_train.iloc[train_idx]
            X_val = X_train.iloc[val_idx]
            y_tr = y_train.iloc[train_idx]
            y_val = y_train.iloc[val_idx]

            model_fold = model.copy()

            model_fold.fit(
                X_tr,
                y_tr,
                cat_features=categorical_features,
                verbose=0
            )

            y_pred = model_fold.predict(X_val)

            metrics = evaluate_regression_model(y_val, y_pred)

            mae_scores.append(metrics["mae_log"])
            rmse_scores.append(metrics["rmse_log"])
            r2_scores.append(metrics["r2_log"])

        return {
            "mae_mean": np.mean(mae_scores),
            "mae_std": np.std(mae_scores),
            "rmse_mean": np.mean(rmse_scores),
            "rmse_std": np.std(rmse_scores),
            "r2_mean": np.mean(r2_scores),
            "r2_std": np.std(r2_scores),
        }

    except Exception as e:
        print(f"Erreur CV CatBoost : {e}")
        return None