"""Phase 7 — every table fragment (paper/tables/*.tex) and empirical figure (paper/figures/*.pdf) is generated here
from machine-readable results. Usage: python -m wp35.visualization.make_all . No number is typed by hand.
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parents[3]
RES = ROOT / "results" / "empirical"; SIM = ROOT / "results" / "simulated"; DIAG = ROOT / "results" / "diagnostics"
TAB = ROOT / "paper" / "tables"; FIG = ROOT / "paper" / "figures"; META = ROOT / "data" / "metadata"
TAB.mkdir(parents=True, exist_ok=True); FIG.mkdir(parents=True, exist_ok=True)


def esc(s: str) -> str:
    return str(s).replace("~", "$\\sim$").replace("&", "\\&").replace("%", "\\%").replace("_", "\\_").replace("≥", "$\\geq$").replace("Δ", "$\\Delta$").replace("→", "$\\to$").replace("·", "$\\cdot$").replace("β", "$\\beta$").replace("²", "$^2$").replace("δ", "$\\delta$")


def fmt(x, d=3):
    if x is None or (isinstance(x, float) and np.isnan(x)):
        return "--"
    return f"{x:.{d}f}"


def stars(p):
    return "" if p is None or np.isnan(p) else ("$^{***}$" if p < 0.01 else "$^{**}$" if p < 0.05 else "$^{*}$" if p < 0.1 else "")


def coef_table(df: pd.DataFrame, path: Path, cols=("model", "term", "coef", "se", "p", "n"), extra=None, label_col="model"):
    lines = []
    for _, r in df.iterrows():
        if "error" in df.columns and isinstance(r.get("error"), str) and r.get("error"):
            lines.append(f"{esc(r[label_col])} & \\multicolumn{{4}}{{l}}{{not estimable: {esc(r['error'][:40])}}} \\\\")
            continue
        cells = [esc(r[label_col]), esc(r.get("term", "")), fmt(r.coef) + stars(r.p), fmt(r.se), str(int(r.n)) if not np.isnan(r.n) else "--"]
        if extra:
            cells += [fmt(r.get(e), 2) if isinstance(r.get(e), float) else esc(r.get(e, "")) for e in extra]
        lines.append(" & ".join(cells) + " \\\\")
    path.write_text("\n".join(lines) + "\n")


def tables():
    reg = json.load(open(META / "source_registry.json"))
    rows = []
    for k, r in reg.items():
        rows.append(f"{esc(r['dataset'][:60])} & {esc(r['release'][:48])} & {esc(r['license'][:40])} & {r['size_bytes']:,} & \\texttt{{{r['sha256'][:12]}}} \\\\")
    (TAB / "tab_sources.tex").write_text("\n".join(rows) + "\n")
    cov = json.load(open(DIAG / "coverage_summary.json")); s, d, m = cov["seshat"], cov["dplace"], cov["modern"]
    rows = [f"Seshat Polaris-2026 polities & {s['polities']} & {s['regions']} regions / {s['macro_regions']} macro-regions \\\\",
            f"Seshat facts (dated) & {s['facts_total']:,} ({s['facts_dated']:,}) & unparsed {s['facts_unparsed']}, inferred {s['facts_inferred']:,} \\\\",
            f"Polities with population / territory / any D & {s['polities_with_population']} / {s['polities_with_territory']} / {s['polities_with_any_D']} & \\\\",
            f"Polities with $\\geq 50\\%$ of K indicators & {s['polities_with_ge50pct_K']} & \\\\",
            f"CrisisDB records (polities) & {s['crisis_records']} ({s['crisis_polities']}) & Y1 coded {s['crisis_y1_coded']}; Y1b coded {s['crisis_y1b_coded']}, positive {s['crisis_y1b_positive']} \\\\",
            f"Y3 succession coded (discontinuous) & {s['y3_coded']} ({s['y3_positive']}) & polities censored at 1900: {s['duration_censored_ge1900']} \\\\",
            f"D-PLACE EA societies & {d['societies']} & EA031 {d['with_EA031']}, EA202 {d['with_EA202']}, EA033 {d['with_EA033']}; {d['families']} language families \\\\",
            f"Modern country-years (countries) & {m['country_years']:,} ({m['countries']}) & {m['years'][0]}--{m['years'][1]}; Maddison population {m['with_maddison_pop']:,} \\\\",
            f"Modern events & civil-war onsets {m['y4_onsets']} / {m['y4_coded']:,} & coups {m['y5_coups']} / {m['y5_coded']:,}; irregular regime ends {m['y6_irregular_ends']} / {m['y6_coded']:,} \\\\"]
    (TAB / "tab_coverage.tex").write_text("\n".join(rows) + "\n")
    md = json.load(open(RES / "measurement_diagnostics.json")); mm = json.load(open(RES / "measurement_diagnostics_modern.json"))
    rows = [f"D: PC1 of log population, log territory, log largest settlement, settlement hierarchy & variance share {md['D_pc1_var_share']:.3f}; loadings " + ", ".join(f"{esc(k)} {v:.2f}" for k, v in md["D_pc1_loadings"].items()) + f"; $n={md['n_D_pc1']}$ \\\\",
            f"K: mean of four blocks (admin, rules/money, communication, records) & inter-block corr. 0.75--0.79; PC1 share {md['K_pc1_var_share']:.3f}; $n={md['n_K_mean']}$ \\\\",
            f"K (2PL IRT, 30 items, MAP) & corr. with block mean {md['corr_zD_zK']['K_mean~K_irt']:.3f}; posterior SD median {md['measurement_error']['K_irt_psd_median']:.2f}; $n={md['n_K_irt']}$ \\\\",
            f"corr$(z(D), z(K))$ & {md['corr_zD_zK']['D_pc1~K_mean']:.3f} (IRT: {md['corr_zD_zK']['D_pc1~K_irt']:.3f}) \\\\",
            f"P = $z(D)-z(K)$ & variance share of $z(D)+z(K)$ retained {md['P_variance_share_of_D_plus_K']:.3f}; corr. with D {md['corr_zD_zK']['P_main~D_pc1']:.2f}, with K {md['corr_zD_zK']['P_main~K_mean']:.2f}, with complexity PC1 {md['corr_zD_zK']['P_main~complexity_pc1']:.2f}; $n={md['n_P_main']}$ polities, {md['n_P_main_crisis']} crisis records \\\\",
            f"IRT invariance (Old World vs Americas/Oceania fits) & difficulty correlation {md['irt_invariance']['corr_difficulty_old_vs_new']:.2f} over {md['irt_invariance']['items_compared']} items \\\\",
            f"Modern K (6 V-Dem items) & PC1 share {mm['K_modern_pc1_var_share']:.3f}; corr$(D,K)$ {mm['corr_D_K_modern']:.3f}; corr$(K,\\text{{polyarchy}})$ {mm['corr_K_modern_polyarchy']:.3f}; $n={mm['n_P_modern']:,}$ \\\\"]
    (TAB / "tab_measurement.tex").write_text("\n".join(rows) + "\n")
    pw1 = pd.read_csv(RES / "power_h2a_binary.csv"); pw2 = pd.read_csv(RES / "power_h2a_fractional.csv")
    (TAB / "tab_power.tex").write_text("\n".join([f"Y1b binary & {r.beta:.2f} & {r.power:.2f} & {int(r.n)} \\\\" for _, r in pw1.iterrows()] + [f"Y1 fractional & {r.beta:.2f} & {r.power:.2f} & {int(r.n)} \\\\" for _, r in pw2.iterrows()]) + "\n")
    coef_table(pd.read_csv(RES / "h1_seshat.csv"), TAB / "tab_h1_seshat.tex", extra=["r2"])
    coef_table(pd.read_csv(RES / "h1_ea.csv"), TAB / "tab_h1_ea.tex", extra=["r2"])
    coef_table(pd.read_csv(RES / "h1_modern.csv"), TAB / "tab_h1_modern.tex", extra=["r2"])
    coef_table(pd.read_csv(RES / "h2_seshat_primary.csv"), TAB / "tab_h2_seshat.tex")
    coef_table(pd.read_csv(RES / "h2_seshat_robustness.csv"), TAB / "tab_h2_seshat_robust.tex")
    w = pd.read_csv(RES / "h2_seshat_cv_wins.csv")
    (TAB / "tab_h2_seshat_cv.tex").write_text("\n".join(f"{esc(r.baseline)} & {esc(r.scheme)} & {int(r.folds)} & {r.win_share:.2f} & {r.mean_diff:+.4f} \\\\" for _, r in w.iterrows()) + "\n")
    b = pd.read_csv(RES / "h2_seshat_cv_binary_summary.csv")
    (TAB / "tab_h2_seshat_cv_binary.tex").write_text("\n".join(f"{esc(r.model)} & {r.oos_loglik:.3f} & {r.auc:.3f} & {r.brier:.3f} \\\\" for _, r in b.iterrows()) + "\n")
    y23 = pd.read_csv(RES / "h2_seshat_y2_y3.csv")
    coef_table(y23, TAB / "tab_h2_seshat_y2y3.tex", extra=["hr", "events"])
    coef_table(pd.read_csv(RES / "h2_modern_primary.csv"), TAB / "tab_h2_modern.tex", extra=["events"])
    rb = pd.read_csv(RES / "h2_modern_robustness.csv")
    coef_table(rb[~rb.model.str.startswith("Sensitivity")], TAB / "tab_h2_modern_robust.tex", extra=["events"])
    sens = rb[rb.model.str.startswith("Sensitivity")].iloc[0]
    (TAB / "tab_sensitivity.tex").write_text(f"LPM $\\beta_P$ with controls & {sens.oster_beta_full:.4f} \\\\\nLPM $\\beta_P$ without controls & {sens.oster_beta_short:.4f} \\\\\n$R^2$ with / without controls & {sens.r2_full:.3f} / {sens.r2_short:.3f} \\\\\nOster bias-adjusted $\\beta_P$ ($R_{{\\max}}=1.3R^2$, $\\delta=1$) & {sens.coef:.4f} \\\\\nPartial $R^2$ of $P$ & {sens.partial_r2_P:.4f} \\\\\nCinelli--Hazlett robustness value $RV_{{q=1}}$ & {sens.robustness_value:.3f} \\\\\n")
    wm = pd.read_csv(RES / "h2_modern_cv_wins.csv")
    (TAB / "tab_h2_modern_cv.tex").write_text("\n".join(f"{esc(r.baseline)} & {esc(r.scheme)} & {esc(r.metric)} & {int(r.folds)} & {r.win_share:.2f} & {r.mean_diff:+.4f} \\\\" for _, r in wm.iterrows()) + "\n")
    sm_ = pd.read_csv(RES / "h2_modern_cv_summary.csv")
    (TAB / "tab_h2_modern_cv_summary.tex").write_text("\n".join(f"{esc(r.scheme)} & {esc(r.model)} & {r.oos_loglik:.3f} & {r.auc:.3f} & {r.brier:.3f} & {r.calibration_slope:.2f} \\\\" for _, r in sm_.iterrows()) + "\n")
    h4s = json.load(open(RES / "h4_seshat_summary.json")); h4m = json.load(open(RES / "h4_modern_summary.json"))
    (TAB / "tab_h4.tex").write_text(f"Seshat Y1 (10$\\times$5-fold) & {h4s['folds']} & {h4s['hinge_beats_linear_share']:.2f} & {h4s['spline_beats_linear_share']:.2f} & {h4s['mean_ll_linear']:.4f} & {h4s['mean_ll_hinge']:.4f} & {h4s['mean_ll_spline']:.4f} \\\\\nModern Y456 (3$\\times$5 country-blocked folds) & {h4m['folds']} & {h4m['hinge_beats_linear_share']:.2f} & {h4m['spline_beats_linear_share']:.2f} & {h4m['mean_ll_linear']:.4f} & {h4m['mean_ll_hinge']:.4f} & {h4m['mean_ll_spline']:.4f} \\\\\n")
    for name in ["h5_seshat_blocks", "h5_modern_blocks"]:
        h5 = pd.read_csv(RES / f"{name}.csv")
        lines = []
        for _, r in h5.iterrows():
            if isinstance(r.get("error"), str) and r.get("error"):
                lines.append(f"{esc(r.model)} & \\multicolumn{{5}}{{l}}{{not estimable}} \\\\"); continue
            wb = f"{fmt(r.get('within_block_coef'))} ({fmt(r.get('within_block_se'))})" if not np.isnan(r.get("within_block_coef", np.nan)) else "--"
            held = int(r["held_out_n"]) if "held_out_n" in r and not np.isnan(r["held_out_n"]) else (int(r["within_block_n"]) if "within_block_n" in r and not np.isnan(r["within_block_n"]) else 0)
            lines.append(f"{esc(r.model.replace('H5 ', ''))} & {fmt(r.coef)}{stars(r.p)} & {fmt(r.se)} & {int(r.n)} & {held} & {wb} \\\\")
        (TAB / f"tab_{name}.tex").write_text("\n".join(lines) + "\n")
    # simulations
    ss = json.load(open(SIM / "simulation_summary.json"))
    pr = pd.DataFrame(ss["paired_vs_none"])
    (TAB / "tab_sim_paired.tex").write_text("\n".join(f"{esc(r.regime)} & {r.admin_capacity_ratio:.1f} & {int(r.n_pairs)} & {r.share_lower_bits_per_success:.2f} & {r.share_lower_bits_incl_setup:.2f} & {r.share_higher_success:.2f} & {r.share_higher_welfare:.2f} & {r.mean_d_bits_per_success:+.1f} & {r.mean_d_success:+.3f} & {r.mean_d_welfare:+.1f} \\\\" for _, r in pr.iterrows()) + "\n")
    ls = pd.DataFrame(ss["institution_loss_shares"])
    (TAB / "tab_sim_loss.tex").write_text("\n".join(f"{esc(r.regime)} & {r.admin_capacity_ratio:.1f} & {int(r.n_cells)} & {r.welfare_loss:.2f} & {r.bits_loss:.2f} & {r.success_loss:.2f} \\\\" for _, r in ls.iterrows()) + "\n")
    sh = pd.DataFrame(ss["shock_effects"])
    (TAB / "tab_sim_shocks.tex").write_text("\n".join(f"{esc(r.shock_type)} & {esc(r.regime)} & {r.admin_capacity_ratio:.1f} & {r.d_success_mean:+.3f} ({r.d_success_std / np.sqrt(r.d_success_count):.3f}) & {r.d_bits_mean:+.3f} & {r.kl_q_p_bits_mean:.3f} \\\\" for _, r in sh.iterrows()) + "\n")
    ab = ss["ablations"]; tc = ss["theorem_checks"]
    lines = [f"Reference (N=200, lattice, no shock): success / bits per success / welfare per agent & " + "; ".join(f"{k}: {v['success_rate']:.3f} / {v['bits_per_success_incl_setup']:.1f} / {v['welfare_per_agent']:.1f}" for k, v in ab["reference"].items()) + " \\\\",
             f"Records without verification & success {ab['records_without_verification']['success_rate']:.3f}, bits per success {ab['records_without_verification']['bits_per_success_incl_setup']:.1f}, welfare {ab['records_without_verification']['welfare_per_agent']:.1f} \\\\",
             "Norms, extra bits per task after a distribution shock by adaptation delay & " + "; ".join(f"delay {k}: {v['d_bits']:+.2f}" for k, v in ab["norms_adaptation"].items()) + f" (KL $= {list(ab['norms_adaptation'].values())[0]['kl_q_p_bits']:.2f}$ bits) \\\\",
             "Break-even (norms cheaper than none incl. set-up) & " + "; ".join(f"{esc(k)}: {'yes' if v else 'no'}" for k, v in ab["breakeven_norms_wins_bits"].items()) + " \\\\",
             "Heterogeneity $\\epsilon$: success none vs norms & " + "; ".join(f"{esc(k)}: {v['success_rate']:.3f}" for k, v in ab["heterogeneity_eps"].items()) + " \\\\",
             "Federated branching $b$: success / messages per task / root load & " + "; ".join(f"b={k}: {v['success_rate']:.2f} / {v['msgs_per_task']:.2f} / {v['root_load_mean']:.0f}" for k, v in ab["federated_branching"].items()) + " \\\\",
             "Coding gain (norms): $H(p)$, $L_{code}$, $L_{raw}$ & " + "; ".join(f"{esc(k)}: {v['H']:.2f}, {v['L_code']:.2f}, {v['L_raw']:.0f}" for k, v in ab["coding_gain"].items()) + " \\\\",
             f"Theorem checks & T1 ($H\\le L<H+1$): {tc['T1_entropy_le_code_lt_entropy_plus_1']}; code shorter than raw whenever $H\\le\\log_2 s-1$: {tc['T1_code_shorter_than_raw_when_H_le_log2s_minus_1']}; T3 extra bits per recipient {tc['T3_post_shock_extra_bits_per_recipient_norms_vs_KL']['mean_extra_bits_per_recipient']:.2f} vs KL {tc['T3_post_shock_extra_bits_per_recipient_norms_vs_KL']['mean_KL']:.2f} \\\\"]
    (TAB / "tab_sim_ablation.tex").write_text("\n".join(lines) + "\n")
    ac = pd.read_csv(SIM / "ablation_admin_capacity.csv")
    (TAB / "tab_sim_capacity.tex").write_text("\n".join(f"{esc(r.regime)} & {r.admin_capacity_ratio:.2f} & {r.success_rate:.3f} & {r.fail_latency_share:.3f} & {r.latency_mean:.2f} & {r.root_load_mean:.0f} \\\\" for _, r in ac.iterrows()) + "\n")
    by = pd.read_csv(SIM / "regime_by_N.csv"); by = by[by.admin_capacity_ratio == 2.0]
    (TAB / "tab_sim_scaling.tex").write_text("\n".join(f"{esc(r.regime)} & {int(r.N)} & {r.success_rate:.3f} & {r.msgs_per_task:.2f} & {r.bits_per_task:.1f} & {r.bits_per_success_incl_setup:.1f} & {r.welfare_per_agent:.1f} & {r.gini_benefit:.3f} \\\\" for _, r in by.sort_values(['regime', 'N']).iterrows()) + "\n")


def figures():
    # coefficient plot: Seshat primary/baselines and modern primary/baselines
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
    for ax, f, title in [(axes[0], "h2_seshat_primary.csv", "Seshat: crisis severity Y1 (fractional logit)"), (axes[1], "h2_modern_primary.csv", "Modern: any instability event in (t, t+5] (logit)")]:
        d = pd.read_csv(RES / f); d = d[d.se.notna()].reset_index(drop=True)
        y = np.arange(len(d))[::-1]
        ax.errorbar(d.coef, y, xerr=1.96 * d.se, fmt="o", ms=3.5, color="#333", capsize=2, lw=1)
        ax.axvline(0, color="k", lw=0.7, ls="--"); ax.set_yticks(y); ax.set_yticklabels([f"{m[:48]} [{t}]" for m, t in zip(d.model, d.term)], fontsize=6)
        ax.set_xlabel("coefficient (per SD) ± 95% CI"); ax.set_title(title, fontsize=9)
    fig.tight_layout(); fig.savefig(FIG / "fig_h2_coefplot.pdf"); plt.close(fig)
    # H5 forest
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
    for ax, f, title in [(axes[0], "h5_seshat_blocks.csv", "Seshat Y1: leave-one-block-out β_P"), (axes[1], "h5_modern_blocks.csv", "Modern Y456: leave-one-block-out β_P")]:
        d = pd.read_csv(RES / f); d = d[d.coef.notna()].reset_index(drop=True); y = np.arange(len(d))[::-1]
        ax.errorbar(d.coef, y, xerr=1.96 * d.se, fmt="o", ms=3.5, color="#333", capsize=2, lw=1, label="pooled, block held out")
        if "within_block_coef" in d:
            ok = d.within_block_coef.notna(); ax.errorbar(d.within_block_coef[ok], y[ok] - 0.25, xerr=1.96 * d.within_block_se[ok], fmt="s", ms=3, color="#1f77b4", capsize=2, lw=0.8, label="within held-out block")
        ax.axvline(0, color="k", lw=0.7, ls="--"); ax.set_yticks(y); ax.set_yticklabels([m.replace("H5 ", "") for m in d.model], fontsize=6); ax.set_title(title, fontsize=9); ax.legend(fontsize=6)
    fig.tight_layout(); fig.savefig(FIG / "fig_h5_forest.pdf"); plt.close(fig)
    # H4 curve (modern)
    c = pd.read_csv(RES / "h4_modern_curve.csv")
    fig, ax = plt.subplots(figsize=(5, 3.6)); ax.plot(c.P_z, c.p_linear, label="linear in P", color="#333"); ax.plot(c.P_z, c.p_spline, label="natural cubic spline (3 df)", color="#d62728")
    ax.set_xlabel("P (standardized)"); ax.set_ylabel("Pr(event in (t, t+5]) at covariate means"); ax.legend(fontsize=8); ax.set_title("Modern panel: shape of the P–instability relation", fontsize=9)
    fig.tight_layout(); fig.savefig(FIG / "fig_h4_curve.pdf"); plt.close(fig)
    # calibration (modern, out-of-sample predictions)
    pp = RES / "h2_modern_cv_predictions.parquet"
    if pp.exists():
        pr = pd.read_parquet(pp); fig, axes = plt.subplots(1, 2, figsize=(9, 3.8))
        for ax, scheme in zip(axes, ["forward", "loro"]):
            for model, col in [("C1 P", "#d62728"), ("B2 K", "#1f77b4"), ("B0 FE+X", "#999999")]:
                s = pr[(pr.scheme == scheme) & (pr.model == model)]
                if len(s) == 0:
                    continue
                bins = pd.qcut(s.p, 10, duplicates="drop"); g = s.groupby(bins, observed=True).agg(p=("p", "mean"), y=("y", "mean"))
                ax.plot(g.p, g.y, marker="o", ms=3, label=model, color=col, lw=1)
            ax.plot([0, 0.6], [0, 0.6], "k--", lw=0.7); ax.set_xlabel("predicted probability (decile mean)"); ax.set_ylabel("observed event share"); ax.set_title(f"Out-of-sample calibration: {scheme}", fontsize=9); ax.legend(fontsize=7)
        fig.tight_layout(); fig.savefig(FIG / "fig_calibration_modern.pdf"); plt.close(fig)
    # historical trajectories: D and K by era, with P distribution
    ind = pd.read_parquet(ROOT / "data/processed/seshat_indices.parquet")
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
    for reg, g in ind.dropna(subset=["D_pc1", "K_mean"]).groupby("macro_region"):
        axes[0].scatter(g.midpoint, g.P_main, s=8, alpha=0.6, label=reg.replace("_", " "))
    axes[0].axhline(0, color="k", lw=0.6); axes[0].set_xlabel("polity midpoint year"); axes[0].set_ylabel("P = z(D) − z(K)"); axes[0].set_title("Coordination pressure across time (Seshat)", fontsize=9); axes[0].legend(fontsize=5, ncol=2)
    cr = pd.read_parquet(ROOT / "data/processed/seshat_h2_sample.parquet").dropna(subset=["P_main_z", "y1_severity"])
    axes[1].scatter(cr.P_main_z, cr.y1_severity, s=12, alpha=0.7, color="#333")
    xs = np.linspace(cr.P_main_z.min(), cr.P_main_z.max(), 50); b = np.polyfit(cr.P_main_z, cr.y1_severity, 1); axes[1].plot(xs, np.polyval(b, xs), color="#d62728", lw=1)
    axes[1].set_xlabel("pre-crisis P (standardized)"); axes[1].set_ylabel("crisis severity (share of adverse consequences)"); axes[1].set_title(f"CrisisDB severity vs P (n = {len(cr)})", fontsize=9)
    fig.tight_layout(); fig.savefig(FIG / "fig_trajectories.pdf"); plt.close(fig)


def main():
    tables(); figures(); print("tables:", len(list(TAB.glob("*.tex"))), "figures:", len(list(FIG.glob("*.pdf"))))


if __name__ == "__main__":
    main()