#!/usr/bin/env python
"""Lab 6: the weight edit is steering in disguise.

    . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
    CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py                 # parts 1 and 2, about a minute
    CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py --parts 3       # part 3, about 2 minutes, needs a finished base run
    CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py --parts 4       # part 4, about a minute, needs part 3's results
    python labs/lab6_steering.py --plot-only                            # redraw part 3's figures, no GPU

Experiments (each writes to its own folder, results/lab6/try-*, and leaves the main results alone):
    CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py --parts 1 --alpha-scale 2      # the same check at alpha 12
    CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py --parts 3 --max-samples 10     # part 3 on 10 loops
    python labs/lab6_steering.py --plot-only --smooth 16                               # less smoothing, no GPU
    TRITON_F32_DEFAULT=ieee TRITON_CACHE_DIR=/nvme/wsp-lab/tmp/triton-ieee \
      CUDA_VISIBLE_DEVICES=7 python labs/lab6_steering.py --parts 1 --try-name ieee    # DeltaNet kernels in full fp32

Part 1. Same math, two ways. In float32, compare the logits of
  (a) the model with down_proj edited in memory:  W' = W - alpha u u^T W[:, cols]
  (b) the unedited model plus a hook that subtracts alpha * s(t) * u from down_proj's output,
      where s(t) = u . (W[:, cols] a[cols](t)) is one number per token.
  They agree to float32 rounding. Then compare the released bf16 file with (b) run in bf16.
Part 2. s(t) on the agent's 8 texts: how hard the 64 neurons push along the loop direction.
Part 3. Real loops. Teacher-force looping answers from the base model and follow, token by token:
  s(t), how much the edit changes the next-token prediction (KL), and the probability of the token that
  actually came next, both straight from the model (temperature 1) and as the chance that the paper's
  sampler settings pick it (wsp/protocol.py, pick_prob). Inside the loop it also works out the chance
  of writing one more full cycle unchanged.
Part 4. Answers that went well. The same s(t) and KL on every answer in the run that finished on its own with
  the right number, side by side with part 3's loops (from part 3's folder for the same flags, or the main results
  if that folder has none). If the edit only reacted to loops, both would stay near zero here.
"""
import argparse
import contextlib
import json
import os
import sys

import numpy as np
import torch

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp import hf, loops, protocol  # noqa: E402
from wsp.edits import projection_edit  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--parts", default="1,2")
ap.add_argument("--run", default=os.path.join(hf.LAB, "runs", "build", "base"))
ap.add_argument("--data", default="build128")
ap.add_argument("--max-samples", type=int, default=48, help="part 3: how many looping answers to teacher-force")
ap.add_argument("--alpha-scale", type=float, default=1.0,
                help="part 1: multiply alpha by this in both the hook and the in-memory edit (2 means alpha 12)")
ap.add_argument("--plot-only", action="store_true", help="redraw the part 3 figures from results/lab6/*.json")
ap.add_argument("--smooth", type=int, default=64, help="moving-average window for the one-loop figure, in tokens")
ap.add_argument("--try-name", help="write this run's results to results/lab6/try-<name>/ instead of the main files")
a = ap.parse_args()
parts = {int(x) for x in a.parts.split(",")}

OUT = os.path.join(ROOT, "results", "lab6")
FIG = os.path.join(ROOT, "results", "figures")
# Experiments write to results/lab6/try-*/, so the main results and figures are never overwritten.
# Changing --alpha-scale, --max-samples or --smooth picks a folder name for you.
name1 = a.try_name or (f"alpha-scale-{a.alpha_scale:g}" if a.alpha_scale != 1.0 else None)
name3 = a.try_name or (f"max-samples-{a.max_samples}" if a.max_samples != 48 else None)
OUT1 = os.path.join(OUT, f"try-{name1}") if name1 else OUT       # part 1
OUT2 = os.path.join(OUT, f"try-{a.try_name}") if a.try_name else OUT  # part 2
OUT3 = os.path.join(OUT, f"try-{name3}") if name3 else OUT       # part 3, and where --plot-only reads it
FIG3 = OUT3 if name3 else FIG
if a.smooth != 64:
    FIG3 = os.path.join(OUT, f"try-{name3 + '-' if name3 else ''}smooth-{a.smooth}")
for d in (OUT, FIG, OUT1, OUT2, OUT3, FIG3):
    os.makedirs(d, exist_ok=True)
LAYERS = (15, 19)
rel = hf.released_edit()
lab3 = json.load(open(os.path.join(ROOT, "results", "lab3", "rebuild.json")))
# the released direction u, signed so that +u points toward looping (Lab 3 measured which sign that is)
u_loop = {}
for L in LAYERS:
    u_rel = rel[L]["u"] / rel[L]["u"].norm()
    u_loop[L] = u_rel if float(u_rel @ torch.tensor(lab3["layers"][str(L)]["u_loop_sign"])) > 0 else -u_rel
pairs = json.load(open(os.path.join(ROOT, "labs", "data", "astra_pairs.json")))


@contextlib.contextmanager
def steered(m, alpha_scale=1.0, record=None):
    """The released edit as hooks at both layers. alpha_scale=0 only records s(t)."""
    with contextlib.ExitStack() as st:
        for L in LAYERS:
            st.enter_context(hf.steer_down(m, L, u_loop[L], alpha_scale * rel[L]["alpha"], rel[L]["cols"],
                                           record=None if record is None else record[L]))
        yield


def kl_rows(lp, lq):  # KL(p || q) per row, from log-probs
    return (lp.exp() * (lp - lq)).sum(-1)


def plot_part3():
    """Draws part 3's two figures from real_loops.json and trace_one_loop.json in OUT3, into FIG3."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    real = json.load(open(os.path.join(OUT3, "real_loops.json")))
    trace = json.load(open(os.path.join(OUT3, "trace_one_loop.json")))
    stats, names = real["samples"], list(real["mean"])

    def smooth(v, w=a.smooth):  # moving average that stays honest at both ends of the answer
        v, k = np.asarray(v, dtype=float), np.ones(w)
        return np.convolve(v, k, mode="same") / np.convolve(np.ones_like(v), k, mode="same")

    tm = trace["meta"]
    xs = np.arange(len(trace["s15"]))  # s(t): one value per answer token
    xk = np.arange(len(trace["kl"]))   # KL and p(next): one fewer, since position t predicts token t+1
    has_pick = "p_pick_edit" in trace  # traces written before the sampler measurement have 3 panels
    fig, ax = plt.subplots(4 if has_pick else 3, 1, figsize=(13, 11.5 if has_pick else 9), sharex=True)
    ax[0].plot(xs, smooth(trace["s15"]), label="layer 15", color="#0969da")
    ax[0].plot(xs, smooth(trace["s19"]), label="layer 19", color="#8250df")
    ax[0].set_ylabel("s(t): push along\nthe loop direction")
    ax[0].legend(loc="upper left")
    ax[1].plot(xk, smooth(trace["kl"]), color="#cf222e")
    ax[1].set_ylabel("KL(base || edited)\nnext-token change")
    ax[2].plot(xk, smooth(trace["p_base"]), label="base model", color="#57606a")
    ax[2].plot(xk, smooth(trace["p_edit"]), label="edited model", color="#1a7f37")
    ax[2].set_ylabel("probability of the\ntoken that came next\n(temperature 1)")
    ax[2].legend(loc="lower left")
    if has_pick:
        ax[3].plot(xk, smooth(trace["p_pick_base"]), label="base model", color="#57606a")
        ax[3].plot(xk, smooth(trace["p_pick_edit"]), label="edited model", color="#1a7f37")
        ax[3].set_ylabel("chance the paper's\nsampler picks it")
        ax[3].legend(loc="lower left")
    ax[-1].set_xlabel("token position in the answer")
    for axi in ax:
        axi.axvline(tm["loop_starts_at"], color="black", ls="--", lw=1)
    ax[0].set_title(f"One real loop from the base model ({tm['id']})\nthe loop starts at token "
                    f"{tm['loop_starts_at']} and repeats every {tm['period']} tokens; lines are "
                    f"{a.smooth}-token moving averages")
    fig.tight_layout()
    fig.savefig(os.path.join(FIG3, "lab6_real_loop.png"), dpi=130)

    has_cycle = "cycle_pick_edit" in stats[0][names[2]]
    fig, ax = plt.subplots(1, 4 if has_cycle else 3, figsize=(18.5 if has_cycle else 14, 3.8))
    short = ["early", "late", "inside loop"]
    for i, (k, lab) in enumerate([("s15", "s(t) at layer 15"), ("kl", "KL(base || edited)"),
                                  ("p_next_base", "p(next token), temperature 1")]):
        vals = np.array([[s[z][k] for z in names] for s in stats])
        ax[i].boxplot([vals[:, j] for j in range(3)], tick_labels=short, showfliers=False)
        if k == "p_next_base":
            ve = np.array([[s[z]["p_next_edit"] for z in names] for s in stats])
            ax[i].plot(range(1, 4), ve.mean(0), "o-", color="#1a7f37", label="edited model (mean)")
            ax[i].plot(range(1, 4), vals.mean(0), "o-", color="#57606a", label="base model (mean)")
            ax[i].legend()
        ax[i].set_title(lab)
    if has_cycle:
        per = [s["period"] for s in stats]
        ax[3].scatter(per, [s[names[2]]["cycle_pick_base"] for s in stats], color="#57606a", s=18, label="base model")
        ax[3].scatter(per, [s[names[2]]["cycle_pick_edit"] for s in stats], color="#1a7f37", s=18,
                      label="edited model")
        ax[3].set_xscale("log")
        ax[3].set_ylim(-0.03, 1.03)
        ax[3].set_xlabel("loop length (tokens per cycle)")
        ax[3].set_title("chance of one more identical cycle\n(paper's sampler settings)")
        ax[3].legend(loc="center left")
    fig.suptitle(f"{len(stats)} real loops from the base model, three stretches of each answer")
    fig.tight_layout()
    fig.savefig(os.path.join(FIG3, "lab6_loops_summary.png"), dpi=130)
    print("wrote lab6_real_loop.png and lab6_loops_summary.png in", FIG3)


if a.plot_only:
    plot_part3()
    sys.exit(0)


# ------------------------------------------------------------------ part 1
if 1 in parts:
    res = {}
    m, tok = hf.load("base", dtype=torch.float32)
    ids = [hf.encode(tok, t) for t in pairs["texts"]]
    with torch.no_grad():
        base = [m(x).logits[0].float() for x in ids]
        with steered(m, a.alpha_scale):
            hook = [m(x).logits[0].float() for x in ids]
        saved = {}
        for L in LAYERS:
            W = hf.down_proj(m, L).weight
            saved[L] = W.data.clone()
            W.data = projection_edit(W.data, u_loop[L].to(W.device), a.alpha_scale * rel[L]["alpha"], rel[L]["cols"])
        edit = [m(x).logits[0].float() for x in ids]
        for L in LAYERS:
            hf.down_proj(m, L).weight.data = saved[L]
    res["fp32_hook_vs_edit_max_abs_logit_diff"] = max(float((h - e).abs().max()) for h, e in zip(hook, edit))
    res["fp32_base_vs_edit_max_abs_logit_diff"] = max(float((b - e).abs().max()) for b, e in zip(base, edit))
    res["fp32_hook_vs_edit_mean_kl"] = float(torch.cat([kl_rows(e.log_softmax(-1), h.log_softmax(-1))
                                                        for h, e in zip(hook, edit)]).mean())
    res["fp32_base_vs_edit_mean_kl"] = float(torch.cat([kl_rows(b.log_softmax(-1), e.log_softmax(-1))
                                                        for b, e in zip(base, edit)]).mean())
    del m
    torch.cuda.empty_cache()

if 1 in parts and a.alpha_scale == 1.0:  # the released file has alpha 6, so this check only makes sense at scale 1
    m16, tok = hf.load("base")
    ids = [hf.encode(tok, t) for t in pairs["texts"]]
    p16, _ = hf.load("qwen35-p50c")
    with torch.no_grad():
        b16 = [m16(x).logits[0].float() for x in ids]
        with steered(m16):
            h16 = [m16(x).logits[0].float() for x in ids]
        r16 = [p16(x).logits[0].float() for x in ids]
    res["bf16_hook_vs_released_file_max_abs_logit_diff"] = max(float((h - r).abs().max()) for h, r in zip(h16, r16))
    res["bf16_hook_vs_released_file_mean_kl"] = float(torch.cat([kl_rows(r.log_softmax(-1), h.log_softmax(-1))
                                                                 for h, r in zip(h16, r16)]).mean())
    res["bf16_base_vs_released_file_mean_kl"] = float(torch.cat([kl_rows(b.log_softmax(-1), r.log_softmax(-1))
                                                                 for b, r in zip(b16, r16)]).mean())
    res["bf16_same_top_token_hook_vs_file"] = float(torch.cat([(h.argmax(-1) == r.argmax(-1)).float()
                                                               for h, r in zip(h16, r16)]).mean())
    # one layer in float64, where no GPU kernel rounds anything: the identity holds to machine precision
    with torch.no_grad(), hf.record_mlp_acts(m16, [15]) as acts:
        m16(ids[1])
    A64 = acts[15][0].double()                               # [tokens, 3584] neuron activations
    W64 = hf.down_proj(m16, 15).weight.detach().double().cpu()
    u64 = u_loop[15].double() / u_loop[15].double().norm()
    al, cols = rel[15]["alpha"], rel[15]["cols"]
    mask = torch.zeros(W64.shape[1], dtype=torch.float64)
    mask[cols] = 1
    W_edit = W64 - al * torch.outer(u64, mask * (u64 @ W64))  # the weight edit
    s = A64[:, cols] @ (W64[:, cols].T @ u64)                 # one number per token
    y_edit = A64 @ W_edit.T
    y_hook = A64 @ W64.T - al * s[:, None] * u64[None, :]
    res["fp64_one_layer_max_abs_diff"] = float((y_edit - y_hook).abs().max())
    res["fp64_one_layer_output_scale"] = float(y_edit.abs().max())
    del p16, m16
    torch.cuda.empty_cache()
if 1 in parts:
    json.dump(res, open(os.path.join(OUT1, "equivalence.json"), "w"), indent=1)
    for k, v in res.items():
        print(f"{k:52s} {v:.3g}")
    if a.alpha_scale != 1.0:
        print("skipped the bf16 and fp64 checks against the released file, which was saved with alpha 6; wrote", OUT1)

# ------------------------------------------------------------------ part 2
if 2 in parts:
    m, tok = hf.load("base")
    rec = {L: [] for L in LAYERS}
    with torch.no_grad(), steered(m, 0.0, rec):
        for t in pairs["texts"]:
            m(hf.encode(tok, t))
    out = {}
    for L in LAYERS:
        tails = [float(s[-36:].mean()) for s in rec[L]]
        out[str(L)] = {"productive_last36": tails[0::2], "looping_last36": tails[1::2],
                       "per_token": [[round(float(x), 4) for x in s] for s in rec[L]]}
        print(f"layer {L}: s over the last 36 tokens, productive {np.round(tails[0::2], 2).tolist()} "
              f"vs looping {np.round(tails[1::2], 2).tolist()}")
    json.dump(out, open(os.path.join(OUT2, "s_on_pairs.json"), "w"), indent=1)
    del m
    torch.cuda.empty_cache()

# ------------------------------------------------------------------ part 3
if 3 in parts:
    m, tok = hf.load("base")
    rows = loops.read_rows(hf.LAB, a.data)
    found = loops.looping_samples(tok, a.run, rows)
    print(f"{len(found)} samples in {a.run} end in a verbatim loop")
    # keep loops with a real run-up before them, spread over the run
    found = [f for f in found if f["onset"] - f["n_prompt"] >= 1024]
    rng = np.random.default_rng(0)
    pick = sorted(rng.choice(len(found), size=min(a.max_samples, len(found)), replace=False).tolist())
    lm_head = m.lm_head
    stats, best, trace = [], None, None
    for k in pick:
        f = found[k]
        x = torch.tensor([f["ids"]], device="cuda")
        rec_b = {L: [] for L in LAYERS}
        rec_e = {L: [] for L in LAYERS}
        n0, on = f["n_prompt"], f["onset"]
        with torch.no_grad():
            with steered(m, 0.0, rec_b):
                hb = m.model(x, use_cache=False).last_hidden_state[0]
            with steered(m, 1.0, rec_e):
                he = m.model(x, use_cache=False).last_hidden_state[0]
            T = x.shape[1]
            # first[v] = where token v first shows up in the answer. The sampler's presence penalty
            # hits token v at position t once first[v] <= t.
            first = torch.full((lm_head.weight.shape[0],), T, device=x.device)
            first.scatter_reduce_(0, x[0, n0:], torch.arange(n0, T, device=x.device), reduce="amin")
            kl, pb, pe, qb, qe = [], [], [], [], []
            nxt = x[0, 1:]
            for c in range(0, T - 1, 1024):
                zb = lm_head(hb[c:min(c + 1024, T - 1)]).float()
                ze = lm_head(he[c:min(c + 1024, T - 1)]).float()
                lb, le = zb.log_softmax(-1), ze.log_softmax(-1)
                kl.append(kl_rows(lb, le).cpu())
                ix = nxt[c:c + lb.shape[0]]
                pb.append(lb.gather(1, ix[:, None])[:, 0].exp().cpu())
                pe.append(le.gather(1, ix[:, None])[:, 0].exp().cpu())
                present = first[None, :] <= torch.arange(c, c + lb.shape[0], device=x.device)[:, None]
                qb.append(protocol.pick_prob(zb, present, ix).cpu())
                qe.append(protocol.pick_prob(ze, present, ix).cpu())
                del zb, ze, lb, le, present
        kl, pb, pe = torch.cat(kl).numpy(), torch.cat(pb).numpy(), torch.cat(pe).numpy()
        qb, qe = torch.cat(qb).numpy(), torch.cat(qe).numpy()
        sb = {L: rec_b[L][0].numpy() for L in LAYERS}
        # positions t predict token t+1; the three stretches of the answer
        zones = {"early (first half before the loop)": slice(n0, n0 + (on - n0) // 2),
                 "late (last 512 tokens before the loop)": slice(on - 512, on),
                 "inside the loop": slice(on + f["period"], T - 1)}
        st = {"id": f["id"], "j": f["j"], "answer_tokens": T - n0, "loop_starts_at": on - n0, "period": f["period"]}
        for z, sl in zones.items():
            st[z] = {"s15": float(sb[15][sl].mean()), "s19": float(sb[19][sl].mean()), "kl": float(kl[sl].mean()),
                     "p_next_base": float(pb[sl].mean()), "p_next_edit": float(pe[sl].mean()),
                     "p_pick_base": float(qb[sl].mean()), "p_pick_edit": float(qe[sl].mean())}
        # Inside the loop the text repeats every `period` tokens. The chance of writing one more full
        # cycle unchanged is the product of `period` token probabilities; average it over every start.
        ins = zones["inside the loop"]
        for n, v in (("base", pb), ("edit", pe), ("pick_base", qb), ("pick_edit", qe)):
            lp = np.concatenate([[0.0], np.cumsum(np.log(np.maximum(v[ins].astype(np.float64), 1e-300)))])
            st["inside the loop"][f"cycle_{n}"] = float(np.exp(lp[f["period"]:] - lp[:-f["period"]]).mean())
        stats.append(st)
        off_middle = abs(st["loop_starts_at"] / st["answer_tokens"] - 0.5)  # plot a loop that starts mid-answer
        if best is None or off_middle < best:
            best = off_middle
            trace = {"meta": st, "s15": sb[15][n0:].tolist(), "s19": sb[19][n0:].tolist(), "kl": kl[n0:].tolist(),
                     "p_base": pb[n0:].tolist(), "p_edit": pe[n0:].tolist(),
                     "p_pick_base": qb[n0:].tolist(), "p_pick_edit": qe[n0:].tolist(), "text": f["text"]}
        print(f"{f['id']} #{f['j']}: loop at token {on - n0} of {T - n0}, period {f['period']}; "
              f"KL early/late/inside {[round(st[z]['kl'], 3) for z in zones]}")
    agg = {z: {k: float(np.mean([s[z][k] for s in stats])) for k in stats[0][z]} for z in zones}
    json.dump({"n": len(stats), "mean": agg, "samples": stats}, open(os.path.join(OUT3, "real_loops.json"), "w"),
              indent=1)
    json.dump(trace, open(os.path.join(OUT3, "trace_one_loop.json"), "w"))
    print(json.dumps(agg, indent=1))

    plot_part3()
    print("wrote", OUT3)
    del m
    torch.cuda.empty_cache()

# ------------------------------------------------------------------ part 4
if 4 in parts:
    m, tok = hf.load("base")
    rows = loops.read_rows(hf.LAB, a.data)
    end_think = tok.convert_tokens_to_ids("</think>")
    done = []
    for rec, j, s, text in loops.iter_samples(a.run):
        if s.get("cut") or not s.get("correct"):
            continue
        prompt = tok(protocol.render(tok, rows[rec["id"]], True), add_special_tokens=False).input_ids
        ans = tok(text, add_special_tokens=False).input_ids
        if end_think in ans:
            done.append({"id": rec["id"], "j": j, "ids": prompt + ans, "n_prompt": len(prompt),
                         "think_end": len(prompt) + ans.index(end_think)})
    print(f"{len(done)} answers in {a.run} finished on their own with the right number")
    if not done:
        sys.exit("part 4 needs at least one answer that finished on its own with the right number; this run has none")
    box64 = np.ones(64) / 64
    stats = []
    for f in done:
        x = torch.tensor([f["ids"]], device="cuda")
        rec_b = {L: [] for L in LAYERS}
        with torch.no_grad():
            with steered(m, 0.0, rec_b):
                hb = m.model(x, use_cache=False).last_hidden_state[0]
            with steered(m, 1.0):
                he = m.model(x, use_cache=False).last_hidden_state[0]
            T = x.shape[1]
            kl = []
            for c in range(0, T - 1, 1024):
                lb = m.lm_head(hb[c:min(c + 1024, T - 1)]).float().log_softmax(-1)
                le = m.lm_head(he[c:min(c + 1024, T - 1)]).float().log_softmax(-1)
                kl.append(kl_rows(lb, le).cpu())
                del lb, le
        kl = torch.cat(kl).numpy()
        sb = {L: rec_b[L][0].numpy() for L in LAYERS}
        n0, te = f["n_prompt"], f["think_end"]
        zones = {"early (first half of the thinking)": slice(n0, n0 + (te - n0) // 2),
                 "late (last 512 tokens of the thinking)": slice(max(n0, te - 512), te),
                 "the written answer after </think>": slice(te, T - 1)}
        st = {"id": f["id"], "j": f["j"], "answer_tokens": T - n0, "thinking_tokens": te - n0}
        for z, sl in zones.items():
            st[z] = {"s15": float(sb[15][sl].mean()), "s19": float(sb[19][sl].mean()), "kl": float(kl[sl].mean())}
        for L in LAYERS:  # the hardest 64-token stretch of the thinking
            st[f"max_s{L}_64_tokens"] = float(np.convolve(sb[L][n0:te], box64, mode="valid").max())
        stats.append(st)
        print(f"{f['id']} #{f['j']}: {te - n0} thinking tokens; s19 early/late/answer "
              f"{[round(st[z]['s19'], 3) for z in zones]}; KL {[round(st[z]['kl'], 4) for z in zones]}")
    agg = {z: {k: float(np.mean([s[z][k] for s in stats])) for k in stats[0][z]} for z in zones}
    out = {"n": len(stats), "mean": agg, "samples": stats}
    # Part 4 writes where part 3 does, and compares with that experiment's loops if part 3 wrote any there.
    OUT4, FIG4 = OUT3, (OUT3 if name3 else FIG)
    loops_file = next((p for p in (os.path.join(OUT3, "real_loops.json"), os.path.join(OUT, "real_loops.json"))
                       if os.path.exists(p)), None)
    if loops_file:  # line the answers up against part 3's loops
        print("comparing with the loops in", os.path.relpath(loops_file, ROOT))
        real = json.load(open(loops_file))
        inside = list(real["mean"])[2]
        out["loops_for_comparison"] = {"n": real["n"], "mean": real["mean"]}
        for L in LAYERS:
            med = float(np.median([s[inside][f"s{L}"] for s in real["samples"]]))
            hit = sum(s[f"max_s{L}_64_tokens"] > med for s in stats)
            out[f"layer {L}"] = {"median loop, mean s inside the loop": med,
                                 "answers whose hardest 64-token stretch beats it": hit}
            print(f"layer {L}: {hit} of {len(stats)} answers have a 64-token stretch with s above the median "
                  f"loop's in-loop mean ({med:.3f})")
    os.makedirs(OUT4, exist_ok=True)
    json.dump(out, open(os.path.join(OUT4, "finished_answers.json"), "w"), indent=1)
    print(json.dumps(agg, indent=1))
    if "loops_for_comparison" in out:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        fz, lz = list(agg), list(real["mean"])
        fig, ax = plt.subplots(1, 3, figsize=(16, 4.2))
        for i, (k, lab) in enumerate((("s15", "s(t) at layer 15"), ("s19", "s(t) at layer 19"),
                                      ("kl", "KL(base || edited)"))):
            data = [[s[z][k] for s in stats] for z in fz] + [[s[z][k] for s in real["samples"]] for z in lz]
            bp = ax[i].boxplot(data, tick_labels=["right answer\nearly", "right answer\nlate",
                                                  "right answer\nwritten", "loop\nearly", "loop\nlate",
                                                  "loop\ninside"], showfliers=False, patch_artist=True)
            for b, col in zip(bp["boxes"], ["#1a7f37"] * 3 + ["#cf222e"] * 3):
                b.set_facecolor(col)
                b.set_alpha(0.35)
            ax[i].set_title(lab)
            ax[i].tick_params(axis="x", labelsize=8)
        fig.suptitle(f"The edit's trigger on {len(stats)} answers that finished with the right number (green) "
                     f"and {real['n']} answers that ended in a loop (red), mean per stretch of each answer")
        fig.tight_layout()
        fig.savefig(os.path.join(FIG4, "lab6_finished_vs_loops.png"), dpi=130)
        print("wrote lab6_finished_vs_loops.png in", FIG4)
    print("wrote", OUT4)