#!/usr/bin/env python3
"""Generate paper/tables/*.tex (LaTeX fragments), paper/tables/numbers.tex (macros) and experiments/EXPERIMENT_MANIFEST.csv
from results/processed/summary.csv, results/processed/boundary_summary.csv, results/raw/ablation_*.csv, results/raw/skew_*.csv
and results/raw/benchmarks.jsonl.  Degrades gracefully when a result set is absent (the fragment then says so)."""
import glob, json, pathlib, sys
import numpy as np, pandas as pd
ROOT = pathlib.Path(__file__).resolve().parents[1]
PROC, RAW, TAB = ROOT / "results" / "processed", ROOT / "results" / "raw", ROOT / "paper" / "tables"
TAB.mkdir(parents=True, exist_ok=True)
LAB = {"r1_q0": "(1,0) plain average", "r1_q1": "(1,1) average + 1 Newton", "r1_q2": "(1,2) average + 2 Newton", "r2_q0": "(2,0) half-block jackknife",
       "r2_q1": "(2,1) half-block + 1 Newton", "r2_q2": "(2,2) half-block + 2 Newton", "r3_q0": "(3,0) quarter-block", "r3_q1": "(3,1) quarter-block + 1 Newton",
       "r3_q2": "(3,2) quarter-block + 2 Newton", "cs_q0": "analytic Cox--Snell, $q=0$", "cs_q1": "analytic Cox--Snell + 1 Newton", "csl1": "CSL, 1 round", "csl2": "CSL, 2 rounds",
       "aee": "AEE (Lin--Xi)", "stale_r1_q1": "(1,1) stale Hessian", "stale_r2_q1": "(2,1) stale Hessian", "stale_r1_q2": "(1,2) stale Hessian", "central": "centralized MLE"}
MAIN = ["r1_q0", "r1_q1", "r1_q2", "r2_q0", "r2_q1", "r2_q2", "r3_q0", "r3_q1", "cs_q1", "csl1", "csl2", "aee", "stale_r2_q1", "central"]


def fmt(x, d=3):
    if x is None or (isinstance(x, float) and not np.isfinite(x)):
        return "--"
    a = abs(x)
    if a == 0:
        return "0"
    if a < 1e-3 or a >= 1e4:
        return f"{x:.2e}".replace("e-0", "e-").replace("e+0", "e+")
    return f"{x:.{d}f}"


def table_cell(s, label, p, N, ms, fname, caption, ests=MAIN):
    d = s[(s.label == label) & (s.p == p) & (s.N == N) & (s.m.isin(ms))]
    if d.empty:
        (TAB / fname).write_text(f"% no data for {label} p={p} N={N}\n"); return False
    ms = sorted(d.m.unique())
    lines = [r"\begin{table}[p]", r"\centering\scriptsize", r"\begin{threeparttable}", f"\\caption{{{caption}}}\\label{{tab:{fname[:-4]}}}",
             r"\begin{tabular}{@{}l" + "r" * (3 * len(ms)) + "@{}}", r"\toprule",
             "Estimator & " + " & ".join(f"\\multicolumn{{3}}{{c}}{{$m = {m}$, $K = {int(N//m):,}$}}" for m in ms) + r" \\",
             " & " + " & ".join("gap & RMSE & cov" for _ in ms) + r" \\", r"\midrule"]
    for e in ests:
        row = [LAB.get(e, e)]
        for m in ms:
            dd = d[(d.estimator == e) & (d.m == m)]
            if dd.empty:
                row += ["--", "--", "--"]; continue
            dd = dd.iloc[0]
            row += [fmt(dd["rms_gap"]), fmt(dd["rmse"]), (f"{dd['cov95']:.3f}" if "cov95" in dd and np.isfinite(dd["cov95"]) else "--")]
        lines.append(" & ".join(row) + r" \\")
    reps = int(d.reps.iloc[0]); cf = d.drop_duplicates("m")[["m", "central_fail", "valid_frac_r2", "valid_frac_r3"]] if "valid_frac_r3" in d else None
    lines += [r"\bottomrule", r"\end{tabular}", r"\begin{tablenotes}\footnotesize",
              f"\\item {reps} paired replications per column. gap $=$ root mean square of $\\sqrt N\\,\\lVert\\tilde\\theta - \\widehat\\theta_N\\rVert$; RMSE $=$ root mean square of $\\sqrt N\\,\\lVert\\tilde\\theta - \\theta_0\\rVert$; cov $=$ empirical coverage of nominal 95\\% Wald intervals (observed information at the estimate; `--' when not computed). Monte Carlo standard errors of the gaps and RMSEs are in \\texttt{{results/processed/summary.csv}}; coverage standard error $\\approx {np.sqrt(0.95*0.05/reps):.3f}$."]
    if cf is not None:
        lines.append("\\item Fraction of workers with all local fits valid (orders 2 / 3): " + "; ".join(f"$m={int(r.m)}$: {r.valid_frac_r2:.3f} / {r.valid_frac_r3:.3f}" for r in cf.itertuples()) + ".")
    lines += [r"\end{tablenotes}", r"\end{threeparttable}", r"\end{table}"]
    (TAB / fname).write_text("\n".join(lines) + "\n"); return True


def table_rates(s, fname="mc_rates.tex"):
    """Log-log slopes of rms_gap vs m at fixed N (largest N per model) for the main estimators."""
    rows = []
    for (label, p), d in s.groupby(["label", "p"]):
        N = d.N.max()
        for e in ["r1_q0", "r1_q1", "r1_q2", "r2_q0", "r2_q1", "r3_q0", "r3_q1", "cs_q1"]:
            dd = d[(d.N == N) & (d.estimator == e) & (d.rms_gap > 1e-9)].sort_values("m")
            if len(dd) >= 3:
                slope = np.polyfit(np.log(dd.m), np.log(dd.rms_gap), 1)[0]
                rows.append((label, p, N, e, slope, len(dd)))
    if not rows:
        (TAB / fname).write_text("% no data\n"); return
    th = {"r1_q0": -1, "r1_q1": -2, "r1_q2": -4, "r2_q0": -2, "r2_q1": -4, "r3_q0": -3, "r3_q1": -6, "cs_q1": -4}
    df = pd.DataFrame(rows, columns=["label", "p", "N", "estimator", "slope", "n"])
    piv = df.pivot_table(index=["label", "p", "N"], columns="estimator", values="slope")
    cols = [c for c in ["r1_q0", "r1_q1", "r1_q2", "r2_q0", "r2_q1", "r3_q0", "r3_q1", "cs_q1"] if c in piv.columns]
    lines = [r"\begin{table}[t]", r"\centering\small", r"\begin{threeparttable}", r"\caption{Empirical log-log slopes of the RMS scaled gap to $\widehat\theta_N$ against $m$ at fixed $N$ (largest $N$ of each design), with the slopes predicted by the leading bias term.}\label{tab:mc_rates}",
             r"\begin{tabular}{@{}llr" + "r" * len(cols) + "@{}}", r"\toprule", "Design & $p$ & $N$ & " + " & ".join(c.replace("_", "\\_") for c in cols) + r" \\",
             "predicted & & & " + " & ".join(str(th[c]) for c in cols) + r" \\", r"\midrule"]
    for idx, r in piv.iterrows():
        lines.append(f"{idx[0].replace('_', ' ')} & {idx[1]} & {int(idx[2]):,} & " + " & ".join(fmt(r[c], 2) for c in cols) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\begin{tablenotes}\footnotesize", r"\item Slopes are least-squares fits over the available $m$ grid; they approach the predicted exponent only where the bias term dominates the random remainder and the sample is in the asymptotic regime. The predicted slope for the Cox--Snell + 1 Newton estimator assumes an exact first-order bias formula ($-4$). Source: \texttt{results/processed/summary.csv}.", r"\end{tablenotes}", r"\end{threeparttable}", r"\end{table}"]
    (TAB / fname).write_text("\n".join(lines) + "\n")


def table_boundary(fname="boundary.tex"):
    f = PROC / "boundary_summary.csv"
    if not f.exists():
        (TAB / fname).write_text("% no boundary data\n"); return
    b = pd.read_csv(f)
    sys.path.insert(0, str(ROOT / "experiments")); from boundary_pred import predicted_mean_gap
    b["pred_finite_m"] = [predicted_mean_gap(int(r.r), int(r.q), int(r.m), int(r.K), int(r.N)) for r in b.itertuples()]
    b.to_csv(PROC / "boundary_summary_with_prediction.csv", index=False)
    lines = [r"\begin{table}[t]", r"\centering\scriptsize", r"\begin{threeparttable}", r"\caption{Boundary study in the exponential model (Theorem~\ref{thm:exponential}(v)): mean of $\sqrt N(\theta^{(q)}_r - \widehat\theta_N)$ when $K \approx cN^{\alpha(r,q)}$, against the limit $-\theta_0(\beta_r c_{\mathrm{eff}}^r)^{2^q}$ ($q\ge1$) or $\beta_r\theta_0 c_{\mathrm{eff}}^r$ ($q=0$), with $c_{\mathrm{eff}} = K/N^{\alpha}$ after rounding $m$ to a multiple of the finest split.}\label{tab:boundary}",
             r"\begin{tabular}{@{}ccrrrrrrrr@{}}", r"\toprule", r"$(r,q)$ & $c$ & $N$ & $K$ & $m$ & reps & mean gap (s.e.) & finite-$m$ prediction & limit at $c_{\mathrm{eff}}$ & gap/limit \\", r"\midrule"]
    for (r, q, c), d in b.groupby(["r", "q", "c"]):
        for row in d.sort_values("N").itertuples():
            ratio = row.mean_gap / row.limit_c_eff if row.limit_c_eff != 0 else np.nan
            lines.append(f"({r},{q}) & {c:g} & $2^{{{int(np.log2(row.N))}}}$ & {row.K:,} & {row.m} & {row.reps} & {row.mean_gap:+.4g} ({row.se:.2g}) & {(f'{row.pred_finite_m:+.4g}' if np.isfinite(row.pred_finite_m) else '--')} & {row.limit_c_eff:+.4g} & {ratio:.2f} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\begin{tablenotes}\footnotesize", r"\item Exact-Gamma simulation of block sufficient statistics; seed 20261028 plus a cell hash. The ratio converges to one only as $m\to\infty$; with $m \asymp N^{1/(2^{q+1}r)}$ the local sample sizes in the table are small and the finite-$m$ bias $\theta_0\sum_j w_js_j/(m-s_j)$ exceeds its leading term. Source: \texttt{results/processed/boundary\_summary.csv}.", r"\end{tablenotes}", r"\end{threeparttable}", r"\end{table}"]
    (TAB / fname).write_text("\n".join(lines) + "\n")


def table_ablation(model, fname):
    files = sorted(glob.glob(str(RAW / f"ablation_{model}_*.csv")))
    if not files:
        (TAB / fname).write_text(f"% no ablation data for {model}\n"); return None
    d = pd.concat([pd.read_csv(f) for f in files], ignore_index=True)
    d = d[np.isfinite(d.gap)]
    g = d.groupby(["N", "m", "variant", "pattern", "r", "q", "hessian", "bits"]).agg(rms=("gap", lambda x: np.sqrt(np.mean(x ** 2))), mean0=("gap0", "mean"), n=("gap", "size"), valid=("valid_frac", "mean")).reset_index()
    g.to_csv(PROC / f"ablation_{model}_summary.csv", index=False)
    # (a) patterns x q at the two largest m of the largest N
    N = g.N.max(); ms = sorted(g[g.N == N].m.unique())
    lines = [r"\begin{table}[p]", r"\centering\scriptsize", r"\begin{threeparttable}", f"\\caption{{Ablation ({model}, $N = {N:,}$): RMS of $\\sqrt N\\,\\lVert\\tilde\\theta - \\widehat\\theta_N\\rVert$ by split pattern, order $r$, Newton rounds $q$ and curvature used (exact / stale averaged local / single machine); quantized exact curvature with $b$ mantissa bits for $q = 1$.}}\\label{{tab:abl_{model}}}",
             r"\begin{tabular}{@{}llr" + "r" * len(ms) + "@{}}", r"\toprule", "Pattern & curvature & $q$ & " + " & ".join(f"$m = {m}$" for m in ms) + r" \\", r"\midrule"]
    for pat in ["1", "1-2", "1-3", "1-2-4", "1-2-3", "1-2-4-8", "1-2-3-4"]:
        for hess in ["exact", "stale", "single"]:
            for q in range(0, 4):
                if hess != "exact" and q == 0:
                    continue
                vals = []
                for m in ms:
                    dd = g[(g.N == N) & (g.m == m) & (g.pattern == pat) & (g.q == q) & (g.hessian == hess) & (g.bits == 64)]
                    vals.append(fmt(float(dd.rms.iloc[0])) if len(dd) else "--")
                if all(v == "--" for v in vals):
                    continue
                lines.append(f"({pat.replace('-', ',')}) & {hess} & {q} & " + " & ".join(vals) + r" \\")
    lines.append(r"\midrule")
    for pat in ["1", "1-2", "1-2-4"]:
        for hess, lab in [("quantizedH", "$b$-bit Hessian"), ("quantizedg", "$b$-bit score"), ("quantizedBoth", "$b$-bit both")]:
            for bits in sorted(g.bits.unique()):
                if bits == 64:
                    continue
                vals = []
                for m in ms:
                    dd = g[(g.N == N) & (g.m == m) & (g.pattern == pat) & (g.q == 1) & (g.hessian == hess) & (g.bits == bits)]
                    vals.append(fmt(float(dd.rms.iloc[0])) if len(dd) else "--")
                if all(v == "--" for v in vals):
                    continue
                lines.append(f"({pat.replace('-', ',')}) & {lab}, $b={bits}$ & 1 & " + " & ".join(vals) + r" \\")
    reps = int(d.groupby(["N", "m"]).rep.nunique().max())
    lines += [r"\bottomrule", r"\end{tabular}", r"\begin{tablenotes}\footnotesize", f"\\item {reps} paired replications per column; $K = N/m$. Rows are omitted when the pattern does not divide $m$. Source: \\texttt{{results/processed/ablation\\_{model}\\_summary.csv}}.", r"\end{tablenotes}", r"\end{threeparttable}", r"\end{table}"]
    (TAB / fname).write_text("\n".join(lines) + "\n"); return g


def table_skew(fname="skew.tex"):
    files = sorted(glob.glob(str(RAW / "skew_*.csv")))
    if not files:
        (TAB / fname).write_text("% no skew data\n"); return
    d = pd.concat([pd.read_csv(f) for f in files], ignore_index=True)
    g = d.groupby(["model", "K", "ratio", "agg", "pilot"]).agg(gap0=("gap_q0", lambda x: np.sqrt(np.mean(x ** 2))), gap1=("gap_q1", lambda x: np.sqrt(np.mean(x ** 2))), gap2=("gap_q2", lambda x: np.sqrt(np.mean(x ** 2))), n=("rep", "size"), N=("N", "first")).reset_index()
    g.to_csv(PROC / "skew_summary.csv", index=False)
    lines = [r"\begin{table}[t]", r"\centering\scriptsize", r"\begin{threeparttable}", r"\caption{Unequal block sizes (Proposition~\ref{prop:robust}): RMS of $\sqrt N\,\lVert\tilde\theta - \widehat\theta_N\rVert$ for size-weighted and equal-weight aggregation of the raw and half-block-jackknife local estimates, with $q = 0, 1, 2$ Newton corrections. Sizes grow geometrically with max/min ratio as shown; $N \approx 2^{18}$.}\label{tab:skew}",
             r"\begin{tabular}{@{}llrrllrrr@{}}", r"\toprule", r"Model & $K$ & ratio & reps & weights & pilot & $q=0$ & $q=1$ & $q=2$ \\", r"\midrule"]
    for r in g.sort_values(["model", "K", "ratio", "agg", "pilot"]).itertuples():
        lines.append(f"{r.model} & {r.K} & {r.ratio:g} & {r.n} & {r.agg} & {r.pilot} & {fmt(r.gap0)} & {fmt(r.gap1)} & {fmt(r.gap2)} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\begin{tablenotes}\footnotesize", r"\item Source: \texttt{results/processed/skew\_summary.csv}.", r"\end{tablenotes}", r"\end{threeparttable}", r"\end{table}"]
    (TAB / fname).write_text("\n".join(lines) + "\n")


def table_bench(fname="bench.tex"):
    f = RAW / "benchmarks.jsonl"
    if not f.exists():
        (TAB / fname).write_text("% no benchmark data\n"); return None
    rows = [json.loads(l) for l in open(f) if l.strip()]
    recs = []
    for r in rows:
        st = r["stats"]; A = st["stages"]["A_B_local_fit_and_upload"]; rounds = [v for k, v in st["stages"].items() if k.startswith("D_F")]
        recs.append({"label": r["label"], "K": r["K"], "hosts": r["n_hosts"], "p": r["p"], "m": r["m"], "r": len(r["splits"]), "q": r["q"], "hessian": r["hessian_mode"],
                     "t_total": r["t_total_after_connect"], "t_local": A["wall"], "fit_mean": A["worker_fit_s_mean"], "fit_max": A["worker_fit_s_max"], "lat99_A": A["latency_p99"],
                     "t_rounds": sum(v["wall"] for v in rounds), "round_wall": np.mean([v["wall"] for v in rounds]) if rounds else 0.0, "der_max": max([v["worker_der_s_max"] for v in rounds]) if rounds else 0.0,
                     "lat99_round": max([v["latency_p99"] for v in rounds]) if rounds else 0.0, "up_B": st["bytes_up"], "down_B": st["bytes_down"], "msgs": st["messages_up"] + st["messages_down"], "grad_final": r["newton"][-1]["grad_norm"] if r["newton"] else np.nan, "t_connect": r["t_connect"]})
    d = pd.DataFrame(recs)
    g = d.groupby(["label", "K", "hosts", "p", "r", "q", "hessian"]).agg(reps=("t_total", "size"), t_total=("t_total", "median"), t_local=("t_local", "median"), fit_mean=("fit_mean", "median"), fit_max=("fit_max", "median"),
                                                                     t_rounds=("t_rounds", "median"), round_wall=("round_wall", "median"), der_max=("der_max", "median"), lat99=("lat99_round", "median"), up_B=("up_B", "first"), down_B=("down_B", "first"), msgs=("msgs", "first"), grad=("grad_final", "median"), t_connect=("t_connect", "median")).reset_index()
    g.to_csv(PROC / "benchmarks_summary.csv", index=False)
    lines = [r"\begin{table}[p]", r"\centering\scriptsize", r"\begin{threeparttable}", r"\caption{Measured multi-node benchmarks of the protocol (logistic regression, $N = 2^{22}$, medians over repeats): wall-clock of the local phase (slowest worker) and of the synchronized derivative phases, worker compute times, 99th-percentile reply latency, and bytes moved.}\label{tab:bench}",
             r"\begin{tabular}{@{}lrrrrrlrrrrrrrr@{}}", r"\toprule", r"Design & $K$ & hosts & $p$ & $r$ & $q$ & curvature & local (s) & fit max (s) & rounds (s) & per round (s) & der.\ max (s) & lat.\ p99 (s) & up (kB) & down (kB) \\", r"\midrule"]
    for r in g.itertuples():
        lines.append(f"{r.label.replace('_', ' ')} & {r.K} & {r.hosts} & {r.p} & {r.r} & {r.q} & {r.hessian} & {r.t_local:.3f} & {r.fit_max:.3f} & {r.t_rounds:.3f} & {r.round_wall:.3f} & {r.der_max:.3f} & {r.lat99:.3f} & {r.up_B/1e3:.1f} & {r.down_B/1e3:.1f} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\begin{tablenotes}\footnotesize", r"\item Python 3.9/3.14 + NumPy over TCP (length-prefixed pickle frames) on the MacLustr LAN (1 GbE/10 GbE mixed); coordinator on a Mac Studio M4 Max. `local' is the wall-clock from the first `fit' broadcast to the last upload, i.e.\ the slowest worker; `rounds' sums the synchronized derivative phases. Bytes include frame headers and pickle overhead. Repeats and raw logs: \texttt{results/raw/benchmarks.jsonl}.", r"\end{tablenotes}", r"\end{threeparttable}", r"\end{table}"]
    (TAB / fname).write_text("\n".join(lines) + "\n"); return g


def manifest():
    rows = []
    for f in sorted(glob.glob(str(RAW / "*.manifest.json"))):
        m = json.load(open(f)); c = m["cell"]
        rows.append({"run_id": m["run_id"], "git_commit": m["git_commit"], "python_version": m["python"], "numpy": m["numpy"], "dependency_lock_sha256": "", "seed": m["seed"], "N": c["N"], "K": c["K"], "m": c["m"], "p": c["p"],
                     "DGP": json.dumps(c["dgp"]), "method": "estimate_all", "optimizer": "batched Newton, step halving, tol 1e-9 (central 1e-12)", "tolerance": "1e-9", "replicates": m["reps"], "parallel_backend": "xargs -P / nohup", "node_count": 1,
                     "hardware": m["host"] + " " + m["machine"], "OS": "macOS", "wall_seconds": round(m["wall_seconds"], 1), "communicated_bytes": 0, "rounds": "in-memory (no network)", "status": m["status"], "raw_output_sha256": m["raw_output_sha256"]})
    pd.DataFrame(rows).to_csv(ROOT / "experiments" / "EXPERIMENT_MANIFEST.csv", index=False)
    return len(rows)


def main():
    s = pd.read_csv(PROC / "summary.csv") if (PROC / "summary.csv").exists() else pd.DataFrame()
    made = {}
    if not s.empty:
        for (label, p), d in s.groupby(["label", "p"]):
            N = d.N.max()
            made[(label, p)] = table_cell(s, label, p, N, sorted(d.m.unique()), f"mc_{label}_p{p}.tex", f"Monte Carlo, {label.replace('_', ' ')} design, $p = {p}$, $N = {N:,}$: scaled gap to $\\widehat\\theta_N$, scaled RMSE and 95\\% coverage by local sample size $m$.")
        table_rates(s)
    table_boundary(); table_ablation("exponential", "abl_exponential.tex"); table_ablation("logistic", "abl_logistic.tex"); table_skew(); table_bench()
    n = manifest()
    # numbers macros
    macros = {"numMCcells": str(s.groupby(["label", "p", "N", "K"]).ngroups if not s.empty else 0), "numManifest": str(n)}
    if not s.empty:
        macros["numMCrepsExp"] = str(int(s[s.label == "exp"].reps.max())) if (s.label == "exp").any() else "0"
        macros["numMCrepsLogit"] = str(int(s[s.label == "logit"].reps.max())) if (s.label == "logit").any() else "0"
    bf = PROC / "benchmarks_summary.csv"
    macros["benchHosts"] = str(int(pd.read_csv(bf).hosts.max())) if bf.exists() else "0"
    with open(TAB / "numbers.tex", "w") as f:
        f.write("% Auto-generated by experiments/make_tables.py\n")
        for k, v in macros.items():
            f.write(f"\\newcommand{{\\{k}}}{{{v}}}\n")
    print("tables written:", sorted(p.name for p in TAB.glob("*.tex")), "manifest rows:", n)


if __name__ == "__main__":
    main()