import time
import joblib
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, 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

class Engine:
    def __init__(self, df, target_col, time_budget=300, metric=None):
        try: 
            import mlflow
            mlflow.autolog(disable=True)
        except: pass
        
        self.audit = AuditTrail()
        # Remove duplicatas logo na entrada
        self.raw_df = df.loc[:, ~df.columns.duplicated()]
        self.target_col = target_col
        self.time_budget = time_budget
        self.metric = metric
        
        self.scout = Scout(self.audit, target_col).analyze(self.raw_df)
        self.df = self.scout.optimize_memory(self.raw_df)
        self.forge = Forge(self.scout)
        self.best_pipeline = None

    def run(self):
        X = self.df.drop(columns=[self.target_col] + self.scout.dropped_cols)
        y = self.df[self.target_col]
        
        if self.scout.problem_type != "regression" and not pd.api.types.is_numeric_dtype(y):
            self.le = LabelEncoder()
            y = self.le.fit_transform(y)
            
        stratify = y if self.scout.problem_type != "regression" else None
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=stratify)
        
        study = optuna.create_study(direction="maximize")
        study.optimize(lambda t: self._objective(t, X_train, y_train), timeout=self.time_budget, show_progress_bar=False)
        
        self._finalize(study.best_params, X_train, y_train)
        self._audit_final(X_test, y_test, X_train, y_train)
        self.audit.generate_report()
        return self.best_pipeline

    def _objective(self, trial, X, y):
        algo = trial.suggest_categorical("algo", ["LGBM", "XGB"])
        if algo == "LGBM":
            params = {
                "n_estimators": trial.suggest_int("n_est", 50, 500),
                "learning_rate": trial.suggest_float("lr", 0.01, 0.3),
                "num_leaves": trial.suggest_int("leaves", 20, 100),
                "verbose": -1
            }
            Model = LGBMRegressor if self.scout.problem_type == "regression" else LGBMClassifier
        else:
            params = {
                "n_estimators": trial.suggest_int("n_est", 50, 500),
                "learning_rate": trial.suggest_float("lr", 0.01, 0.3),
                "max_depth": trial.suggest_int("depth", 3, 10),
                "verbosity": 0
            }
            Model = XGBRegressor if self.scout.problem_type == "regression" else XGBClassifier
        
        pipe = Pipeline([("pre", self.forge.build("tree")), ("model", Model(**params))])
        
        cv = KFold(3) if self.scout.problem_type == "regression" else StratifiedKFold(3)
        scr = self.metric if self.metric else ("r2" if self.scout.problem_type == "regression" else "roc_auc")
        
        scores = cross_validate(pipe, X, y, cv=cv, scoring=scr, n_jobs=1)
        mean = scores["test_score"].mean()
        self.audit.add_model_result(algo, params, {"cv_mean": mean, "metric": scr}, 0)
        return mean

    def _finalize(self, params, X, y):
        algo = params.pop("algo")
        if algo == "LGBM":
            Model = LGBMRegressor if self.scout.problem_type == "regression" else LGBMClassifier
            p = {"n_estimators": params["n_est"], "learning_rate": params["lr"], "num_leaves": params["leaves"], "verbose": -1}
        else:
            Model = XGBRegressor if self.scout.problem_type == "regression" else XGBClassifier
            p = {"n_estimators": params["n_est"], "learning_rate": params["lr"], "max_depth": params["depth"], "verbosity": 0}
            
        self.best_pipeline = Pipeline([("pre", self.forge.build("tree")), ("model", Model(**p))])
        self.best_pipeline.fit(X, y)
        joblib.dump(self.best_pipeline, f"{self.audit.exp_name}.pkl")

    def _audit_final(self, X_test, y_test, X_train, y_train):
        try:
            if self.scout.problem_type != "regression":
                self.audit.generate_risk_report(self.best_pipeline, X_train, y_train, X_test, y_test)
            
            # SHAP
            model = self.best_pipeline.named_steps["model"]
            pre = self.best_pipeline.named_steps["pre"]
            bg = shap.utils.sample(pre.transform(X_test), 50)
            explainer = shap.TreeExplainer(model)
            shap_vals = explainer.shap_values(bg)
            plt.figure()
            vals = shap_vals[1] if isinstance(shap_vals, list) else shap_vals
            shap.summary_plot(vals, bg, show=False)
            plt.savefig(f"{self.audit.exp_name}_shap.png")
            plt.close()
        except Exception as e: print(f"Audit error: {e}")
