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

Reads the step-02/03 outputs and writes eight PNG figures to ``figures/``.
Same conventions as WP10: no in-figure titles on single panels, bold
"Panel A/B" headers on multi-panel figures, one accent hue doubled by a
marker/linestyle difference, direct labels over legend boxes.

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 wp11 import config, sample  # noqa: E402
from wp11.plotstyle import (BLUE, BLUES, GREY, INK, LIGHT, RED, TEXTWIDTH,  # noqa: E402
                            apply_style, panel_label, ygrid)

import matplotlib.pyplot as plt  # 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)


def _forest(ax, res, metric="mdape", xlabel="Median absolute error (%)"):
    y = np.arange(len(res))[::-1].astype(float)
    ax.barh(y, res[metric], height=0.62, color=LIGHT, edgecolor=INK,
            linewidth=0.4)
    for yy, v in zip(y, res[metric]):
        ax.text(v + 0.25, yy, f"{v:.1f}", va="center", fontsize=7.5)
    ax.set_yticks(y)
    ax.set_yticklabels(res["name"], fontsize=8)
    ax.set_xlabel(xlabel)
    ax.spines["left"].set_visible(False)
    ax.tick_params(axis="y", length=0)


# ---------------------------------------------------------------- fig 1-4
def fig_groups():
    """One figure per design axis: MdAPE bars, random holdout."""
    r = pd.read_csv(RES / "horserace_random.csv")
    for gkey, fname in (("A.", "fig_forms.png"), ("B.", "fig_time.png"),
                        ("C.", "fig_space.png"), ("D.", "fig_methods.png")):
        res = r[r["group"].str.startswith(gkey)].reset_index(drop=True)
        fig, ax = plt.subplots(figsize=(0.85 * TEXTWIDTH,
                                        0.55 + 0.42 * len(res)))
        _forest(ax, res)
        ax.set_xlim(0, res["mdape"].max() * 1.15)
        _save(fig, fname)


# ---------------------------------------------------------------- fig 5
def fig_random_vs_temporal():
    """The generalization gap: random vs forward-in-time holdout."""
    a = pd.read_csv(RES / "horserace_random.csv")
    b = pd.read_csv(RES / "horserace_temporal.csv")
    m = a.merge(b, on=["name", "group"], suffixes=("_r", "_t"))
    keep = m["name"].str.match(r"(A5|B4|C4|D)")
    m = m[keep].reset_index(drop=True)
    fig, ax = plt.subplots(figsize=(0.85 * TEXTWIDTH, 3.7))
    y = np.arange(len(m))[::-1].astype(float)
    ax.scatter(m["mdape_r"], y + 0.14, s=22, color=BLUE, zorder=3,
               label="Random 80/20 holdout")
    ax.scatter(m["mdape_t"], y - 0.14, s=24, facecolor="white",
               edgecolor=RED, marker="s", zorder=3,
               label="Train $<$ 2025, test 2025–26")
    for yy, r0, t0 in zip(y, m["mdape_r"], m["mdape_t"]):
        ax.plot([r0, t0], [yy + 0.14, yy - 0.14], color=GREY, lw=0.6,
                zorder=2)
    ax.set_yticks(y)
    ax.set_yticklabels(m["name"], fontsize=8)
    ax.set_xlabel("Median absolute error (%)")
    ax.legend(loc="upper left", fontsize=8, bbox_to_anchor=(0.02, 0.35))
    ax.spines["left"].set_visible(False)
    ax.tick_params(axis="y", length=0)
    _save(fig, "fig_generalization.png")


# ---------------------------------------------------------------- fig 6
def fig_index():
    """Constant-quality monthly indices across specifications."""
    ix = pd.read_csv(RES / "index.csv")
    ix["date"] = pd.PeriodIndex(ix["month"], freq="M").to_timestamp()
    fig, ax = plt.subplots(figsize=(TEXTWIDTH, 3.2))
    forms = ["Semi-log", "Log-log", "Semi-log + quadratics",
             "Semi-log + splines", "Gradient boosting"]
    styles = [(BLUES[2], "-"), (BLUES[3], (0, (5, 2))), (BLUES[4], "-"),
              (BLUES[5], (0, (1, 1.2))), (RED, "-")]
    for f, (c, ls) in zip(forms, styles):
        g = ix[ix["form"] == f].sort_values("date")
        lw = 1.5 if f == "Gradient boosting" else 1.0
        ax.plot(g["date"], g["index"], color=c, ls=ls, lw=lw)
    # the four linear forms end within ~2 points — one bundle label
    lin_end = (ix[ix["form"] != "Gradient boosting"]
               .sort_values("date").groupby("form")["index"].last())
    gb = ix[ix["form"] == "Gradient boosting"].sort_values("date")
    last_date = gb["date"].iloc[-1]
    ax.annotate("Four linear forms", xy=(last_date, lin_end.mean()),
                xytext=(6, 4), textcoords="offset points", fontsize=7.5,
                color=BLUES[4], va="center")
    ax.annotate("Gradient boosting", xy=(last_date, gb["index"].iloc[-1]),
                xytext=(6, -4), textcoords="offset points", fontsize=7.5,
                color=RED, va="center")
    ax.set_ylabel("Constant-quality index (Jan 2021 = 100)")
    ax.margins(x=0.14)
    ygrid(ax)
    _save(fig, "fig_index.png")


# ---------------------------------------------------------------- fig 7
def fig_profiles():
    """Implicit ln-price profiles: OLS quadratic vs gradient boosting."""
    p = pd.read_csv(RES / "profiles.csv")
    fig, axes = plt.subplots(1, 2, figsize=(TEXTWIDTH, 3.0))
    for ax, var, xlabel, ptitle in (
            (axes[0], "age", "Building age (years)", "Panel A. Age profile"),
            (axes[1], "area", "Floor area (m$^2$)",
             "Panel B. Floor-area profile")):
        for model, color, ls in (("OLS quadratic", BLUE, "-"),
                                 ("Gradient boosting", RED, (0, (5, 2)))):
            g = p[(p["var"] == var) & (p["model"] == model)]
            ax.plot(g["x"], g["y"], color=color, ls=ls, lw=1.3, label=model)
        ax.set_xlabel(xlabel)
        ax.set_ylabel("ln price, relative to leftmost point")
        panel_label(ax, ptitle)
    axes[0].legend(fontsize=8)
    fig.subplots_adjust(wspace=0.28)
    _save(fig, "fig_profiles.png")


# ---------------------------------------------------------------- fig 8
def fig_learning():
    """Accuracy versus training-set size."""
    l = pd.read_csv(RES / "learning.csv")
    fig, ax = plt.subplots(figsize=(0.72 * TEXTWIDTH, 3.2))
    for model, color, mk in (("OLS quadratic (muni+quarter FE)", BLUE, "o"),
                             ("Gradient boosting", RED, "s")):
        g = l[l["model"] == model].sort_values("n_train")
        ax.plot(g["n_train"], g["mdape"], marker=mk, ms=4.5, color=color,
                lw=1.2, mfc="white" if mk == "s" else color, label=model)
    ax.set_xscale("log")
    ax.set_xlabel("Training sales (log scale)")
    ax.set_ylabel("Median absolute error (%)")
    ax.legend(fontsize=8)
    ygrid(ax)
    _save(fig, "fig_learning.png")


# ---------------------------------------------------------------- fig 9
def fig_segments():
    """MdAPE by market segment, OLS vs GB."""
    s = pd.read_csv(RES / "segments.csv")
    order = [x for x in ["Class: single_family", "Class: condo",
                         "Class: plex", "Class: cottage",
                         "Muni < 1k sales", "Muni 1k–10k sales",
                         "Muni > 10k sales"] if x in set(s["segment"])]
    labels = {"Class: single_family": "Single-family", "Class: condo":
              "Condominium", "Class: plex": "Plex", "Class: cottage":
              "Cottage", "Muni < 1k sales": "Muni $<$ 1k sales",
              "Muni 1k–10k sales": "Muni 1k–10k sales",
              "Muni > 10k sales": "Muni $>$ 10k sales"}
    y = np.arange(len(order))[::-1].astype(float)
    fig, ax = plt.subplots(figsize=(0.8 * TEXTWIDTH, 3.4))
    for model, color, mk, dy in (("OLS quadratic", BLUE, "o", 0.14),
                                 ("Gradient boosting", RED, "s", -0.14)):
        g = s[s["model"] == model].set_index("segment").loc[order]
        ax.scatter(g["mdape"], y + dy, s=22, color=color, marker=mk,
                   facecolor="white" if mk == "s" else color,
                   edgecolor=color, label=model, zorder=3)
    ax.set_yticks(y)
    ax.set_yticklabels([labels[o] for o in order], fontsize=8)
    ax.set_xlabel("Median absolute error (%)")
    ax.legend(fontsize=8, loc="lower right")
    ax.spines["left"].set_visible(False)
    ax.tick_params(axis="y", length=0)
    _save(fig, "fig_segments.png")


def main() -> None:
    config.ensure_dirs()
    apply_style()
    fig_groups()
    fig_random_vs_temporal()
    fig_index()
    fig_profiles()
    fig_learning()
    fig_segments()
    made = sorted(p.name for p in OUT.glob("fig_*.png"))
    print(f"{len(made)} figures written:", ", ".join(made))


if __name__ == "__main__":
    main()