#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Step 04 — Publication figures (print-journal calibre).

Reads the analysis sample and the step-02/03 outputs, writes ten PNG figures
to ``figures/``. No titles are drawn inside single-panel figures — captions
in the manuscript carry the message; multi-panel figures use bold
"Panel A/B" headers. One accent hue per figure, doubled by linestyle or
marker so nothing is encoded by colour alone.

Usage:  python scripts/04_make_figures.py
"""
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from wp10 import config, sample  # noqa: E402
from wp10.plotstyle import (BLUE, BLUES, GREY, INK, LIGHT, RED, TEXTWIDTH,  # noqa: E402
                            apply_style, panel_label, ygrid)

import matplotlib.pyplot as plt  # noqa: E402
import matplotlib.dates as mdates  # noqa: E402

OUT = config.FIGURES
RES = config.REPRODUCED


def _save(fig, name):
    fig.savefig(OUT / name, bbox_inches="tight", pad_inches=0.02)
    plt.close(fig)


# ---------------------------------------------------------------- fig 1
def fig_ratio_dist(df):
    """Panel A: ratio histogram. Panel B: densities by roll lag."""
    fig, axes = plt.subplots(1, 2, figsize=(TEXTWIDTH, 2.9))

    ax = axes[0]
    ax.hist(df["ratio"], bins=120, range=(0, 2.5), color=LIGHT,
            edgecolor=INK, linewidth=0.25)
    med = df["ratio"].median()
    ax.axvline(med, color=INK, lw=0.9)
    ax.axvline(1.0, color=GREY, lw=0.8, ls=(0, (4, 3)))
    ax.text(med - 0.05, ax.get_ylim()[1] * 0.97, f"median = {med:.2f}",
            ha="right", va="top", fontsize=8)
    ax.text(1.04, ax.get_ylim()[1] * 0.75, "AV = SP", fontsize=8, color=GREY)
    ax.set_xlabel("Assessment ratio $AV/SP$")
    ax.set_ylabel("Sales")
    ax.set_yticks([])
    ax.spines["left"].set_visible(False)
    panel_label(ax, "Panel A. All sales")

    ax = axes[1]
    specs = [("Roll lag $<$ 24 m", df["lag_months"] < 24, BLUES[2], "-"),
             ("24–48 m", df["lag_months"].between(24, 48), BLUES[4], (0, (5, 2))),
             ("$>$ 48 m", df["lag_months"] > 48, BLUES[5], (0, (1, 1.2)))]
    for label, m, color, ls in specs:
        ax.hist(df.loc[m, "ratio"], bins=120, range=(0, 2.5), density=True,
                histtype="step", lw=1.2, color=color, ls=ls, label=label)
    ax.set_xlabel("Assessment ratio $AV/SP$")
    ax.set_ylabel("Density")
    ax.legend(loc="upper right", handlelength=2.4)
    panel_label(ax, "Panel B. By roll lag at sale")
    fig.subplots_adjust(wspace=0.25)
    _save(fig, "fig_ratio_dist.png")


# ---------------------------------------------------------------- fig 2
def fig_binscatter():
    """Within-cell binned scatter of ln ratio on ln price — the core fact."""
    b = pd.read_csv(RES / "binscatter.csv")
    fig, ax = plt.subplots(figsize=(0.72 * TEXTWIDTH, 3.3))
    ax.axhline(0, color=GREY, lw=0.7, ls=(0, (4, 3)))
    slope = np.polyfit(b["x"], b["y"], 1, w=b["n"])[0]
    icept = np.average(b["y"] - slope * b["x"], weights=b["n"])
    xs = np.linspace(b["x"].min(), b["x"].max(), 50)
    ax.plot(xs, slope * xs + icept, color=BLUE, lw=1.2, zorder=2)
    ax.errorbar(b["x"], b["y"], yerr=1.96 * b["se"], fmt="o", ms=4,
                mfc=INK, mec=INK, ecolor=INK, elinewidth=0.6, capsize=0,
                lw=0, zorder=3)
    ax.annotate(f"slope = ${slope:.3f}$", xy=(0.55, slope * 0.55 + icept),
                xytext=(0.32, 0.12), fontsize=8.5, color=BLUE,
                arrowprops=dict(arrowstyle="-", color=BLUE, lw=0.6,
                                shrinkA=2, shrinkB=2))
    ax.set_xlabel("ln sale price (demeaned within municipality × roll × year)")
    ax.set_ylabel("ln assessment ratio (demeaned)")
    _save(fig, "fig_binscatter.png")


# ---------------------------------------------------------------- fig 3
def fig_time(df):
    """Median ratio by sale month; sequential blues by roll vintage,
    each segment labelled directly (no legend)."""
    d = df.copy()
    d["month"] = d["sale_date"].dt.to_period("M").dt.to_timestamp()
    fig, ax = plt.subplots(figsize=(TEXTWIDTH, 2.9))
    rolls = sorted(d["roll"].unique())
    for i, roll in enumerate(rolls):
        gg = d[d["roll"] == roll].groupby("month")["ratio"]
        g = gg.median()[gg.size() >= 100]  # drop thin months (spurious spikes)
        g = g[g.index.notna()].sort_index()
        if len(g) < 3:
            continue
        color = BLUES[min(i, len(BLUES) - 1)]
        ax.plot(g.index, g.values, lw=1.3, color=color)
        ax.annotate(roll, xy=(g.index[-1], g.values[-1]),
                    xytext=(3, 0), textcoords="offset points",
                    fontsize=7.5, color=color, va="center")
    ax.axhline(1.0, color=GREY, lw=0.7, ls=(0, (4, 3)))
    ax.text(pd.Timestamp("2021-02-01"), 1.012, "parity", fontsize=7.5,
            color=GREY)
    ax.set_ylabel("Median assessment ratio")
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y"))
    ygrid(ax)
    _save(fig, "fig_time.png")


# ---------------------------------------------------------------- fig 4
def fig_prb_muni():
    """Distribution of municipality-level PRB."""
    m = pd.read_csv(RES / "iaao_muni.csv")
    fig, ax = plt.subplots(figsize=(0.72 * TEXTWIDTH, 3.1))
    ax.hist(m["prb"], bins=40, color=LIGHT, edgecolor=INK, linewidth=0.3)
    lo, hi = config.IAAO_PRB_RANGE
    ax.axvspan(lo, hi, color="0.92", zorder=0)
    ax.axvline(0, color=GREY, lw=0.7, ls=(0, (4, 3)))
    med = m["prb"].median()
    ax.axvline(med, color=RED, lw=1.0)
    ymax = ax.get_ylim()[1]
    ax.text(med - 0.008, ymax * 0.97, f"median = {med:.2f}", ha="right",
            va="top", fontsize=8, color=RED)
    ax.text((lo + hi) / 2, ymax * 0.55, "IAAO\nband", ha="center",
            fontsize=7.5, color="0.35")
    share = (m["prb"] < 0).mean()
    ax.text(0.02, 0.97, f"{share:.0%} of municipalities\nhave PRB $<$ 0",
            transform=ax.transAxes, fontsize=8, va="top")
    ax.set_xlabel("Municipality-level PRB (median across sale years)")
    ax.set_ylabel("Municipalities")
    ygrid(ax)
    _save(fig, "fig_prb_muni.png")


# ---------------------------------------------------------------- fig 5
def fig_map():
    """Municipal PRB across the province (diverging hue, neutral midpoint)."""
    m = pd.read_csv(RES / "iaao_muni.csv")
    fig, ax = plt.subplots(figsize=(TEXTWIDTH, 4.4))
    v = m["prb"].clip(-0.30, 0.10)
    sc = ax.scatter(m["lng"], m["lat"], c=v, s=np.sqrt(m["n"]) * 0.9,
                    cmap="RdBu", vmin=-0.30, vmax=0.30,
                    edgecolor=INK, linewidth=0.25, alpha=0.9)
    cb = fig.colorbar(sc, ax=ax, shrink=0.75, pad=0.02)
    cb.set_label("PRB (negative = regressive)", fontsize=8)
    cb.ax.tick_params(labelsize=7.5)
    cb.outline.set_linewidth(0.5)
    offsets = {"Montréal": (6, -14), "Québec": (8, 4), "Gatineau": (-8, -14),
               "Sherbrooke": (8, -10), "Saguenay": (8, 4)}
    for _, r in m.nlargest(12, "n").iterrows():
        if r["name"] in offsets:
            ax.annotate(r["name"], (r["lng"], r["lat"]),
                        xytext=offsets[r["name"]], textcoords="offset points",
                        fontsize=7.5)
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.set_xlim(-80, -63)
    ax.set_ylim(44.9, 49.6)
    _save(fig, "fig_map.png")


# ---------------------------------------------------------------- fig 6
def fig_quantile():
    """β(τ) with a shaded 95% band; FE/IV benchmarks labelled directly."""
    q = pd.read_csv(RES / "quantile.csv")
    v = pd.read_csv(RES / "vertical.csv")
    fe = v.loc[v["estimator"].str.startswith("Cheng FE"), "beta"].iloc[0]
    iv = v.loc[v["estimator"].str.startswith("Clapp"), "beta"].iloc[0]
    fig, ax = plt.subplots(figsize=(0.72 * TEXTWIDTH, 3.3))
    ax.fill_between(q["tau"], q["beta"] - 1.96 * q["se"],
                    q["beta"] + 1.96 * q["se"], color=BLUE, alpha=0.18, lw=0)
    ax.plot(q["tau"], q["beta"], "o-", color=BLUE, ms=4, lw=1.3)
    x1, x0 = q["tau"].max(), q["tau"].min()
    for yv, lab, ls, xa, ha in [
            (1.0, r"$\beta = 1$ (proportional)", (0, (4, 3)), x1, "right"),
            (iv, f"Clapp IV ({iv:.2f})", (0, (1, 1.2)), x1, "right"),
            (fe, f"Cheng FE ({fe:.2f})", (0, (6, 2)), x0, "left")]:
        ax.axhline(yv, color=GREY, lw=0.8, ls=ls)
        ax.annotate(lab, xy=(xa, yv), xytext=(0, 3), textcoords="offset points",
                    ha=ha, fontsize=7.5, color="0.35")
    ax.set_xlabel(r"Quantile $\tau$ of the conditional $\ln AV$ distribution")
    ax.set_ylabel(r"$\beta(\tau)$")
    ax.set_xticks(q["tau"])
    _save(fig, "fig_quantile.png")


# ---------------------------------------------------------------- fig 7
def fig_heterogeneity():
    """Forest plot of Cheng-FE γ by subgroup, with panel groupings."""
    h = pd.read_csv(RES / "heterogeneity.csv").set_index("group")
    panels = [
        ("Property class", ["Single-family", "Condominium", "Plex (2–5 units)",
                            "Cottage"]),
        ("Building age", ["Age < 20 y", "Age 20–60 y", "Age > 60 y"]),
        ("Assessed land share", ["Land share < 0.2", "Land share 0.2–0.4",
                                 "Land share > 0.4"]),
        ("Roll lag at sale", ["Roll lag < 24 m", "Roll lag 24–48 m",
                              "Roll lag > 48 m"]),
        ("Municipality size", ["Muni < 1k sales", "Muni 1k–10k sales",
                               "Muni > 10k sales"]),
        ("Sale year", [f"Sales {y}" for y in range(2021, 2027)]),
    ]
    rows, ypos, headers, seps = [], [], [], []
    y = 0
    for name, keys in panels:
        headers.append((y, name))
        y -= 1
        for k in keys:
            if k in h.index:
                rows.append((y, k, h.loc[k]))
                ypos.append(y)
                y -= 1
        seps.append(y + 0.5)
        y -= 0.6
    fig, ax = plt.subplots(figsize=(0.85 * TEXTWIDTH, 5.6))
    for yy, k, r in rows:
        ax.errorbar(r["gamma"], yy, xerr=1.96 * r["se"], fmt="o", ms=3.8,
                    mfc=INK, mec=INK, ecolor=INK, elinewidth=0.8, capsize=1.5)
    ax.axvline(0, color=GREY, lw=0.7, ls=(0, (4, 3)))
    for yy, name in headers:
        ax.text(-0.72, yy, name, fontsize=8.5, fontweight="bold", va="center")
    labels = {yy: k.replace("<", "$<$").replace(">", "$>$")
              for yy, k, _ in rows}
    ax.set_yticks(list(labels.keys()))
    ax.set_yticklabels(labels.values(), fontsize=8)
    ax.set_ylim(y + 0.4, 1.0)
    ax.set_xlim(-0.72, 0.12)
    ax.set_xlabel(r"$\gamma$ = elasticity of the assessment ratio with respect"
                  " to price (negative = regressive)")
    ax.spines["left"].set_visible(False)
    ax.tick_params(axis="y", length=0)
    _save(fig, "fig_heterogeneity.png")


# ---------------------------------------------------------------- fig 8
def fig_taxshift():
    """Median excess tax burden by within-market price decile (polarity)."""
    t = pd.read_csv(RES / "taxshift.csv")
    fig, ax = plt.subplots(figsize=(0.85 * TEXTWIDTH, 3.1))
    vals = 100 * t["median_rel"]
    colors = [RED if v > 0 else BLUE for v in vals]
    ax.bar(t["decile"], vals, width=0.72, color=colors, edgecolor=INK,
           linewidth=0.4)
    ax.axhline(0, color=INK, lw=0.7)
    for d, v in zip(t["decile"], vals):
        va = "bottom" if v > 0 else "top"
        off = 1.2 if v > 0 else -1.2
        ax.text(d, v + off, f"{v:+.1f}", ha="center", va=va, fontsize=7.5)
    ax.text(2.5, 45, "over-taxed", fontsize=8, color=RED, ha="center")
    ax.text(8.5, 14, "under-taxed", fontsize=8, color=BLUE, ha="center")
    ax.set_xticks(t["decile"])
    ax.set_xlabel("Within-market sale-price decile")
    ax.set_ylabel("Median excess tax burden (%)")
    ax.set_ylim(min(vals) - 8, max(vals) + 9)
    _save(fig, "fig_taxshift.png")


# ---------------------------------------------------------------- fig 9
def fig_cod():
    """Panel A: COD distribution. Panel B: COD vs market size."""
    m = pd.read_csv(RES / "iaao_muni.csv")
    fig, axes = plt.subplots(1, 2, figsize=(TEXTWIDTH, 2.9))

    ax = axes[0]
    ax.hist(m["cod"], bins=40, color=LIGHT, edgecolor=INK, linewidth=0.3)
    ax.axvline(config.IAAO_COD_MAX_SF, color=RED, lw=1.0)
    ymax = ax.get_ylim()[1]
    ax.text(config.IAAO_COD_MAX_SF + 0.6, ymax * 0.95,
            f"IAAO ceiling ({config.IAAO_COD_MAX_SF:.0f})", fontsize=7.5,
            color=RED, va="top")
    med = m["cod"].median()
    ax.axvline(med, color=INK, lw=0.9, ls=(0, (5, 2)))
    ax.text(med + 0.6, ymax * 0.72, f"median = {med:.1f}", fontsize=7.5)
    ax.set_xlabel("Municipality COD (median across sale years)")
    ax.set_ylabel("Municipalities")
    panel_label(ax, "Panel A. Distribution")

    ax = axes[1]
    ax.scatter(m["n"], m["cod"], s=7, facecolor="none", edgecolor=INK,
               linewidth=0.5, alpha=0.65)
    ax.axhline(config.IAAO_COD_MAX_SF, color=RED, lw=0.9)
    ax.set_xscale("log")
    ax.set_xlabel("Sales in municipality, 2021–2026 (log scale)")
    ax.set_ylabel("COD")
    panel_label(ax, "Panel B. COD and market size")
    fig.subplots_adjust(wspace=0.25)
    _save(fig, "fig_cod.png")


# ---------------------------------------------------------------- fig 10
def fig_robustness():
    """γ across sample variants — FE (filled circles) vs IV (open squares)."""
    r = pd.read_csv(RES / "robustness.csv")
    fig, ax = plt.subplots(figsize=(0.85 * TEXTWIDTH, 3.6))
    y = np.arange(len(r))[::-1].astype(float)
    ax.errorbar(r["gamma_fe"], y + 0.16, xerr=1.96 * r["se_fe"], fmt="o",
                ms=4, mfc=BLUE, mec=BLUE, ecolor=BLUE, elinewidth=0.8,
                capsize=1.5, lw=0, label="Cheng FE")
    ax.errorbar(r["gamma_iv"], y - 0.16, xerr=1.96 * r["se_iv"], fmt="s",
                ms=4, mfc="white", mec=RED, ecolor=RED, elinewidth=0.8,
                capsize=1.5, lw=0, label="Clapp IV")
    ax.axvline(0, color=GREY, lw=0.7, ls=(0, (4, 3)))
    ax.set_yticks(y)
    labels = [str(v).replace("<=", "$\\leq$").replace(">=", "$\\geq$")
              .replace("<", "$<$").replace("$100k", "\\$100k")
              for v in r["variant"]]
    ax.set_yticklabels(labels, fontsize=8)
    ax.set_xlabel(r"$\gamma$ (negative = regressive)")
    ax.legend(loc="lower left", markerscale=1.1)
    ax.spines["left"].set_visible(False)
    ax.tick_params(axis="y", length=0)
    _save(fig, "fig_robustness.png")


def main() -> None:
    config.ensure_dirs()
    apply_style()
    df = sample.load()
    fig_ratio_dist(df)
    fig_binscatter()
    fig_time(df)
    fig_prb_muni()
    fig_map()
    fig_quantile()
    fig_heterogeneity()
    fig_taxshift()
    fig_cod()
    fig_robustness()
    made = sorted(p.name for p in OUT.glob("fig_*.png"))
    print(f"{len(made)} figures written:", ", ".join(made))


if __name__ == "__main__":
    main()