#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Step 04 — Generate every figure used in the paper (17 PNG files).

Number-bearing summary figures (R^2 ladder, variance decomposition, OOS
accuracy, FSA premia, heterogeneity, quantile) are drawn from the results
tier selected with ``--results`` so the paper's published numbers are used by
default; distribution/scatter/map figures are drawn from the micro sample.

Usage:  python scripts/04_make_figures.py [--results reference|reproduced]
"""
import argparse
import json
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 wp9 import models, sample  # noqa: E402
from wp9.config import FIGURES, REFERENCE, REPRODUCED, RESULTS, ensure_dirs  # noqa: E402
from wp9.plotstyle import (ACCENT, ACCENT2, ACCENT3, GREEN, GREY, LIGHT,  # noqa: E402
                           NEUTRAL, ORANGE, apply_style)

apply_style()
import matplotlib.pyplot as plt  # noqa: E402
import statsmodels.api as sm  # noqa: E402

PROVINCE_NAMES = {"ON": "Ontario", "QC": "Quebec", "BC": "British Columbia",
                  "AB": "Alberta", "SK": "Saskatchewan", "MB": "Manitoba",
                  "NS": "Nova Scotia", "NL": "Nfld. & Labrador", "NB": "New Brunswick"}


def save(fig, name):
    fig.tight_layout()
    fig.savefig(FIGURES / name, bbox_inches="tight")
    plt.close(fig)
    print(" fig", name)


# ---------------------------------------------------------------- micro-data
def fig_price_dist(s):
    fig, ax = plt.subplots(1, 2, figsize=(11, 4))
    shown = s[s["price_cad"] <= 3e6]["price_cad"]
    ax[0].hist(shown / 1e3, bins=80, color=ACCENT, alpha=0.9)
    ax[0].axvline(s["price_cad"].median() / 1e3, color=ACCENT2, ls="--", lw=1.2,
                  label=f"median \\${s['price_cad'].median():,.0f}")
    ax[0].set_xlabel("List price (thousand CAD, truncated at \\$3M)")
    ax[0].set_ylabel("Listings")
    ax[0].set_title("(a) Raw list price")
    ax[0].legend()
    ax[1].hist(s["ln_price"], bins=80, color=ACCENT, alpha=0.9)
    ax[1].set_xlabel("ln(price)")
    ax[1].set_title("(b) Log price (dependent variable)")
    save(fig, "fig_price_dist.png")


def fig_province_ppm2(s):
    g = (s.groupby("prov")["ppm2"].agg(["median", "mean"])
         .sort_values("median"))
    fig, ax = plt.subplots(figsize=(7.5, 4.2))
    y = np.arange(len(g))
    ax.barh(y, g["median"], color=ACCENT, alpha=0.9, label="median")
    ax.barh(y, (g["mean"] - g["median"]).clip(lower=0), left=g["median"],
            color=GREY, alpha=0.55, label="mean$-$median gap")
    ax.set_yticks(y)
    ax.set_yticklabels([PROVINCE_NAMES.get(p, p) for p in g.index])
    ax.set_xlabel("Price per m$^2$ of living area (CAD)")
    ax.set_title("Price per square metre by province")
    ax.legend()
    save(fig, "fig_province_ppm2.png")


def fig_size_gradient(s):
    bins = np.exp(np.linspace(np.log(45), np.log(470), 14))
    banded = (s.assign(band=pd.cut(s["living_m2"], bins))
              .groupby(["band", "cat"], observed=True)
              .agg(x=("living_m2", "median"), price=("price_cad", "median"),
                   n=("price_cad", "size"))
              .reset_index())
    fig, ax = plt.subplots(figsize=(7, 4.4))
    for cat, colour, label in (("house", ACCENT, "Houses"), ("condo", ACCENT2, "Condominiums")):
        g = banded[(banded["cat"] == cat) & (banded["n"] >= 30)]
        ax.plot(g["x"], g["price"] / 1e3, "o-", color=colour, lw=2, ms=5, label=label)
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel("Living area (m$^2$, log scale)")
    ax.set_ylabel("Median list price (thousand CAD, log scale)")
    ax.set_title("Median price by living-area bin and dwelling type")
    ax.legend()
    save(fig, "fig_size_gradient.png")


def _canada_axes(ax, s):
    ax.set_xlim(-140, -50)
    ax.set_ylim(41, 61)
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    for prov, g in s.groupby("prov"):
        if len(g) < 400:
            continue
        ax.annotate(prov, (g["lon"].median(), g["lat"].quantile(0.9) + 1.2),
                    fontsize=8, fontweight="bold", color="#333333", ha="center")


def fig_maps(s):
    colour = np.log(s["ppm2"])
    fig, ax = plt.subplots(figsize=(7, 5))
    sc = ax.scatter(s["lon"], s["lat"], c=colour, s=2, alpha=0.35, cmap="viridis")
    _canada_axes(ax, s)
    ax.set_title("Listing locations, coloured by log price per m$^2$")
    fig.colorbar(sc, ax=ax, shrink=0.75, label="ln(price per m$^2$)")
    save(fig, "fig_map.png")

    g = (s.groupby("fsa")
         .agg(lat=("lat", "mean"), lon=("lon", "mean"),
              ppm2=("ppm2", "median"), n=("ppm2", "size")))
    g = g[g["n"] >= 25]
    fig, ax = plt.subplots(figsize=(7, 5))
    sc = ax.scatter(g["lon"], g["lat"], c=np.log(g["ppm2"]),
                    s=np.sqrt(g["n"]) * 1.8, alpha=0.75, cmap="viridis",
                    edgecolors="white", linewidths=0.2)
    _canada_axes(ax, s)
    ax.set_title("FSA neighbourhood medians (bubble area $\\propto\\sqrt{\\mathrm{listings}}$)")
    fig.colorbar(sc, ax=ax, shrink=0.75, label="ln(median price per m$^2$)")
    save(fig, "fig_fsa_map.png")


def fig_fit_resid():
    grand = pd.read_parquet(REPRODUCED / "grand_model.parquet")
    fig, ax = plt.subplots(figsize=(5.4, 5.2))
    ax.hexbin(grand["pred_grand"], grand["ln_price"], gridsize=90, cmap="Blues",
              mincnt=1, bins="log")
    lims = [grand["ln_price"].min(), grand["ln_price"].max()]
    ax.plot(lims, lims, color=ACCENT2, lw=1.4, ls="--")
    ax.set_xlabel("Predicted ln(price)")
    ax.set_ylabel("Actual ln(price)")
    ax.set_title("Grand model: predicted vs. actual")
    save(fig, "fig_fit.png")

    resid = grand["resid_grand"]
    fig, ax = plt.subplots(1, 2, figsize=(9.6, 4.1))
    ax[0].hist(resid, bins=100, color=ACCENT, alpha=0.9, density=True)
    grid = np.linspace(resid.quantile(0.001), resid.quantile(0.999), 200)
    ax[0].plot(grid, np.exp(-0.5 * ((grid - resid.mean()) / resid.std()) ** 2)
               / (resid.std() * np.sqrt(2 * np.pi)), color=ACCENT2, lw=1.5,
               label="Normal density")
    ax[0].set_xlabel("Residual")
    ax[0].set_title("(a) Residual distribution")
    ax[0].legend()
    sm.qqplot(resid, line="45", fit=True, ax=ax[1], markerfacecolor=ACCENT,
              markeredgecolor=ACCENT, markersize=2, alpha=0.4)
    ax[1].set_title("(b) Normal Q--Q plot")
    save(fig, "fig_resid.png")


def fig_moran(res_dir):
    ext2 = json.load(open(res_dir / "ext2.json"))
    coords = np.load(REPRODUCED / "moran_coords.npy")
    panels = [("Structural-only residuals", np.load(REPRODUCED / "moran_resid_struct.npy"),
               ext2["moran"]["struct_I"]),
              ("Grand model (FSA FE) residuals", np.load(REPRODUCED / "moran_resid_grand.npy"),
               ext2["moran"]["grand_I"])]
    from sklearn.neighbors import NearestNeighbors
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4))
    for j, (title, resid, moran) in enumerate(panels):
        nn = NearestNeighbors(n_neighbors=11).fit(coords)
        _, idx = nn.kneighbors(coords)
        z = resid - resid.mean()
        lag = z[idx[:, 1:]].mean(axis=1)
        ax[j].scatter(z, lag, s=3, alpha=0.15, color=ACCENT)
        slope, intercept = np.polyfit(z, lag, 1)
        xs = np.linspace(z.min(), z.max(), 10)
        ax[j].plot(xs, slope * xs + intercept, color=ACCENT2, lw=1.8)
        ax[j].axhline(0, color="k", lw=0.5)
        ax[j].axvline(0, color="k", lw=0.5)
        ax[j].set_xlabel("Residual ($z$)")
        ax[j].set_ylabel("Spatial lag of residual")
        ax[j].set_title(f"{title}\nMoran's I = {moran:.3f}")
    save(fig, "fig_moran.png")


def fig_gradient(res_dir):
    banded = pd.read_csv(REPRODUCED / "gradient_bins.csv")
    fig, ax = plt.subplots(figsize=(7, 4.2))
    ax.plot(banded["x"], (np.exp(banded["prem"]) - 1) * 100, "o-", color=ACCENT, lw=2)
    ax.axhline(0, color="k", lw=0.7, ls=":")
    ax.set_xscale("symlog")
    ax.set_xlabel("Distance to nearest major metro (km, symlog)")
    ax.set_ylabel("Location premium vs structure-only (%)")
    ax.set_title("The urban price gradient: value falls with distance to metro")
    save(fig, "fig_gradient.png")


def fig_nonlinear(res_dir):
    band = pd.read_csv(REPRODUCED / "nonlinear_band.csv")
    baseline = pd.read_csv(res_dir / "robustness.csv").iloc[0]["ln_living"]
    fig, ax = plt.subplots(figsize=(7, 4.2))
    area = np.exp(band["ln_area"])
    ax.plot(area, band["elasticity"], color=ACCENT, lw=2)
    ax.fill_between(area, band["elasticity"] - 1.96 * band["se"],
                    band["elasticity"] + 1.96 * band["se"], color=ACCENT, alpha=0.18)
    ax.axhline(baseline, color=ACCENT2, ls="--", lw=1,
               label=f"linear-model elasticity ({baseline:.2f})")
    ax.set_xscale("log")
    ax.set_xlabel("Living area (m$^2$)")
    ax.set_ylabel("Marginal elasticity of price w.r.t. area")
    ax.set_title("Diminishing returns to floor space (quadratic spec., grand model)")
    ax.legend()
    save(fig, "fig_nonlinear.png")


# ----------------------------------------------------------- results-driven
def fig_r2(res_dir):
    fit = json.load(open(res_dir / "fit.json"))
    names = ["M1", "M2", "M3", "M4", "M5"]
    labels = ["M1\nStructural", "M2\n+Type/Own.", "M3\n+Province",
              "M4\nHouses+FSA", "M5\nGrand+FSA"]
    values = [fit[m]["r2"] for m in names]
    fig, ax = plt.subplots(figsize=(7.5, 4))
    bars = ax.bar(labels, values, color=[ACCENT, ACCENT, ACCENT, ACCENT3, ACCENT2], alpha=0.92)
    for bar, val in zip(bars, values):
        ax.text(bar.get_x() + bar.get_width() / 2, val + 0.012, f"{val:.3f}",
                ha="center", fontsize=9)
    ax.set_ylim(0, 0.9)
    ax.set_ylabel("$R^2$ (share of log-price variance explained)")
    ax.set_title("Explanatory power across the specification ladder")
    save(fig, "fig_r2.png")


def fig_decomp(res_dir):
    fit = json.load(open(res_dir / "fit.json"))
    blocks = [("Structure", fit["M1"]["r2"]),
              ("Dwelling type & ownership", fit["M2"]["r2"] - fit["M1"]["r2"]),
              ("Province", fit["M3"]["r2"] - fit["M2"]["r2"]),
              ("Neighbourhood (FSA)", fit["M5"]["r2"] - fit["M3"]["r2"]),
              ("Unexplained", 1 - fit["M5"]["r2"])]
    colours = [ACCENT, ACCENT3, GREEN, ACCENT2, NEUTRAL]
    fig, ax = plt.subplots(figsize=(8, 2.4))
    left = 0.0
    for (label, share), colour in zip(blocks, colours):
        ax.barh(0, share, left=left, color=colour, edgecolor="white")
        if share > 0.03:
            ax.text(left + share / 2, 0, f"{label}\n{share * 100:.1f}%",
                    ha="center", va="center", fontsize=8,
                    color="black" if colour == NEUTRAL else "white")
        left += share
    ax.set_xlim(0, 1)
    ax.set_ylim(-0.5, 0.5)
    ax.set_yticks([])
    ax.set_xlabel("Share of log-price variance")
    ax.set_title("Variance decomposition of Canadian house prices")
    save(fig, "fig_decomp.png")


def fig_forest(res_dir):
    coefs = pd.read_csv(res_dir / "coef_M3.csv", index_col=0)
    order = [("ln_living", "$\\ln$ living area"), ("bathrooms", "Full bathrooms"),
             ("half_baths", "Half bathrooms"), ("bedrooms", "Bedrooms"),
             ("parking_n", "Parking spaces"), ("stories_n", "Storeys"),
             ("ln_lot", "$\\ln(1+$lot m$^2)$")]
    rows = [(label, coefs.loc[key, "coef"], coefs.loc[key, "se"])
            for key, label in order if key in coefs.index]
    fig, ax = plt.subplots(figsize=(7, 4))
    y = np.arange(len(rows))[::-1]
    ax.errorbar([r[1] for r in rows], y, xerr=[1.96 * r[2] for r in rows],
                fmt="o", color=ACCENT, ecolor=GREY, capsize=3, ms=6)
    ax.axvline(0, color=ACCENT2, ls="--", lw=1)
    ax.set_yticks(y)
    ax.set_yticklabels([r[0] for r in rows])
    ax.set_xlabel("Coefficient on ln(price), 95% cluster-robust CI")
    ax.set_title("Structural implicit prices (houses, province-FE model M3)")
    save(fig, "fig_forest.png")


def fig_heterogeneity(res_dir):
    het = pd.read_csv(res_dir / "heterogeneity.csv").sort_values("elast")
    fig, ax = plt.subplots(figsize=(7, 4))
    y = np.arange(len(het))
    ax.errorbar(het["elast"], y, xerr=1.96 * het["se_el"], fmt="o",
                color=ACCENT, ecolor=GREY, capsize=3, ms=6)
    ax.axvline(het["elast"].mean(), color=ACCENT2, ls="--", lw=1,
               label="cross-province mean")
    ax.set_yticks(y)
    ax.set_yticklabels([f"{p} (n={n:,})" for p, n in zip(het["prov"], het["n"])])
    ax.set_xlabel("Living-area elasticity (within-FSA)")
    ax.set_title("Heterogeneity in the size elasticity of price across provinces")
    ax.legend()
    save(fig, "fig_heterogeneity.png")


def fig_premia(res_dir):
    premia = pd.read_csv(res_dir / "fsa_premia.csv")
    premia = premia.rename(columns={premia.columns[0]: "fsa"})
    shown = pd.concat([premia.nsmallest(12, "premium_pct"),
                       premia.nlargest(12, "premium_pct")]).sort_values("premium_pct")
    colours = [ACCENT2 if v < 0 else ACCENT for v in shown["premium_pct"]]
    fig, ax = plt.subplots(figsize=(7.5, 5))
    ax.barh(np.arange(len(shown)), shown["premium_pct"], color=colours, alpha=0.9)
    ax.set_yticks(np.arange(len(shown)))
    ax.set_yticklabels([f"{f} ({p})" for f, p in zip(shown["fsa"], shown["prov"])],
                       fontsize=8)
    ax.axvline(0, color="k", lw=0.8)
    ax.set_xlabel("Neighbourhood (FSA) price premium vs. national median (%), net of structure")
    ax.set_title("Highest- and lowest-valued neighbourhoods in Canada")
    save(fig, "fig_premia.png")


def fig_quantile(res_dir):
    q = pd.read_csv(res_dir / "quantile.csv")
    taus = q[q["tau"].notna()]
    ols = q[q["tau"].isna()].iloc[0]
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    for j, (var, se_var, title) in enumerate(
            [("ln_living", "se_living", "$\\ln$ living-area elasticity"),
             ("bathrooms", "se_bath", "Full-bathroom premium")]):
        ax[j].errorbar(taus["tau"], taus[var], yerr=1.96 * taus[se_var],
                       fmt="o-", color=ACCENT, capsize=3, label="Quantile")
        ax[j].axhline(ols[var], color=ACCENT2, ls="--", label="OLS")
        ax[j].set_xlabel("Quantile of price ($\\tau$)")
        ax[j].set_title(title)
        ax[j].legend()
    save(fig, "fig_quantile.png")


def fig_oos(res_dir):
    oos = json.load(open(res_dir / "oos.json"))
    within10 = oos["within10"]
    mid = oos["within20"] - oos["within10"]
    beyond = 100 - oos["within20"]
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar([5, 15, 30], [within10, mid, beyond], width=8,
           color=[ACCENT, ACCENT3, LIGHT], alpha=0.9)
    for x, v, label in ((5, within10, "≤10%"), (15, mid, "10–20%"), (30, beyond, ">20%")):
        ax.text(x, v + 1, f"{label}\n{v:.0f}%", ha="center", fontsize=9)
    ax.axvline(oos["median_ape"], color=ACCENT2, ls="--",
               label=f"median APE {oos['median_ape']:.1f}%")
    ax.set_xticks([5, 15, 30])
    ax.set_xticklabels(["≤10%", "10–20%", ">20%"])
    ax.set_xlabel("Absolute percentage error (out-of-sample)")
    ax.set_ylabel("Share of test listings (%)")
    ax.set_ylim(0, 52)
    ax.set_title(f"Out-of-sample valuation accuracy (OOS $R^2$={oos['oos_r2']:.3f})", pad=12)
    ax.legend(loc="upper right", fontsize=8)
    save(fig, "fig_oos.png")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results", choices=["reference", "reproduced"],
                        default="reference",
                        help="results tier for number-bearing figures (default: reference)")
    args = parser.parse_args()
    res_dir = REFERENCE if args.results == "reference" else REPRODUCED
    ensure_dirs()
    s = sample.load_sample()

    fig_price_dist(s)
    fig_province_ppm2(s)
    fig_size_gradient(s)
    fig_maps(s)
    fig_fit_resid()
    fig_r2(res_dir)
    fig_decomp(res_dir)
    fig_forest(res_dir)
    fig_heterogeneity(res_dir)
    fig_premia(res_dir)
    fig_quantile(res_dir)
    fig_oos(res_dir)
    fig_nonlinear(res_dir)
    fig_gradient(res_dir)
    fig_moran(res_dir)
    print("all figures written to", FIGURES)


if __name__ == "__main__":
    main()