#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Interpretability (protocol §11): SHAP values of the Task A LightGBM at three origins (importance and its stability over
time), accumulated local effects of the market features, and the DSGFN attention weights / RAREF regime probabilities for Task B.
Runs in the tree family process (no torch) for SHAP; the DSGFN/RAREF diagnostics are produced by 36_ablation.py."""
from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
import numpy as np
import polars as pl

from wp23 import config as C
from wp23 import data as D
from wp23.models_a import LGBM

OUT = C.RESULTS / "interpret"; OUT.mkdir(parents=True, exist_ok=True)


def main() -> None:
    import shap
    sales = D.load_sales(); rows_imp, rows_ale = [], []
    for o in ["2023-03", "2024-03", "2025-03"]:
        tr, te = D.task_a_sets(sales, o); Xtr, ytr, cols, cats = D.design(tr); Xte, yte, _, _ = D.design(te, cats)
        m = LGBM(quantiles=False).fit(Xtr, ytr); rng = np.random.default_rng(0); idx = rng.choice(len(Xte), min(4000, len(Xte)), replace=False)
        ex = shap.TreeExplainer(m.m); sv = ex.shap_values(Xte[idx])
        imp = np.abs(sv).mean(axis=0)
        for c, v in zip(cols, imp):
            rows_imp.append({"origin": o, "feature": c, "mean_abs_shap": float(v)})
        # ALE for selected continuous market features (10 bins)
        for feat in ["mun_mom_3_12", "reg_mom_3_12", "knn_q", "mun_q_3", "mun_vol_ratio", "lfa", "age_f"]:
            j = cols.index(feat); x = Xte[idx][:, j]; q = np.quantile(x, np.linspace(0, 1, 11)); q[-1] += 1e-9
            ale = []
            for b in range(10):
                sel = (x >= q[b]) & (x < q[b + 1])
                if sel.sum() < 20: ale.append(np.nan); continue
                lo = Xte[idx][sel].copy(); hi = lo.copy(); lo[:, j] = q[b]; hi[:, j] = q[b + 1]; ale.append(float(np.mean(m.m.predict(hi) - m.m.predict(lo))))
            acc = np.nancumsum(np.nan_to_num(ale)); acc -= np.nanmean(acc)
            for b in range(10):
                rows_ale.append({"origin": o, "feature": feat, "bin": b, "x_lo": float(q[b]), "x_hi": float(q[b + 1]), "ale": float(acc[b])})
        print(o, "done", flush=True)
    pl.DataFrame(rows_imp).write_csv(OUT / "shap_importance.csv"); pl.DataFrame(rows_ale).write_csv(OUT / "ale.csv")
    imp = pl.DataFrame(rows_imp).pivot(on="origin", index="feature", values="mean_abs_shap")
    cols_o = [c for c in imp.columns if c != "feature"]; imp = imp.with_columns(pl.mean_horizontal(cols_o).alias("mean")).sort("mean", descending=True)
    imp.write_csv(OUT / "shap_importance_wide.csv"); print(imp.head(20))


if __name__ == "__main__":
    main()