"""SIMULATED — summaries, phase diagrams, scaling curves and ablations from results/simulated/runs_*.parquet.
Usage: python -m wp35.simulations.analyze . Writes results/simulated/*.csv, simulation_summary.json and paper/figures/fig_sim_*.pdf
"""
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]; SIM = ROOT / "results" / "simulated"; FIG = ROOT / "paper" / "figures"
METRICS = ["success_rate", "msgs_per_task", "bits_per_task", "bits_per_success", "bits_per_success_incl_setup", "tasks_per_kbit", "latency_mean",
           "fail_latency_share", "fail_decode_share", "admin_overhead_share", "welfare_per_agent", "gini_benefit", "root_load_mean", "success_pre", "success_post", "bits_per_task_pre", "bits_per_task_post"]
COL = {"none": "#8c8c8c", "norms": "#1f77b4", "records": "#2ca02c", "central": "#d62728", "federated": "#9467bd"}


def ci(x):
    x = np.asarray(x, float); x = x[np.isfinite(x)]
    if len(x) < 2:
        return (np.nan, np.nan)
    m = x.mean(); h = 1.96 * x.std(ddof=1) / np.sqrt(len(x)); return (m - h, m + h)


def main():
    m = pd.read_parquet(SIM / "runs_main.parquet"); a = pd.read_parquet(SIM / "runs_ablation.parquet")
    factors = ["N", "network", "bandwidth", "zipf", "shock_type", "regime", "admin_capacity_ratio"]
    # 1. cell means with CIs
    g = m.groupby(factors)[METRICS].agg(["mean", "std", "count"]); g.columns = ["_".join(c) for c in g.columns]; g = g.reset_index()
    g.to_csv(SIM / "cells_main.csv", index=False)
    # 2. regime comparison pooled over cells (H3): bits per success and success rate, by N
    by = m.groupby(["regime", "N", "admin_capacity_ratio"])[["success_rate", "bits_per_task", "bits_per_success", "bits_per_success_incl_setup", "msgs_per_task", "welfare_per_agent", "gini_benefit", "fail_latency_share", "fail_decode_share"]].mean().reset_index()
    by.to_csv(SIM / "regime_by_N.csv", index=False)
    # 3. paired comparison within cell: institution regime vs none (same N/network/bandwidth/zipf/shock/seed)
    key = ["N", "network", "bandwidth", "zipf", "shock_type", "seed"]
    none = m[m.regime == "none"].set_index(key)
    rows = []
    for reg in ["norms", "records", "central", "federated"]:
        for ratio in sorted(m[m.regime == reg].admin_capacity_ratio.unique()):
            sub = m[(m.regime == reg) & (m.admin_capacity_ratio == ratio)].set_index(key)
            j = sub.join(none, rsuffix="_none", how="inner")
            d_bits = j.bits_per_success - j.bits_per_success_none; d_bits_setup = j.bits_per_success_incl_setup - j.bits_per_success_incl_setup_none
            d_succ = j.success_rate - j.success_rate_none; d_welf = j.welfare_per_agent - j.welfare_per_agent_none
            rows.append({"regime": reg, "admin_capacity_ratio": ratio, "n_pairs": len(j),
                         "share_lower_bits_per_success": float((d_bits < 0).mean()), "share_lower_bits_incl_setup": float((d_bits_setup < 0).mean()),
                         "share_higher_success": float((d_succ > 0).mean()), "share_higher_welfare": float((d_welf > 0).mean()),
                         "mean_d_bits_per_success": float(d_bits.mean()), "mean_d_success": float(d_succ.mean()), "mean_d_welfare": float(d_welf.mean()),
                         "ci_d_bits": ci(d_bits), "ci_d_success": ci(d_succ), "ci_d_welfare": ci(d_welf)})
    paired = pd.DataFrame(rows); paired.to_csv(SIM / "paired_vs_none.csv", index=False)
    # 4. cells where institutions LOSE (welfare or bits-per-success worse than none) — fragility search
    lose = []
    for reg in ["norms", "records", "central", "federated"]:
        sub = m[m.regime == reg]
        cm = sub.groupby(["N", "network", "bandwidth", "zipf", "shock_type", "admin_capacity_ratio"])[["welfare_per_agent", "bits_per_success_incl_setup", "success_rate"]].mean()
        nm = none.groupby(["N", "network", "bandwidth", "zipf", "shock_type"])[["welfare_per_agent", "bits_per_success_incl_setup", "success_rate"]].mean()
        cm = cm.reset_index().merge(nm.reset_index(), on=["N", "network", "bandwidth", "zipf", "shock_type"], suffixes=("", "_none"))
        cm["regime"] = reg; cm["welfare_loss"] = cm.welfare_per_agent < cm.welfare_per_agent_none; cm["bits_loss"] = cm.bits_per_success_incl_setup > cm.bits_per_success_incl_setup_none; cm["success_loss"] = cm.success_rate < cm.success_rate_none
        lose.append(cm)
    lose = pd.concat(lose); lose.to_csv(SIM / "cells_institution_vs_none.csv", index=False)
    lose_summary = lose.groupby(["regime", "admin_capacity_ratio"])[["welfare_loss", "bits_loss", "success_loss"]].mean().reset_index(); lose_summary["n_cells"] = lose.groupby(["regime", "admin_capacity_ratio"]).size().values
    lose_summary.to_csv(SIM / "institution_loss_shares.csv", index=False)
    # 5. shock effects: Δ success and Δ bits (post − pre) by regime and shock type
    # difference-in-differences: (post - pre) in the shocked run minus (post - pre) in the unshocked run of the same cell/seed
    m["d_success_raw"] = m.success_post - m.success_pre; m["d_bits_raw"] = m.bits_per_task_post - m.bits_per_task_pre
    ctrl = m[m.shock_type == "none"].set_index(["N", "network", "bandwidth", "zipf", "regime", "admin_capacity_ratio", "seed"])[["d_success_raw", "d_bits_raw"]]
    sh = m[m.shock_type != "none"].join(ctrl, on=["N", "network", "bandwidth", "zipf", "regime", "admin_capacity_ratio", "seed"], rsuffix="_ctrl")
    sh["d_success"] = sh.d_success_raw - sh.d_success_raw_ctrl; sh["d_bits"] = sh.d_bits_raw - sh.d_bits_raw_ctrl
    sh = sh[(sh.shock_type != "distribution") | (sh.zipf > 0)]       # distribution shocks are a no-op in uniform cells (KL = 0)
    shock = sh.groupby(["shock_type", "regime", "admin_capacity_ratio"])[["d_success", "d_bits", "kl_q_p_bits"]].agg(["mean", "std", "count"]); shock.columns = ["_".join(c) for c in shock.columns]; shock = shock.reset_index(); shock.to_csv(SIM / "shock_effects.csv", index=False)
    shock_bw = sh.groupby(["shock_type", "bandwidth", "regime", "admin_capacity_ratio"])[["d_success", "d_bits", "success_pre"]].mean().reset_index(); shock_bw.to_csv(SIM / "shock_effects_by_bandwidth.csv", index=False)
    # 6. ablations
    abl = {}
    ref = a[(a.shock_type == "none") & (a["T"] == 40) & (a.learn_bits == 200.0) & (a.eps == 0.25) & (a.s == 64) & (a.zipf == 1.0) & (a.verify_records) & (a.branching == 8) & (a.admin_capacity_ratio == 2.0)]
    abl["reference"] = ref.groupby("regime")[["success_rate", "bits_per_success_incl_setup", "welfare_per_agent"]].mean().round(4).to_dict("index")
    nv = a[(a.regime == "records") & (~a.verify_records)]; abl["records_without_verification"] = nv[["success_rate", "bits_per_success_incl_setup", "welfare_per_agent"]].mean().round(4).to_dict()
    ad = a[(a.regime == "norms") & (a.shock_type == "distribution")].copy(); ad["d_bits"] = ad.bits_per_task_post - ad.bits_per_task_pre
    abl["norms_adaptation"] = ad.groupby("adapt_delay")[["d_bits", "success_post", "kl_q_p_bits"]].mean().round(4).to_dict("index")
    be = a[(a.regime.isin(["none", "norms"])) & (a.shock_type == "none") & (a.eps == 0.25) & (a.verify_records)]
    be = be.groupby(["s", "zipf", "T", "learn_bits", "regime"])[["bits_per_success_incl_setup", "welfare_per_agent"]].mean().unstack("regime")
    be.columns = ["_".join(c) for c in be.columns]; be["norms_wins_bits"] = be.bits_per_success_incl_setup_norms < be.bits_per_success_incl_setup_none; be = be.reset_index(); be.to_csv(SIM / "ablation_breakeven.csv", index=False)
    abl["breakeven_norms_wins_bits"] = {f"s={int(r.s)},zipf={r.zipf},T={int(r['T'])},learn={r.learn_bits}": bool(r.norms_wins_bits) for _, r in be.iterrows()}
    he = a[(a.regime.isin(["none", "norms"])) & (a.shock_type == "none") & (a["T"] == 40) & (a.learn_bits == 200.0) & (a.s == 64) & (a.zipf == 1.0)]
    hh = he.groupby(["eps", "regime"])[["success_rate", "fail_decode_share"]].mean().round(4).reset_index()
    abl["heterogeneity_eps"] = {f"eps={r.eps},{r.regime}": {"success_rate": r.success_rate, "fail_decode_share": r.fail_decode_share} for _, r in hh.iterrows()}
    pc = a[(a.regime.isin(["central", "federated"])) & (a.branching == 8) & (a.shock_type == "none")]
    phase = pc.groupby(["regime", "admin_capacity_ratio"])[["success_rate", "fail_latency_share", "latency_mean", "root_load_mean", "bits_per_success_incl_setup"]].mean().reset_index(); phase.to_csv(SIM / "ablation_admin_capacity.csv", index=False)
    br = a[(a.regime == "federated") & (a.admin_capacity_ratio == 1.0)]; abl["federated_branching"] = br.groupby("branching")[["success_rate", "msgs_per_task", "root_load_mean"]].mean().round(4).to_dict("index")
    ce = a[(a.regime.isin(["none", "norms"])) & (a.shock_type == "none") & (a["T"] == 40) & (a.learn_bits == 200.0) & (a.eps == 0.25)]
    cg = ce.groupby(["s", "zipf", "regime"])[["H_p_bits", "L_code_bits", "L_raw_bits", "bits_per_task"]].mean().reset_index(); cg.to_csv(SIM / "ablation_coding_gain.csv", index=False)
    abl["coding_gain"] = {f"s={int(r.s)},zipf={r.zipf}": {"H": round(r.H_p_bits, 3), "L_code": round(r.L_code_bits, 3), "L_raw": r.L_raw_bits, "bits_per_task": round(r.bits_per_task, 3)} for _, r in cg[cg.regime == "norms"].iterrows()}
    summary = {"runs_main": int(len(m)), "runs_ablation": int(len(a)), "cells_main": int(len(g)), "seeds": int(m.seed.nunique()),
               "paired_vs_none": paired.drop(columns=["ci_d_bits", "ci_d_success", "ci_d_welfare"]).round(4).to_dict("records"),
               "institution_loss_shares": lose_summary.round(4).to_dict("records"), "shock_effects": shock.round(4).to_dict("records"), "ablations": abl,
               "theorem_checks": {"T1_code_shorter_than_raw_when_H_le_log2s_minus_1": bool((cg[(cg.regime == "norms") & (cg.H_p_bits <= cg.L_raw_bits - 1)].L_code_bits <= cg[(cg.regime == "norms") & (cg.H_p_bits <= cg.L_raw_bits - 1)].L_raw_bits).all()),
                                   "T1_cells_where_code_longer_than_raw": cg[(cg.regime == "norms") & (cg.L_code_bits > cg.L_raw_bits)][["s", "zipf", "H_p_bits", "L_code_bits", "L_raw_bits"]].round(3).to_dict("records"),
                                   "T1_entropy_le_code_lt_entropy_plus_1": bool(((cg.H_p_bits <= cg.L_code_bits + 1e-9) & (cg.L_code_bits < cg.H_p_bits + 1)).all()),
                                   "T3_post_shock_extra_bits_per_recipient_norms_vs_KL": {"mean_extra_bits_per_recipient": float(ad[ad.adapt_delay == 1000].d_bits.mean() / 2), "mean_KL": float(ad.kl_q_p_bits.mean()), "within_KL_pm_1": bool(abs(ad[ad.adapt_delay == 1000].d_bits.mean() / 2 - ad.kl_q_p_bits.mean()) < 1)}}}
    (SIM / "simulation_summary.json").write_text(json.dumps(summary, indent=2, default=float))
    # ---- figures ----
    FIG.mkdir(parents=True, exist_ok=True)
    fig, axes = plt.subplots(1, 3, figsize=(12.5, 3.8))
    for reg in ["none", "norms", "records", "central", "federated"]:
        sub = m[(m.regime == reg) & (m.shock_type == "none") & ((m.admin_capacity_ratio == 2.0))]
        s = sub.groupby("N")[["bits_per_success_incl_setup", "success_rate", "msgs_per_task"]].agg(["mean", "std", "count"])
        for ax, met in zip(axes, ["bits_per_success_incl_setup", "success_rate", "msgs_per_task"]):
            mu, sd, n = s[(met, "mean")], s[(met, "std")], s[(met, "count")]
            ax.errorbar(mu.index, mu, yerr=1.96 * sd / np.sqrt(n), marker="o", ms=3, label=reg, color=COL[reg], lw=1.2, capsize=2)
    axes[0].set_ylabel("bits per successful task (incl. set-up)"); axes[0].set_yscale("log"); axes[1].set_ylabel("task success rate"); axes[2].set_ylabel("messages per task")
    for ax in axes:
        ax.set_xscale("log"); ax.set_xlabel("N agents"); ax.set_xticks([50, 200, 800]); ax.set_xticklabels([50, 200, 800])
    axes[0].legend(fontsize=7, title="regime (budget multiple 2)", title_fontsize=7); fig.suptitle("SIMULATED — scaling by regime, no shock (mean ± 95% CI over cells and seeds)", fontsize=9)
    fig.tight_layout(); fig.savefig(FIG / "fig_sim_scaling.pdf"); plt.close(fig)
    # phase diagram: success vs admin capacity ratio
    fig, ax = plt.subplots(figsize=(5, 3.6))
    for reg in ["central", "federated"]:
        s = phase[phase.regime == reg]; ax.plot(s.admin_capacity_ratio, s.success_rate, marker="o", label=reg, color=COL[reg])
    ax.axvline(1.0, ls="--", color="k", lw=0.8); ax.set_xscale("log"); ax.set_xlabel("administrative capacity / expected relayed load  (ρ⁻¹)"); ax.set_ylabel("success rate"); ax.legend(fontsize=8)
    ax.set_title("SIMULATED — hierarchy saturation (N = 200, lattice)", fontsize=9); fig.tight_layout(); fig.savefig(FIG / "fig_sim_phase.pdf"); plt.close(fig)
    # shock effects figure
    fig, axes = plt.subplots(1, 2, figsize=(9, 3.6))
    for i, st in enumerate(["distribution", "network"]):
        s = shock[(shock.shock_type == st) & (shock.admin_capacity_ratio == 2.0)]
        axes[i].bar(s.regime, s.d_success_mean, yerr=1.96 * s.d_success_std / np.sqrt(s.d_success_count), color=[COL[r] for r in s.regime], capsize=3)
        axes[i].axhline(0, color="k", lw=0.6); axes[i].set_title(f"{st} shock: Δ success (post − pre)", fontsize=9); axes[i].tick_params(axis="x", labelsize=8)
    fig.suptitle("SIMULATED — shock responses by regime", fontsize=9); fig.tight_layout(); fig.savefig(FIG / "fig_sim_shocks.pdf"); plt.close(fig)
    # break-even figure
    fig, ax = plt.subplots(figsize=(5, 3.6))
    for (s_, z, lb), s in be.groupby(["s", "zipf", "learn_bits"]):
        ax.plot(s["T"], s.bits_per_success_incl_setup_norms / s.bits_per_success_incl_setup_none, marker="o", ms=3, label=f"s={int(s_)}, zipf={z}, learn={int(lb)}")
    ax.axhline(1, ls="--", color="k", lw=0.8); ax.set_xscale("log"); ax.set_xlabel("horizon T (periods)"); ax.set_ylabel("bits per success, norms / none"); ax.legend(fontsize=7)
    ax.set_title("SIMULATED — break-even of a shared code (Theorem 2)", fontsize=9); fig.tight_layout(); fig.savefig(FIG / "fig_sim_breakeven.pdf"); plt.close(fig)
    print(json.dumps({k: v for k, v in summary.items() if k in ("runs_main", "runs_ablation", "cells_main", "institution_loss_shares", "theorem_checks")}, indent=1, default=float))
    print(paired.drop(columns=["ci_d_bits", "ci_d_success", "ci_d_welfare"]).round(3).to_string(index=False))
    print(shock[["shock_type", "regime", "admin_capacity_ratio", "d_success_mean", "d_bits_mean", "kl_q_p_bits_mean"]].round(3).to_string(index=False))
    print(json.dumps(abl, indent=1, default=float)[:3000])


if __name__ == "__main__":
    main()