#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
#
"""Generate the LaTeX table fragments (paper/tables/*.tex) and the three new
figures (figures/fig12-fig14) of the October 2026 revision from the CSV outputs
of scripts 05-07 in results/revision/. Nothing is typed by hand: every number
in these tables comes from a stored CSV.

Usage:  python scripts/08_revision_tables_figures.py
"""

import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from wp3 import config  # noqa: E402
from wp3.plotting import ACCENT, NEUTRAL, PRIMARY, PRIMARY_LIGHT, apply_paper_style, label, panel_title  # noqa: E402

REV = config.RESULTS_DIR / "revision"
TAB = config.PROJECT_ROOT / "paper" / "tables"
FIG = config.FIGURES_DIR
TAB.mkdir(exist_ok=True, parents=True)

HOLDOUT = set(config.HOLDOUT_STATES)


def f(x, d=3, plus=False):
    if pd.isna(x):
        return "---"
    s = f"{x:+.{d}f}" if plus else f"{x:.{d}f}"
    return s.replace("-", "$-$")


def n(x):
    return f"{int(x):,}"


def write(name, body):
    r"""Write a table body. The final row is written without its trailing
    line break: the \\ is supplied by the LaTeX source after \input, because a
    row terminator at the very end of an \input file leaves tokens between
    \cr and the following \bottomrule (``Misplaced \noalign'')."""
    body = body.rstrip()
    if body.endswith("\\\\"):
        body = body[:-2].rstrip()
    (TAB / name).write_text(body + "\n")
    print("wrote", TAB / name)


def tex_label(v):
    return label(v).replace("Log(", "$\\ln$(")


# --------------------------------------------------------------------------- tables
def table_missingness():
    m = pd.read_csv(REV / "missingness_full.csv")
    imp = pd.read_csv(REV / "implausible_values.csv", index_col=0)["n"]
    nice = {"avg_school_rating": "Average school rating", "school_count": "School count",
            "nearest_school_distance": "Nearest-school distance", "transit_score": "Transit Score",
            "hoa_fee_value (ln HOA)": "HOA fee ($\\ln$ HOA)", "has_garage": "Garage indicator",
            "roof_type / construction_material / foundation_type": "Roof, construction, foundation type",
            "flooring (central air / forced air / hardwood dummies)": "Flooring (hardwood dummy)",
            "stories": "Stories", "parking_spaces": "Parking spaces",
            "year_built (age, age^2, interactions)": "Year built (age, age$^2$, interactions)",
            "lot_size_sqft (ln lot)": "Lot size ($\\ln$ lot)",
            "cooling_type (central air / forced air / hardwood dummies)": "Cooling type (central-air dummy)",
            "heating_type (central air / forced air / hardwood dummies)": "Heating type (forced-air dummy)",
            "bike_score": "Bike Score", "walk_score": "Walk Score"}
    rows = []
    for _, r in m.iterrows():
        if r["missing_share"] < 0.0005:
            continue
        rows.append(f"{nice.get(r['input'], r['input']).replace('_', '\\_')} & "
                    f"{r['missing_share'] * 100:.1f} & {r['imputation'].replace('_', '\\_')} \\\\")
    body = "\n".join(rows)
    write("missingness.tex", body)
    rows = [f"{k.replace('_', '\\_').replace('>', '$>$')} & {n(v)} & {v / 788842 * 100:.2f} \\\\"
            for k, v in imp.items()]
    write("implausible.tex", "\n".join(rows))


def table_cluster_se():
    c = pd.read_csv(REV / "cluster_se.csv", index_col=0)
    key = ["ln_sqft", "bedrooms", "bathrooms", "age", "age_sq", "ln_lot", "has_pool",
           "has_garage", "on_waterfront", "has_central_air", "has_hardwood", "luxury_score",
           "walk_score", "bike_score", "transit_score", "avg_school_rating",
           "property_tax_rate", "is_condo", "has_hoa", "tag_new_construction",
           "tag_foreclosure", "region_Northeast", "region_South", "region_West"]
    rows = []
    for v in key:
        r = c.loc[v]
        rows.append(f"{tex_label(v)} & {f(r.coef, 4)} & {f(r.se_hc3, 4)} & {f(r.se_cluster_zip3, 4)} & "
                    f"{f(r.se_cluster_state, 4)} & {f(r.ratio_zip3_hc3, 1)} & {f(r.ratio_state_hc3, 1)} & "
                    f"{f(r.t_cluster_state, 2)} \\\\")
    sig = {k: int((c[f"t_{k}"].abs() > 1.96).sum()) for k in ["hc3", "cluster_zip3", "cluster_state"]}
    rows.append("\\midrule")
    rows.append(f"Coefficients significant at 5\\% (of {len(c)}) & & {sig['hc3']} & {sig['cluster_zip3']} & "
                f"{sig['cluster_state']} & & & \\\\")
    rows.append(f"Median SE ratio across all coefficients & & & & & {c.ratio_zip3_hc3.median():.1f} & "
                f"{c.ratio_state_hc3.median():.1f} & \\\\")
    write("cluster_se.tex", "\n".join(rows))
    return c, sig


def table_lot_recoding():
    d = pd.read_csv(REV / "lot_recoding.csv")
    labels = ["Published: $\\ln(\\text{lot}) = 0$ if missing or zero",
              "Recoded: log of state-median lot if missing or zero",
              "Recoded + missing-lot indicator",
              "Complete cases (lot reported, positive)",
              "Recoded, lot winsorised at p99.5"]
    rows = []
    for lab, (_, r) in zip(labels, d.iterrows()):
        ind = f(r.beta_lot_missing_indicator, 3) if not pd.isna(r.beta_lot_missing_indicator) else "---"
        rows.append(f"{lab} & {n(r.N)} & {f(r.R2, 3)} & {f(r.beta_ln_lot, 3)} & ({f(r.se_ln_lot, 4)}) & "
                    f"{f(r.beta_ln_sqft, 3)} & {ind} \\\\")
    write("lot_recoding.tex", "\n".join(rows))
    return d


def table_state_heterogeneity():
    s = pd.read_csv(REV / "state_heterogeneity_summary.csv")
    p = pd.read_csv(REV / "pooled_vs_interacted.csv")
    rows = []
    for _, r in s.iterrows():
        rows.append(f"{tex_label(r.variable)} & {f(r.pooled_state_fe, 3)} & {f(r.p10, 3)} & {f(r["median"], 3)} & "
                    f"{f(r.p90, 3)} & {f(r['min'], 3)} & {f(r['max'], 3)} & {int(r.n_states_sign_differs_from_pooled)} \\\\")
    write("state_heterogeneity.tex", "\n".join(rows))
    rows = [f"{r.model} & {f(r.r2_test, 3)} \\\\" for _, r in p.iterrows()]
    write("pooled_vs_interacted.tex", "\n".join(rows))
    return s, p


def table_r2_bootstrap():
    b = pd.read_csv(REV / "r2_bootstrap.csv")
    for scheme, fname in [("observation bootstrap", "r2_bootstrap_obs.tex"),
                          ("ZIP3-cluster bootstrap", "r2_bootstrap_zip3.tex")]:
        sub = b[b.scheme == scheme]
        rows = []
        for _, r in sub.iterrows():
            q = r.quantity.replace("R2", "$R^2$")
            rows.append(f"\\quad {q} & {f(r.point, 4)} & [{f(r.ci_low, 4)}, {f(r.ci_high, 4)}] & {f(r.boot_se, 4)} \\\\")
        write(fname, "\n".join(rows))
    return b


def table_moran_models():
    m = pd.read_csv(REV / "moran_xgb_vs_ols.csv")
    rows = []
    for model in ["OLS", "LightGBM", "XGBoost"]:
        sub = m[m.model == model].sort_values("subsample")
        vals = " & ".join(f(v, 4) for v in sub["I"])
        rows.append(f"{model} & {vals} & {f(sub['I'].mean(), 4)} & {f(sub['I'].std(ddof=0), 4)} & "
                    f"{f(sub['z'].min(), 1)}--{f(sub['z'].max(), 1)} \\\\")
    write("moran_models.tex", "\n".join(rows))
    return m


def table_designs():
    sm = pd.read_csv(REV / "size_matched_split.csv")
    rows = []
    short = {"random 80/20 (631,073 train)": "Random 80/20 split",
             "random, geographic training size": "Random split, geographic training size",
             "10-state geographic holdout": "Ten-state geographic holdout"}
    for _, r in sm.iterrows():
        r.design = short.get(r.design, r.design)
        within = f(r.r2_within, 3) if not pd.isna(r.r2_within) else "---"
        rows.append(f"{r.design} & {n(r.n_train)} & {n(r.n_test)} & {f(r.r2, 3)} & {within} & "
                    f"{f(r.rmse, 3)} & {r.mdape:.1f} \\\\")
    write("designs.tex", "\n".join(rows))
    return sm


def table_variants():
    v = pd.read_csv(REV / "variants_two_designs.csv")
    short = {"Full model (62 features)": "Full model",
             "Without school variables (59)": "No school variables",
             "Without region dummies (59)": "No region dummies",
             "Without school variables and region dummies (56)": "No school variables, no region dummies",
             "Full model + 6 missingness indicators (68)": "Full model + 6 missingness indicators"}
    rows = []
    for _, r in v.iterrows():
        rows.append(f"{short.get(r.variant, r.variant)} & {int(r.n_features)} & {f(r.r2_random, 3)} & {f(r.r2_geo, 3)} & "
                    f"{f(r.r2_geo_100_trees, 3)} & {f(r.r2_geo_within_state, 3)} & {r.mdape_geo:.1f} \\\\")
    write("variants.tex", "\n".join(rows))
    return v


def table_ols_geo():
    o = pd.read_csv(REV / "ols_geo_holdout.csv")
    x = pd.read_csv(REV / "xgb_geo_trimmed.csv")
    rows = []
    short = {"OLS, 62 regressors, all observations": "OLS, all observations",
             "OLS, 62 regressors, implausible values removed": "OLS, implausible records removed",
             "OLS, no region dummies, implausible values removed": "OLS, no region dummies, implausible removed",
             "XGBoost, 100 trees, trimmed test set": "XGBoost, 100 rounds, same trimmed test set",
             "XGBoost, 1000 trees, trimmed test set": "XGBoost, 1,000 rounds, same trimmed test set"}
    for _, r in o.iterrows():
        r.specification = short.get(r.specification, r.specification)
        r2 = f(r.r2, 3) if r.r2 > -10 else f"$-{abs(r.r2) / 1e6:.1f}\\times 10^{{6}}$"
        rw = f(r.r2_within, 3) if r.r2_within > -10 else "---"
        rows.append(f"{r.specification} & {n(r.n_test)} & {r2} & {rw} & {f(r.rmse, 3) if r.rmse < 10 else f(r.rmse, 0)} & {r.mdape:.1f} \\\\")
    for _, r in x.iterrows():
        rows.append(f"{short.get(r.model, r.model)} & {n(r.n_test)} & {f(r.r2, 3)} & --- & {f(r.rmse, 3)} & {r.mdape:.1f} \\\\")
    write("ols_geo.tex", "\n".join(rows))
    return o, x


def table_grouped_cv():
    g = pd.read_csv(REV / "grouped_cv_by_state.csv")
    for variant, fname in [("Full model (62 features)", "grouped_cv_full.tex"),
                           ("Without school variables (59)", "grouped_cv_noschool.tex")]:
        sub = g[g.variant == variant]
        rows = []
        for _, r in sub.iterrows():
            lab = f"Fold {r.fold}" if r.fold != "pooled" else "Pooled over the five folds"
            rows.append(f"\\quad {lab} & {int(r.n_states)} & {n(r.n_test)} & {f(r.r2, 3)} & "
                        f"{f(r.r2_within, 3)} & {f(r.rmse, 3)} & {r.mdape:.1f} \\\\")
        write(fname, "\n".join(rows))
    return g


def table_loso():
    p = REV / "loso_by_state.csv"
    if not p.exists():
        print("loso_by_state.csv not available yet"); return None
    l = pd.read_csv(p)
    rows = []
    summ = {}
    for variant in l.variant.unique():
        sub = l[l.variant == variant]
        w = sub.n_test.values
        summ[variant] = {
            "median": sub.r2.median(), "p10": sub.r2.quantile(0.1), "p90": sub.r2.quantile(0.9),
            "min": sub.r2.min(), "max": sub.r2.max(), "wmean": np.average(sub.r2, weights=w),
            "neg": int((sub.r2 < 0).sum()), "bias_med": sub.mean_error.abs().median(),
            "mdape_med": sub.mdape.median(), "within_pooled": np.nan}
        s = summ[variant]
        short = {"Full model (62 features)": "Full model", "Without school variables (59)": "No school variables"}
        rows.append(f"{short.get(variant, variant)} & {f(s['median'], 3)} & {f(s['p10'], 3)} & {f(s['p90'], 3)} & "
                    f"{f(s['min'], 3)} & {f(s['max'], 3)} & {f(s['wmean'], 3)} & {s['neg']} & "
                    f"{f(s['bias_med'], 3)} & {s['mdape_med']:.1f} \\\\")
    write("loso_summary.tex", "\n".join(rows))
    # per-state table (full model) for the appendix
    full = l[l.variant.str.startswith("Full")].sort_values("r2", ascending=False)
    rows = []
    for _, r in full.iterrows():
        rows.append(f"{r.state} & {n(r.n_test)} & {f(r.r2, 3)} & {f(r.r2_100_trees, 3)} & "
                    f"{f(r.mean_error, 3, plus=True)} & {r.mdape:.1f} \\\\")
    write("loso_states.tex", "\n".join(rows))
    return l, summ


def table_iqr_bootstrap():
    p = REV / "iqr_wald_bootstrap.csv"
    if not p.exists():
        print("iqr_wald_bootstrap.csv not available yet"); return None
    d = pd.read_csv(p)
    d["absz"] = d.z_boot_zip3.abs()
    d = d.sort_values("absz", ascending=False)
    rows = []
    for _, r in d.iterrows():
        star = "***" if abs(r.z_boot_zip3) > 3.29 else "**" if abs(r.z_boot_zip3) > 2.576 else "*" if abs(r.z_boot_zip3) > 1.96 else ""
        rows.append(f"{tex_label(r.variable)} & {f(r.difference, 3, plus=True)} & {f(r.z_independent, 2)} & "
                    f"{f(r.z_boot_obs, 2)} & {f(r.z_boot_zip3, 2)} & [{f(r.ci_low_zip3, 3)}, {f(r.ci_high_zip3, 3)}] & {star} \\\\")
    write("iqr_bootstrap.tex", "\n".join(rows))
    return d


def table_zestimate():
    z = pd.read_csv(REV / "list_to_zestimate.csv", index_col=0)
    rows = []
    for q in ["0.1", "0.25", "0.5", "0.75", "0.9"]:
        r = z.loc[float(q)] if float(q) in z.index else z.loc[q]
        rows.append(f"$p_{{{int(float(q) * 100)}}}$ & " + " & ".join(f(r[c], 3) for c in z.columns) + " \\\\")
    rows.append("Share above 1 & " + " & ".join(f"{z.loc['share_above_1', c] * 100:.1f}\\%" for c in z.columns) + " \\\\")
    rows.append("$N$ & " + " & ".join(n(z.loc['n', c]) for c in z.columns) + " \\\\")
    write("zestimate.tex", "\n".join(rows))
    return z


def table_trees_curve():
    c = pd.read_csv(REV / "geo_trees_curve.csv")
    rows = [f"{int(r.n_trees)} & {f(r.r2_random_test, 3)} & {f(r.r2_eval_within_training_states, 3)} & {f(r.r2_geo, 3)} \\\\"
            for _, r in c.iterrows() if int(r.n_trees) in (25, 50, 75, 100, 200, 300, 500, 700, 1000)]
    write("trees_curve.tex", "\n".join(rows))
    return c


# --------------------------------------------------------------------------- figures
def fig12_trees_curve(c):
    apply_paper_style()
    fig, ax = plt.subplots(figsize=(8.5, 5))
    ax.plot(c.n_trees, c.r2_random_test, color=PRIMARY, lw=2, marker="o", ms=5,
            label="Random 20% test set")
    ax.plot(c.n_trees, c.r2_eval_within_training_states, color=NEUTRAL, lw=2, ls="--", marker="s", ms=4,
            label="Early-stopping set (within training states)")
    ax.plot(c.n_trees, c.r2_geo, color=ACCENT, lw=2, marker="D", ms=5,
            label="Ten held-out states")
    k = c.loc[c.r2_geo.idxmax()]
    ax.annotate(f"peak {k.r2_geo:.3f} at {int(k.n_trees)} rounds", xy=(k.n_trees, k.r2_geo),
                xytext=(k.n_trees + 120, k.r2_geo + 0.06), fontsize=11, color=ACCENT,
                arrowprops=dict(arrowstyle="-", color=ACCENT, lw=0.8))
    last = c.iloc[-1]
    ax.annotate(f"{last.r2_geo:.3f}", xy=(last.n_trees, last.r2_geo), xytext=(8, -4),
                textcoords="offset points", fontsize=11, color=ACCENT)
    ax.set_xlabel("Boosting rounds")
    ax.set_ylabel("Out-of-sample $R^2$")
    ax.set_ylim(0.3, 0.9)
    ax.legend(loc="lower right")
    fig.savefig(FIG / "fig12_geo_trees_curve.png")
    plt.close(fig)


def fig13_loso(l):
    apply_paper_style()
    full = l[l.variant.str.startswith("Full")].sort_values("r2")
    fig, ax = plt.subplots(figsize=(8.5, 10))
    colors = [ACCENT if s in HOLDOUT else PRIMARY for s in full.state]
    ax.barh(full.state, full.r2, color=colors, height=0.72)
    ax.axvline(0, color=NEUTRAL, lw=0.9)
    ax.axvline(full.r2.median(), color=NEUTRAL, lw=1, ls=":")
    ax.text(full.r2.median(), len(full) - 0.2, f"median {full.r2.median():.2f}", fontsize=11,
            color=NEUTRAL, ha="center", va="bottom")
    ax.set_xlabel("Leave-one-state-out $R^2$ (XGBoost, 62 features, 1,000 rounds)")
    ax.grid(axis="x", linestyle=":", alpha=0.45); ax.grid(axis="y", visible=False)
    ax.tick_params(axis="y", labelsize=9)
    from matplotlib.patches import Patch
    ax.legend(handles=[Patch(color=ACCENT, label="State in the original ten-state holdout"),
                       Patch(color=PRIMARY, label="Other states and DC")], loc="lower right")
    fig.savefig(FIG / "fig13_loso_by_state.png")
    plt.close(fig)


def fig14_state_heterogeneity(summ_row_pooled):
    apply_paper_style()
    h = pd.read_csv(REV / "state_heterogeneity.csv").sort_values("ln_sqft")
    fig, axes = plt.subplots(1, 2, figsize=(12, 9), sharey=False)
    for ax, var, pooled, title in [
            (axes[0], "ln_sqft", summ_row_pooled["ln_sqft"], "(a) Living-area elasticity"),
            (axes[1], "has_garage", summ_row_pooled["has_garage"], "(b) Garage coefficient")]:
        hh = h.sort_values(var)
        sizes = 10 + 60 * np.sqrt(hh.n / hh.n.max())
        ax.scatter(hh[var], hh.state, s=sizes, color=PRIMARY, alpha=0.85, edgecolor="white", lw=0.6)
        ax.axvline(pooled, color=ACCENT, lw=1.5, ls="--")
        ax.text(pooled, -1.2, f"pooled, state FE: {pooled:.3f}", color=ACCENT, fontsize=10.5,
                ha="center", va="top")
        ax.axvline(0, color=NEUTRAL, lw=0.8)
        ax.tick_params(axis="y", labelsize=8.5)
        ax.grid(axis="x", linestyle=":", alpha=0.45); ax.grid(axis="y", visible=False)
        panel_title(ax, title)
        ax.set_xlabel("State-by-state OLS coefficient (natural units)")
    fig.tight_layout()
    fig.savefig(FIG / "fig14_state_heterogeneity.png")
    plt.close(fig)


def main():
    table_missingness()
    table_cluster_se()
    table_lot_recoding()
    s, p = table_state_heterogeneity()
    table_r2_bootstrap()
    table_moran_models()
    table_designs()
    table_variants()
    table_ols_geo()
    table_grouped_cv()
    table_zestimate()
    c = table_trees_curve()
    fig12_trees_curve(c)
    loso = table_loso()
    if loso is not None:
        fig13_loso(loso[0])
    pooled = dict(zip(s.variable, s.pooled_state_fe))
    fig14_state_heterogeneity(pooled)
    table_iqr_bootstrap()


if __name__ == "__main__":
    main()