"""Build every empirical table and figure from raw records. Each output gets a provenance entry (source run files + script)."""
from __future__ import annotations
import json, sys, glob
from pathlib import Path
import numpy as np, pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from ..util import ROOT, read_json, write_json, read_jsonl, utc_now
from .load import load_experiments, LABELS
from .stats import cluster_bootstrap_mean, paired_difference, holm, summarize_latency

# fixed hue per method family (colorblind-validated palette; see results/figures/PALETTE.txt)
PAL = {"A": "#2a78d6", "B_quant": "#eb6834", "B_sel": "#1baf7a", "B_ref": "#6250d6", "C": "#e87ba4", "A_struct": "#eda100", "neg": "#e34948", "hosted": "#008300"}
METHOD_COLOR = {"A0": "#7a7a78", "A1": "#2a78d6", "A2": "#2a78d6", "A3": "#2a78d6", "A4": "#2a78d6", "A5": "#eda100", "A6": "#eda100", "A7": "#2a78d6", "A8": "#2a78d6",
                "B0": "#6250d6", "B1": "#6250d6", "B2": "#eb6834", "B3": "#1baf7a", "B4": "#e34948", "B5": "#1baf7a", "B6": "#1baf7a", "B9": "#eb6834", "C0": "#e34948", "C2": "#e87ba4", "C3": "#e87ba4"}
MARKER = {"A0": "x", "A1": "o", "A2": "v", "A3": "^", "A4": "s", "A5": "D", "A6": "P", "A7": "o", "A8": "<", "B0": "o", "B1": "o", "B2": "s", "B3": "^", "B4": "x", "B5": "D", "B6": "P", "B9": "*", "C0": "x", "C2": "s", "C3": "D"}
TAB = ROOT / "results/tables"; FIG = ROOT / "results/figures"
PROV: dict = {}

def _prov(path: Path, sources: list[str], note: str = ""):
    PROV[str(path.relative_to(ROOT))] = {"sources": sources, "script": "src/wp30/analysis/analyze.py", "generated_utc": utc_now(), "note": note}

def _save_fig(fig, name: str, sources: list[str], note: str = ""):
    FIG.mkdir(parents=True, exist_ok=True)
    for ext in ("pdf", "png"):
        p = FIG / f"{name}.{ext}"; fig.savefig(p, bbox_inches="tight", dpi=160 if ext == "png" else None)
        _prov(p, sources, note)
    plt.close(fig)

def _style():
    plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False, "axes.grid": True, "grid.alpha": 0.25, "grid.linewidth": 0.5,
                         "legend.frameon": False, "figure.dpi": 110, "axes.titlesize": 9.5, "axes.labelsize": 9})

def fmt_bytes(b):
    return f"{b/1e6:.1f} MB" if b >= 1e6 else (f"{b/1e3:.1f} kB" if b >= 1e3 else f"{b:.0f} B")

# ------------------------------------------------------------------------------------------------------------- tables
def table_primary(df: pd.DataFrame, exp: str, sources: list[str]) -> pd.DataFrame:
    rows = []
    for (fam, pair, cond), g in df.groupby(["task_family", "pair", "condition"]):
        q = cluster_bootstrap_mean(g, "quality_operational"); ok = (g["status"] == "success").mean()
        rows.append({"family": fam, "pair": pair, "condition": cond, "method": g["method_id"].iloc[0], "access": g["question_access"].iloc[0], "n": q["n"], "n_docs": q["n_clusters"],
                     "quality": q["mean"], "q_lo": q["lo"], "q_hi": q["hi"], "success_rate": ok, "payload_bytes_median": g["serialized_payload_bytes"].median(),
                     "cache_bytes_median": g["cache_tensor_bytes"].median(), "cold_ms_p50": g["cost_cold_ms"].median(), "warm_ms_p50": g["cost_warm_query_ms"].median(),
                     "amortized_ms_p50": g["cost_amortized_ms"].median(), "ttft_ms_p50": g["ttft_ms"].median(), "peak_decode_mem_median": g["peak_decode_memory"].median(), "peak_cold_mem_median": g["peak_memory_cold"].median(),
                     "fetched_bytes_mean": g["fetched_bytes"].mean()})
    t = pd.DataFrame(rows).sort_values(["family", "pair", "condition"])
    TAB.mkdir(parents=True, exist_ok=True); p = TAB / f"{exp}_primary.csv"; t.to_csv(p, index=False); _prov(p, sources, "per-condition quality (cluster bootstrap CI) and resource medians")
    return t

def table_failures(df: pd.DataFrame, exp: str, sources):
    t = df.groupby(["method_id", "status"]).size().unstack(fill_value=0)
    p = TAB / f"{exp}_failures.csv"; t.to_csv(p); _prov(p, sources, "attempt status counts by method"); return t

def table_qtype(df: pd.DataFrame, exp: str, sources):
    d = df[df["task_family"] == "synthetic"]
    rows = []
    for (cond, qt), g in d.groupby(["condition", "qtype"]):
        q = cluster_bootstrap_mean(g, "quality_operational", n_boot=500); rows.append({"condition": cond, "qtype": qt, "quality": q["mean"], "lo": q["lo"], "hi": q["hi"], "n": q["n"]})
    t = pd.DataFrame(rows); p = TAB / f"{exp}_by_question_type.csv"; t.to_csv(p, index=False); _prov(p, sources, "synthetic accuracy by question type"); return t

def primary_tests(df: pd.DataFrame, exp: str, sources, tests: list[dict]):
    rows = []
    for t in tests:
        sub = df[df["task_family"].isin(t["families"])] if t.get("families") else df
        if t.get("pair"): sub = sub[sub["pair"] == t["pair"]]
        r = paired_difference(sub, t.get("value", "quality_operational"), "condition", t["a"], t["b"], keys=("document_id", "question_id", "pair"))
        r.update({"name": t["name"], "hypothesis": t.get("hypothesis", "")}); rows.append(r)
    out = pd.DataFrame(rows)
    if len(out):
        out["p_holm"] = holm([p if not np.isnan(p) else 1.0 for p in out["p"]])
    p = TAB / f"{exp}_primary_tests.csv"; out.to_csv(p, index=False); _prov(p, sources, "registered primary comparisons, paired, cluster bootstrap, Holm"); return out

# ------------------------------------------------------------------------------------------------------------ figures
def _short_pair(p: str) -> str:
    return p.replace("Qwen2.5-3B-Instruct-bf16->Qwen2.5-3B-Instruct-bf16", "Qwen2.5-3B").replace("Llama-3.2-3B-Instruct-bf16->Llama-3.2-3B-Instruct-bf16", "Llama-3.2-3B").replace("Qwen2.5-1.5B-Instruct-bf16->Qwen2.5-3B-Instruct-bf16", "Qwen1.5B->3B")

def _mk(m: str, bits) -> str:
    if m == "B2" and bits == 4: return "p"
    if m == "A7": return "h"
    return MARKER.get(m, "o")

def fig_quality_vs_bytes(df: pd.DataFrame, exp: str, sources, access_filter=None, name="fig2_quality_vs_bytes"):
    _style()
    fams = [f for f in ["synthetic", "exec", "public_qasper", "public_2wikimqa"] if f in set(df["task_family"])]
    pairs = sorted(df["pair"].unique())
    fig, axes = plt.subplots(len(pairs), len(fams), figsize=(3.1 * len(fams), 2.7 * len(pairs)), squeeze=False, sharey="row")
    for i, pair in enumerate(pairs):
        for j, fam in enumerate(fams):
            ax = axes[i][j]; d = df[(df["pair"] == pair) & (df["task_family"] == fam)]
            if access_filter: d = d[d["question_access"].isin(access_filter)]
            for cond, g in d.groupby("condition"):
                m = g["method_id"].iloc[0]; q = cluster_bootstrap_mean(g, "quality_operational", n_boot=300)
                x = g["serialized_payload_bytes"].median()
                if m == "A0" or x is None or np.isnan(x): continue
                ax.errorbar(max(x, 1), q["mean"], yerr=[[q["mean"] - q["lo"]], [q["hi"] - q["mean"]]], fmt=_mk(m, g["bits"].iloc[0]), color=METHOD_COLOR.get(m, "#444"), ms=4.5, lw=0.9, capsize=1.5, mfc="white" if m.startswith("B") else None)
            ax.set_xscale("log"); ax.set_title(f"{fam.replace('public_', '')} - {_short_pair(pair)}", fontsize=8.5); ax.set_ylim(-0.03, 1.03)
            if i == len(pairs) - 1: ax.set_xlabel("serialized payload bytes (log)")
            if j == 0: ax.set_ylabel("quality (acc. or F1)")
    present = set(df["method_id"]); has4 = ((df["method_id"] == "B2") & (df["bits"] == 4)).any()
    handles = [plt.Line2D([], [], marker=_mk(m, None), color=METHOD_COLOR[m], ls="", ms=5, mfc="white" if m.startswith("B") else None, label=LABELS[m] + (" 8-bit" if m == "B2" else "")) for m in ["A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8", "B0", "B1", "B2", "B3", "B4", "B5", "B6", "B9", "C0", "C2", "C3"] if m in present]
    if has4: handles.append(plt.Line2D([], [], marker="p", color=METHOD_COLOR["B2"], ls="", ms=5, mfc="white", label="B2 packed quant 4-bit"))
    fig.legend(handles=handles, loc="lower center", ncol=5, fontsize=7, bbox_to_anchor=(0.5, -0.06 - 0.02 * (len(handles) // 5)))
    fig.subplots_adjust(hspace=0.5); fig.suptitle("Quality vs measured serialized payload bytes (medians; 95% cluster-bootstrap CIs)", fontsize=9.5, y=1.0)
    _save_fig(fig, name, sources, "points = conditions; x = median payload bytes; y = mean quality")

def fig_quality_vs_latency(df: pd.DataFrame, exp: str, sources, name="fig3_quality_vs_latency_memory"):
    _style()
    from matplotlib.ticker import NullFormatter
    df = df[(df["hardware_id"] == "laptop-m5max-48g") & (df["pair"].str.startswith("Qwen"))]   # one machine, one pair: comparable timings
    fams = [f for f in ["synthetic", "exec", "public_qasper", "public_2wikimqa"] if f in set(df["task_family"])]
    fig, axes = plt.subplots(2, len(fams), figsize=(3.1 * len(fams), 5.6), squeeze=False)
    for j, fam in enumerate(fams):
        d = df[df["task_family"] == fam]
        for cond, g in d.groupby("condition"):
            m = g["method_id"].iloc[0]; q = g["quality_operational"].mean()
            if m == "A0": continue
            axes[0][j].scatter(g["cost_cold_ms"].median(), q, marker=_mk(m, g["bits"].iloc[0]), color=METHOD_COLOR.get(m, "#444"), s=22, facecolors="white" if m.startswith("B") else METHOD_COLOR.get(m, "#444"), lw=0.9)
            axes[1][j].scatter(g["peak_memory_cold"].median() / 1e9, q, marker=_mk(m, g["bits"].iloc[0]), color=METHOD_COLOR.get(m, "#444"), s=22, facecolors="white" if m.startswith("B") else METHOD_COLOR.get(m, "#444"), lw=0.9)
        axes[0][j].set_xscale("log"); axes[0][j].xaxis.set_minor_formatter(NullFormatter()); axes[0][j].set_title(fam.replace("public_", ""), fontsize=8.5); axes[0][j].set_xlabel("cold latency, ms (log)"); axes[1][j].set_xlabel("peak MLX memory incl. build, GB")
        axes[0][j].set_ylim(-0.03, 1.03); axes[1][j].set_ylim(-0.03, 1.03)
    axes[0][0].set_ylabel("quality"); axes[1][0].set_ylabel("quality")
    present = set(df["method_id"])
    handles = [plt.Line2D([], [], marker=_mk(m, None), color=METHOD_COLOR[m], ls="", ms=5, mfc="white" if m.startswith("B") else None, label=LABELS[m] + (" 8-bit" if m == "B2" else "")) for m in ["A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8", "B0", "B1", "B2", "B3", "B4", "B5", "B6", "B9"] if m in present]
    if ((df["method_id"] == "B2") & (df["bits"] == 4)).any(): handles.append(plt.Line2D([], [], marker="p", color=METHOD_COLOR["B2"], ls="", ms=5, mfc="white", label="B2 packed quant 4-bit"))
    fig.legend(handles=handles, loc="lower center", ncol=6, fontsize=7, bbox_to_anchor=(0.5, -0.07))
    fig.subplots_adjust(hspace=0.45)
    fig.suptitle("Quality vs cold latency (top) and peak MLX memory (bottom) - Qwen2.5-3B pair on the laptop (models resident: ~6.5 GB baseline)", fontsize=9)
    _save_fig(fig, name, sources)

def fig_latency_decomposition(df: pd.DataFrame, exp: str, sources, name="fig4_latency_decomposition"):
    _style()
    d = df[(df["task_family"] == "synthetic") & (df["status"] == "success") & (df["hardware_id"] == "laptop-m5max-48g") & (df["pair"].str.startswith("Qwen"))]
    conds = [c for c in ["A1", "A3|budget_tokens=512|aware", "A4|budget_tokens=512|blind", "A5|budget_tokens=512|blind", "A7|codec=zstd", "A8|budget_tokens=512|blind", "B0", "B1", "B2|bits=8", "B3|keep_frac=0.25", "B5|keep_frac=0.25|aware", "B6|keep_frac=0.25|aware"] if c in set(d["condition"])]
    comps = [("span_encode_ms", "sender encode (text)"), ("span_doc_prefill_ms", "document prefill (cache)"), ("span_compress_ms", "compress/select"), ("span_serialize_ms", "serialize"), ("span_load_ms", "deserialize/load"), ("span_receiver_prefill_ms", "receiver prefill"), ("ttft_ms", "query prefill→first token"), ("span_decode_ms", "decode")]
    cols = ["#2a78d6", "#6250d6", "#1baf7a", "#eda100", "#e87ba4", "#eb6834", "#008300", "#7a7a78"]
    fig, ax = plt.subplots(figsize=(7.2, 3.8))
    for i, c in enumerate(conds):
        g = d[d["condition"] == c]; left = 0
        for (col, lab), color in zip(comps, cols):
            if col == "span_encode_ms" and c.startswith("B"):
                continue   # for cache methods the encode span is the sum of the components plotted separately
            v = g[col].median() if col in g and g[col].notna().any() else 0.0
            if v and v > 0:
                ax.barh(i, v, left=left, color=color, height=0.6, edgecolor="white", linewidth=0.8); left += v
    ax.set_yticks(range(len(conds))); ax.set_yticklabels([c.replace("|", " ") for c in conds], fontsize=7.5); ax.invert_yaxis(); ax.set_xlabel("median span duration, ms (serial accounting; one question per document build)")
    ax.legend(handles=[plt.Rectangle((0, 0), 1, 1, color=c) for c in cols], labels=[lab for _, lab in comps], fontsize=7, ncol=4, loc="upper center", bbox_to_anchor=(0.5, -0.22)); ax.set_title("Latency decomposition - synthetic family, Qwen2.5-3B on the laptop, cold scenario (document build counted once)", fontsize=8.5)
    _save_fig(fig, name, sources)

def fig_handoffs(df: pd.DataFrame, exp: str, sources, name="fig5_handoffs"):
    _style()
    d = df[df["topology"].astype(str).str.startswith("chain")]
    if d.empty: return
    fig, axes = plt.subplots(1, 3, figsize=(9.6, 3.4))
    d["chain"] = d["run_group_id"].str.split(":").str[1]
    for chain, g in d.groupby("chain"):
        rep = g["method_id"].iloc[0]; mode = g["topology"].iloc[0].replace("chain_", ""); agents = chain.split("|")[-1]
        alt = "-" in agents
        color = "#eda100" if rep == "A5" else "#2a78d6"; ls = "-" if mode == "closed" else "--"; mk = "s" if alt else "o"
        hop = g.groupby("handoff_index")
        ret = hop["retention_rate"].mean(); uns = hop["unsupported"].mean(); acc = hop["quality_operational"].mean()
        lab = f"{rep} {mode}" + (" (Qwen/Llama alternating)" if alt else " (Qwen)")
        axes[0].plot(ret.index, ret.values, ls, color=color, marker=mk, ms=4, label=lab, mfc="white" if alt else None)
        axes[1].plot(uns.index, uns.values, ls, color=color, marker=mk, ms=4, mfc="white" if alt else None)
        axes[2].plot(acc.index, acc.values, ls, color=color, marker=mk, ms=4, mfc="white" if alt else None)
    for ax, t in zip(axes, ["retained ledger facts (share of 62)", "unsupported numbers/ids per message", "answer accuracy (15 question types)"]):
        ax.set_xlabel("handoffs"); ax.set_title(t, fontsize=8.5); ax.set_xticks(sorted(d["handoff_index"].unique()))
    axes[0].set_ylim(0, 1.02); axes[2].set_ylim(0, 1.02); axes[1].set_ylim(0, 0.3)
    fig.subplots_adjust(bottom=0.3, wspace=0.3)
    fig.legend(*axes[0].get_legend_handles_labels(), loc="lower center", ncol=3, fontsize=7, bbox_to_anchor=(0.5, 0.0))
    _save_fig(fig, name, sources, "means over documents; chains defined in configs/chains*.yaml")

def fig_breakeven(net_rows: list[dict], exp: str, sources, name="fig6_breakeven"):
    """Left: measured two-node transfer vs remote text prefill (medians with p25-p75 bars per target length).
    Right: analytic transfer-time model from the measured bandwidth and the measured remote prefill slope."""
    _style()
    if not net_rows: return
    n = pd.DataFrame(net_rows); n = n[n["label"] != "loopback"]
    fig, axes = plt.subplots(1, 2, figsize=(8.6, 3.1))
    ax = axes[0]
    series = [("text_remote_prefill", None, "none", "text -> remote prefill", "#2a78d6", "o", "-"),
              ("B1", None, "none", "B1 bf16 cache", "#6250d6", "D", "-"), ("B1", None, "zstd", "B1 bf16 + zstd", "#6250d6", "D", ":"),
              ("B2", 8.0, "none", "B2 8-bit cache", "#eb6834", "s", "-"), ("B2", 4.0, "none", "B2 4-bit cache (Qwen: quality 0)", "#e34948", "x", "--")]
    for meth, bits, codec, lab, color, mk, ls in series:
        g = n[(n["method"] == meth) & (n["codec"].fillna("none") == codec)]
        g = g[g["bits"].isna()] if bits is None else g[g["bits"] == bits]
        if g.empty: continue
        agg = g.groupby("target_tokens").agg(x=("prefix_tokens", "median"), y=("roundtrip_ms", "median"), lo=("roundtrip_ms", lambda s: s.quantile(0.25)), hi=("roundtrip_ms", lambda s: s.quantile(0.75))).reset_index()
        ax.errorbar(agg["x"], agg["y"], yerr=[agg["y"] - agg["lo"], agg["hi"] - agg["y"]], fmt=mk + ls, color=color, ms=4.5, lw=1.1, capsize=2, label=lab, mfc="white" if meth == "B1" else None)
    loc = n[n["method"] == "local_prefill"]
    if len(loc):
        agg = loc.groupby("target_tokens").agg(x=("prefix_tokens", "median"), y=("local_prefill_ms", "median")).reset_index()
        ax.plot(agg["x"], agg["y"], "^-.", color="#7a7a78", ms=4.5, lw=1, label="local prefill only (no transfer, no decode)")
    ax.set_xlabel("context tokens (median per length class)"); ax.set_ylabel("ms"); ax.set_title("Measured: M3U96a -> M3U96b, 1 GbE (18 runs per point)", fontsize=8.5)
    ax.legend(fontsize=6.5, loc="upper left")
    ax = axes[1]
    bw = n[n["method"].isin(["B1", "B2"])]["effective_gbps"].median()
    b1 = n[(n["method"] == "B1") & (n["codec"].fillna("none") == "none")]
    bpt = (b1["serialized_bytes"] / b1["prefix_tokens"]).median() if len(b1) else np.nan
    tr = n[n["method"] == "text_remote_prefill"]
    if len(tr) and not np.isnan(bpt):
        slope = (tr["remote_prefill_ms"] / tr["prefix_tokens"]).median()
        L = np.linspace(500, 32000, 50)
        for frac, lab, color in [(1.0, "bf16 cache", "#6250d6"), (0.53, "8-bit cache", "#eb6834")]:
            for bw_g, ls, alpha in [(bw, "-", 0.95), (10.0, "--", 0.5)]:
                ax.plot(L, L * bpt * frac * 8 / (bw_g * 1e9) * 1000, ls, color=color, lw=1.1, alpha=alpha, label=f"{lab} transfer @ {bw_g:.1f} Gb/s" + (" (hypothetical)" if bw_g == 10.0 else " (measured)"))
        ax.plot(L, L * slope, "-", color="#2a78d6", lw=1.6, label="remote text prefill (measured slope, M3 Ultra)")
        ax.set_xlabel("context tokens"); ax.set_ylabel("ms"); ax.set_title("Analytic model (Qwen2.5-3B cache, measured slopes)", fontsize=8.5)
        ax.legend(fontsize=6.3, loc="upper left"); ax.set_ylim(0, None)
    _save_fig(fig, name, sources, "left = measured (medians, p25-p75); right = analytic model from measured bandwidth and prefill slope; dashed 10 Gb/s is hypothetical")

# ------------------------------------------------------------------------------------------------------------ main
def main(protocol_path: str) -> int:
    _style()
    prot = read_json(ROOT / protocol_path) if (ROOT / protocol_path).exists() else {}
    core_exp = prot.get("core_experiment_ids", ["E1_core"]) if prot else ["E1_core"]
    exps = {p.stem: p.name for p in (ROOT / "runs/raw").glob("*.jsonl")}
    (FIG).mkdir(parents=True, exist_ok=True); (TAB).mkdir(parents=True, exist_ok=True)
    (FIG / "PALETTE.txt").write_text("categorical palette validated with dataviz validate_palette.js (light): #2a78d6,#eb6834,#1baf7a,#eda100,#e87ba4,#008300,#6250d6,#e34948 — ALL CHECKS PASS\n")
    # choose the experiment set to analyze: core if present, else pilot
    main_exp = "E1_core" if "E1_core" in exps else ("E1_pilot" if "E1_pilot" in exps else None)
    if main_exp:
        src = [f"runs/raw/{main_exp}.jsonl"]
        df = load_experiments([f"{main_exp}.jsonl"]); df_all = load_experiments([f"{main_exp}.jsonl"], last_attempt_only=False)
        table_primary(df, main_exp, src); table_failures(df_all, main_exp, src); table_qtype(df, main_exp, src)
        tests = [
            {"name": "T1 A5 vs A4 (512, blind, synthetic)", "a": "A5|budget_tokens=512|blind", "b": "A4|budget_tokens=512|blind", "families": ["synthetic"], "hypothesis": "RQ5/H1"},
            {"name": "T2 B2-8bit vs B0", "a": "B2|bits=8", "b": "B0", "hypothesis": "H2"},
            {"name": "T3 B2-4bit vs B0", "a": "B2|bits=4", "b": "B0", "hypothesis": "H2"},
            {"name": "T4 B5-aware vs B3 (1/4)", "a": "B5|keep_frac=0.25|aware", "b": "B3|keep_frac=0.25", "hypothesis": "H3"},
            {"name": "T6 A1 vs B0 (text vs native cache, same content)", "a": "A1", "b": "B0", "hypothesis": "consistency"},
        ]
        primary_tests(df, main_exp, src, tests)
        fig_quality_vs_bytes(df, main_exp, src)
        fig_quality_vs_bytes(df[df["question_access"].isin(["blind", "none"])], main_exp, src, name="figA_quality_vs_bytes_blind")
        fig_quality_vs_bytes(df[df["question_access"] == "aware"], main_exp, src, name="figA_quality_vs_bytes_aware")
        fig_quality_vs_latency(df, main_exp, src); fig_latency_decomposition(df, main_exp, src)
    chain_exp = "E4_core" if "E4_core" in exps else ("E4_pilot" if "E4_pilot" in exps else None)
    if chain_exp:
        dfc = load_experiments([f"{chain_exp}.jsonl"]); src = [f"runs/raw/{chain_exp}.jsonl"]
        from .chains_recompute import recompute
        dfc = recompute(dfc, chain_exp)   # ledger metrics recomputed from stored hop messages (span-offset fix)
        fig_handoffs(dfc, chain_exp, src)
        rows = []
        for (grp, h), g in dfc.groupby(["run_group_id", "handoff_index"]):
            q = cluster_bootstrap_mean(g, "quality_operational", n_boot=500); r = cluster_bootstrap_mean(g.drop_duplicates("document_id"), "retention_rate", n_boot=500)
            rows.append({"chain": grp.split(":")[1], "handoffs": h, "accuracy": q["mean"], "acc_lo": q["lo"], "acc_hi": q["hi"], "retention": r["mean"], "ret_lo": r["lo"], "ret_hi": r["hi"],
                         "unsupported_mean": g.drop_duplicates("document_id")["unsupported"].mean(), "fetched_bytes_mean": g["fetched_bytes"].mean(), "msg_bytes_mean": g["msg_bytes_hop"].mean(), "n": q["n"], "n_docs": q["n_clusters"]})
        t = pd.DataFrame(rows); p = TAB / f"{chain_exp}_handoffs.csv"; t.to_csv(p, index=False); _prov(p, src, "repeated handoffs")
        primary_tests(dfc[dfc["handoff_index"] == 4].assign(condition=dfc["run_group_id"].str.split(":").str[1]), chain_exp + "_h4", src,
                      [{"name": "T5 A5 chain vs A4 chain at 4 handoffs (closed)", "a": "A5|closed|qwen3b", "b": "A4|closed|qwen3b", "hypothesis": "H5"}])
    fan_exp = "E5_core" if "E5_core" in exps else ("E5_pilot" if "E5_pilot" in exps else None)
    if fan_exp:
        dff = load_experiments([f"{fan_exp}.jsonl"]); src = [f"runs/raw/{fan_exp}.jsonl"]
        rows = []
        for (cond, qt), g in dff.groupby(["condition", "qtype"]):
            q = cluster_bootstrap_mean(g, "quality_operational", n_boot=500); rows.append({"condition": cond, "qtype": qt, "quality": q["mean"], "lo": q["lo"], "hi": q["hi"], "n": q["n"], "payload_bytes_median": g["serialized_payload_bytes"].median(), "fetched_bytes_mean": g["fetched_bytes"].mean()})
        t = pd.DataFrame(rows); p = TAB / f"{fan_exp}_fanin.csv"; t.to_csv(p, index=False); _prov(p, src, "fan-in by method and question type")
    net_files = [p for p in (ROOT / "runs/raw").glob("E6_*.jsonl") if "smoke" not in p.name] or list((ROOT / "runs/raw").glob("E6_*.jsonl"))
    if net_files:
        rows = [r for p in net_files for r in read_jsonl(p)]
        fig_breakeven(rows, "E6", [str(p.relative_to(ROOT)) for p in net_files])
        n = pd.DataFrame(rows)
        agg = n.groupby(["label", "method", "bits", "codec", "target_tokens"], dropna=False).agg(n=("repeat", "size"), bytes=("serialized_bytes", "median"), roundtrip_ms_p50=("roundtrip_ms", "median"), send_ms_p50=("send_ms", "median"),
                                                                                              remote_load_ms_p50=("remote_load_ms", "median"), remote_ttft_p50=("remote_ttft_ms", "median"), remote_prefill_p50=("remote_prefill_ms", "median"), eff_gbps=("effective_gbps", "median"), correct=("correct", "mean"),
                                                                                              local_prefill_p50=("local_prefill_ms", "median"), local_ttft_p50=("local_ttft_ms", "median"), remote_decode_p50=("remote_decode_ms", "median")).reset_index()
        p = TAB / "E6_network.csv"; agg.to_csv(p, index=False); _prov(p, [str(x.relative_to(ROOT)) for x in net_files], "two-node transfer measurements")
    hosted = [p for p in (ROOT / "runs/raw").glob("ED_hosted_core.jsonl")] or [p for p in (ROOT / "runs/raw").glob("ED_hosted_pilot.jsonl")]
    if hosted:
        dh = load_experiments([p.name for p in hosted]); src = [str(p.relative_to(ROOT)) for p in hosted]
        rows = []
        for (fam, recv, cond), g in dh.groupby(["task_family", "receiver", "condition"]):
            q = cluster_bootstrap_mean(g, "quality_operational", n_boot=500)
            rows.append({"family": fam, "provider_model": recv, "condition": cond, "quality": q["mean"], "lo": q["lo"], "hi": q["hi"], "n": q["n"], "cost_usd_mean": g["cost_usd"].mean(), "ttft_ms_p50": g["ttft_ms"].median(), "total_ms_p50": g["total_wall_ms"].median(), "success_rate": (g["status"] == "success").mean()})
        t = pd.DataFrame(rows); p = TAB / "ED_hosted.csv"; t.to_csv(p, index=False); _prov(p, src, "hosted receivers; cost estimated from provider price tables")
    if "E3_core" in exps:
        d3 = load_experiments(["E3_core.jsonl"]); s3 = load_experiments(["E3_core_sender.jsonl"]) if "E3_core_sender" in exps else pd.DataFrame()
        src3 = ["runs/raw/E3_core.jsonl", "runs/raw/E3_core_sender.jsonl"]
        d3["cond_s"] = d3["condition"].str.replace("|mapper=runs/checkpoints/mapper_v2_qwen1b5_to_qwen3b.npz", "", regex=False).str.replace("budget_tokens=512", "512", regex=False).str.replace("keep_frac=0.25", "1/4", regex=False)
        if len(s3): s3["cond_s"] = "sender alone " + s3["condition"]
        allr = pd.concat([d3, s3], ignore_index=True)
        rows = []
        for (fam, cond), g in allr.groupby(["task_family", "cond_s"]):
            q = cluster_bootstrap_mean(g, "quality_operational"); byq = g.groupby("qtype")["quality_operational"].mean().to_dict()
            rows.append({"family": fam, "condition": cond, "quality": q["mean"], "lo": q["lo"], "hi": q["hi"], "n": q["n"], "n_docs": q["n_clusters"], "payload_bytes_median": g["serialized_payload_bytes"].median(),
                         "cold_ms_p50": g["cost_cold_ms"].median(), "warm_ms_p50": g["cost_warm_query_ms"].median(), "unknown_share": g["prediction_artifact"].fillna("").str.lower().str.contains("unknown").mean(), **{f"q_{k}": v for k, v in byq.items()}})
        t3 = pd.DataFrame(rows); p = TAB / "E3_core_transfer.csv"; t3.to_csv(p, index=False); _prov(p, src3, "same-family mapper transfer vs text handoffs and references")
        primary_tests(d3.assign(condition=d3["cond_s"]), "E3_core", src3, [{"name": "C2 mapper vs A4 summary (synthetic)", "a": "C2", "b": "A4|512|blind", "families": ["synthetic"], "hypothesis": "H10"},
                                                                          {"name": "C2 mapper vs A1 receiver full context (synthetic)", "a": "C2", "b": "A1", "families": ["synthetic"], "hypothesis": "H10"}])
    if "E7_qtypes" in exps:
        d7 = load_experiments(["E7_qtypes.jsonl"]); src7 = ["runs/raw/E7_qtypes.jsonl"]
        rows = []
        for (cond, qt, pair), g in d7.groupby(["condition", "qtype", "pair"]):
            q = cluster_bootstrap_mean(g, "quality_operational", n_boot=500); rows.append({"condition": cond, "qtype": qt, "pair": _short_pair(pair), "quality": q["mean"], "lo": q["lo"], "hi": q["hi"], "n": q["n"], "abstained": g["abstained"].mean()})
        t7 = pd.DataFrame(rows); p = TAB / "E7_qtypes.csv"; t7.to_csv(p, index=False); _prov(p, src7, "EXPLORATORY: question types 6-15 on synthetic test documents")
        _style(); fig, axes = plt.subplots(1, 2, figsize=(9.4, 3.4), sharey=True)
        qorder = ["negation", "revision", "temporal", "aggregation", "unit_conversion", "unanswerable"]
        for ax, pair in zip(axes, sorted(t7["pair"].unique())):
            d = t7[t7["pair"] == pair]
            conds = [c for c in ["A1", "A3|budget_tokens=512|blind", "A4|budget_tokens=512|blind", "A5|budget_tokens=512|blind", "B2|bits=8", "B3|keep_frac=0.25", "B5|keep_frac=0.25|aware"] if c in set(d["condition"])]
            w = 0.8 / max(1, len(conds)); x = np.arange(len(qorder))
            for i, c in enumerate(conds):
                m = c.split("|")[0]; vals = [d[(d["condition"] == c) & (d["qtype"] == qt)]["quality"].mean() if ((d["condition"] == c) & (d["qtype"] == qt)).any() else np.nan for qt in qorder]
                ax.bar(x + (i - len(conds) / 2 + 0.5) * w, vals, w, color=METHOD_COLOR.get(m, "#444"), alpha=0.95 if c in ("A1", "A4|budget_tokens=512|blind", "A5|budget_tokens=512|blind") else 0.6, edgecolor="white", linewidth=0.5, label=LABELS.get(m, m) + (" 8-bit" if m == "B2" else "") + (" (aware)" if "aware" in c else ""))
            ax.set_xticks(x); ax.set_xticklabels([q.replace("_", "\n") for q in qorder], fontsize=7.5); ax.set_title(f"{pair}", fontsize=8.5); ax.set_ylim(0, 1.02)
        axes[0].set_ylabel("accuracy"); axes[0].legend(fontsize=6.5, ncol=2, loc="upper center", bbox_to_anchor=(1.05, -0.2))
        fig.suptitle("EXPLORATORY E7: question types not covered by the core (synthetic test, 40 documents)", fontsize=9)
        _save_fig(fig, "figA_qtypes_exploratory", src7, "exploratory extension A-004")
    r8 = TAB / "E8_router_summary.csv"
    if r8.exists():
        _style(); S = pd.read_csv(r8)
        fig, axes = plt.subplots(1, 2, figsize=(8.4, 3.0), sharey=True)
        order = ["router", "best_static_dev", "full_context", "native_cache", "attention_cache_aware", "structured", "oracle_infeasible"]
        col = {"router": "#e87ba4", "best_static_dev": "#2a78d6", "full_context": "#2a78d6", "native_cache": "#6250d6", "attention_cache_aware": "#1baf7a", "structured": "#eda100", "oracle_infeasible": "#7a7a78"}
        for ax, acc in zip(axes, ["blind", "aware"]):
            d = S[S["access"] == acc]
            seen_pts = {}
            for _, r in d.iterrows():
                key = (round(r["cost_ms"]), round(r["quality"], 4))
                lab = r["policy"].replace("_", " ").replace("best static dev", "best static (dev)")
                if key in seen_pts:   # identical policies (e.g. dev-best static == full context): one label
                    seen_pts[key].set_text(seen_pts[key].get_text() + " = " + lab); continue
                mk = "*" if r["policy"] == "router" else ("x" if r["policy"] == "oracle_infeasible" else "o")
                ax.errorbar(r["cost_ms"], r["quality"], xerr=[[r["cost_ms"] - r["c_lo"]], [r["c_hi"] - r["cost_ms"]]], yerr=[[r["quality"] - r["q_lo"]], [r["q_hi"] - r["quality"]]], fmt=mk, color=col.get(r["policy"], "#444"), ms=8 if mk == "*" else 5, capsize=2, lw=0.9)
                seen_pts[key] = ax.annotate(lab, (r["cost_ms"], r["quality"]), fontsize=6.5, xytext=(4, -9 if r["policy"] == "router" else 3), textcoords="offset points")
            ax.set_xscale("log"); ax.set_xlabel("mean amortized latency per query, ms (log)"); ax.set_title(f"{acc} regime (held-out test, n={int(d['n'].max())})", fontsize=8.5)
        axes[0].set_ylabel("mean quality (all families, both pairs)")
        fig.suptitle("E8: dev-fitted router vs static policies on held-out test instances (95% cluster-bootstrap bars)", fontsize=9)
        _save_fig(fig, "fig7_router", ["results/tables/E8_router_summary.csv", "runs/raw/E1_core.jsonl", "runs/checkpoints/router_policy.json"], "oracle = infeasible upper reference")
    write_json(ROOT / "results/provenance.json", PROV)
    print(f"analysis done: {len(PROV)} artifacts with provenance")
    return 0