#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""WP23 figures from results/ (missing inputs skipped)."""
from __future__ import annotations

import sys
from pathlib import Path

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

from wp23 import config as C

sys.path.insert(0, str(C.WP22_SRC))
from sharp_hpi.plotting import INK2, SERIES, style  # noqa: E402

style(); FIG = C.FIGURES; FIG.mkdir(exist_ok=True); PFIG = C.ROOT / "paper" / "figures"; PFIG.mkdir(parents=True, exist_ok=True); R = C.RESULTS
LAB = {"naive_mun_median": "Naive municipal median", "hedonic_ols": "Hedonic OLS", "lightgbm": "LightGBM", "xgboost": "XGBoost", "random_forest": "Random forest", "mlp": "Quantile MLP", "hshn_partial": "HSHN partial", "hshn_full": "HSHN full", "hshn_none": "HSHN none",
       "no_change": "No change", "momentum": "Momentum", "ar2_cell": "AR(2)", "ridge_panel": "Ridge panel", "lightgbm_panel": "LightGBM panel", "lstm": "LSTM", "tft_lite": "TFT-lite", "dsgfn": "DSGFN", "raref": "RAREF", "bhpf": "BHPF", "qanss": "QANSS", "qrefa": "QREFA", "mean_growth": "Mean growth", "local_linear_trend": "Local linear trend", "tcn": "TCN", "gru": "GRU", "mlp_panel": "MLP panel"}
COL = {m: SERIES[i % 8] for i, m in enumerate(["no_change", "momentum", "ar2_cell", "lightgbm_panel", "tft_lite", "bhpf", "qanss", "dsgfn"])}; COL.update({"raref": "#7a7a7a", "qrefa": "#0b0b0b", "ridge_panel": "#a0a0a0", "lstm": "#c98500"})
COLA = {m: SERIES[i % 8] for i, m in enumerate(["naive_mun_median", "hedonic_ols", "lightgbm", "xgboost", "random_forest", "mlp", "hshn_partial", "hshn_full"])}; COLA["hshn_none"] = "#7a7a7a"


def save(fig, name):
    fig.savefig(FIG / f"{name}.pdf", dpi=200); fig.savefig(PFIG / f"{name}.pdf", dpi=200); fig.savefig(FIG / f"{name}.png", dpi=150); plt.close(fig); print("fig", name)


def ok(*p):
    miss = [str(x) for x in p if not Path(x).exists()]
    if miss: print("skip:", miss)
    return not miss


def fig_density():
    if not ok(C.WP22_SDA, C.WP22_ANALYTIC): return
    import geopandas as gpd
    s = pl.read_parquet(C.WP22_ANALYTIC, columns=["mun_code"]).group_by("mun_code").agg(pl.len().alias("n")).to_pandas()
    mun = gpd.read_file(C.WP22_SDA, layer="munic_s")[["MUS_CO_GEO", "geometry"]].rename(columns={"MUS_CO_GEO": "mun_code"}).to_crs("EPSG:32198")
    g = mun.merge(s, on="mun_code", how="left"); g["geometry"] = g.geometry.simplify(150); g["n"] = g["n"].fillna(0)
    fig, ax = plt.subplots(figsize=(7.2, 5)); g.plot(ax=ax, column=np.log10(g["n"] + 1), cmap="Blues", edgecolor="white", linewidth=0.1, legend=True, legend_kwds={"shrink": 0.5, "label": "log10(1 + sales 2021–2026)"})
    g[g["n"] >= C.MIN_CELL_SALES].boundary.plot(ax=ax, color="#e34948", linewidth=0.5)
    for coll in ax.collections: coll.set_rasterized(True)
    ax.set_xlim(-650000, 500000); ax.set_ylim(130000, 700000); ax.axis("off"); save(fig, "fig_map_density")


def fig_perf(tag, order, col, xlab):
    f = R / f"eval_{tag}" / "scores.csv"
    if not ok(f): return
    d = pl.read_csv(f, infer_schema_length=None).filter(pl.col("group") == "all")
    def _k(x):
        try: return (0, int(x))
        except (TypeError, ValueError): return (1, str(x))
    hs = sorted(d["h"].unique().to_list(), key=_k)
    fig, axes = plt.subplots(1, 2, figsize=(9, 3.2))
    for m in order:
        s = d.filter(pl.col("model") == m)
        if s.height == 0: continue
        v = {str(r["h"]): r for r in s.iter_rows(named=True)}
        axes[0].plot(range(len(hs)), [100 * v[str(h)]["mae"] if str(h) in v else np.nan for h in hs], marker="o", ms=3, color=col.get(m, "#999"), label=LAB.get(m, m), lw=1.6 if m.startswith(("hshn", "dsgfn", "qanss", "bhpf", "raref", "qrefa")) else 1.0)
        axes[1].plot(range(len(hs)), [100 * v[str(h)]["crps"] if str(h) in v and v[str(h)]["crps"] == v[str(h)]["crps"] else np.nan for h in hs], marker="o", ms=3, color=col.get(m, "#999"), lw=1.0)
    for ax, yl in zip(axes, ["MAE (log points ×100)", "CRPS (×100)"]):
        ax.set_xticks(range(len(hs))); ax.set_xticklabels([str(h) for h in hs]); ax.set_xlabel(xlab); ax.set_ylabel(yl)
    axes[0].legend(fontsize=6, ncol=2); save(fig, f"fig_perf_{tag}")


def fig_calibration(tag, order, col):
    f = R / f"eval_{tag}" / "calibration.csv"
    if not ok(f): return
    d = pl.read_csv(f, infer_schema_length=None)
    def _k(x):
        try: return (0, int(x))
        except (TypeError, ValueError): return (1, str(x))
    hs = sorted(d["h"].unique().to_list(), key=_k); h = hs[min(1, len(hs) - 1)]
    fig, ax = plt.subplots(figsize=(4.2, 3.6)); ax.plot([0, 1], [0, 1], color=INK2, lw=0.6, ls="--")
    for m in order:
        s = d.filter((pl.col("model") == m) & (pl.col("h").cast(pl.Utf8) == str(h))).sort("tau")
        if s.height: ax.plot(s["tau"], s["share_below"], marker="o", ms=3, color=col.get(m, "#999"), label=LAB.get(m, m))
    ax.set_xlabel("Nominal quantile level"); ax.set_ylabel("Share of outcomes below the quantile"); ax.set_title(f"Calibration, horizon {h}", fontsize=9); ax.legend(fontsize=6); save(fig, f"fig_calibration_{tag}")


def fig_trajectories():
    f = R / "task_b" 
    files = sorted(f.glob("predictions_*.parquet"))
    if not files: print("skip trajectories"); return
    p = pl.concat([pl.read_parquet(x) for x in files], how="diagonal").filter((pl.col("h") == 3) & (pl.col("class") == "house"))
    cells = [("province", "QC", "Province"), ("region", "06", "Montréal"), ("region", "03", "Capitale-Nationale"), ("municipality", "81017", "Gatineau")]
    fig, axes = plt.subplots(2, 2, figsize=(9, 5.4))
    for ax, (lvl, gid, nm) in zip(axes.ravel(), cells):
        s = p.filter((pl.col("level") == lvl) & (pl.col("geo_id") == gid)).sort("origin")
        if s.height == 0: continue
        tgt = s.filter(pl.col("model") == "no_change"); x = np.arange(tgt.height); ax.plot(x, 100 * tgt["target"], color="#0b0b0b", lw=1.6, label="Realised 3-month change")
        for m in ["momentum", "qanss", "dsgfn", "lightgbm_panel", "bhpf"]:
            t = s.filter(pl.col("model") == m)
            if t.height: ax.plot(x[:t.height], 100 * t["pred"], color=COL.get(m), lw=0.9, label=LAB.get(m, m))
        ax.set_title(f"{nm}, houses", fontsize=9); idx = [i for i, o in enumerate(tgt["origin"].to_list()) if o.endswith("-01")]; ax.set_xticks(idx); ax.set_xticklabels([tgt["origin"][i][:4] for i in idx]); ax.axhline(0, color=INK2, lw=0.5)
    axes[0, 0].set_ylabel("Log points ×100"); axes[1, 0].set_ylabel("Log points ×100"); axes[0, 0].legend(fontsize=6); save(fig, "fig_trajectories_b")


def fig_error_map():
    f = R / "eval_b" / "scores.csv"
    if not ok(f, C.WP22_SDA): return
    import geopandas as gpd
    d = pl.read_csv(f, infer_schema_length=None).filter((pl.col("group") == "region_code") & (pl.col("h").cast(pl.Utf8) == "3"))
    best = d.group_by("region_code").agg(pl.col("mae").min().alias("best"), pl.col("mae").filter(pl.col("model") == "no_change").first().alias("nc")).with_columns((pl.col("best") / pl.col("nc")).alias("rel"))
    reg = gpd.read_file(C.WP22_SDA, layer="regio_s")[["RES_CO_REG", "geometry"]].rename(columns={"RES_CO_REG": "region_code"}).to_crs("EPSG:32198"); reg["geometry"] = reg.geometry.simplify(300)
    g = reg.merge(best.with_columns(pl.col("region_code").cast(pl.Utf8)).to_pandas(), on="region_code", how="left")
    fig, ax = plt.subplots(figsize=(6.4, 4.6)); g.plot(ax=ax, column="rel", cmap="Blues_r", vmin=0.5, vmax=1.0, edgecolor="white", linewidth=0.3, legend=True, legend_kwds={"shrink": 0.5, "label": "Best model MAE / no-change MAE, h = 3"})
    for coll in ax.collections: coll.set_rasterized(True)
    ax.set_xlim(-650000, 500000); ax.set_ylim(130000, 700000); ax.axis("off"); save(fig, "fig_error_map")


def fig_simulation():
    f = R / "simulation" / "summary.csv"
    if not ok(f): return
    d = pl.read_csv(f, infer_schema_length=None).filter((pl.col("level") == "municipality") & (pl.col("h") == 3)); scen = [s for s in ["rw", "momentum", "regimes", "spillover", "sparse"] if s in d["scenario"].to_list()]
    models = [m for m in ["no_change", "momentum", "ar2_cell", "lightgbm_panel", "tft_lite", "bhpf", "raref", "qanss", "dsgfn"] if m in d["model"].to_list()]
    fig, ax = plt.subplots(figsize=(9, 3.2)); w = 0.8 / len(models)
    for j, m in enumerate(models):
        v = {r["scenario"]: r for r in d.filter(pl.col("model") == m).iter_rows(named=True)}
        ax.bar(np.arange(len(scen)) + j * w - 0.4 + w / 2, [v[s]["rel_mae"] if s in v else np.nan for s in scen], width=w * 0.9, color=COL.get(m, "#999"), label=LAB.get(m, m), yerr=[v[s]["rel_mae_mcse"] if s in v else 0 for s in scen], error_kw=dict(lw=0.5), lw=0)
    ax.axhline(1, color=INK2, lw=0.6, ls="--"); ax.set_xticks(range(len(scen))); ax.set_xticklabels(scen); ax.set_ylabel("MAE relative to no-change (municipalities, h = 3)"); ax.legend(fontsize=6, ncol=3); save(fig, "fig_simulation")


def fig_mcs():
    for tag in ["a", "b"]:
        f = R / f"eval_{tag}" / "mcs.csv"
        if not ok(f): continue
        d = pl.read_csv(f, infer_schema_length=None); hs = sorted(d["h"].unique().to_list(), key=str); models = sorted(d["model"].unique().to_list(), key=lambda m: -d.filter(pl.col("model") == m)["mcs_p"].mean())
        M = np.array([[d.filter((pl.col("model") == m) & (pl.col("h").cast(pl.Utf8) == str(h)))["mcs_p"].to_list()[:1] or [np.nan] for h in hs] for m in models]).squeeze(-1)
        fig, ax = plt.subplots(figsize=(4.5, 0.3 * len(models) + 1)); im = ax.imshow(M, cmap="Blues", vmin=0, vmax=1, aspect="auto")
        ax.set_yticks(range(len(models))); ax.set_yticklabels([LAB.get(m, m) for m in models], fontsize=7); ax.set_xticks(range(len(hs))); ax.set_xticklabels([str(h) for h in hs]); ax.set_xlabel("Horizon")
        for i in range(len(models)):
            for j in range(len(hs)):
                if M[i, j] == M[i, j]: ax.text(j, i, f"{M[i, j]:.2f}", ha="center", va="center", fontsize=6, color="white" if M[i, j] > 0.6 else "#0b0b0b")
        ax.set_title(f"MCS p-values, Task {tag.upper()}", fontsize=9); ax.grid(False); save(fig, f"fig_mcs_{tag}")


if __name__ == "__main__":
    for fn in [fig_density, lambda: fig_perf("a", ["naive_mun_median", "hedonic_ols", "random_forest", "xgboost", "lightgbm", "mlp", "hshn_partial", "hshn_full", "hshn_none"], COLA, "Horizon band (months)"),
               lambda: fig_perf("b", ["no_change", "momentum", "ar2_cell", "ridge_panel", "lightgbm_panel", "lstm", "tft_lite", "bhpf", "raref", "qanss", "dsgfn", "qrefa"], COL, "Horizon (months)"),
               lambda: fig_calibration("a", ["lightgbm", "mlp", "hshn_partial", "hedonic_ols"], COLA), lambda: fig_calibration("b", ["no_change", "lightgbm_panel", "tft_lite", "bhpf", "qanss", "dsgfn"], COL), fig_trajectories, fig_error_map, fig_simulation, fig_mcs]:
        try: fn()
        except Exception as ex:
            import traceback; traceback.print_exc()


def fig_error_time():
    X = R / "extra"
    if ok(X / "error_by_origin_b.csv"):
        d = pl.read_csv(X / "error_by_origin_b.csv", infer_schema_length=None).filter(pl.col("h") == 3)
        fig, ax = plt.subplots(figsize=(8.4, 3.0))
        for m in ["no_change", "ar2_cell", "ridge_panel", "lightgbm_panel", "qanss", "qrefa", "dsgfn"]:
            s = d.filter(pl.col("model") == m).sort("origin")
            if s.height: ax.plot(range(s.height), 100 * s["mae"], color=COL.get(m, "#999"), lw=1.6 if m in ("lightgbm_panel", "no_change") else 1.0, label=LAB.get(m, m))
        orig = d.filter(pl.col("model") == "no_change").sort("origin")["origin"].to_list(); idx = [i for i, o in enumerate(orig) if o.endswith(("-01", "-07"))]
        ax.set_xticks(idx); ax.set_xticklabels([orig[i] for i in idx], rotation=45, fontsize=7); ax.set_ylabel("MAE (log points ×100), h = 3"); ax.legend(fontsize=6.5, ncol=4); save(fig, "fig_error_time_b")
    if ok(X / "error_by_origin_a.csv"):
        d = pl.read_csv(X / "error_by_origin_a.csv", infer_schema_length=None)
        fig, axes = plt.subplots(1, 2, figsize=(8.4, 2.8))
        for m in ["naive_mun_median", "hedonic_ols", "lightgbm", "hshn_partial", "mlp"]:
            s = d.filter(pl.col("model") == m).sort("origin")
            if s.height:
                axes[0].plot(range(s.height), 100 * s["mae"], marker="o", ms=3, color=COLA.get(m, "#999"), label=LAB.get(m, m)); axes[1].plot(range(s.height), 100 * s["bias"], marker="o", ms=3, color=COLA.get(m, "#999"))
        orig = d.filter(pl.col("model") == "lightgbm").sort("origin")["origin"].to_list()
        for ax, yl in zip(axes, ["MAE (log ×100), all horizons", "Bias (log ×100)"]):
            ax.set_xticks(range(len(orig))); ax.set_xticklabels(orig, rotation=45, fontsize=6.5); ax.set_ylabel(yl)
        axes[1].axhline(0, color=INK2, lw=0.5); axes[0].legend(fontsize=6.5); save(fig, "fig_error_time_a")


def fig_regimes():
    X = R / "extra"
    if not ok(X / "raref_regimes.csv"): return
    d = pl.read_csv(X / "raref_regimes.csv", infer_schema_length=None).filter(pl.col("class") == "house").sort("origin")
    fig, ax = plt.subplots(figsize=(8.4, 2.8)); x = np.arange(d.height)
    ax.stackplot(x, d["p_low"], d["p_mid"], d["p_high"], colors=[SERIES[7], "#d9d9d6", SERIES[2]], alpha=0.75, labels=["low-growth state", "middle state", "high-growth state"], lw=0)
    ax2 = ax.twinx() if False else None  # single axis rule: growth shown as markers scaled to [0,1]
    g = d["prov_growth_last"].to_numpy(); ax.plot(x, 0.5 + 10 * g, color="#0b0b0b", lw=1.2, label="provincial monthly growth (×10, centred at 0.5)")
    idx = [i for i, o in enumerate(d["origin"].to_list()) if o.endswith(("-01", "-07"))]; ax.set_xticks(idx); ax.set_xticklabels([d["origin"][i] for i in idx], rotation=45, fontsize=7)
    ax.set_ylim(0, 1); ax.set_ylabel("Filtered regime probability (RAREF), houses"); ax.legend(fontsize=6.5, loc="upper center", bbox_to_anchor=(0.5, -0.32), ncol=4, frameon=False); save(fig, "fig_regimes")


def fig_decline_reliability():
    X = R / "extra"
    if not ok(X / "decline_reliability_b.csv"): return
    d = pl.read_csv(X / "decline_reliability_b.csv", infer_schema_length=None).filter(pl.col("h") == 3)
    fig, ax = plt.subplots(figsize=(4.2, 3.6)); ax.plot([0, 1], [0, 1], color=INK2, lw=0.6, ls="--")
    for m in ["lightgbm_panel", "qrefa", "bhpf", "qanss", "ridge_panel", "dsgfn"]:
        s = d.filter(pl.col("model") == m).sort("p_mean")
        if s.height: ax.plot(s["p_mean"], s["freq"], marker="o", ms=4, color=COL.get(m, "#999"), label=LAB.get(m, m))
    ax.set_xlabel("Forecast probability of a decline (h = 3)"); ax.set_ylabel("Observed frequency of declines"); ax.legend(fontsize=6.5); save(fig, "fig_decline_reliability")


def fig_mun_distribution():
    X = R / "extra"
    if not ok(X / "municipal_mase_b.csv"): return
    d = pl.read_csv(X / "municipal_mase_b.csv", infer_schema_length=None)
    fig, ax = plt.subplots(figsize=(5.2, 3.0))
    for m in ["lightgbm_panel", "qanss", "ridge_panel", "ar2_cell"]:
        s = d.filter(pl.col("model") == m)["mase"].to_numpy(); ax.hist(s, bins=np.linspace(0.4, 1.6, 31), histtype="step", lw=1.4, color=COL.get(m, "#999"), label=LAB.get(m, m))
    ax.axvline(1, color=INK2, lw=0.6, ls="--"); ax.set_xlabel("MASE relative to no change, h = 3 (192 municipal cells)"); ax.set_ylabel("Cells"); ax.legend(fontsize=6.5); save(fig, "fig_mun_distribution")


if __name__ == "__main__":
    for fn in [fig_error_time, fig_regimes, fig_decline_reliability, fig_mun_distribution]:
        try: fn()
        except Exception as ex:
            import traceback; traceback.print_exc()