import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import logging
import os
import datetime
import json
from sklearn.metrics import confusion_matrix, classification_report

# Estilo dos gráficos
plt.style.use("ggplot")


class AuditTrail:
    def __init__(self, exp_name=None):
        # Define o nome do experimento
        if exp_name:
            self.exp_name = exp_name
        else:
            self.exp_name = f"run_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"

        # --- ESTRUTURAS LEGADAS (Mantidas para compatibilidade) ---
        self.log = []
        self.metadata = {}  # Metadados do dataset
        self.reports = {}
        self.confusion_matrix_data = None
        self.dropped_columns = []
        self.transformations = (
            []
        )  # Legado, mas mantido caso o Engine antigo escreva aqui
        self.data_health = {}

        # --- NOVAS ESTRUTURAS V2.0 (DataFrames Ricos) ---
        # Catálogo de features com estatísticas avançadas (Skew, Kurtosis, Missing) do Scout V2
        self.feature_catalog = pd.DataFrame()

        # Log de decisões estruturadas do pipeline (Ex: Por que RobustScaler?) do Forge V2
        self.pipeline_decisions = []

        # Leaderboard detalhado para cada trial do Optuna
        self.model_results = []

    # --- MÉTODOS DE LOGGING BÁSICOS (Compatibilidade com Engine v1) ---

    def log_event(self, step, message, details=None):
        """Registra um evento no log geral de texto."""
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        if details:
            entry = f"[{timestamp}] [{step}] {message} | Details: {details}"
        else:
            entry = f"[{timestamp}] [{step}] {message}"
        self.log.append(entry)
        # Opcional: printar no console para debug em tempo real
        # print(entry)

    def log_action(self, step, message, details=None):
        """Alias para log_event (usado extensivamente no código legado)."""
        self.log_event(step, message, details)

    # --- NOVOS MÉTODOS V2.0 (Integração com Scout/Forge aprimorados) ---

    def register_feature_stats(self, df_stats: pd.DataFrame):
        """
        Recebe o DataFrame de estatísticas detalhadas do Scout V2.
        Espera colunas como: column, missing_pct, skewness, kurtosis, outliers_pct, recommendation
        """
        self.feature_catalog = df_stats
        self.log_event(
            "Scout_V2", "Catálogo de Features estatísticas registrado com sucesso."
        )

    def log_pipeline_decision(self, feature, action, reason, parameters=None):
        """
        Registra decisões lógicas do Forge V2 de forma estruturada.
        """
        decision_entry = {
            "timestamp": datetime.datetime.now().strftime("%H:%M:%S"),
            "feature": str(feature),
            "action": action,
            "reason": reason,
            "parameters": json.dumps(parameters or {}),
        }
        self.pipeline_decisions.append(decision_entry)
        # Também loga no texto para garantir redundância
        self.log_event(
            "Forge_V2", f"Action: {action} on {feature}", f"Reason: {reason}"
        )

    # --- MÉTODOS DE RASTREABILIDADE LEGADOS (Mantidos para não quebrar Engine) ---

    def log_transformation(self, trans_type, columns, method, details=""):
        """Legado: Registra transformações simples."""
        self.transformations.append(
            {
                "step": trans_type,
                "columns": str(columns),
                "method": method,
                "details": details,
            }
        )
        self.log_event("Transformation", f"{trans_type} aplicado via {method}", details)

    def log_dropped_column(self, column, reason):
        """Legado: Registra colunas removidas."""
        self.dropped_columns.append({"column": column, "reason": reason})
        self.log_event("Drop", f"Coluna removida: {column}", reason)

    def log_data_health(self, initial_shape, missing_stats, duplicates):
        """Legado: Registra saúde inicial."""
        self.data_health = {
            "initial_rows": initial_shape[0],
            "initial_cols": initial_shape[1],
            "duplicates_removed": duplicates,
            "missing_values_summary": missing_stats,
        }

    # --- INTEGRAÇÃO COM MODELAGEM (Engine) ---

    def add_model_result(self, model_name, params, metrics, train_time):
        """
        Registra o resultado de uma rodada de treino.
        Atualizado para garantir que salva tudo no self.model_results
        """
        self.model_results.append(
            {
                "algo": model_name,
                "params": str(params),  # String para não quebrar Markdown
                "cv_mean": metrics.get("cv_mean", 0),
                "cv_std": metrics.get("cv_std", 0),
                "metric_name": metrics.get("metric", "score"),
                "duration": round(train_time, 4),
                "timestamp": datetime.datetime.now().strftime("%H:%M:%S"),
            }
        )

    def log_confusion_matrix(self, y_true, y_pred):
        """Salva e plota a matriz de confusão."""
        cm = confusion_matrix(y_true, y_pred)
        self.confusion_matrix_data = cm

        plt.figure(figsize=(6, 5))
        sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", cbar=False)
        plt.title(f"Confusion Matrix - {self.exp_name}")
        plt.ylabel("Real")
        plt.xlabel("Previsto")
        plt.tight_layout()
        plt.show()
        filename = f"{self.exp_name}_confusion_matrix.png"
        plt.savefig(filename)
        plt.close()
        print(f"📉 Matriz de Confusão salva: {filename}")

    # --- RELATÓRIO FINAL (O Grande Consolidador) ---

    # No arquivo scoutml/audit.py, substitua APENAS o método generate_report por este:

    def generate_report(self, filename="scout_audit_report.md"):
        """
        Gera o relatório final em Markdown combinando V1, V2 e Auditoria de Risco.
        """
        filename = f"{self.exp_name}_audit_report.md"

        with open(filename, "w", encoding="utf-8") as f:
            # --- CABEÇALHO ---
            f.write(f"# 📊 ScoutML Enterprise Audit Report\n")
            f.write(f"**Experiment Name:** `{self.exp_name}`\n")
            f.write(
                f"**Generated at:** {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
            )
            f.write(f"**Target Column:** `{self.metadata.get('target', 'N/A')}`\n\n")

            # --- SEÇÃO 1: Data Profiling (Scout V2) ---
            f.write("## 1. Data Integrity & Profiling (Scout V2) 🕵️\n")
            if not self.feature_catalog.empty:
                f.write(
                    "Estatísticas avançadas calculadas (Skewness, Kurtosis, Outliers via IQR):\n\n"
                )
                cols_to_show = [
                    "column",
                    "missing_pct",
                    "skewness",
                    "outliers_pct",
                    "recommendation",
                ]
                valid_cols = [
                    c for c in cols_to_show if c in self.feature_catalog.columns
                ]
                f.write(self.feature_catalog[valid_cols].to_markdown(index=False))
            elif self.data_health:
                f.write(f"- Initial Rows: {self.data_health.get('initial_rows')}\n")
            else:
                f.write("*Nenhuma estatística detalhada disponível.*\n")
            f.write("\n\n")

            # --- SEÇÃO 2: Preprocessing Decisions (Forge V2) ---
            f.write("## 2. Pipeline Decisions (Forge V2) 🔧\n")
            if self.pipeline_decisions:
                df_decisions = pd.DataFrame(self.pipeline_decisions)
                f.write(df_decisions.to_markdown(index=False))
            else:
                f.write("*Nenhuma transformação complexa registrada.*\n")
            f.write("\n\n")

            # --- SEÇÃO 3: Features Removidas ---
            f.write("## 3. Dropped Features 🗑️\n")
            if self.dropped_columns:
                df_drop = pd.DataFrame(self.dropped_columns)
                f.write(df_drop.to_markdown(index=False))
            else:
                f.write("*Nenhuma coluna removida explicitamente.*\n")
            f.write("\n\n")

            # --- SEÇÃO 4: Model Leaderboard ---
            f.write("## 4. Model Leaderboard 🏆\n")
            if self.model_results:
                df_results = pd.DataFrame(self.model_results)
                if "cv_mean" in df_results.columns:
                    df_results = df_results.sort_values("cv_mean", ascending=False)

                cols = [
                    "algo",
                    "cv_mean",
                    "cv_std",
                    "metric_name",
                    "duration",
                    "params",
                ]
                valid_cols = [c for c in cols if c in df_results.columns]
                f.write(df_results[valid_cols].to_markdown(index=False))
            else:
                f.write("*Nenhum modelo treinado.*\n")
            f.write("\n\n")

            # --- SEÇÃO 5: Final Model Audit (Holdout) ---
            f.write("## 5. Final Model Performance (Holdout) 🧪\n")
            metrics = self.metadata.get("final_metrics", {})
            if metrics:
                for k, v in metrics.items():
                    if isinstance(v, dict):
                        continue  # Pula dicts aninhados como classification_report
                    f.write(f"- **{k}:** {v:.4f}\n")

                # Classification Report formatado
                if "classification_report" in metrics:
                    f.write("\n**Classification Report:**\n```json\n")
                    f.write(json.dumps(metrics["classification_report"], indent=2))
                    f.write("\n```\n")

            # Imagem do SHAP
            f.write(f"\n![SHAP Summary]({self.exp_name}_shap_summary.png)\n\n")

            # --- SEÇÃO 6: RISK AUDIT (A parte que faltava!) ---
            if "risk_metrics" in self.metadata:
                risk = self.metadata["risk_metrics"]
                f.write("## 6. Risk & Stability Audit ⚖️\n")
                f.write(
                    "> Análise crítica de estabilidade (Train vs Test) e discriminação.\n\n"
                )

                f.write(
                    f"- **KS Score:** `{risk.get('ks_score', 0):.2f}` (Ideal: > 40.0)\n"
                )
                f.write(
                    f"- **PSI Score:** `{risk.get('psi_score', 0):.4f}` (Ideal: < 0.1)\n\n"
                )

                f.write("### 6.1 Stability by Risk Decile\n")
                if "audit_table" in risk:
                    # Tabela de faixas
                    f.write(risk["audit_table"].to_markdown(index=False))

                f.write("\n\n### 6.2 Visual Evidence\n")
                # Incorpora as imagens geradas no Markdown
                f.write(f"**Kolmogorov-Smirnov (KS) Curve:**\n\n")
                f.write(f"![KS Curve]({self.exp_name}_ks_curve.png)\n\n")

                f.write(f"**PSI / Stability Plot:**\n\n")
                f.write(f"![PSI Plot]({self.exp_name}_risk_audit.png)\n\n")

                f.write(f"**Confusion Matrix:**\n\n")
                f.write(
                    f"![Confusion Matrix]({self.exp_name}_confusion_matrix.png)\n\n"
                )

            # --- SEÇÃO 7: Logs ---
            f.write("## 7. System Logs 📜\n")
            for entry in self.log:
                f.write(f"- {entry}\n")

        print(f"📄 Relatório COMPLETO (com Risco e Imagens) salvo em: {filename}")

    # --- MÉTODOS DE AUDITORIA DE RISCO (Mantidos Integrais) ---

    def _calculate_psi(self, expected, actual):
        """Cálculo interno do PSI (Mantido para compatibilidade)."""
        actual = np.where(actual == 0, 0.0001, actual)
        expected = np.where(expected == 0, 0.0001, expected)
        psi_val = np.sum((actual - expected) * np.log(actual / expected))
        return psi_val

    def _plot_ks_curve(self, y_true, y_prob, ks_score):
        """Gera e salva o gráfico da curva KS (Mantido)."""
        df = pd.DataFrame({"label": y_true, "prob": y_prob})
        df = df.sort_values("prob", ascending=False).reset_index(drop=True)

        df["is_bad"] = df["label"]
        df["is_good"] = 1 - df["label"]

        cum_bad = df["is_bad"].cumsum() / df["is_bad"].sum()
        cum_good = df["is_good"].cumsum() / df["is_good"].sum()
        x_axis = np.arange(len(df)) / len(df)

        diff = np.abs(cum_bad - cum_good)
        max_idx = np.argmax(diff)
        max_x = x_axis[max_idx]

        plt.figure(figsize=(10, 6))
        plt.plot(x_axis, cum_bad, label="Captura (Target=1)", color="red", linewidth=2)
        plt.plot(
            x_axis, cum_good, label="Acumulado (Target=0)", color="green", linewidth=2
        )
        plt.plot(x_axis, diff, label="KS", color="grey", linestyle="--", alpha=0.5)

        plt.vlines(
            max_x,
            cum_good[max_idx],
            cum_bad[max_idx],
            color="blue",
            linestyle="--",
            label=f"KS Max: {ks_score:.2f}",
        )

        plt.title(f"Curva KS - Score: {ks_score:.2f}")
        plt.xlabel("População Ordenada por Score")
        plt.ylabel("% Acumulada")
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.show()
        filename = f"{self.exp_name}_ks_curve.png"
        plt.savefig(filename)
        plt.close()
        print(f"📈 Gráfico KS salvo: {filename}")

    def generate_risk_report(
        self, model, X_train, y_train, X_test, y_test, n_bins=5, save_plot=True
    ):
        """
        Gera relatório de Risco (PSI, KS, Decis).
        IMPORTANTE: Mantido integralmente pois é chamado externamente.
        """
        print(f"\n🔍 Iniciando Auditoria de Risco...")

        # Verifica proba vs predict
        if hasattr(model, "predict_proba"):
            train_probs = model.predict_proba(X_train)[:, 1]
            test_probs = model.predict_proba(X_test)[:, 1]
        else:
            # Fallback para regressão ou modelos sem proba
            train_probs = model.predict(X_train)
            test_probs = model.predict(X_test)

        # Gera Matriz de Confusão
        y_pred = model.predict(X_test)
        self.log_confusion_matrix(y_test, y_pred)

        # Treino
        df_train = pd.DataFrame({"prob": train_probs, "actual": y_train})
        df_train["bin"] = pd.qcut(
            df_train["prob"].rank(method="first"), n_bins, labels=False
        )
        thresholds = df_train.groupby("bin")["prob"].max().values

        # Teste
        df_test = pd.DataFrame({"prob": test_probs, "actual": y_test})
        df_test["bin"] = np.digitize(test_probs, thresholds, right=True)
        df_test["bin"] = df_test["bin"].clip(0, n_bins - 1)

        # PSI
        train_dist = (
            df_train["bin"]
            .value_counts(normalize=True)
            .reindex(range(n_bins), fill_value=0)
            .sort_index()
            .values
        )
        test_dist = (
            df_test["bin"]
            .value_counts(normalize=True)
            .reindex(range(n_bins), fill_value=0)
            .sort_index()
            .values
        )
        psi_score = self._calculate_psi(train_dist, test_dist)

        # KS
        df_ks = df_test.sort_values("prob", ascending=False).reset_index(drop=True)
        df_ks["is_fraud"] = df_ks["actual"]
        df_ks["is_not_fraud"] = 1 - df_ks["actual"]
        cum_fraud = df_ks["is_fraud"].cumsum() / df_ks["is_fraud"].sum()
        cum_not_fraud = df_ks["is_not_fraud"].cumsum() / df_ks["is_not_fraud"].sum()
        ks_score = np.max(np.abs(cum_fraud - cum_not_fraud)) * 100

        if save_plot:
            self._plot_ks_curve(y_test, test_probs, ks_score)

        # Tabela Comparativa
        train_res = (
            df_train.groupby("bin", observed=True)["actual"]
            .mean()
            .reindex(range(n_bins), fill_value=0)
            .reset_index()
            .rename(columns={"actual": "Treino"})
        )
        test_res = (
            df_test.groupby("bin", observed=True)["actual"]
            .mean()
            .reindex(range(n_bins), fill_value=0)
            .reset_index()
            .rename(columns={"actual": "Teste"})
        )

        comparison = pd.merge(train_res, test_res, on="bin")
        comparison["bin"] = comparison["bin"] + 1

        print("\n" + "=" * 50)
        print(f"📊 RELATÓRIO DE RISCO (ScoutML Enterprise)")
        print("=" * 50)
        print(f"✅ KS Score:  {ks_score:.2f}")
        print(f"✅ PSI Score: {psi_score:.4f}")
        print("-" * 50)
        # Tenta usar tabulate se instalado, senão string padrão
        try:
            print(comparison.to_markdown(index=False))
        except:
            print(comparison.to_string(index=False))
        print("=" * 50)

        if save_plot:
            plt.figure(figsize=(10, 5))
            plt.plot(
                comparison["bin"],
                comparison["Treino"],
                "o-",
                label="Treino",
                color="#1f77b4",
            )
            plt.plot(
                comparison["bin"],
                comparison["Teste"],
                "s--",
                label="Teste",
                color="#ff7f0e",
            )
            plt.axhline(
                y=y_train.mean(), color="black", linestyle=":", label=f"Média Global"
            )
            plt.title(f"Estabilidade por Faixas | PSI: {psi_score:.3f}")
            plt.xlabel(f"Faixas de Risco (1=Menor | {n_bins}=Maior)")
            plt.ylabel("Taxa do Evento")
            plt.xticks(range(1, n_bins + 1))
            plt.legend()
            plt.grid(True, alpha=0.3)
            plt.show()
            filename = f"{self.exp_name}_risk_audit.png"
            plt.savefig(filename)
            plt.close()
            print(f"📈 Gráfico Estabilidade salvo: {filename}")

        return {
            "ks_score": ks_score,
            "psi_score": psi_score,
            "thresholds": thresholds.tolist(),
            "audit_table": comparison,
        }
