#!/usr/bin/env python
"""Lab 8 (capstone): build your own edit from the model's own loops, then test it on problems it never saw.

    . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
    CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py                          # about 5 minutes on one GPU
    CUDA_VISIBLE_DEVICES="" python scripts/build_models.py --variants configs/capstone.json   # CPU, about a minute
    python scripts/run_jobs.py plans/extra.json --gpus 6,7 --only 'dev/cap*'    # about 30 minutes, one run per GPU
    python scripts/summarize.py /nvme/wsp-lab/runs                               # writes /nvme/wsp-lab/summary.md

Experiments (each writes to its own folder, results/lab8/try-*, and leaves the main results, the figures and
configs/capstone.json alone; the script prints the commands that build and test the try's two edits):
    CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --max-samples 16    # u from only 16 loops
    CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --min-before 2048   # only loops that start late
    CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --alpha 8           # same u and neurons, stronger edit
    CUDA_VISIBLE_DEVICES=7 python labs/lab8_capstone.py --split-by problem  # stability checks with no problem in both halves

The paper's agent built its edit from 8 short texts that it wrote and labeled itself. Here the data is the
unedited model's own answers to 128 practice problems (runs/build/base; none of them is used for testing).
For every answer that ends in a verbatim loop, take three 36-token windows:
  in the loop    starts two periods after the loop starts, when the model is surely repeating
  just before    ends 32 tokens before the loop starts                          (recipe A)
  well before    the middle of the text before the loop                         (recipe B)
u = mean over answers of (stream in the loop - stream before), after blocks 15 and 19, scaled to length 1.
Neuron score = (how much neuron j writes along u) x (how much more neuron j fires in the loop; last 24 tokens of
each window). The 64 highest-scoring neurons get the edit with alpha 6: the agent's recipe, your data.

Before you trust any direction you extract, measure:
  agreement     cos with the released u, and how many of the released 64 neurons you picked too
  stability     split the answers in half at random: do the two halves give the same u?
  data needed   how many answers you need before u stops changing (the agent used 4 pairs)
  transfer      does your u, built from the model's loops, separate the agent's 8 hand-written texts?
and one question about the loops themselves: lined up at the token where each loop starts, when does the stream
move along the released u, and when do the released 64 neurons start pushing (s(t), the edit's trigger)?
"""
import argparse
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  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--run", default=os.path.join(hf.LAB, "runs", "build", "base"))
ap.add_argument("--data", default="build128")
ap.add_argument("--min-before", type=int, default=512, help="skip loops that start fewer answer tokens in")
ap.add_argument("--max-samples", type=int, default=0, help="use at most this many loops (0 = all)")
ap.add_argument("--alpha", type=float, default=6)
ap.add_argument("--split-by", choices=("answer", "problem"), default="answer",
                help="stability checks: split the loops at random by answer, or keep both answers to a problem "
                     "in the same half (the edits themselves do not change)")
ap.add_argument("--try-name", help="write to results/lab8/try-<name>/ (a name is picked for you when a flag changes)")
a = ap.parse_args()

# Experiments write to results/lab8/try-*/ and name their edits capA-<try>/capB-<try>, so the main results, the
# figures, configs/capstone.json and the models/qwen35-cap*-a6 folders are never overwritten.
changed = [f"{k.replace('_', '-')}-{getattr(a, k):g}" for k in ("min_before", "max_samples", "alpha")
           if getattr(a, k) != ap.get_default(k)]
if a.run != ap.get_default("run") or a.data != ap.get_default("data"):
    changed.append(f"{os.path.basename(os.path.normpath(a.run))}-{a.data}")
# --split-by on its own changes only the stability checks, so those edits are the main run's. Every other experiment,
# including a --try-name run after you change the code (GAP, say), writes edits you can build.
new_edits = bool(changed) or a.split_by == "answer"
if a.split_by != "answer":
    changed.append(f"split-by-{a.split_by}")
TRY = a.try_name or "-".join(changed) or None
OUT = os.path.join(ROOT, "results", "lab8", f"try-{TRY}") if TRY else os.path.join(ROOT, "results", "lab8")
FIG = OUT if TRY else os.path.join(ROOT, "results", "figures")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIG, exist_ok=True)
LAYERS = (15, 19)
WIN, AWIN, GAP = 36, 24, 32          # window sizes as in the agent's recipe; GAP = distance of window A from onset
PRE, POST = 2048, 256                # the onset profile covers this many tokens before and after each loop starts
RECIPES = {"A": "just before the loop", "B": "well before the loop"}
rng = np.random.default_rng(0)


def unit(v):
    return v / v.norm()


rel = hf.released_edit()
lab3 = json.load(open(os.path.join(ROOT, "results", "lab3", "rebuild.json")))
# the released u, signed to point toward looping (Lab 3 found which way it points)
u_rel = {L: unit(rel[L]["u"].float()) * float(np.sign(lab3["layers"][str(L)]["cos_with_released_u"]))
         for L in LAYERS}

# 1. find the loops
rows = loops.read_rows(hf.LAB, a.data)
m, tok = hf.load("base")
n_total = n_cut = 0
for rec, j, s, text in loops.iter_samples(a.run):
    n_total += 1
    n_cut += bool(s.get("cut"))
found = loops.looping_samples(tok, a.run, rows)
print(f"{n_total} answers in {a.run}: {n_cut} hit the token cap, {len(found)} of those end in a verbatim loop")


def windows(f):
    n0, on, p, T = f["n_prompt"], f["onset"], f["period"], len(f["ids"])
    r0 = min(on + 2 * p, T - 1 - WIN)   # the last token is not part of the checked loop
    if on - n0 < a.min_before or r0 < on:
        return None
    b0 = n0 + (on - n0) // 2 - WIN // 2
    return {"loop": (r0, r0 + WIN), "A": (on - GAP - WIN, on - GAP), "B": (b0, b0 + WIN)}


used = [(f, w) for f in found if (w := windows(f))]
if a.max_samples:
    used = used[:a.max_samples]
N = len(used)
print(f"using {N} loops (loop starts at least {a.min_before} answer tokens in)")
assert N >= 8, "too few loops to build a direction"
Wd = {L: hf.down_proj(m, L).weight.detach().float().cpu() for L in LAYERS}
r_rel = {L: Wd[L][:, rel[L]["cols"]].T @ u_rel[L] for L in LAYERS}   # per released neuron: its push along u

# 2. one forward pass per answer, cut where nothing later is needed (the model is causal)
dh = {k: {L: [] for L in LAYERS} for k in RECIPES}
da = {k: {L: [] for L in LAYERS} for k in RECIPES}
prof = {q: {L: [] for L in LAYERS} for q in ("stream", "s")}
meta, example = [], None
order_by_onset = sorted(range(N), key=lambda i: used[i][0]["onset"] - used[i][0]["n_prompt"])
example_i = order_by_onset[N // 2]
for i, (f, w) in enumerate(used):
    n0, on = f["n_prompt"], f["onset"]
    end = max(w["loop"][1], min(on + POST, len(f["ids"])))
    x = torch.tensor([f["ids"][:end]], device="cuda")
    with torch.no_grad(), hf.record_outputs(m, LAYERS) as H, hf.record_mlp_acts(m, LAYERS) as A:
        m.model(x, use_cache=False)
    pos = np.arange(end) - on
    keep = (pos >= -PRE) & (pos < POST) & (np.arange(end) >= n0)
    ex = {} if i == example_i else None
    for L in LAYERS:
        h, act = H[L][0], A[L][0]
        hw = {k: h[s:e].mean(0) for k, (s, e) in w.items()}
        aw = {k: act[e - AWIN:e].mean(0) for k, (s, e) in w.items()}
        for k in RECIPES:
            dh[k][L].append(hw["loop"] - hw[k])
            da[k][L].append(aw["loop"] - aw[k])
        for q, v in (("stream", h @ u_rel[L]), ("s", act[:, rel[L]["cols"]] @ r_rel[L])):
            row = np.full(PRE + POST, np.nan, dtype=np.float32)
            row[pos[keep] + PRE] = v.numpy()[keep]
            prof[q][L].append(row)
        if ex is not None:
            ex[f"stream{L}"] = (h[n0:] @ u_rel[L]).tolist()
            ex[f"s{L}"] = (act[n0:, rel[L]["cols"]] @ r_rel[L]).tolist()
    meta.append({"id": f["id"], "j": f["j"], "loop_starts_at": on - n0, "period": f["period"],
                 "answer_tokens": len(f["ids"]) - n0, "windows": {k: [s - n0, e - n0] for k, (s, e) in w.items()}})
    if ex is not None:
        example = {"meta": meta[-1], **ex}
    del H, A
    if (i + 1) % 25 == 0:
        print(f"  {i + 1}/{N} answers")

# 3. directions, neurons, and the checks
res = {"answers": n_total, "cut": n_cut, "verbatim_loops": len(found), "used": N, "min_before": a.min_before,
       "periods": np.percentile([x["period"] for x in meta], [10, 25, 50, 75, 90]).tolist(),
       "loop_starts_at": np.percentile([x["loop_starts_at"] for x in meta], [10, 25, 50, 75, 90]).tolist(),
       "recipes": {}, "loops": meta}
print(f"loop period, 10/25/50/75/90th percentile: {res['periods']} tokens")
print(f"loop starts at answer token, same percentiles: {res['loop_starts_at']}")
pairs = json.load(open(os.path.join(ROOT, "labs", "data", "astra_pairs.json")))
with torch.no_grad(), hf.record_outputs(m, LAYERS) as H:
    for t in pairs["texts"]:
        m.model(hf.encode(tok, t), use_cache=False)
astra_tail = {L: torch.stack([h[-WIN:].mean(0) for h in H[L]]) for L in LAYERS}   # [8 texts, 1024]
SIZES = sorted({n for n in (1, 2, 4, 8, 16, 32, 64, 128) if n < N // 2} | {N // 2})
problem_of = np.unique([f["id"] for f, _ in used], return_inverse=True)[1]   # which problem each loop came from
n_problems = int(problem_of.max()) + 1
half_sizes = []


def split():
    """Two halves of the loops for the stability checks."""
    if a.split_by == "answer":   # two answers to the same problem can land in opposite halves
        perm = rng.permutation(N)
        h1, h2 = perm[:N // 2], perm[N // 2:]
    else:                        # every problem's loops stay together, so the halves share no problem
        first = rng.permutation(n_problems)[:n_problems // 2]
        in1 = np.isin(problem_of, first)
        h1, h2 = np.flatnonzero(in1), np.flatnonzero(~in1)
    half_sizes.extend((len(h1), len(h2)))
    return h1, h2


capstone = {}
for k, desc in RECIPES.items():
    res["recipes"][k] = {"windows": desc, "layers": {}}
    spec = {}
    for L in LAYERS:
        D, DA = torch.stack(dh[k][L]), torch.stack(da[k][L])
        u = unit(D.mean(0))
        score = (u @ Wd[L]) * DA.mean(0)
        top_order = torch.argsort(score, descending=True)
        top = sorted(top_order[:64].tolist())
        halves, curve = [], {n: {"held_out_half": [], "released": [], "used": []} for n in SIZES}
        for _ in range(20):
            h1, h2 = split()
            u2 = unit(D[h2].mean(0))
            halves.append(float(unit(D[h1].mean(0)) @ u2))
            for n in SIZES:   # a problem split can leave fewer than N // 2 loops in h1; then the last point uses all
                used_n = min(n, len(h1))
                un = unit(D[rng.choice(h1, used_n, replace=False)].mean(0))
                curve[n]["held_out_half"].append(float(un @ u2))
                curve[n]["released"].append(float(un @ u_rel[L]))
                curve[n]["used"].append(used_n)
        proj = astra_tail[L] @ u
        gaps = (proj[1::2] - proj[0::2]).tolist()   # looping minus productive, one per pair
        r = {"cos_with_released_u": float(u @ u_rel[L]),
             "overlap_with_released_64": len(set(top) & set(rel[L]["cols"])),
             "split_half_cos_mean_sd": [float(np.mean(halves)), float(np.std(halves))],
             "data_needed": [{"answers": n, "cos_with_held_out_half": float(np.mean(c["held_out_half"])),
                              "cos_with_released_u": float(np.mean(c["released"])),
                              **({"answers_used_min_max": [min(c["used"]), max(c["used"])]} if min(c["used"]) < n
                                 else {})} for n, c in curve.items()],
             "astra_pairs_looping_minus_productive": gaps,
             "gap_size": float(D.mean(0).norm()),
             "top12": top_order[:12].tolist(), "cols": top,
             "scores_top200": [round(float(x), 6) for x in score[top_order[:200]]]}
        res["recipes"][k]["layers"][str(L)] = r
        spec[str(L)] = {"alpha": a.alpha, "dir": [round(float(x), 8) for x in u], "cols": top}
        print(f"recipe {k} ({desc}), layer {L}: cos with released u {r['cos_with_released_u']:+.3f}; "
              f"released neurons picked {r['overlap_with_released_64']}/64; split-half cos "
              f"{r['split_half_cos_mean_sd'][0]:.3f}; agent's 4 pairs, looping minus productive along u: "
              f"{np.round(gaps, 2).tolist()}")
    capstone[f"cap{k}-{TRY}" if TRY else f"cap{k}-a{a.alpha:g}"] = {
        "about": f"Lab 8 recipe {k}: u and 64 neurons from the base model's own loops (windows {desc}), "
                 f"alpha {a.alpha:g}",
        "layers": spec}
    print("   answers needed:", [(d["answers"], round(d["cos_with_held_out_half"], 3))
                                 for d in res["recipes"][k]["layers"]["15"]["data_needed"]])
res["cos_A_B"] = {}
for L in LAYERS:
    uA, uB = (torch.tensor(spec_k["layers"][str(L)]["dir"]) for spec_k in capstone.values())
    res["cos_A_B"][str(L)] = float(unit(uA) @ unit(uB))
print("cos between recipe A's u and recipe B's u:", {L: round(v, 3) for L, v in res["cos_A_B"].items()})
res["split_by"] = a.split_by
res["problems"] = n_problems
res["problems_with_two_loops"] = int((np.bincount(problem_of) > 1).sum())
res["half_sizes_min_max"] = [min(half_sizes), max(half_sizes)]
print(f"{N} loops from {n_problems} problems ({res['problems_with_two_loops']} with both answers looping); "
      f"split by {a.split_by}: each half held {min(half_sizes)} to {max(half_sizes)} loops")

# 4. the onset profile: mean over loops, lined up at the loop's first token
res["onset_profile"] = {}
for q in prof:
    for L in LAYERS:
        P = np.stack(prof[q][L])
        mean = np.nanmean(P, 0)

        def avg(lo, hi):
            return float(np.nanmean(mean[lo + PRE:hi + PRE]))
        res["onset_profile"][f"{q}{L}"] = {"-2048..-1024": avg(-2048, -1024), "-512..-128": avg(-512, -128),
                                          "-128..0": avg(-128, 0), "0..128": avg(0, 128), "128..256": avg(128, 256),
                                          "mean": [round(float(x), 4) for x in mean]}
        print(f"{'stream along u' if q == 'stream' else 's(t), the trigger'} at layer {L}, mean over loops: "
              + ", ".join(f"{z} {v:.2f}" for z, v in res["onset_profile"][f"{q}{L}"].items() if z != "mean"))
json.dump(res, open(os.path.join(OUT, "capstone.json"), "w"), indent=1)
json.dump(example, open(os.path.join(OUT, "example_answer.json"), "w"))
if TRY and not new_edits:
    variants = plan = None
elif TRY:
    variants, plan = os.path.join(OUT, "variant.json"), os.path.join(OUT, "plan.json")
    json.dump(capstone, open(variants, "w"), indent=1)
    json.dump({"about": f"Lab 8 experiment {TRY}: both recipes' edits on the 256 unseen practice problems",
               "runs": [{"name": f"dev/{v}", "model": f"qwen35-{v}", "data": "dev256", "n": 2, "shards": 1, "est": 8}
                        for v in capstone]}, open(plan, "w"), indent=1)
else:
    variants, plan = os.path.join(ROOT, "configs", "capstone.json"), os.path.join(ROOT, "plans", "extra.json")
    json.dump(capstone, open(variants, "w"), indent=1)

# figures
import matplotlib  # noqa: E402
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402


def smooth(v, w=32):
    v = np.asarray(v, dtype=float)
    return np.convolve(v, np.ones(w) / w, mode="valid")


fig, ax = plt.subplots(2, 2, figsize=(13, 6.5), sharex=True)
xs = np.arange(-PRE, POST)
for c, L in enumerate(LAYERS):
    for r_, (q, lab) in enumerate((("stream", "stream along the released u"), ("s", "s(t): the 64 neurons' push"))):
        P = np.stack(prof[q][L])
        lo, mid, hi = np.nanpercentile(P, [25, 50, 75], axis=0)
        ax[r_, c].fill_between(xs, lo, hi, color="#0969da", alpha=0.15, label="middle half of loops")
        ax[r_, c].plot(xs, mid, color="#0969da", lw=1, label="median")
        ax[r_, c].axvline(0, color="black", ls="--", lw=1)
        ax[r_, c].set_title(f"layer {L}: {lab}")
        ax[r_, c].set_xlabel("tokens from the start of the loop")
ax[0, 0].legend(loc="upper left")
fig.suptitle(f"{N} loops from the unedited model, lined up at the token where each loop starts")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab8_loop_onset.png"), dpi=130)

fig, ax = plt.subplots(1, 2, figsize=(13, 4))
for c, L in enumerate(LAYERS):
    for k, col in (("A", "#cf222e"), ("B", "#8250df")):
        dn = res["recipes"][k]["layers"][str(L)]["data_needed"]
        n = [d["answers"] for d in dn]
        ax[c].plot(n, [d["cos_with_held_out_half"] for d in dn], "o-", color=col,
                   label=f"recipe {k} vs the other half of the answers")
        ax[c].plot(n, [d["cos_with_released_u"] for d in dn], "o--", color=col, alpha=0.6,
                   label=f"recipe {k} vs the released u")
    ax[c].axvline(4, color="#57606a", lw=0.8, ls=":")
    ax[c].annotate("the agent used 4 pairs", (4, 0.05), xytext=(4, 0), textcoords="offset points", fontsize=8)
    ax[c].set_xscale("log", base=2)
    ax[c].set_ylim(0, 1)
    ax[c].set_xlabel("answers used to build u")
    ax[c].set_ylabel("cos")
    ax[c].set_title(f"layer {L}: how many loops does a direction need?")
ax[0].legend(fontsize=8, loc="lower right")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab8_data_needed.png"), dpi=130)

em = example["meta"]
fig, ax = plt.subplots(2, 1, figsize=(13, 6), sharex=True)
for r_, (q, lab) in enumerate((("stream", "stream along the released u"), ("s", "s(t): the 64 neurons' push"))):
    for L, col in zip(LAYERS, ("#0969da", "#8250df")):
        v = smooth(example[f"{q}{L}"])
        ax[r_].plot(np.arange(len(v)) + 16, v, color=col, label=f"layer {L}")
    for k, col in (("B", "#8250df"), ("A", "#cf222e"), ("loop", "#1a7f37")):
        ax[r_].axvspan(*em["windows"][k], color=col, alpha=0.25)
    ax[r_].axvline(em["loop_starts_at"], color="black", ls="--", lw=1)
    ax[r_].set_ylabel(lab)
ax[0].legend(loc="upper left")
ax[1].set_xlabel("token position in the answer (32-token moving average)")
ax[0].set_title(f"One answer ({em['id']}): the loop starts at token {em['loop_starts_at']} and repeats every "
                f"{em['period']} tokens. Shaded: window B, window A, loop window")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab8_one_answer.png"), dpi=130)
rp = lambda p: os.path.relpath(p, ROOT)  # noqa: E731
if variants is None:
    print("wrote", rp(OUT), "and three figures in", rp(FIG))
    print("--split-by changes only the stability checks, so the edits match the main run's (configs/capstone.json) and "
          "there is nothing new to build or test (if you also changed the code, run again without --split-by)")
    sys.exit(0)
print("wrote", rp(OUT), rp(variants), "and three figures in", rp(FIG))
print("next, build both edits and test them on the 256 unseen practice problems (one GPU per run):")
print(f"  CUDA_VISIBLE_DEVICES=\"\" python scripts/build_models.py --variants {rp(variants)}")
print(f"  python scripts/run_jobs.py {rp(plan)} --gpus 6,7 --only " + " ".join(f"dev/{v}" for v in capstone))