#!/usr/bin/env python
"""Lab 3: rebuild the released edit from the agent's 8 sentences, using only hooks and arithmetic.

    . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
    CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py      # about a minute on one GPU

Experiments (each writes to its own folder, results/lab3/try-*, and leaves the main results alone; the script
prints the commands that build the try's edit and test it on the 256 unseen practice problems):
    CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py --pairs 2           # only the agent's first 2 pairs
    CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py --windows 36 36     # neurons averaged over 36 final tokens, not 24
    CUDA_VISIBLE_DEVICES=7 python labs/lab3_rebuild.py --layers 15 19 21   # also edit layer 21, the third-strongest MLP
(plans/lab3-tries.json tests exactly these three: runs dev/rebuilt-pairs2, dev/rebuilt-windows36-36 and
dev/rebuilt-layers15-19-21.)

What it does:
  1. Runs the 8 texts in labs/data/astra_pairs.json (4 productive, 4 looping) through the unedited model and
     records every block's output (the residual stream) and every MLP's neuron activations.
  2. For every layer, measures how far apart looping and productive text sit (averaged over the last 36 tokens)
     and how much of that gap each block adds, split into the MLP's share and the mixer's share. This is how the
     agent picked layers 15 and 19.
  3. Direction u at layer L = the looping-minus-productive gap after block L, scaled to length 1.
  4. Neuron score = (how much neuron j writes along u) x (how much more neuron j fires in looping text, last 24
     tokens). The 64 highest are the neurons the edit touches.
  5. Compares all of it with the released edit and writes configs/rebuilt.json, an edit you can build and test.
"""
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  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--pairs", type=int, default=4, choices=[1, 2, 3, 4], help="use only the agent's first N pairs")
ap.add_argument("--windows", type=int, nargs=2, default=[36, 24], metavar=("STREAM", "NEURONS"),
                help="final tokens to average for the stream and for the neurons (the agent used 36 and 24)")
ap.add_argument("--layers", type=int, nargs="+", default=[15, 19], help="layers to build the edit for")
a = ap.parse_args()
LAYERS = tuple(a.layers)
H_WIN, A_WIN = a.windows  # the agent's windows: last 36 tokens for the stream, last 24 for neurons
TRY = "try-" + "-".join(p for p in (f"pairs{a.pairs}" if a.pairs != 4 else "",
                                     f"windows{H_WIN}-{A_WIN}" if a.windows != [36, 24] else "",
                                     "layers" + "-".join(map(str, LAYERS)) if list(LAYERS) != [15, 19] else "") if p)
TRY = None if TRY == "try-" else TRY
OUT = os.path.join(ROOT, "results", "lab3", *([TRY] if TRY else []))
FIG = OUT if TRY else os.path.join(ROOT, "results", "figures")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIG, exist_ok=True)

pairs = json.load(open(os.path.join(ROOT, "labs", "data", "astra_pairs.json")))
texts = pairs["texts"][:2 * a.pairs]
N = len(texts)
m, tok = hf.load("base")
n_layers = len(hf.blocks(m))
# hidden_states[-1] has the final norm applied, so the last block's raw output is not available to measure
assert all(0 <= L < n_layers - 1 for L in LAYERS), f"--layers must be between 0 and {n_layers - 2}"

# 1. record. hidden[l] is the stream entering block l (hidden[0] = token embeddings, hidden[l+1] = block l's output)
hidden, acts, lens = [], [], []
with torch.no_grad(), hf.record_outputs(m, range(n_layers)) as outs, hf.record_mlp_acts(m, range(n_layers)) as A:
    for t in texts:
        ids = hf.encode(tok, t)
        lens.append(ids.shape[1])
        o = m(ids, output_hidden_states=True)
        hidden.append([h[0].float().cpu() for h in o.hidden_states])
    # the hook copy of block l's output must equal transformers' hidden_states[l+1]
    hook_check = max(float((outs[l][i] - hidden[i][l + 1]).abs().max()) for i in range(N) for l in range(n_layers - 1))
assert lens == pairs["agent_token_lengths"][:N], (lens, pairs["agent_token_lengths"])
print(f"token counts match the agent's: {lens}; hooks vs hidden_states max diff {hook_check:.2e}")


def pair_gap(per_text):  # per_text: [N, d] -> mean over the pairs of (looping - productive)
    x = torch.stack(per_text)
    return (x[1::2] - x[0::2]).mean(0)


gapH = [pair_gap([hidden[i][l][-H_WIN:].mean(0) for i in range(N)]) for l in range(n_layers + 1)]
gapA36 = [pair_gap([A[l][i][-H_WIN:].mean(0) for i in range(N)]) for l in range(n_layers)]
gapA24 = {L: pair_gap([A[L][i][-A_WIN:].mean(0) for i in range(N)]) for L in LAYERS}
Wd = {l: hf.down_proj(m, l).weight.detach().float().cpu() for l in range(n_layers)}

# 2. where does the gap grow? (the agent's table, pass 3)
growth = []
for l in range(n_layers - 1):
    d = gapH[l + 1] / gapH[l + 1].norm()
    total = float((gapH[l + 1] - gapH[l]) @ d)
    mlp = float((Wd[l] @ gapA36[l]) @ d)
    growth.append({"layer": l, "gap_size_after": float(gapH[l + 1].norm()), "growth": total, "mlp": mlp,
                   "mixer": total - mlp})
top_mlp = sorted(growth, key=lambda g: -g["mlp"])[:4]
print("layers whose MLP adds most to the gap:", [(g["layer"], round(g["mlp"], 3)) for g in top_mlp])

# 3-4. direction and neurons, then compare with the released edit
rel = hf.released_edit()
rebuilt, report = {}, {"token_lengths": lens, "hook_vs_hidden_states_max_diff": hook_check, "growth": growth,
                       "layers": {}}
for L in LAYERS:
    u = gapH[L + 1] / gapH[L + 1].norm()  # points from productive toward looping
    write = u @ Wd[L]                       # [3584]: how much each neuron writes along u per unit of activation
    score = write * gapA24[L]
    order = torch.argsort(score, descending=True)
    top = sorted(order[:64].tolist())
    pos = score.clamp_min(0).sum()
    cos = float(u @ rel[L]["u"] / rel[L]["u"].norm()) if L in rel else None  # the released edit: layers 15, 19
    report["layers"][str(L)] = {
        "cos_with_released_u": cos,
        "released_u_points_to": None if cos is None else "looping" if cos > 0 else "productive",
        "overlap_with_released_64": len(set(top) & set(rel[L]["cols"])) if L in rel else None,
        "top12": order[:12].tolist(),
        "positive_score_share_top8_32_64": [float(score[order[:k]].clamp_min(0).sum() / pos) for k in (8, 32, 64)],
        "gap_size": float(gapH[L + 1].norm()),
        "u_loop_sign": [round(float(x), 8) for x in u],
        "cols": top,
        "scores_sorted": [round(float(x), 6) for x in score[order]],
    }
    rebuilt[str(L)] = {"alpha": 6, "dir": [round(float(x), 8) for x in u], "cols": top}
    vs = (f"cos(rebuilt u, released u) = {cos:+.5f}; neurons shared with the released 64: "
          f"{report['layers'][str(L)]['overlap_with_released_64']}/64; " if L in rel
          else "the released edit leaves this layer alone; ")
    print(f"layer {L}: {vs}top-8/32/64 share of positive score: "
          f"{[round(x, 3) for x in report['layers'][str(L)]['positive_score_share_top8_32_64']]}")
if 15 in LAYERS and 19 in LAYERS:
    u15, u19 = (torch.tensor(report["layers"][str(L)]["u_loop_sign"]) for L in (15, 19))
    report["cos_u15_u19_loop_sign"] = float(u15 @ u19)
    print(f"cos(u15, u19) with both pointing toward looping: {report['cos_u15_u19_loop_sign']:+.3f}")

json.dump(report, open(os.path.join(OUT, "rebuild.json"), "w"), indent=1)
if TRY:  # its own variant name, so building it never replaces models/qwen35-rebuilt
    name = f"rebuilt-{TRY[4:]}"
    variants, plan = os.path.join(OUT, "variant.json"), os.path.join(OUT, "plan.json")
    json.dump({name: {"about": f"Lab 3 experiment: {TRY[4:]}", "layers": rebuilt}}, open(variants, "w"), indent=1)
    json.dump({"about": f"Lab 3 experiment {TRY[4:]}: the rebuilt edit on the 256 unseen practice problems",
               "runs": [{"name": f"dev/{name}", "model": f"qwen35-{name}", "data": "dev256", "n": 2, "shards": 1,
                         "est": 8}]}, open(plan, "w"), indent=1)
else:
    name = "rebuilt"
    variants, plan = os.path.join(ROOT, "configs", "rebuilt.json"), os.path.join(ROOT, "plans", "extra.json")
    json.dump({"rebuilt": {"about": "Lab 3: the edit rebuilt from the agent's 8 texts (alpha 6, my u, my top 64)",
                           "layers": rebuilt}}, open(variants, "w"), indent=1)

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

fig, ax = plt.subplots(1, 2, figsize=(13, 4.2))
xs = [g["layer"] for g in growth]
# leave out the last position: hidden_states[-1] has gone through the final norm, which rescales everything
ax[0].plot(range(n_layers), [float(g.norm()) for g in gapH[:n_layers]], "o-", color="#0969da")
ax[0].text(0.02, 0.95, f"not shown: after the last block and the final norm\nthe gap is {float(gapH[-1].norm()):.0f}",
           transform=ax[0].transAxes, va="top", fontsize=8, color="#57606a")
ax[0].set_xlabel("position in the stream (0 = embeddings, l+1 = after block l)")
ax[0].set_ylabel("size of looping - productive gap")
ax[0].set_title("The loop signal builds up through the network")
ax[1].bar([x - 0.2 for x in xs], [g["mlp"] for g in growth], 0.4, label="added by the MLP", color="#cf222e")
ax[1].bar([x + 0.2 for x in xs], [g["mixer"] for g in growth], 0.4, label="added by attention / DeltaNet",
          color="#8c959f")
for L in LAYERS:
    ax[1].annotate(f"layer {L}", (L - 0.2, growth[L]["mlp"]), xytext=(0, 6), textcoords="offset points",
                   ha="center", fontsize=9)
ax[1].set_xlabel("block")
ax[1].set_ylabel("growth of the gap along its own direction")
ax[1].set_title("Which blocks push text toward the loop")
ax[1].legend()
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab3_where_the_loop_grows.png"), dpi=130)

fig, ax = plt.subplots(1, len(LAYERS), figsize=(6.5 * len(LAYERS), 3.8), squeeze=False)
ax = ax[0]
for k, L in enumerate(LAYERS):
    s = np.array(report["layers"][str(L)]["scores_sorted"])
    ax[k].plot(np.arange(len(s)), s, color="#8c959f", lw=1)
    ax[k].plot(np.arange(64), s[:64], color="#cf222e", lw=2, label="the 64 neurons the edit touches")
    ax[k].set_xscale("symlog", linthresh=10)
    ax[k].axhline(0, color="black", lw=0.5)
    ax[k].set_title(f"layer {L}: neuron scores, highest first")
    ax[k].set_xlabel("neuron rank (of 3,584)")
    ax[k].set_ylabel("writes along u x fires more in loops")
    ax[k].legend()
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab3_neuron_scores.png"), dpi=130)
print("wrote", OUT, "(an experiment: configs/rebuilt.json and results/figures are untouched)" if TRY
      else "and configs/rebuilt.json and two figures")
print("next, build the edit and test it on the 256 unseen practice problems (about 20 minutes on one GPU):")
print(f"  CUDA_VISIBLE_DEVICES=\"\" python scripts/build_models.py --variants {os.path.relpath(variants, ROOT)}")
print(f"  python scripts/run_jobs.py {os.path.relpath(plan, ROOT)} --gpus 7 --only dev/{name}")