#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 7 — figures generated exclusively from results/ and data/processed/. Missing inputs are skipped with a note."""
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
import matplotlib.pyplot as plt

from sharp_hpi import config as C
from sharp_hpi.plotting import CLASS_LABELS, METHOD_COLORS, METHOD_LABELS, SERIES, INK2, month_ticks, style

FIG = C.FIGURES; FIG.mkdir(exist_ok=True); PFIG = C.ROOT / "paper" / "figures"; PFIG.mkdir(exist_ok=True, parents=True)
style()
R = C.RESULTS


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=160); plt.close(fig); print("fig", name)


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


def rebase(li: np.ndarray, months: list[str]) -> np.ndarray:
    sel = [i for i, m in enumerate(months) if m.startswith("2021") and not np.isnan(li[i])]
    return 100 * np.exp(li - np.mean(li[sel])) if sel else np.full_like(li, np.nan)


# 1. coverage ---------------------------------------------------------------
def fig_coverage():
    p = R / "profile" / "coverage_month_type.csv"; h = R / "cleaning" / "headline_share_by_year.csv"
    if not exists(p):
        return
    d = pl.read_csv(p).group_by("month").agg(pl.col("n").sum()).sort("month")
    months = d["month"].to_list(); n = d["n"].to_numpy()
    fig, ax = plt.subplots(figsize=(6.4, 2.6))
    ax.bar(range(len(months)), n, color=SERIES[0], width=0.8, linewidth=0)
    ax.set_ylabel("Registered sales per month"); month_ticks(ax, months)
    ax.annotate("feed start-up", xy=(0, n[0]), xytext=(4, n.max() * 0.85), fontsize=8, color=INK2, arrowprops=dict(arrowstyle="-", color=INK2, lw=0.6))
    ax.annotate("partial month", xy=(len(months) - 1, n[-1]), xytext=(len(months) - 14, n.max() * 0.7), fontsize=8, color=INK2, arrowprops=dict(arrowstyle="-", color=INK2, lw=0.6))
    save(fig, "fig_coverage_monthly")


def fig_arrival():
    p = R / "vintages" / "arrival_profile_by_lag.csv"
    if not exists(p):
        return
    d = pl.read_csv(p).filter(pl.col("lag_bin") <= 70)
    fig, ax = plt.subplots(figsize=(4.2, 2.6))
    for i, (v, sub) in enumerate(d.group_by("vintage", maintain_order=True)):
        sub = sub.sort("lag_bin"); ax.plot(sub["lag_bin"], 100 * sub["share_visible"], color=SERIES[i % 8], marker="o", ms=3, label=f"vintage {v[0]}")
    ax.set_xlabel("Days between transaction date and extraction"); ax.set_ylabel("% of eventual transactions visible"); ax.legend(); ax.set_ylim(50, 101)
    save(fig, "fig_arrival_profile")


# 2. province methods ------------------------------------------------------
def fig_province_methods():
    p = R / "benchmarks" / "series.parquet"
    if not exists(p):
        return
    s = pl.read_parquet(p)
    if (R / "benchmarks" / "series_hmts.parquet").exists():
        s = pl.concat([s, pl.read_parquet(R / "benchmarks" / "series_hmts.parquet")])
    methods = ["S_sharp", "C_pooled_td", "D_rtd_w13_mean", "E_hi_fisher_chained", "F_geks_w25_movement", "G_hmts_w13", "H_repeat_sales_cs", "K_qhpi", "A_median"]
    fig, axes = plt.subplots(1, 3, figsize=(9.6, 3.0), sharey=True)
    for ax, cls in zip(axes, ["house", "condo", "plex"]):
        for m in methods:
            sub = s.filter((pl.col("class") == cls) & (pl.col("level") == "province") & (pl.col("method") == m)).sort("month")
            if sub.height == 0:
                continue
            sub = sub.filter(pl.col("month") < "2026-10"); months = sub["month"].to_list(); li = rebase(sub["log_index"].to_numpy(), months)
            ax.plot(range(len(months)), li, color=METHOD_COLORS.get(m, "#999"), lw=1.6 if m == "S_sharp" else 1.0, label=METHOD_LABELS.get(m, m), alpha=1 if m == "S_sharp" else 0.9)
        ax.set_title(CLASS_LABELS[cls]); month_ticks(ax, months); ax.axhline(100, color=INK2, lw=0.5)
    axes[0].set_ylabel("Index (2021 average = 100)"); axes[0].legend(loc="upper left", fontsize=6.5)
    save(fig, "fig_province_methods")


def fig_region_methods():
    p = R / "benchmarks" / "series.parquet"
    if not exists(p):
        return
    s = pl.read_parquet(p)
    regs = [("06", "Montréal"), ("03", "Capitale-Nationale"), ("07", "Outaouais"), ("01", "Bas-Saint-Laurent"), ("09", "Côte-Nord"), ("11", "Gaspésie–Îles-de-la-Madeleine")]
    methods = ["S_sharp", "S_sharp_filtered", "C_pooled_td", "H_repeat_sales_cs", "A_median", "K_qhpi_direct_local_td", "K_qhpi_hierarchical"]
    fig, axes = plt.subplots(2, 3, figsize=(9.6, 5.4), sharey=True)
    for ax, (reg, nm) in zip(axes.ravel(), regs):
        for m in methods:
            sub = s.filter((pl.col("class") == "house") & (pl.col("level") == "region") & (pl.col("geo_id") == reg) & (pl.col("method") == m)).sort("month")
            if sub.height == 0:
                continue
            sub = sub.filter(pl.col("month") < "2026-10"); months = sub["month"].to_list(); li = rebase(sub["log_index"].to_numpy(), months)
            col = METHOD_COLORS.get(m, METHOD_COLORS["K_qhpi"] if m.startswith("K_") else "#999")
            if m == "S_sharp_filtered":
                col = "#0b0b0b"
            ax.plot(range(len(months)), li, color=col, lw=1.6 if m == "S_sharp" else 1.0, ls="--" if m == "S_sharp_filtered" else "-", label=METHOD_LABELS.get(m, "QHPI v2.1 architecture") if not m.startswith("K_") else "QHPI v2.1 architecture")
        n = s.filter((pl.col("class") == "house") & (pl.col("level") == "region") & (pl.col("geo_id") == reg) & (pl.col("method") == "A_median")).sort("month")["n"].to_numpy()
        ax.set_title(f"{nm}\nhouses, median {np.nanmedian(n[:-1]):.0f} sales/month", fontsize=8); month_ticks(ax, months); ax.axhline(100, color=INK2, lw=0.5)
    axes[0, 0].set_ylabel("Index (2021 = 100)"); axes[1, 0].set_ylabel("Index (2021 = 100)")
    h, l = axes[0, 0].get_legend_handles_labels(); uniq = dict(zip(l, h)); fig.legend(uniq.values(), uniq.keys(), loc="lower center", ncol=4, fontsize=7.5, bbox_to_anchor=(0.5, -0.04))
    save(fig, "fig_region_methods")


# 3. municipal examples ------------------------------------------------------
def fig_municipal_examples():
    p = C.PROCESSED / "index_panel.parquet"
    if not exists(p):
        return
    panel = pl.read_parquet(p).filter((pl.col("geography_level") == "municipality") & (pl.col("property_type") == "house"))
    sizes = panel.group_by(["geography_id", "geography_name"]).agg(pl.col("actual_transaction_count").sum().alias("n")).sort("n", descending=True)
    picks = []
    for lo, hi in [(20000, 10**9), (3000, 8000), (800, 1500), (250, 400), (80, 130), (20, 40)]:
        c = sizes.filter((pl.col("n") >= lo) & (pl.col("n") < hi))
        if c.height:
            picks.append(c.row(c.height // 2))
    fig, axes = plt.subplots(2, 3, figsize=(9.6, 5.2))
    for ax, (gid, nm, n) in zip(axes.ravel(), picks):
        sub = panel.filter(pl.col("geography_id") == gid).sort("month"); months = sub["month"].to_list(); x = np.arange(len(months))
        full = ~sub["partial_month_flag"].to_numpy()
        ax.fill_between(x[full], sub["lower_95"].to_numpy()[full], sub["upper_95"].to_numpy()[full], color=SERIES[0], alpha=0.18, lw=0)
        ax.plot(x[full], sub["index_value"].to_numpy()[full], color=SERIES[0], lw=1.5, label="SHARP-HPI (smoothed)")
        ax.plot(x[~full], sub["index_value"].to_numpy()[~full], marker="o", mfc="white", mec=SERIES[0], ms=4, ls="none", label="partial month (Oct. 2026)")
        ax2 = ax.inset_axes([0, 0, 1, 0.22]); ax2.bar(x, sub["actual_transaction_count"], color=INK2, alpha=0.35, width=0.8, lw=0); ax2.axis("off")
        g = sub["reliability_grade"].to_list(); ax.set_title(f"{nm}: {int(n)} sales, grades {'/'.join(sorted(set(g)))}", fontsize=8)
        month_ticks(ax, months); ax.set_ylim(80, max(180, float(sub['upper_95'].max()) + 5))
    axes[0, 0].set_ylabel("Index (2021 = 100)"); axes[1, 0].set_ylabel("Index (2021 = 100)"); axes[0, 0].legend(loc="upper left", fontsize=7)
    fig.text(0.01, 0.005, "Shaded band: 95 % interval. Grey bars (bottom of each panel): monthly sales count.", fontsize=7, color=INK2)
    save(fig, "fig_municipal_examples")


def fig_gain_vs_n():
    p = R / "panel" / "reliability_cells.csv"
    if not exists(p):
        return
    d = pl.read_csv(p).filter(pl.col("geography_level") == "municipality")
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))
    for i, cls in enumerate(["house", "condo", "plex"]):
        s = d.filter(pl.col("property_type") == cls)
        axes[0].scatter(s["median_monthly_n"], s["mean_gain"], s=7, color=SERIES[i], alpha=0.6, label=CLASS_LABELS[cls], lw=0)
        axes[1].scatter(s["median_monthly_n"], s["mean_ess"], s=7, color=SERIES[i], alpha=0.6, lw=0)
    axes[0].set_xscale("log"); axes[0].set_xlabel("Median monthly sales (log scale)"); axes[0].set_ylabel("Mean local data weight (Kalman gain)"); axes[0].legend(fontsize=7)
    axes[1].set_xscale("log"); axes[1].set_yscale("log"); axes[1].set_xlabel("Median monthly sales (log scale)"); axes[1].set_ylabel("Mean effective sample size")
    save(fig, "fig_gain_vs_n")


def fig_hedonic():
    import json
    p = R / "panel" / "stage1_meta.json"
    if not exists(p):
        return
    meta = json.loads(Path(p).read_text())
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))
    labels = ["0–1", "2–5", "6–10", "11–20", "21–40", "41–60", "61–80", "81–120", "120+"]
    for i, cls in enumerate(["house", "condo", "plex"]):
        c = meta[cls]["coef"]
        age = [0.0] + [c.get(f"C(age_bin)[T.a0{k}_{lab.replace('–', '-')}]", np.nan) for k, lab in enumerate(labels[1:], start=1)]
        axes[0].plot(range(len(labels)), 100 * (np.exp(age) - 1), marker="o", ms=3, color=SERIES[i], label=CLASS_LABELS[cls])
        x = np.linspace(-1.0, 1.0, 50); y = c["lfa"] * x + c["lfa2"] * x ** 2 + c["lfa3"] * x ** 3
        axes[1].plot(x, 100 * (np.exp(y) - 1), color=SERIES[i], label=CLASS_LABELS[cls])
    axes[0].set_xticks(range(len(labels))); axes[0].set_xticklabels(labels, rotation=45, fontsize=7); axes[0].set_xlabel("Age (years)"); axes[0].set_ylabel("% price effect vs 0–1 years"); axes[0].legend(fontsize=7)
    axes[1].set_xlabel("log floor area, deviation from class median"); axes[1].set_ylabel("% price effect"); axes[1].axhline(0, color=INK2, lw=0.5)
    save(fig, "fig_hedonic_profiles")


# 4. map ---------------------------------------------------------------------
def fig_map():
    p = C.PROCESSED / "index_panel.parquet"
    if not exists(p, C.SDA_GPKG):
        return
    import geopandas as gpd
    panel = pl.read_parquet(p).filter((pl.col("geography_level") == "municipality") & (pl.col("property_type") == "house"))
    last = panel.filter(pl.col("month") == "2026-09").select(["geography_id", "index_value", "reliability_grade", "suppression_flag"])
    mun = gpd.read_file(C.SDA_GPKG, layer="munic_s")[["MUS_CO_GEO", "geometry"]].rename(columns={"MUS_CO_GEO": "geography_id"})
    reg = gpd.read_file(C.SDA_GPKG, layer="regio_s")
    g = mun.merge(last.to_pandas(), on="geography_id", how="left").to_crs("EPSG:32198")
    g["geometry"] = g.geometry.simplify(150, preserve_topology=True)
    g = g.cx[-900000:900000, 100000:900000]  # southern Quebec extent (where sales exist)
    fig, ax = plt.subplots(figsize=(7.2, 5.2))
    g[g["index_value"].isna() | g["suppression_flag"].fillna(True)].plot(ax=ax, color="#eeeeea", edgecolor="white", linewidth=0.1)
    gg = g[~g["index_value"].isna() & ~g["suppression_flag"].fillna(True)]
    gg.plot(ax=ax, column="index_value", cmap="Blues", vmin=110, vmax=200, edgecolor="white", linewidth=0.1, legend=True, legend_kwds={"shrink": 0.5, "label": "House index, Sept. 2026 (2021 = 100)"})
    gg[gg["reliability_grade"].isin(["D", "E"])].plot(ax=ax, facecolor="none", edgecolor="#e34948", hatch="////", linewidth=0.0, alpha=0.5)
    rb = reg.to_crs("EPSG:32198"); rb["geometry"] = rb.geometry.simplify(300, preserve_topology=True); rb.boundary.plot(ax=ax, color=INK2, linewidth=0.3)
    for coll in ax.collections:
        coll.set_rasterized(True)
    ax.set_xlim(-650000, 500000); ax.set_ylim(130000, 700000); ax.axis("off")
    ax.set_title("SHARP-HPI municipal house index, September 2026 (hatched: reliability grade D/E; grey: suppressed or no sales)", fontsize=8)
    save(fig, "fig_map_municipal_house")


# 5. Monte Carlo ------------------------------------------------------------
def fig_mc():
    p = R / "simulation" / "summary.csv"
    if not exists(p):
        return
    d = pl.read_csv(p)
    methods = ["S_sharp", "S_sharp_filtered", "S_noEB", "S_LLT", "S_coarseloc", "S_norobust", "K_qhpi", "C_pooled_td", "D_rtd13", "B_stratified", "A_median"]
    cols = {m: METHOD_COLORS.get(m.replace("D_rtd13", "D_rtd_w13_mean").replace("B_stratified", "B_stratified_median"), SERIES[i % 8]) for i, m in enumerate(methods)}
    scen = d["scenario"].unique().sort().to_list()
    fig, axes = plt.subplots(2, 1, figsize=(9.0, 5.6), sharex=True)
    for ax, lvl in zip(axes, ["municipality", "region"]):
        sub = d.filter(pl.col("level") == lvl)
        w = 0.8 / len(methods)
        for j, m in enumerate(methods):
            s = sub.filter(pl.col("method") == m)
            vals = {r["scenario"]: (r["rmse_ret"], r["rmse_ret_mcse"]) for r in s.iter_rows(named=True)}
            ys = [vals.get(sc, (np.nan, 0))[0] for sc in scen]; es = [vals.get(sc, (0, 0))[1] for sc in scen]
            ax.bar(np.arange(len(scen)) + j * w - 0.4 + w / 2, 100 * np.array(ys), width=w * 0.92, color=cols[m], label=METHOD_LABELS.get(m, m), yerr=100 * np.array(es), error_kw=dict(lw=0.5, capsize=0), lw=0)
        ax.set_ylabel(f"RMSE of monthly log change (%), {lvl}"); ax.set_ylim(0, None)
    axes[1].set_xticks(range(len(scen))); axes[1].set_xticklabels(scen, rotation=20, fontsize=7.5)
    axes[0].legend(ncol=4, fontsize=6.5, loc="upper left")
    save(fig, "fig_mc_rmse")
    fig, ax = plt.subplots(figsize=(9.0, 2.8))
    for j, m in enumerate(["S_sharp", "S_sharp_filtered", "S_noEB", "S_LLT", "K_qhpi", "C_pooled_td"]):
        s = d.filter((pl.col("level") == "municipality") & (pl.col("method") == m)); vals = {r["scenario"]: r["cov95"] for r in s.iter_rows(named=True)}
        ax.plot(range(len(scen)), [100 * vals.get(sc, np.nan) for sc in scen], marker="o", ms=4, color=cols[m], label=METHOD_LABELS.get(m, m))
    ax.axhline(95, color=INK2, lw=0.6, ls="--"); ax.set_ylabel("Coverage of 95 % intervals (%), municipalities"); ax.set_xticks(range(len(scen))); ax.set_xticklabels(scen, rotation=20, fontsize=7.5); ax.legend(ncol=3, fontsize=6.5)
    save(fig, "fig_mc_coverage")


# 6. revisions ---------------------------------------------------------------
def fig_revisions():
    p = R / "vintages" / "revisions_by_release.csv"; q = C.PROCESSED / "index_vintages.parquet"
    if not exists(p, q):
        return
    d = pl.read_csv(p)
    fig, axes = plt.subplots(1, 2, figsize=(8.4, 2.8))
    df_f = pl.read_csv(R / "vintages" / "revisions_by_release_filtered.csv") if (R / "vintages" / "revisions_by_release_filtered.csv").exists() else None
    for i, lvl in enumerate(["province", "region", "municipality"]):
        s = d.filter((pl.col("class") == "house") & (pl.col("level") == lvl)).sort("release")
        axes[0].plot(s["release"], 100 * s["mae_rev_m1_vs_final"], marker="o", ms=3, color=SERIES[i], label=f"{lvl} (full re-estimation)")
        if df_f is not None:
            sf = df_f.filter((pl.col("class") == "house") & (pl.col("level") == lvl)).sort("release")
            axes[0].plot(sf["release"], 100 * sf["mae_rev_m1_vs_final"], ls="--", color=SERIES[i], lw=0.9, label=f"{lvl} (frozen first release)")
    axes[0].set_xlabel("Release number of a month"); axes[0].set_ylabel("Mean |revision| of monthly change vs final (log pts × 100)"); axes[0].legend(fontsize=6)
    v = pl.read_parquet(q).filter((pl.col("class") == "house") & (pl.col("level") == "province"))
    if "mode" in v.columns:
        v = v.filter(pl.col("mode") == "smoothed")
    vints = sorted(v["vintage"].unique().to_list())
    for k, vm in enumerate(vints):
        s = v.filter(pl.col("vintage") == vm).sort("month"); months = s["month"].to_list(); li = rebase(s["log_index"].to_numpy(), months)
        axes[1].plot(range(len(months)), li, color=SERIES[0], alpha=0.12 + 0.85 * (k == len(vints) - 1), lw=0.8 + 0.8 * (k == len(vints) - 1))
    axes[1].set_ylabel("Province house index, every vintage (2021 = 100)"); month_ticks(axes[1], months)
    save(fig, "fig_revisions")


# 7. validation ---------------------------------------------------------------
def fig_validation():
    p = R / "validation" / "placebo_cells.csv"; q = C.PROCESSED / "index_panel.parquet"
    if exists(p, q):
        d = pl.read_csv(p).filter(pl.col("level") == "municipality")
        panel = pl.read_parquet(q).filter((pl.col("geography_level") == "municipality") & (pl.col("property_type") == "house")).sort(["geography_id", "month"])
        real = panel.with_columns((np.log(pl.col("index_value")) - np.log(pl.col("index_value")).shift(12)).over("geography_id").alias("d12")).filter(pl.col("month") == "2026-06")
        fig, ax = plt.subplots(figsize=(4.6, 2.8))
        ax.hist(100 * d.filter(pl.col("class") == "house")["max_abs_12m_change"], bins=40, color=SERIES[2], alpha=0.8, label="Placebo: max |12-month change| of municipal deviation", lw=0)
        ax.hist(100 * real["d12"].drop_nulls().abs(), bins=40, color=SERIES[0], alpha=0.6, label="Actual: |12-month change| to June 2026", lw=0)
        ax.set_xlabel("Log points × 100"); ax.set_ylabel("Municipalities"); ax.legend(fontsize=7)
        save(fig, "fig_placebo")
    p = R / "validation" / "composition_shock.csv"
    if exists(p):
        d = pl.read_csv(p)
        fig, ax = plt.subplots(figsize=(4.6, 2.6))
        regs = d["region"].unique().sort().to_list(); w = 0.26
        for j, m in enumerate(["A_median", "C_pooled_td", "S_sharp"]):
            s = d.filter(pl.col("method") == m); vals = {str(r["region"]).zfill(2): r["shock_effect"] for r in s.iter_rows(named=True)}
            ax.bar(np.arange(len(regs)) + (j - 1) * w, [100 * vals.get(str(r).zfill(2), np.nan) for r in regs], width=w * 0.9, color=METHOD_COLORS[m], label=METHOD_LABELS[m], lw=0)
        ax.set_xticks(range(len(regs))); ax.set_xticklabels([f"region {str(r).zfill(2)}" for r in regs]); ax.set_ylabel("Shift of 2023→2024 change (log pts × 100)"); ax.axhline(0, color=INK2, lw=0.5); ax.legend(fontsize=7)
        save(fig, "fig_composition_shock")
    p = R / "validation" / "aggregation_sensitivity.csv"
    if exists(p):
        d = pl.read_csv(p)
        fig, ax = plt.subplots(figsize=(4.6, 2.6)); ws = ["stock_units", "stock_value", "tx", "eq"]; lab = {"stock_units": "stock (units)", "stock_value": "stock (value)", "tx": "transactions", "eq": "equal"}
        for i, cls in enumerate(["house", "condo", "plex"]):
            s = d.filter(pl.col("class") == cls); vals = {r["weights"]: 100 * (r["total_change_bottomup"] - r["total_change_direct"]) for r in s.iter_rows(named=True)}
            ax.bar(np.arange(len(ws)) + (i - 1) * 0.26, [vals.get(w, np.nan) for w in ws], width=0.24, color=SERIES[i], label=CLASS_LABELS[cls], lw=0)
        ax.set_xticks(range(len(ws))); ax.set_xticklabels([lab[w] for w in ws]); ax.set_ylabel("Bottom-up − direct, 2021→2026 (log pts × 100)"); ax.axhline(0, color=INK2, lw=0.5); ax.legend(fontsize=7)
        save(fig, "fig_aggregation")
    p = R / "validation" / "moran_residuals.csv"
    if exists(p):
        d = pl.read_csv(p)
        fig, ax = plt.subplots(figsize=(4.6, 2.6))
        for i, sc in enumerate(["province", "Montréal", "Québec", "Gatineau"]):
            s = d.filter((pl.col("class") == "house") & (pl.col("scope") == sc)).sort("month"); ax.plot(range(s.height), s["moran_I_knn10"], marker="o", ms=3, color=SERIES[i], label=sc)
            ax.set_xticks(range(s.height)); ax.set_xticklabels(s["month"].to_list(), fontsize=7)
        ax.set_ylabel("Moran's I of stage-1 residuals (10-NN)"); ax.axhline(0, color=INK2, lw=0.5); ax.legend(fontsize=7)
        save(fig, "fig_moran")


if __name__ == "__main__":
    for fn in [fig_coverage, fig_arrival, fig_province_methods, fig_region_methods, fig_municipal_examples, fig_gain_vs_n, fig_hedonic, fig_map, fig_mc, fig_revisions, fig_validation]:
        try:
            fn()
        except Exception as ex:
            print("FAILED", fn.__name__, repr(ex))