"""E2 stage A (item 21): dense short-CoT supervision through the carry
whiteboard on GSM8K.

train_carry.py skeleton, one change that matters: CE targets are the
model's own VERIFIED terse scratchpad + answer (gsm_cot_data.json,
~30-60 tokens) instead of the ~3-token bare answer — the dense-output
ingredient the latent GSM arms structurally lacked.
Control: --feedforward = same supervision, adapter(e,e), no recurrence.
"""

import argparse
import json
import os
import math
import random
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F

from carry_common import PAUSE_ID, carry_logits
from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import load_model  # noqa: E402

OUT = Path(os.environ.get("LOOP_OUT",
    Path(__file__).resolve().parent.parent / "results-loop"))
STEPS = 600
BATCH = 4
LR = 1e-3
WARMUP = 20
K_PREFILL = 2
P_BY_LABEL = {"easy": 2, "hard": 6}

ap = argparse.ArgumentParser()
ap.add_argument("--feedforward", action="store_true")
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--lensnoise", default=None, metavar="RANK,SCALE",
                help="E2-N1: inject noise into the carried state during "
                     "training, shaped by the top-RANK sensitivity "
                     "directions of jbar at the band entrance; noise norm "
                     "= SCALE * per-position state norm (e.g. 32,0.05)")
ap.add_argument("--drop-steps", type=int, default=0, metavar="D",
                help="E2-L rung B: delete the first D scratchpad steps, "
                     "each replaced by --pause-per-step extra pauses")
ap.add_argument("--pause-per-step", type=int, default=10,
                help="pauses per deleted step (median step = 10 tokens)")
ap.add_argument("--warm-start", default=None, metavar="ADAPTER_PT")
ap.add_argument("--steps", type=int, default=STEPS)
ap.add_argument("--lr", type=float, default=LR)
ap.add_argument("--tag-suffix", default="",
                help="appended to TAG (distinguish control variants)")
ap.add_argument("--bandlora", type=int, default=0, metavar="RANK",
                help="item 24: loop-only LoRA (lora_band.LoopLoRA) on every "
                     "band layer, uniform scale 1.0 — active only during "
                     "band re-runs, k=0 stays bit-exact")
ap.add_argument("--lora-lr", type=float, default=1e-3)
ap.add_argument("--lensteach", type=float, default=0.0, metavar="LAMBDA",
                help="item 25: latent process supervision — lens-CE at the "
                     "replacement pauses against the DELETED step's tokens "
                     "(1:1 pause j <-> step token j), mixed at LAMBDA into "
                     "the output CE")
ap.add_argument("--lensteach-gen", type=float, default=0.0, metavar="LAMBDA",
                help="item 26: result-staging supervision during generation "
                     "— at each visible scratchpad line's pre-'=' positions "
                     "(result not yet in causal context), lens-CE the "
                     "carried state against that line's result tokens")
ap.add_argument("--inner-iters", type=int, default=0, metavar="M",
                help="item 27: M in-place band iterations at the anchor "
                     "(internal looping, no extra tokens); with --lensteach "
                     "the lens-CE aligns iteration i <-> deleted-step token "
                     "i instead of pause positions")
ap.add_argument("--base-pauses", type=int, default=-1, metavar="P",
                help="override P_BY_LABEL with a fixed pause count; 0 = NO "
                     "pause tokens at all (inner iterations anchor on the "
                     "last prompt position)")
ap.add_argument("--traj-tf", type=float, default=0.0, metavar="LAMBDA",
                help="item 29: teacher-forced TRANSITION loss — iteration i "
                     "gets teacher waypoint T[i-1] as merge input, its band "
                     "output is pulled onto T[i] (cosine); 10 independent "
                     "supervised one-step regressions")
ap.add_argument("--traj-fr", type=float, default=0.0, metavar="LAMBDA",
                help="item 29: free-running trajectory loss — the burst's "
                     "own state s_i pulled onto waypoint T[i] (exposure-gap "
                     "term)")
ap.add_argument("--traj-span", choices=("step", "full"), default="step",
                help="waypoints sampled evenly across the deleted step "
                     "(job 1, d=1) or the full CoT (job 2, answer-only)")
ap.add_argument("--kvmem", type=int, default=0, metavar="CODE",
                help="item 31: KVMemoryAdapter (code width) — burst states "
                     "become per-band-layer KV prefix entries readable by "
                     "frozen attention during the answer scan")
ap.add_argument("--freeze-merge", action="store_true",
                help="freeze the (warm-started) merge adapter; train only "
                     "the kvmem adapter")
ap.add_argument("--symchain", choices=("tf", "st"), default=None,
                help="item 32: discrete latent chain — each burst tick "
                     "feeds back the lens-snapped token embedding through "
                     "a zero-init projector. tf: teacher-forced symbols "
                     "(ground-truth deleted-step tokens); st: free-running "
                     "straight-through snaps")
ap.add_argument("--teachstate", type=float, default=0.0, metavar="LAMBDA",
                help="item 28 (Nils's variant): teacher-state distillation "
                     "— frozen warm-start adapter runs the FULL cot (step "
                     "visible); its band-exit state at the deleted step's "
                     "last token becomes the target; burst final iterate "
                     "trained to it by cosine, weighted LAMBDA")
ARGS = ap.parse_args()
STEPS, LR = ARGS.steps, ARGS.lr
TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
    f"_b{ARGS.drop_steps}" if ARGS.drop_steps else "") + (
    f"_blr{ARGS.bandlora}" if ARGS.bandlora else "") + (
    f"_lt{str(ARGS.lensteach).replace('.', '')}" if ARGS.lensteach else "") + (
    f"_lg{str(ARGS.lensteach_gen).replace('.', '')}"
    if ARGS.lensteach_gen else "") + (
    f"_ii{ARGS.inner_iters}" if ARGS.inner_iters else "") + (
    f"_ts{str(ARGS.teachstate).replace('.', '')}"
    if ARGS.teachstate else "") + (
    f"_tj{ARGS.traj_span[0]}{str(ARGS.traj_tf).replace('.', '')}"
    f"_{str(ARGS.traj_fr).replace('.', '')}"
    if (ARGS.traj_tf or ARGS.traj_fr) else "") + (
    f"_kvm{ARGS.kvmem}" if ARGS.kvmem else "") + (
    f"_sc{ARGS.symchain}" if ARGS.symchain else "") + (
    "_fm" if ARGS.freeze_merge else "") + (
    f"_p{ARGS.base_pauses}" if ARGS.base_pauses >= 0 else "") + (
    f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
    f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix


def drop_cot_steps(cot, d):
    """Delete the first d scratchpad lines (front-first: the deleted
    computation must ride the pause-chain before the visible remainder).
    Returns (new_cot, n_deleted, deleted_text); unparseable cots pass
    through intact."""
    lines = [l for l in cot.split("\n") if l.strip()]
    ans = [l for l in lines if l.startswith("Answer")]
    steps = [l for l in lines if not l.startswith("Answer")]
    if len(ans) != 1 or not steps:
        return cot, 0, ""
    n = min(d, len(steps))
    return "\n".join(steps[n:] + ans), n, "\n".join(steps[:n])


class LensNoiseWrapper(torch.nn.Module):
    """Perturb the carried state s (not the anchor e) before the merge,
    within the span of the lens's top-r readout-sensitive directions.
    Train-time only (noise_on flag); eval and checkpoints use .base."""

    def __init__(self, base, jbar_layer, rank, scale):
        super().__init__()
        self.base = base
        self.scale = scale
        self.noise_on = True
        J = jbar_layer.float()
        _, _, Vt = torch.linalg.svd(J, full_matrices=False)
        self.register_buffer("V", Vt[:rank].T.contiguous())  # (d, r)

    def forward(self, e, s):
        if self.noise_on and self.scale > 0:
            z = torch.randn(*s.shape[:-1], self.V.shape[1],
                            device=s.device, dtype=torch.float32)
            n = z @ self.V.T
            n = n * (s.float().norm(dim=-1, keepdim=True) * self.scale
                     / (n.norm(dim=-1, keepdim=True) + 1e-6))
            s = (s.float() + n).to(s.dtype)
        return self.base(e, s)


def gen_staging_targets(tok, cot):
    """Item 26: per-line result-staging spans, computed in TOKEN space.
    For each cot line with '=', the pre-'=' positions (result not yet in
    causal context) target the line's result tokens; the Answer line's
    'Answer:' positions target the answer tokens.
    Returns [(rel_positions, result_token_ids), ...] relative to tok(cot)."""
    ids = tok(cot, add_special_tokens=False)["input_ids"]
    decoded = [tok.decode([t]) for t in ids]
    out, line_start = [], 0
    for i, d in enumerate(decoded + ["\n"]):
        if "\n" not in d and i < len(ids):
            continue
        line = list(range(line_start, min(i, len(ids))))
        line_start = i + 1
        if not line:
            continue
        text = "".join(decoded[j] for j in line)
        seps = [j for j in line if decoded[j].strip() == "="]
        if seps:
            sep = seps[-1]
        elif text.strip().startswith("Answer"):
            colons = [j for j in line if decoded[j].strip() == ":"]
            if not colons:
                continue
            sep = colons[-1]
        else:
            continue
        result = [ids[j] for j in line if j > sep
                  and decoded[j].strip()]
        pre = [j for j in line if j <= sep]
        if result and pre:
            out.append((pre, result))
    return out


def sym_burst_forward(looper, adapter, proj, ids, msk, plens, m, sym_tf,
                      k, lens_fn, embed_w, start_id):
    """Item 32 forward: discrete latent chain at the anchor (teacher-forced
    or straight-through symbols), then the visible-token carry scan."""
    from carry_common import (prompt_prefill, build_step_updates,
                              carry_steps, sym_iterate)
    dev = ids.device
    calls, _ = looper.capture(ids, msk, logits_to_keep=1)
    e = looper._hin[looper.l0].detach()
    B = ids.shape[0]
    ar = torch.arange(ids.shape[1], device=dev)
    pmask = ar[None, :] < plens[:, None].to(dev)
    S, X = prompt_prefill(looper, adapter, e, calls, pmask, k)
    rows = torch.arange(B, device=dev)
    anchor = (plens - 1).to(dev)
    states = []
    S, X = sym_iterate(looper, adapter, proj, e, calls, S, X, rows, anchor,
                       m, lens_fn, embed_w, sym_tf=sym_tf,
                       start_id=start_id, use_checkpoint=True,
                       iter_states=states)
    total = msk.sum(-1)
    updates = build_step_updates(plens.to(dev), total.to(dev), dev)
    S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
                       use_checkpoint=True)
    return looper.suffix_logits(S, calls), states


def traj_burst_forward(looper, adapter, ids, msk, plens, m, T, tf_on, k,
                       mem_adapter=None):
    """Item 29 forward: optional teacher-forced transition predictions,
    then the free-running burst (whose final state seeds the answer scan,
    matching inference), then the visible-token carry. Returns
    (logits, tf_preds, fr_states)."""
    from torch.utils.checkpoint import checkpoint
    from carry_common import prompt_prefill, build_step_updates, carry_steps
    dev = ids.device
    calls, _ = looper.capture(ids, msk, logits_to_keep=1)
    e = looper._hin[looper.l0].detach()
    B = ids.shape[0]
    ar = torch.arange(ids.shape[1], device=dev)
    pmask = ar[None, :] < plens[:, None].to(dev)
    S, X = prompt_prefill(looper, adapter, e, calls, pmask, k)
    rows = torch.arange(B, device=dev)
    anchor = (plens - 1).to(dev)
    s0 = S[rows, anchor]

    def upd(S, X, seed):
        x_new = adapter(e[rows, anchor], seed)
        X = X.clone()
        X[rows, anchor] = x_new.to(X.dtype)
        S = checkpoint(lambda X_: looper.band(X_, calls), X,
                       use_reentrant=False)
        return S, X

    tf_preds, fr_states = [], []
    if tf_on:
        for i in range(m):
            seed = s0 if i == 0 else T[:, i - 1].to(s0.dtype)
            S, X = upd(S, X, seed)
            tf_preds.append(S[rows, anchor])
    for i in range(m):
        seed = s0 if i == 0 else fr_states[-1]
        S, X = upd(S, X, seed)
        fr_states.append(S[rows, anchor])
    if mem_adapter is not None:
        from kv_memory import arm_memory
        arm_memory(mem_adapter(torch.stack(fr_states, 1)))
    total = msk.sum(-1)
    updates = build_step_updates(plens.to(dev), total.to(dev), dev)
    S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
                       use_checkpoint=True)
    out = looper.suffix_logits(S, calls)
    # NOTE: memory stays armed through backward (checkpoint recompute
    # must see the same graph); the caller disarms after opt.step().
    return out, tf_preds, fr_states


def lr_at(step):
    if step < WARMUP:
        return LR * (step + 1) / WARMUP
    t = (step - WARMUP) / max(1, STEPS - WARMUP)
    return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t))


def build_batch(tok, items, p, device="cuda"):
    seqs, labs, plens = [], [], []
    for it in items:
        pr = tok(chat_prompt(tok, it["question"], DIRECT_SUFFIX),
                 add_special_tokens=False)["input_ids"]
        a = tok(it["cot"] + "<end_of_turn>",
                add_special_tokens=False)["input_ids"]
        pi = p + it.get("extra_pauses", 0)
        seqs.append(pr + [PAUSE_ID] * pi + a)
        labs.append([-100] * (len(pr) + pi) + a)
        plens.append(len(pr))
    T = max(len(s) for s in seqs)
    pad = tok.pad_token_id or 0
    ids = torch.full((len(seqs), T), pad, dtype=torch.long)
    lab = torch.full((len(seqs), T), -100, dtype=torch.long)
    msk = torch.zeros((len(seqs), T), dtype=torch.long)
    for i, (s, l) in enumerate(zip(seqs, labs)):
        ids[i, : len(s)] = torch.tensor(s)
        lab[i, : len(s)] = torch.tensor(l)
        msk[i, : len(s)] = 1
    return (ids.to(device), msk.to(device), lab.to(device),
            torch.tensor(plens, device=device))


@torch.no_grad()
def val_loss(looper, adapter, tok, items, p):
    tot, n = 0.0, 0
    for i in range(0, len(items), BATCH):
        chunk = items[i : i + BATCH]
        ids, msk, lab, plens = build_batch(tok, chunk, p)
        ckw = {}
        if ARGS.inner_iters:
            extras = torch.tensor([c.get("extra_pauses", 0) for c in chunk],
                                  device=plens.device)
            ckw = dict(inner_iters=ARGS.inner_iters,
                       inner_at=plens + p + extras - 1)
        logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL,
                              feedforward=ARGS.feedforward, **ckw)
        loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
                               lab[:, 1:].flatten(), ignore_index=-100)
        tot += loss.item() * len(ids)
        n += len(ids)
    return tot / n


def main():
    rng = random.Random(ARGS.seed)
    torch.manual_seed(ARGS.seed)
    star = {it["idx"]: it for it in json.load(open(OUT / "star_data.json"))
            if it["split"] == "train"}
    cots = json.load(open(OUT / "gsm_cot_data.json"))
    data = []
    n_dropped = 0
    for r in cots:
        it = star.get(r["idx"])
        if it is None:
            continue
        cot, ndel, deleted = (drop_cot_steps(r["cot"], ARGS.drop_steps)
                              if ARGS.drop_steps else (r["cot"], 0, ""))
        n_dropped += ndel
        data.append({"question": it["question"], "label": r["label"],
                     "cot": cot, "deleted": deleted,
                     "extra_pauses": ndel * ARGS.pause_per_step})
    if ARGS.drop_steps:
        print(f"rung B d={ARGS.drop_steps}: {n_dropped} steps deleted "
              f"across {len(data)} items "
              f"({ARGS.pause_per_step} pauses each)", flush=True)

    model, tok = load_model(dtype=torch.bfloat16)
    for pp in model.parameters():
        pp.requires_grad_(False)
    looper = BandLooper(model)
    adapter = MergeAdapter(
        d=model.config.get_text_config().hidden_size).cuda()
    if ARGS.lensnoise:
        r, sc = ARGS.lensnoise.split(",")
        jbar = torch.load(Path(__file__).resolve().parent.parent
                          / "results/jbar.pt", map_location="cpu")["Jbar"]
        from loop_common import BAND
        adapter = LensNoiseWrapper(adapter, jbar[BAND[0]], int(r),
                                   float(sc)).cuda()
        print(f"lens-noise: rank={r} scale={sc} on jbar L{BAND[0]}",
              flush=True)
    if ARGS.warm_start:
        (adapter.base if ARGS.lensnoise else adapter).load_state_dict(
            torch.load(ARGS.warm_start, map_location="cuda"))
        print(f"warm-started from {ARGS.warm_start}", flush=True)
    lora_params, lora_band_layers = [], []
    if ARGS.bandlora:
        from lora_band import inject_band_lora
        from loop_common import BAND
        lora_band_layers = list(range(BAND[0], BAND[1] + 1))
        scales = {l: 1.0 for l in lora_band_layers}
        lora_params = inject_band_lora(looper.tm, BAND[0], scales,
                                       rank=ARGS.bandlora)
        for p in lora_params:
            p.data = p.data.cuda()
        print(f"band-lora r={ARGS.bandlora}: "
              f"{sum(p.numel() for p in lora_params)/1e6:.1f}M params, "
              f"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, "
              f"lr={ARGS.lora_lr}", flush=True)
    mem_adapter = None
    if ARGS.kvmem:
        from kv_memory import install, KVMemoryAdapter
        from loop_common import BAND
        install(model)
        mem_adapter = KVMemoryAdapter(model, band=BAND,
                                      code=ARGS.kvmem).cuda()
        print(f"kv-memory adapter: code={ARGS.kvmem}, "
              f"{sum(p.numel() for p in mem_adapter.parameters())/1e6:.1f}M "
              f"params, band {BAND}", flush=True)
    if ARGS.freeze_merge:
        for p_ in adapter.parameters():
            p_.requires_grad_(False)
        print("merge adapter FROZEN", flush=True)
    lens_teach = None
    if ARGS.lensteach or ARGS.lensteach_gen:
        from loop_common import BAND
        J30 = torch.load(Path(__file__).resolve().parent.parent
                         / "results/jbar.pt",
                         map_location="cuda")["Jbar"][BAND[1]].float()
        tm = looper.tm
        softcap = model.config.get_text_config().final_logit_softcapping

        def lens_teach(h):
            proj = h.float() @ J30.T
            x = tm.norm(proj.to(tm.norm.weight.dtype))
            lg = model.lm_head(x)
            if softcap:
                lg = softcap * torch.tanh(lg / softcap)
            return lg

        if ARGS.lensteach:
            n_tgt = 0
            for it in data:
                if it["deleted"]:
                    cap = ARGS.inner_iters or it["extra_pauses"]
                    it["lens_targets"] = tok(
                        it["deleted"], add_special_tokens=False
                    )["input_ids"][:cap]
                    n_tgt += bool(it["lens_targets"])
            print(f"lens-teach λ={ARGS.lensteach}: targets on {n_tgt} "
                  f"items (pause j <-> deleted-step token j, lens at "
                  f"L{BAND[1]})", flush=True)
        if ARGS.lensteach_gen:
            n_gt, n_spans = 0, 0
            for it in data:
                gt = gen_staging_targets(tok, it["cot"])
                if gt:
                    it["gen_targets"] = gt
                    n_gt += 1
                    n_spans += len(gt)
            print(f"lens-teach-gen λ={ARGS.lensteach_gen}: staging targets "
                  f"on {n_gt} items ({n_spans} line-spans; pre-'=' "
                  f"positions target the line result)", flush=True)
    if ARGS.traj_tf or ARGS.traj_fr:
        assert ARGS.warm_start and ARGS.inner_iters \
            and ARGS.base_pauses == 0, \
            "traj mode needs --warm-start, --inner-iters, --base-pauses 0"
        data = [it for it in data if it["deleted"]]
        m = ARGS.inner_iters
        t0_ = time.time()
        with torch.no_grad():
            for i in range(0, len(data), BATCH):
                chunk = data[i:i + BATCH]
                full_items = [{"question": c["question"],
                               "cot": c["deleted"] + "\n" + c["cot"],
                               "extra_pauses": 0} for c in chunk]
                ids_, msk_, _, plens_ = build_batch(tok, full_items, 0)
                _, S_ = carry_logits(looper, adapter, ids_, msk_, plens_,
                                     K_PREFILL, return_states=True)
                for b, c in enumerate(chunk):
                    if ARGS.traj_span == "step":
                        span = len(tok(c["deleted"],
                                       add_special_tokens=False)["input_ids"])
                    else:
                        span = len(tok(c["deleted"] + "\n" + c["cot"],
                                       add_special_tokens=False)["input_ids"])
                    idxs = [round(j * (span - 1) / (m - 1)) for j in range(m)]
                    pos = [int(plens_[b]) + jj for jj in idxs]
                    c["traj_states"] = S_[b, pos].float().clone()
        print(f"trajectory waypoints: {len(data)} items x {m} states "
              f"({ARGS.traj_span} span, {time.time()-t0_:.0f}s; frozen "
              f"warm-start teacher, full cot)", flush=True)
    sym_proj, EMBED_W, START_ID = None, None, None
    if ARGS.symchain:
        assert ARGS.warm_start and ARGS.inner_iters and ARGS.lensteach, \
            "symchain needs --warm-start, --inner-iters, --lensteach"
        d_ = model.config.get_text_config().hidden_size
        sym_proj = torch.nn.Linear(d_, d_).cuda()
        torch.nn.init.zeros_(sym_proj.weight)
        torch.nn.init.zeros_(sym_proj.bias)
        EMBED_W = model.get_input_embeddings().weight.detach()
        START_ID = tok("\n", add_special_tokens=False)["input_ids"][0]
        print(f"symchain [{ARGS.symchain}]: zero-init proj "
              f"({d_}x{d_}), start_id={START_ID}", flush=True)
    if ARGS.teachstate:
        assert ARGS.warm_start and ARGS.inner_iters, \
            "teachstate needs --warm-start (frozen teacher) + --inner-iters"
        t0_ = time.time()
        todo = [it for it in data if it["deleted"]]
        pt = ARGS.base_pauses if ARGS.base_pauses >= 0 else 0
        with torch.no_grad():
            for i in range(0, len(todo), BATCH):
                chunk = todo[i:i + BATCH]
                full_items = [{"question": c["question"],
                               "cot": c["deleted"] + "\n" + c["cot"],
                               "extra_pauses": 0} for c in chunk]
                ids_, msk_, _, plens_ = build_batch(tok, full_items, pt)
                _, S_ = carry_logits(looper, adapter, ids_, msk_, plens_,
                                     K_PREFILL, return_states=True)
                for b, c in enumerate(chunk):
                    nstep = len(tok(c["deleted"],
                                    add_special_tokens=False)["input_ids"])
                    pos = int(plens_[b]) + pt + nstep - 1
                    c["teacher_state"] = S_[b, pos].float().clone()
        print(f"teacher states: {len(todo)} captured ({time.time()-t0_:.0f}s;"
              f" frozen warm-start adapter, full cot, band-exit at the "
              f"deleted step's last token)", flush=True)
    groups = ([] if ARGS.freeze_merge else
              [{"params": list(adapter.parameters()), "lr": LR,
                "base": LR}])
    if mem_adapter is not None:
        groups.append({"params": list(mem_adapter.parameters()), "lr": LR,
                       "base": LR})
    if sym_proj is not None:
        groups.append({"params": list(sym_proj.parameters()), "lr": LR,
                       "base": LR})
    if lora_params:
        groups.append({"params": lora_params, "lr": ARGS.lora_lr,
                       "base": ARGS.lora_lr})
    opt = torch.optim.AdamW(groups, weight_decay=0.01)

    keep = [it for it in data
            if len(tok(it["question"])["input_ids"])
            + len(tok(it["cot"])["input_ids"])
            + it.get("extra_pauses", 0) + 16 <= 460]
    rng.shuffle(keep)
    pool = {l: [it for it in keep if it["label"] == l]
            for l in ("easy", "hard")}
    val = {l: pool[l][:12] for l in pool}
    pool = {l: pool[l][12:] for l in pool}
    print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}",
          flush=True)
    if min(len(v) for v in pool.values()) < BATCH:
        print("INSUFFICIENT POOL — aborting", flush=True)
        return

    log = []
    t0 = time.time()
    for step in range(STEPS):
        lbl = ("easy", "hard")[step % 2]
        batch = rng.sample(pool[lbl], BATCH)
        p = (ARGS.base_pauses if ARGS.base_pauses >= 0
             else P_BY_LABEL[lbl])
        ids, msk, lab, plens = build_batch(tok, batch, p)
        for g in opt.param_groups:
            g["lr"] = g["base"] * lr_at(step) / LR
        ckw, itstates, tfp, frs = {}, None, None, None
        if ARGS.symchain:
            sym_tf = None
            if ARGS.symchain == "tf":
                mat = torch.full((len(batch), ARGS.inner_iters), START_ID,
                                 dtype=torch.long)
                for b_i, it in enumerate(batch):
                    tgt = it.get("lens_targets") or []
                    row = [START_ID] + tgt[:-1]
                    mat[b_i, :len(row)] = torch.tensor(row)
                sym_tf = mat.cuda()
            logits, itstates = sym_burst_forward(
                looper, adapter, sym_proj, ids, msk, plens,
                ARGS.inner_iters, sym_tf, K_PREFILL, lens_teach, EMBED_W,
                START_ID)
        elif ARGS.traj_tf or ARGS.traj_fr or ARGS.kvmem:
            T = (torch.stack([torch.as_tensor(b_["traj_states"])
                              for b_ in batch]).cuda()
                 if (ARGS.traj_tf or ARGS.traj_fr) else None)
            logits, tfp, frs = traj_burst_forward(
                looper, adapter, ids, msk, plens, ARGS.inner_iters, T,
                bool(ARGS.traj_tf), K_PREFILL, mem_adapter=mem_adapter)
        else:
            if ARGS.inner_iters:
                extras = torch.tensor([b_.get("extra_pauses", 0)
                                       for b_ in batch],
                                      device=plens.device)
                assert (extras == extras[0]).all(), \
                    "inner-iters needs a batch-uniform pause block"
                ckw = dict(inner_iters=ARGS.inner_iters,
                           inner_at=plens + p + extras - 1)
                if ARGS.lensteach or ARGS.teachstate:
                    itstates = []
                    ckw["iter_states"] = itstates
            if ARGS.lensteach or ARGS.lensteach_gen or ARGS.inner_iters:
                logits, S = carry_logits(looper, adapter, ids, msk, plens,
                                         K_PREFILL, use_checkpoint=True,
                                         feedforward=ARGS.feedforward,
                                         return_states=True, **ckw)
            else:
                logits = carry_logits(looper, adapter, ids, msk, plens,
                                      K_PREFILL, use_checkpoint=True,
                                      feedforward=ARGS.feedforward)
        loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
                               lab[:, 1:].flatten(), ignore_index=-100)
        lce_val, lgen_val = 0.0, 0.0
        if ARGS.lensteach and ARGS.inner_iters:
            terms = []
            for b, it in enumerate(batch):
                tgt = it.get("lens_targets")
                if not tgt or not itstates:
                    continue
                m = min(len(tgt), len(itstates))
                hs = torch.stack([itstates[i][b] for i in range(m)])
                terms.append(F.cross_entropy(
                    lens_teach(hs).float(),
                    torch.tensor(tgt[:m], device=hs.device)))
            if terms:
                lce = torch.stack(terms).mean()
                lce_val = lce.item()
                loss = loss + ARGS.lensteach * lce
        elif ARGS.lensteach:
            terms = []
            for b, it in enumerate(batch):
                tgt = it.get("lens_targets")
                if not tgt:
                    continue
                s0 = int(plens[b]) + p
                hs = S[b, s0 : s0 + len(tgt)]
                terms.append(F.cross_entropy(
                    lens_teach(hs).float(),
                    torch.tensor(tgt, device=hs.device)))
            if terms:
                lce = torch.stack(terms).mean()
                lce_val = lce.item()
                loss = loss + ARGS.lensteach * lce
        if ARGS.lensteach_gen:
            gterms = []
            for b, it in enumerate(batch):
                base = int(plens[b]) + p + it.get("extra_pauses", 0)
                for pre, res in it.get("gen_targets", []):
                    posl = torch.tensor([base + r for r in pre],
                                        device=S.device)
                    lg = lens_teach(S[b, posl]).float()
                    gterms.append(torch.stack([
                        F.cross_entropy(
                            lg, torch.full((len(pre),), r, dtype=torch.long,
                                           device=S.device))
                        for r in res]).mean())
            if gterms:
                gl = torch.stack(gterms).mean()
                lgen_val = gl.item()
                loss = loss + ARGS.lensteach_gen * gl
        lts_val = 0.0
        if ARGS.teachstate and itstates:
            tterms = []
            for b, it in enumerate(batch):
                T = it.get("teacher_state")
                if T is None:
                    continue
                tterms.append(1 - F.cosine_similarity(
                    itstates[-1][b].float(), T, dim=0))
            if tterms:
                lt = torch.stack(tterms).mean()
                lts_val = lt.item()
                loss = loss + ARGS.teachstate * lt
        ltf_val, lfr_val = 0.0, 0.0
        if ARGS.traj_tf and tfp:
            ltf = torch.stack([
                (1 - F.cosine_similarity(p_.float(), T[:, i], dim=-1)).mean()
                for i, p_ in enumerate(tfp)]).mean()
            ltf_val = ltf.item()
            loss = loss + ARGS.traj_tf * ltf
        if ARGS.traj_fr and frs:
            lfr = torch.stack([
                (1 - F.cosine_similarity(s_.float(), T[:, i], dim=-1)).mean()
                for i, s_ in enumerate(frs)]).mean()
            lfr_val = lfr.item()
            loss = loss + ARGS.traj_fr * lfr
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            [p_ for p_ in adapter.parameters() if p_.requires_grad]
            + lora_params
            + (list(mem_adapter.parameters()) if mem_adapter is not None
               else [])
            + (list(sym_proj.parameters()) if sym_proj is not None
               else []), 1.0)
        opt.step()
        if ARGS.kvmem:
            from kv_memory import arm_memory
            arm_memory(None)
        log.append({"step": step, "loss": loss.item(), "lce": lce_val,
                    "lgen": lgen_val, "lts": lts_val, "ltf": ltf_val,
                    "lfr": lfr_val})
        if step % 10 == 0:
            print(f"step {step:4d}  {lbl:4s} loss={loss.item():.4f}  "
                  f"lce={lce_val:.3f}  lgen={lgen_val:.3f}  "
                  f"lts={lts_val:.3f}  ltf={ltf_val:.3f}  "
                  f"lfr={lfr_val:.3f}  "
                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
        if step % 200 == 199 or step == STEPS - 1:
            if ARGS.lensnoise:
                adapter.noise_on = False
            for l in ("easy", "hard"):
                v = val_loss(looper, adapter, tok, val[l],
                             ARGS.base_pauses
                             if ARGS.base_pauses >= 0
                             else P_BY_LABEL[l])
                print(f"  val@{step}: {l}={v:.3f}", flush=True)
            if ARGS.lensnoise:
                adapter.noise_on = True
            sd = (adapter.base if ARGS.lensnoise else adapter).state_dict()
            torch.save(sd, OUT / f"adapter_{TAG}_e{step+1}.pt")
            if mem_adapter is not None:
                torch.save(mem_adapter.state_dict(),
                           OUT / f"kvmem_{TAG}_e{step+1}.pt")
            if sym_proj is not None:
                torch.save(sym_proj.state_dict(),
                           OUT / f"symproj_{TAG}_e{step+1}.pt")
            if lora_params:
                torch.save({"rank": ARGS.bandlora, "band": lora_band_layers,
                            "tensors": [p.detach().cpu()
                                        for p in lora_params]},
                           OUT / f"lora_{TAG}_e{step+1}.pt")
            json.dump(log, open(OUT / f"train_{TAG}_log.json", "w"))
    print("done", flush=True)


if __name__ == "__main__":
    main()
