#!/usr/bin/env python3
"""Sprint 1-2 pilot (Tier 0, micro-world A / OPSEQ).

What this script establishes, with artifacts:
  1. Donor/recipient capability gap on the target operator (Gate B), repeated over seeds.
  2. Hook no-op guarantee on trained models (Gate C): instrumented == uninstrumented, bitwise.
  3. Within-model causal localization of the target operator in the donor (activation patching,
     zero-ablation, head patching), with a sham baseline and reproducibility checks (RQ1).

No transplant is performed here. Nothing in this script claims transfer.

Usage: python scripts/run_pilot.py [--seeds 0 1 2] [--steps 3000] [--device auto] [--quick] [--dry-run]
"""
from __future__ import annotations

import argparse
import gzip
import json
import sys
import time
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from wp_ai_0001.config import ModelRoleConfig, OpSeqTaskConfig, PilotConfig, ToyModelConfig, TrainConfig, config_hash  # noqa: E402
from wp_ai_0001.datasets import opseq  # noqa: E402
from wp_ai_0001.datasets.common import Split, assert_disjoint, encode_split  # noqa: E402
from wp_ai_0001.evaluation.metrics import chance_level, evaluate  # noqa: E402
from wp_ai_0001.instrumentation.provenance import Manifest, make_run_id, utc_now  # noqa: E402
from wp_ai_0001.localization.patching import ablation_sweep, head_patch_sweep, patch_sweep  # noqa: E402
from wp_ai_0001.models.hooks import Capture, Zero  # noqa: E402
from wp_ai_0001.models.toy_transformer import resolve_device  # noqa: E402
from wp_ai_0001.orchestration import registry  # noqa: E402
from wp_ai_0001.statistics.bootstrap import cluster_bootstrap_mean, cluster_bootstrap_paired_diff, seeds_summary  # noqa: E402
from wp_ai_0001.training.train_toy import load as load_model  # noqa: E402
from wp_ai_0001.training.train_toy import train  # noqa: E402

OMEGA = "qua"
MAX_SEQ = 24


def build_splits(cfg_donor: OpSeqTaskConfig, cfg_recip: OpSeqTaskConfig, n_train: int, n_dev: int, n_eval: int, base_seed: int) -> dict[str, Split]:
    """All splits from distinct seeds; evaluation splits are disjoint from every training split."""
    s = base_seed
    sp: dict[str, Split] = {}
    sp["donor_train"] = opseq.generate(cfg_donor, n_train, s + 1, "donor_train", mode="train")
    sp["recipient_train"] = opseq.generate(cfg_recip, n_train, s + 2, "recipient_train", mode="train")
    used = sp["donor_train"].item_hashes() | sp["recipient_train"].item_hashes()
    sp["donor_dev"] = opseq.generate(cfg_donor, n_dev, s + 3, "donor_dev", mode="train", exclude_hashes=used)
    used |= sp["donor_dev"].item_hashes()
    sp["recipient_dev"] = opseq.generate(cfg_recip, n_dev, s + 4, "recipient_dev", mode="train", exclude_hashes=used)
    used |= sp["recipient_dev"].item_hashes()
    # evaluation strata (shared by all models)
    sp["target_iid"] = opseq.generate(cfg_donor, n_eval, s + 10, "target_iid", mode="train", require_operator=OMEGA, exclude_hashes=used)
    used |= sp["target_iid"].item_hashes()
    sp["control_iid"] = opseq.generate(cfg_donor, n_eval, s + 11, "control_iid", mode="train", forbid_operator=OMEGA, exclude_hashes=used)
    used |= sp["control_iid"].item_hashes()
    sp["comp_heldout"] = opseq.generate(cfg_donor, n_eval, s + 12, "comp_heldout", mode="comp", min_len=2, exclude_hashes=used)
    used |= sp["comp_heldout"].item_hashes()
    sp["len_ood"] = opseq.generate(cfg_donor, max(200, n_eval // 2), s + 13, "len_ood", mode="any", min_len=6, max_len=7, exclude_hashes=used)
    for k in ("donor_train", "donor_dev", "target_iid", "control_iid", "comp_heldout", "len_ood"):
        opseq.validate(sp[k], cfg_donor)
    opseq.validate(sp["recipient_train"], cfg_recip)
    opseq.validate(sp["recipient_dev"], cfg_recip)
    assert_disjoint(sp["donor_train"], sp["donor_dev"], sp["target_iid"], sp["control_iid"], sp["comp_heldout"], sp["len_ood"])
    assert_disjoint(sp["recipient_train"], sp["recipient_dev"], sp["target_iid"], sp["control_iid"], sp["comp_heldout"], sp["len_ood"])
    return sp


def omega_step_accuracy(split: Split, predictions: list[list[str]]) -> np.ndarray:
    """Per item: 1 if the scratchpad token at the first Ω step is correct (items without Ω: NaN)."""
    out = np.full(len(split.examples), np.nan)
    for i, (e, pred) in enumerate(zip(split.examples, predictions)):
        ops = list(e.prompt[2:-1])
        if OMEGA in ops:
            j = ops.index(OMEGA)
            out[i] = float(j < len(pred) and pred[j] == e.answers[j])
    return out


def localization_set(cfg: OpSeqTaskConfig, corrupt_op: str = "zib") -> tuple[Split, Split]:
    """Length-3 programs with Ω at index 1; corrupted twin replaces Ω by ``corrupt_op`` so that r_2 differs."""
    others = [o for o in cfg.operators if o != OMEGA]
    clean, corrupt = [], []
    for a in others:
        for b in others:
            for v0 in range(cfg.modulus):
                ec = opseq.make_example(v0, [a, OMEGA, b], cfg.modulus)
                ex = opseq.make_example(v0, [a, corrupt_op, b], cfg.modulus)
                if ec.answers[1] != ex.answers[1]:
                    clean.append(ec)
                    corrupt.append(ex)
    return (Split("loc_clean", clean, "opseq", opseq.GENERATOR_VERSION, 0, cfg.model_dump()),
            Split("loc_corrupt", corrupt, "opseq", opseq.GENERATOR_VERSION, 0, cfg.model_dump()))


def noop_check(model, enc, device) -> dict:
    ids = torch.from_numpy(enc["ids"][:256]).to(device)
    pos = torch.from_numpy(enc["prompt_len"][:256] - 1).to(device)
    model.hooks.clear()
    with torch.no_grad():
        ref = model(ids, pos)
        again = model(ids, pos)
        cap = Capture()
        for n in model.hook_names():
            model.hooks.add(n, cap)
        with_capture = model(ids, pos)
        model.hooks.add("resid_post.0", Zero())
        with model.hooks.disabled():
            disabled = model(ids, pos)
        enabled = model(ids, pos)
    model.hooks.clear()
    return {"deterministic_repeat": bool(torch.equal(ref, again)), "capture_all_bitwise_identity": bool(torch.equal(ref, with_capture)),
            "disabled_modifying_hook_bitwise_identity": bool(torch.equal(ref, disabled)),
            "enabled_zero_hook_changes_output": bool(not torch.equal(ref, enabled)), "n_hook_points": len(model.hook_names()), "n_items": int(ids.shape[0])}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2])
    ap.add_argument("--steps", type=int, default=3000)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--quick", action="store_true", help="tiny settings for a CPU smoke test")
    ap.add_argument("--dry-run", action="store_true", help="build splits and manifest only")
    ap.add_argument("--n-boot", type=int, default=2000)
    args = ap.parse_args()

    n_train, n_dev, n_eval, steps = (3000, 300, 300, 200) if args.quick else (40_000, 1_000, 1_000, args.steps)
    task_d = OpSeqTaskConfig()
    task_r = OpSeqTaskConfig(excluded_operators=[OMEGA])
    model_cfg = ToyModelConfig(max_seq_len=MAX_SEQ) if not args.quick else ToyModelConfig(d_model=64, n_layers=3, n_heads=2, d_ff=128, max_seq_len=MAX_SEQ)
    cfg = PilotConfig(
        experiment_id="pilot_tier0_opseq" + ("_quick" if args.quick else ""), seeds=args.seeds, target_operator=OMEGA, n_eval=n_eval, n_bootstrap=args.n_boot,
        donor=ModelRoleConfig(role="donor", name="donor", model=model_cfg, task=task_d, train=TrainConfig(steps=steps, device=args.device, n_train=n_train, n_dev=n_dev)),
        recipient=ModelRoleConfig(role="recipient", name="recipient", model=model_cfg, task=task_r, train=TrainConfig(steps=steps, device=args.device, n_train=n_train, n_dev=n_dev)),
    )
    chash = config_hash(cfg)
    run_id = make_run_id(cfg.experiment_id, chash)
    out_dir = ROOT / "experiments" / "summaries" / run_id
    art_dir = ROOT / "artifacts" / "runs" / run_id
    models_dir = ROOT / "artifacts" / "models" / run_id
    fig_dir = ROOT / "artifacts" / "figures" / run_id
    for d in (out_dir, art_dir, fig_dir):
        d.mkdir(parents=True, exist_ok=True)
    man = Manifest(run_id, cfg.experiment_id, cfg.model_dump(), ROOT / "experiments" / "manifests")
    registry.record(run_id, cfg.experiment_id, "planned", config_hash=chash)
    device = resolve_device(args.device)
    t_start = time.time()
    print(f"[{utc_now()}] run_id={run_id} device={device}")

    # ---------------- data ----------------
    splits = build_splits(task_d, task_r, n_train, n_dev, n_eval, base_seed=20261109)
    data_dir = art_dir / "data"
    for name, sp in splits.items():
        sp.save(data_dir / f"{name}.json")
        man.data["data_hashes"][name] = {"hash": sp.hash, "n": len(sp), "seed": sp.seed, "generator": f"{sp.generator}@{sp.generator_version}"}
    man.save()
    vocab = opseq.vocab_for(task_d)
    chance = {k: chance_level(splits[k]) for k in ("target_iid", "control_iid", "comp_heldout", "len_ood")}
    print("splits:", {k: (len(v), v.hash) for k, v in splits.items()})
    if args.dry_run:
        man.set_status("planned", "dry run: splits and manifest only")
        print("dry run complete")
        return

    man.set_status("running")
    registry.record(run_id, cfg.experiment_id, "running")
    results: dict = {"run_id": run_id, "config_hash": chash, "device": str(device), "chance_level": chance, "seeds": {}, "noop": {}, "localization": {}}
    try:
        # ---------------- training + evaluation ----------------
        eval_names = ["target_iid", "control_iid", "comp_heldout", "len_ood"]
        per_item: dict[str, dict[str, np.ndarray]] = {}
        for seed in cfg.seeds:
            results["seeds"][seed] = {}
            for role in (cfg.donor, cfg.recipient):
                tcfg = role.train.model_copy(update={"seed": seed})
                mdir = models_dir / f"{role.name}_s{seed}"
                t0 = time.time()
                if (mdir / "model.pt").exists():
                    model, _ = load_model(mdir, device)
                    hist = json.loads((mdir / "history.json").read_text())
                    print(f"  loaded {role.name} s{seed}")
                else:
                    hist = []
                    model = train(role.model, tcfg, splits[f"{role.name}_train"], splits[f"{role.name}_dev"], vocab, mdir, hist)
                    print(f"  trained {role.name} s{seed} in {time.time()-t0:.0f}s dev_acc={hist[-1]['dev_acc']:.3f}")
                model.eval()
                rec = {"n_params": model.n_params, "train_seconds": time.time() - t0, "final_dev": hist[-1], "eval": {}}
                for en in eval_names:
                    r = evaluate(model, splits[en], vocab, device)
                    ci = cluster_bootstrap_mean(r.correct.astype(float), r.clusters, cfg.n_bootstrap, seed=seed)
                    om = omega_step_accuracy(splits[en], r.predictions)
                    om_ci = cluster_bootstrap_mean(om[~np.isnan(om)], [c for c, o in zip(r.clusters, om) if not np.isnan(o)], cfg.n_bootstrap, seed=seed) if (~np.isnan(om)).any() else None
                    rec["eval"][en] = {**r.to_summary(), "ci95": ci.to_dict(), "omega_step_acc": om_ci.to_dict() if om_ci else None, "chance": chance[en]}
                    per_item[f"{role.name}_s{seed}_{en}"] = {"correct": r.correct, "trace_correct": r.trace_correct, "omega": om}
                    with gzip.open(art_dir / f"items_{role.name}_s{seed}_{en}.json.gz", "wt") as f:
                        json.dump({"split_hash": r.split_hash, "predictions": r.predictions, "answers": r.answers, "clusters": r.clusters}, f)
                enc_dev = encode_split(splits["donor_dev"], vocab, MAX_SEQ)
                results["noop"][f"{role.name}_s{seed}"] = noop_check(model, enc_dev, device)
                results["seeds"][seed][role.name] = rec
            # paired donor-recipient difference on identical items
            results["seeds"][seed]["gap"] = {}
            for en in eval_names:
                d, r_ = per_item[f"donor_s{seed}_{en}"], per_item[f"recipient_s{seed}_{en}"]
                ci = cluster_bootstrap_paired_diff(d["correct"].astype(float), r_["correct"].astype(float), splits[en].examples and [e.cluster for e in splits[en].examples], cfg.n_bootstrap, seed=seed)
                results["seeds"][seed]["gap"][en] = ci.to_dict()
            man.save()

        results["across_seeds"] = {en: {role: seeds_summary([results["seeds"][s][role]["eval"][en]["accuracy"] for s in cfg.seeds]) for role in ("donor", "recipient")} for en in eval_names}
        results["across_seeds_omega_step"] = {role: seeds_summary([results["seeds"][s][role]["eval"]["target_iid"]["omega_step_acc"]["estimate"] for s in cfg.seeds]) for role in ("donor", "recipient")}

        # ---------------- localization (donor models) ----------------
        clean, corrupt = localization_set(task_d)
        enc_c, enc_x = encode_split(clean, vocab, MAX_SEQ), encode_split(corrupt, vocab, MAX_SEQ)
        read_pos = int(enc_c["prompt_len"][0])  # index of r_1; predicts r_2 (the Ω step)
        assert (enc_c["prompt_len"] == read_pos).all() and (enc_x["prompt_len"] == read_pos).all()
        ids_c = torch.from_numpy(enc_c["ids"]).to(device)
        ids_x = torch.from_numpy(enc_x["ids"]).to(device)
        pos = torch.full((len(clean),), read_pos, dtype=torch.long, device=device)
        ans_c = torch.from_numpy(enc_c["answers"][:, 1]).to(device)
        ans_x = torch.from_numpy(enc_x["answers"][:, 1]).to(device)
        positions = list(range(read_pos + 1))
        pos_labels = ["<bos>", "v0", "op1", "Ω", "op3", "=", "r1"]
        loc: dict = {"n_items": len(clean), "read_position": read_pos, "positions": pos_labels, "families": {}, "ablation": {}, "heads": {}, "reproducibility": {}}
        heat_by_seed: dict[str, list[np.ndarray]] = {}
        for seed in cfg.seeds:
            model, _ = load_model(models_dir / f"donor_s{seed}", device)
            with torch.no_grad():
                base_c = (model.answer_logits(ids_c, pos).argmax(-1) == ans_c).float().mean().item()
                base_x = (model.answer_logits(ids_x, pos).argmax(-1) == ans_c).float().mean().item()
            loc.setdefault("baselines", {})[f"s{seed}"] = {"clean_acc_on_clean_answer": base_c, "corrupt_run_acc_on_clean_answer(sham)": base_x}
            for fam in ("resid_post", "attn_out", "mlp_out"):
                pr = patch_sweep(model, ids_c, ids_x, pos, ans_c, ans_x, family=fam, positions=positions)
                pr2 = patch_sweep(model, ids_c, ids_x, pos, ans_c, ans_x, family=fam, positions=positions)
                loc["families"].setdefault(fam, {})[f"s{seed}"] = {"flip_rate": pr.flip_rate.tolist(), "logit_recovery": pr.logit_recovery.tolist()}
                loc["reproducibility"][f"{fam}_s{seed}_repeat_max_abs_diff"] = float(np.abs(pr.flip_rate - pr2.flip_rate).max())
                heat_by_seed.setdefault(fam, []).append(pr.flip_rate)
                abl = ablation_sweep(model, ids_c, pos, ans_c, family=fam, positions=positions, mode="zero")
                loc["ablation"].setdefault(fam, {})[f"s{seed}"] = abl.tolist()
            for p_name, p_idx in (("omega_token", 3), ("read_position", read_pos)):
                loc["heads"].setdefault(p_name, {})[f"s{seed}"] = head_patch_sweep(model, ids_c, ids_x, pos, ans_c, ans_x, p_idx).tolist()
            print(f"  localization donor s{seed} done")
        for fam, mats in heat_by_seed.items():
            if len(mats) > 1:
                flat = np.stack([m.ravel() for m in mats])
                cc = np.corrcoef(flat)
                loc["reproducibility"][f"{fam}_across_seed_pearson_min"] = float(cc[np.triu_indices(len(mats), 1)].min())
            loc["families"][fam]["mean_flip_rate"] = np.mean(mats, axis=0).tolist()
        results["localization"] = loc

        # ---------------- figures ----------------
        make_figures(results, fig_dir, pos_labels, cfg)
        results["elapsed_seconds"] = time.time() - t_start
        (out_dir / "results.json").write_text(json.dumps(results, indent=2, default=float))
        write_summary(results, out_dir, cfg, fig_dir)
        man.data["outputs"] = {"results": str(out_dir.relative_to(ROOT) / "results.json"), "figures": str(fig_dir.relative_to(ROOT)), "models": str(models_dir.relative_to(ROOT)), "items": str(art_dir.relative_to(ROOT))}
        man.set_status("completed")
        registry.record(run_id, cfg.experiment_id, "completed", results=str(out_dir.relative_to(ROOT) / "results.json"))
        print(f"[{utc_now()}] completed in {results['elapsed_seconds']:.0f}s -> {out_dir}")
    except Exception as exc:  # noqa: BLE001
        man.set_status("failed", repr(exc))
        registry.record(run_id, cfg.experiment_id, "failed", error=repr(exc))
        raise


def make_figures(results: dict, fig_dir: Path, pos_labels: list[str], cfg: PilotConfig) -> None:
    import matplotlib.pyplot as plt

    from wp_ai_0001.visualization import DIVERGING_CMAP, ENTITY_COLORS, SEQUENTIAL_CMAP, apply_style, heatmap
    from wp_ai_0001.visualization.style import DOUBLE_WIDTH_IN, TEXT_SECONDARY

    apply_style()
    seeds = {str(k): v for k, v in results["seeds"].items()}  # JSON round-trips keys as strings
    eval_names = ["target_iid", "control_iid", "comp_heldout", "len_ood"]
    # Figure 1: capability gap, per split, donor vs recipient, mean over seeds with per-seed points and CI of seed 0
    fig, ax = plt.subplots(figsize=(DOUBLE_WIDTH_IN * 0.6, 2.2))
    x = np.arange(len(eval_names))
    w = 0.36
    for k, role in enumerate(("donor", "recipient")):
        means = [results["across_seeds"][en][role]["mean"] for en in eval_names]
        ax.bar(x + (k - 0.5) * w, means, width=w - 0.04, color=ENTITY_COLORS[role], label=role, linewidth=0)
        for s in cfg.seeds:
            ax.scatter(x + (k - 0.5) * w, [seeds[str(s)][role]["eval"][en]["accuracy"] for en in eval_names], s=8, color="#ffffff", edgecolor=ENTITY_COLORS[role], linewidth=0.6, zorder=3)
    ax.scatter(x, [results["chance_level"][en] for en in eval_names], marker="_", s=120, color=ENTITY_COLORS["chance"], label="majority-class chance", zorder=4)
    ax.set_xticks(x)
    ax.set_xticklabels(["target (Ω)", "control (no Ω)", "held-out comp.", "length OOD"])
    ax.set_ylabel("final-answer exact match")
    ax.set_ylim(0, 1.02)
    ax.legend(loc="upper right", ncol=1)
    fig.savefig(fig_dir / "fig_gap_tier0.pdf")
    plt.close(fig)
    # Figure 2: localization heatmaps (mean flip rate over seeds) for three families
    fams = ["resid_post", "attn_out", "mlp_out"]
    fig, axes = plt.subplots(1, 3, figsize=(DOUBLE_WIDTH_IN, 2.1), sharey=True)
    for ax, fam in zip(axes, fams):
        m = np.array(results["localization"]["families"][fam]["mean_flip_rate"])
        heatmap(ax, m, [f"L{l}" for l in range(m.shape[0])], pos_labels, cmap=SEQUENTIAL_CMAP, vmin=0, vmax=1, cbar_label="restored-answer rate", title=fam, annotate=True)
        ax.set_xlabel("patched token position")
    axes[0].set_ylabel("layer")
    fig.savefig(fig_dir / "fig_localization_patching_tier0.pdf")
    plt.close(fig)
    # Figure 3: zero-ablation (mean over seeds) accuracy
    fig, axes = plt.subplots(1, 3, figsize=(DOUBLE_WIDTH_IN, 2.1), sharey=True)
    for ax, fam in zip(axes, fams):
        m = np.mean([np.array(v) for v in results["localization"]["ablation"][fam].values()], axis=0)
        heatmap(ax, m, [f"L{l}" for l in range(m.shape[0])], pos_labels, cmap=SEQUENTIAL_CMAP, vmin=0, vmax=1, cbar_label="accuracy after zero-ablation", title=fam, annotate=True)
        ax.set_xlabel("ablated token position")
    axes[0].set_ylabel("layer")
    fig.savefig(fig_dir / "fig_localization_ablation_tier0.pdf")
    plt.close(fig)
    # Figure 4: head-level patching at the Ω token and at the read position
    fig, axes = plt.subplots(1, 2, figsize=(DOUBLE_WIDTH_IN * 0.75, 2.1), sharey=True)
    for ax, key in zip(axes, ("omega_token", "read_position")):
        m = np.mean([np.array(v) for v in results["localization"]["heads"][key].values()], axis=0)
        heatmap(ax, m, [f"L{l}" for l in range(m.shape[0])], [f"h{h}" for h in range(m.shape[1])], cmap=SEQUENTIAL_CMAP, vmin=0, vmax=1, cbar_label="restored-answer rate", title=f"head patch @ {key.replace('_', ' ')}", annotate=True)
    axes[0].set_ylabel("layer")
    fig.savefig(fig_dir / "fig_localization_heads_tier0.pdf")
    plt.close(fig)
    fig.text(0, 0, "", color=TEXT_SECONDARY)


def write_summary(results: dict, out_dir: Path, cfg: PilotConfig, fig_dir: Path) -> None:
    seeds = {str(k): v for k, v in results["seeds"].items()}
    L = ["# Pilot summary — " + results["run_id"], "", f"Generated {utc_now()} on device `{results['device']}`; config hash `{results['config_hash']}`; seeds {cfg.seeds}.", "",
         "## Gate B — donor/recipient gap (final-answer exact match, mean over seeds; CI = cluster bootstrap on programs)", "",
         "| split | chance | donor mean (sd) | recipient mean (sd) | paired gap seed 0 [95% CI] |", "|---|---|---|---|---|"]
    for en in ("target_iid", "control_iid", "comp_heldout", "len_ood"):
        d, r = results["across_seeds"][en]["donor"], results["across_seeds"][en]["recipient"]
        g = seeds[str(cfg.seeds[0])]["gap"][en]
        L.append(f"| {en} | {results['chance_level'][en]:.3f} | {d['mean']:.3f} ({(d['sd'] or 0):.3f}) | {r['mean']:.3f} ({(r['sd'] or 0):.3f}) | {g['estimate']:+.3f} [{g['low']:+.3f}, {g['high']:+.3f}] |")
    om = results["across_seeds_omega_step"]
    L += ["", f"Ω-step accuracy on target_iid (token emitted at the first Ω step): donor {om['donor']['mean']:.3f}, recipient {om['recipient']['mean']:.3f}.", "",
          "## Gate C — hook no-op checks (bitwise equality of logits on 256 dev items)", "", "| model | repeat | capture-all | disabled hook | enabled zero hook changes output |", "|---|---|---|---|---|"]
    for k, v in results["noop"].items():
        L.append(f"| {k} | {v['deterministic_repeat']} | {v['capture_all_bitwise_identity']} | {v['disabled_modifying_hook_bitwise_identity']} | {v['enabled_zero_hook_changes_output']} |")
    loc = results["localization"]
    L += ["", f"## RQ1 — within-donor localization of Ω (n = {loc['n_items']} clean/corrupt pairs, read position = r1, predicting r2)", "",
          "Baselines: " + "; ".join(f"s{k[1:]}: clean {v['clean_acc_on_clean_answer']:.3f}, sham (corrupt run scored on clean answer) {v['corrupt_run_acc_on_clean_answer(sham)']:.3f}" for k, v in loc["baselines"].items()), "",
          "Mean restored-answer rate after patching the clean activation into the corrupted run (rows = layers, columns = positions " + ", ".join(loc["positions"]) + "):", ""]
    for fam in ("resid_post", "attn_out", "mlp_out"):
        m = np.array(loc["families"][fam]["mean_flip_rate"])
        L.append(f"- `{fam}`: " + " | ".join(f"L{i}: " + " ".join(f"{v:.2f}" for v in row) for i, row in enumerate(m)))
    L += ["", "Reproducibility: " + "; ".join(f"{k} = {v:.4f}" for k, v in loc["reproducibility"].items()), "",
          "Figures: " + ", ".join(f"`{p.name}`" for p in sorted(fig_dir.glob("*.pdf"))), "",
          "## Interpretation (observed vs. speculative)", "",
          "Observed: the numbers above are produced by the artifacts listed in the manifest. Any statement about *transfer* is out of scope for this run: no transplant was performed.",
          "Speculative: whether the localized positions/layers constitute a transferable 'organ' is the object of the next experiment (activation graft with sham/random controls)."]
    (out_dir / "SUMMARY.md").write_text("\n".join(L) + "\n")


if __name__ == "__main__":
    main()