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

plt.style.use("ggplot")

class AuditTrail:
    def __init__(self, exp_name=None):
        self.exp_name = exp_name or f"run_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.log = []
        self.metadata = {}
        self.feature_catalog = pd.DataFrame()
        self.pipeline_decisions = []
        self.model_results = []

    def log_event(self, step, message, details=None):
        entry = f"[{datetime.datetime.now().strftime('%H:%M:%S')}] [{step}] {message}" + (f" | {details}" if details else "")
        self.log.append(entry)

    def log_action(self, step, message, details=None):
        self.log_event(step, message, details)

    def register_feature_stats(self, df_stats):
        self.feature_catalog = df_stats
        self.log_event("Scout", "Stats registered.")

    def log_pipeline_decision(self, feature, action, reason):
        self.pipeline_decisions.append({
            "timestamp": datetime.datetime.now().strftime("%H:%M:%S"),
            "feature": str(feature),
            "action": action,
            "reason": reason
        })
        self.log_event("Forge", f"{action} on {feature}", reason)

    def add_model_result(self, model, params, metrics, time):
        self.model_results.append({
            "algo": model, "params": str(params),
            "cv_mean": metrics.get("cv_mean",0), "metric": metrics.get("metric",""),
            "duration": round(time,4)
        })

    def generate_risk_report(self, model, X_train, y_train, X_test, y_test):
        try:
            if hasattr(model, "predict_proba"):
                train_p = model.predict_proba(X_train)[:,1]
                test_p = model.predict_proba(X_test)[:,1]
            else:
                train_p = model.predict(X_train)
                test_p = model.predict(X_test)
            
            cm = confusion_matrix(y_test, model.predict(X_test))
            plt.figure(figsize=(5,4))
            sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
            plt.title("Confusion Matrix")
            plt.savefig(f"{self.exp_name}_cm.png")
            plt.close()

            df_tr = pd.DataFrame({"p": train_p}); df_te = pd.DataFrame({"p": test_p, "y": y_test})
            df_tr["bin"] = pd.qcut(df_tr["p"].rank(method="first"), 10, labels=False)
            thresh = df_tr.groupby("bin")["p"].max().values
            df_te["bin"] = np.digitize(test_p, thresh, right=True).clip(0,9)
            
            tr_d = df_tr["bin"].value_counts(normalize=True).reindex(range(10), fill_value=0.0001).sort_index().values
            te_d = df_te["bin"].value_counts(normalize=True).reindex(range(10), fill_value=0.0001).sort_index().values
            psi = np.sum((tr_d - te_d) * np.log(tr_d/te_d))
            
            df_ks = df_te.sort_values("p", ascending=False).reset_index(drop=True)
            df_ks["bad"] = df_ks["y"]; df_ks["good"] = 1-df_ks["y"]
            ks = np.max(np.abs(df_ks["bad"].cumsum()/df_ks["bad"].sum() - df_ks["good"].cumsum()/df_ks["good"].sum())) * 100
            
            return {"ks_score": ks, "psi_score": psi}
        except: return {}

    def generate_report(self):
        with open(f"{self.exp_name}_report.md", "w", encoding="utf-8") as f:
            f.write(f"# ScoutML Report: {self.exp_name}\n\n")
            if not self.feature_catalog.empty:
                f.write("## 1. Data Stats\n" + self.feature_catalog.to_markdown(index=False) + "\n\n")
            if self.pipeline_decisions:
                f.write("## 2. Pipeline\n" + pd.DataFrame(self.pipeline_decisions).to_markdown(index=False) + "\n\n")
            if self.model_results:
                f.write("## 3. Models\n" + pd.DataFrame(self.model_results).to_markdown(index=False) + "\n\n")
            if "risk_metrics" in self.metadata:
                r = self.metadata["risk_metrics"]
                f.write(f"## 4. Risk\nKS: {r.get('ks_score',0):.2f} | PSI: {r.get('psi_score',0):.4f}\n")
                f.write(f"![CM]({self.exp_name}_cm.png)")
