import time
import joblib
import logging
import optuna
import shap
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import (
    cross_validate,
    StratifiedKFold,
    KFold,
    train_test_split,
)
from sklearn.metrics import (
    roc_auc_score,
    mean_squared_error,
    accuracy_score,
    classification_report,
    confusion_matrix,
    r2_score,
)
from sklearn.dummy import DummyClassifier, DummyRegressor
from lightgbm import LGBMClassifier, LGBMRegressor
from xgboost import XGBClassifier, XGBRegressor

from .audit import AuditTrail
from .scout import Scout
from .forge import Forge

logger = logging.getLogger("ScoutML.Engine")


class Engine:
    def __init__(
        self,
        df,
        target_col,
        time_budget=300,
        correlation_threshold=0.95,
        metric=None,  # <--- NOVO: O usuário agora decide a métrica
    ):
        self.audit = AuditTrail()
        self.raw_df = df
        self.target_col = target_col
        self.time_budget = time_budget
        self.correlation_threshold = correlation_threshold
        self.custom_metric = metric  # Guarda a métrica escolhida

        # Fluxo de Inicialização
        self.scout = Scout(
            self.audit, target_col, threshold=self.correlation_threshold
        ).analyze(self.raw_df)

        self.df = self.scout.optimize_memory(self.raw_df)
        self.forge = Forge(self.scout)
        self.device_params = self._detect_hardware()

        self.best_pipeline = None
        self.label_encoder = None

    def _detect_hardware(self):
        """Verifica disponibilidade de GPU para acelerar XGBoost/LightGBM."""
        try:
            import torch

            if torch.cuda.is_available():
                logger.info("🚀 GPU Detectada. Configurando Boosters para modo GPU...")
                return {"device": "gpu", "n_jobs": -1}
        except ImportError:
            pass
        return {"device": "cpu", "n_jobs": -1}

    def _get_cv_strategy(self):
        if self.scout.problem_type == "regression":
            return KFold(n_splits=5, shuffle=True, random_state=42)
        else:
            return StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

    def _get_scoring_metric(self):
        """Define a métrica de otimização."""
        # 1. Se o usuário definiu (ex: 'recall', 'f1', 'precision'), usa ela!
        if self.custom_metric:
            return self.custom_metric

        # 2. Caso contrário, usa o padrão robusto
        if self.scout.problem_type == "regression":
            return "neg_root_mean_squared_error"
        elif self.scout.problem_type == "binary":
            return "roc_auc"
        else:
            return "f1_weighted"

    def run(self):
        # 1. Preparação dos Dados
        X = self.df.drop(columns=[self.target_col] + self.scout.dropped_cols)
        y = self.df[self.target_col]

        # Tratamento de Target Categórico (Ex: 'Sim'/'Não' -> 1/0)
        if self.scout.problem_type in ["binary", "multiclass"]:
            if not pd.api.types.is_numeric_dtype(y):
                logger.info(f"🔤 Codificando Target não-numérico ({y.dtype})...")
                self.label_encoder = LabelEncoder()
                y = self.label_encoder.fit_transform(y)

                # Log do mapeamento (Seguro para JSON)
                try:
                    mapping = {
                        str(k): int(v)
                        for k, v in zip(
                            self.label_encoder.classes_,
                            range(len(self.label_encoder.classes_)),
                        )
                    }
                    self.audit.log_action(
                        "transformation", f"Target Encoded: {mapping}"
                    )
                except:
                    pass

        # Split Treino/Teste (Holdout)
        stratify = y if self.scout.problem_type != "regression" else None

        # CORREÇÃO: Precisamos guardar X_train e y_train para o relatório de PSI/KS
        X_train, X_holdout, y_train, y_holdout = train_test_split(
            X, y, test_size=0.2, random_state=42, stratify=stratify
        )

        # 2. Executa Baseline
        self._run_baseline(X_train, y_train)

        # 3. Definição da Função Objetivo (Optuna)
        def objective(trial):
            algo = trial.suggest_categorical("algo", ["LightGBM", "XGBoost"])

            # Espaço de Hiperparâmetros
            if algo == "LightGBM":
                params = {
                    "n_estimators": trial.suggest_int("n_estimators", 100, 1000),
                    "learning_rate": trial.suggest_float(
                        "learning_rate", 1e-3, 0.3, log=True
                    ),
                    "num_leaves": trial.suggest_int("num_leaves", 20, 100),
                    "reg_alpha": trial.suggest_float("reg_alpha", 1e-3, 10.0, log=True),
                    "reg_lambda": trial.suggest_float(
                        "reg_lambda", 1e-3, 10.0, log=True
                    ),
                    "n_jobs": 4,
                    "verbose": -1,
                }
                if self.scout.problem_type == "regression":
                    model = LGBMRegressor(**params)
                else:
                    if self.scout.is_imbalanced:
                        params["class_weight"] = "balanced"
                    model = LGBMClassifier(**params)

            elif algo == "XGBoost":
                params = {
                    "n_estimators": trial.suggest_int("n_estimators", 100, 1000),
                    "learning_rate": trial.suggest_float(
                        "learning_rate", 1e-3, 0.3, log=True
                    ),
                    "max_depth": trial.suggest_int("max_depth", 3, 12),
                    "subsample": trial.suggest_float("subsample", 0.5, 1.0),
                    "colsample_bytree": trial.suggest_float(
                        "colsample_bytree", 0.5, 1.0
                    ),
                    "n_jobs": 4,
                    "verbosity": 0,
                }
                if self.device_params["device"] == "gpu":
                    params["tree_method"] = "gpu_hist"

                if self.scout.problem_type == "regression":
                    model = XGBRegressor(**params)
                else:
                    model = XGBClassifier(**params)

            # Pipeline
            preprocessor = self.forge.build(model_type="tree")
            pipeline = Pipeline([("preprocessor", preprocessor), ("model", model)])

            cv = self._get_cv_strategy()
            scoring_metric = self._get_scoring_metric()  # Usa a métrica escolhida

            start_time = time.time()
            scores = cross_validate(
                pipeline,
                X_train,
                y_train,
                cv=cv,
                scoring=scoring_metric,
                n_jobs=1,
                return_estimator=False,
            )
            elapsed = time.time() - start_time

            mean_score = scores["test_score"].mean()
            std_score = scores["test_score"].std()

            self.audit.add_model_result(
                model_name=algo,
                params=trial.params,
                metrics={
                    "cv_mean": mean_score,
                    "cv_std": std_score,
                    "metric": scoring_metric,
                },
                train_time=elapsed,
            )
            return mean_score

        logger.info(
            f"🔥 [Engine] Iniciando Otimização (Métrica: {self._get_scoring_metric()})"
        )

        # Otimização
        study = optuna.create_study(
            direction="maximize"
        )  # Scikit-learn scores são "quanto maior melhor" (exceto erros negativos, que também sobem para 0)
        study.optimize(objective, timeout=self.time_budget, show_progress_bar=False)

        logger.info("🏆 [Engine] Treinando modelo final...")
        self._finalize_model(study.best_params, X_train, y_train)

        # --- CORREÇÃO CRÍTICA: Passar X_train/y_train para gerar PSI e KS ---
        self._audit_final_model(X_holdout, y_holdout, X_train, y_train)

        self.audit.generate_report()
        return self.best_pipeline

    def _run_baseline(self, X, y):
        metric = self._get_scoring_metric()
        if self.scout.problem_type == "regression":
            dummy = DummyRegressor(strategy="mean")
        else:
            dummy = DummyClassifier(strategy="most_frequent")

        cv = self._get_cv_strategy()
        scores = cross_validate(dummy, X, y, cv=cv, scoring=metric)
        self.audit.add_model_result(
            "Baseline (Dummy)",
            {},
            {
                "cv_mean": scores["test_score"].mean(),
                "cv_std": scores["test_score"].std(),
                "metric": metric,
            },
            0.1,
        )

    def _finalize_model(self, best_params, X, y):
        algo = best_params.pop("algo")
        # Recria o modelo com os melhores parâmetros
        if algo == "LightGBM":
            ModelClass = (
                LGBMRegressor
                if self.scout.problem_type == "regression"
                else LGBMClassifier
            )
            best_params["verbose"] = -1
            if self.scout.problem_type != "regression" and self.scout.is_imbalanced:
                best_params["class_weight"] = "balanced"
        else:
            ModelClass = (
                XGBRegressor
                if self.scout.problem_type == "regression"
                else XGBClassifier
            )
            if self.device_params["device"] == "gpu":
                best_params["tree_method"] = "gpu_hist"

        model = ModelClass(**best_params)
        preprocessor = self.forge.build(model_type="tree")

        self.best_pipeline = Pipeline(
            [("preprocessor", preprocessor), ("model", model)]
        )
        self.best_pipeline.fit(X, y)

        # Serialização
        joblib.dump(self.best_pipeline, f"{self.audit.exp_name}_model.pkl")
        if self.label_encoder:
            joblib.dump(self.label_encoder, f"{self.audit.exp_name}_label_encoder.pkl")

    def _audit_final_model(self, X_test, y_test, X_train=None, y_train=None):
        """
        Executa auditoria completa.
        AGORA: Recebe X_train/y_train para calcular Estabilidade (PSI) e KS.
        """
        logger.info("🧪 [Audit] Gerando métricas de Holdout, SHAP e Risco...")

        y_pred = self.best_pipeline.predict(X_test)

        # 1. Métricas Padrão (Dependem do tipo de problema)
        if self.scout.problem_type != "regression":
            # Classificação
            if self.scout.problem_type == "binary":
                if hasattr(self.best_pipeline, "predict_proba"):
                    y_proba = self.best_pipeline.predict_proba(X_test)[:, 1]
                else:
                    y_proba = y_pred
                auc = roc_auc_score(y_test, y_proba)
            else:
                # Multiclass
                if hasattr(self.best_pipeline, "predict_proba"):
                    y_proba = self.best_pipeline.predict_proba(X_test)
                    auc = roc_auc_score(y_test, y_proba, multi_class="ovr")
                else:
                    auc = 0

            acc = accuracy_score(y_test, y_pred)
            report = classification_report(y_test, y_pred, output_dict=True)

            self.audit.metadata["final_metrics"] = {
                "holdout_auc": auc,
                "holdout_accuracy": acc,
                "classification_report": report,
            }

            # --- 2. AUDITORIA DE RISCO (KS, PSI, Matriz Confusão) ---
            # Só executamos se tivermos os dados de treino para comparação
            if X_train is not None and y_train is not None:
                try:
                    logger.info("📊 Calculando PSI e KS (Train vs Test)...")
                    # Chama o método do audit.py que gera os gráficos e tabelas
                    risk_metrics = self.audit.generate_risk_report(
                        model=self.best_pipeline,
                        X_train=X_train,
                        y_train=y_train,
                        X_test=X_test,
                        y_test=y_test,
                        save_plot=True,  # Garante que salva os PNGs
                    )
                    # IMPORTANTE: Salva no metadados para o relatório MD pegar depois
                    self.audit.metadata["risk_metrics"] = risk_metrics
                except Exception as e:
                    logger.warning(f"⚠️ Falha ao gerar relatório de Risco: {e}")
                    import traceback

                    traceback.print_exc()

        else:
            # Regressão
            mse = mean_squared_error(y_test, y_pred)
            rmse = np.sqrt(mse)
            r2 = r2_score(y_test, y_pred)
            self.audit.metadata["final_metrics"] = {
                "holdout_rmse": rmse,
                "holdout_r2": r2,
            }

        # 3. Explicabilidade (SHAP)
        try:
            model = self.best_pipeline.named_steps["model"]
            preprocessor = self.best_pipeline.named_steps["preprocessor"]
            X_transformed = preprocessor.transform(X_test)

            # Amostra para performance
            background = shap.utils.sample(X_transformed, 100)

            # TreeExplainer é o ideal para XGB/LGBM
            explainer = shap.TreeExplainer(model)
            try:
                shap_values = explainer.shap_values(background)
            except:
                shap_values = explainer(background).values

            plt.figure(figsize=(10, 6))

            # Ajuste para formato do SHAP em classificação binária
            if isinstance(shap_values, list) and len(shap_values) == 2:
                vals = shap_values[1]  # Classe positiva
            else:
                vals = shap_values

            shap.summary_plot(vals, background, show=False)
            plt.savefig(f"{self.audit.exp_name}_shap_summary.png", bbox_inches="tight")
            plt.close()
            logger.info("🧠 SHAP Plot salvo.")

        except Exception as e:
            logger.warning(f"⚠️ Aviso SHAP: {e}")
