"""Train ONLY the merge adapter at the L13->L14 boundary (WORKSPACE_LOOPING.md §3).

Frozen base model; band L14-30 unrolled k times through the adapter; loss =
cross-entropy on answer tokens of the *direct* (no-CoT) prompt, supervised by
gold answers on items the frozen model can reach (STaR filter from
prep_star_data.py). Difficulty->depth curriculum:
    k=1: easy items    k=2: easy+hard    k=4: hard items
so hard items are only ever seen at depth, forcing the loop to be used.
"""

import json
import random
import time
from pathlib import Path

import torch
import torch.nn.functional as F

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

OUT = Path(__file__).resolve().parent.parent / "results-loop"
STEPS = 800
BATCH = 8
LR = 1e-3
WARMUP = 20
MAX_TOK = 400          # skip overlong items
K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
SEED = 0


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


@torch.no_grad()
def val_answer_acc(looper, adapter, tok, items, k):
    """Teacher-forced exact match: argmax at every answer position."""
    ok = 0
    for i in range(0, len(items), BATCH):
        ids, msk, lab = build_train_batch(tok, items[i : i + BATCH])
        logits = looper.loop_logits(adapter, ids, k, attention_mask=msk)
        pred = logits[:, :-1].argmax(-1)
        tgt = lab[:, 1:]
        m = tgt != -100
        ok += (((pred == tgt) | ~m).all(-1)).sum().item()
    return ok / len(items)


def main():
    rng = random.Random(SEED)
    torch.manual_seed(SEED)
    data = json.load(open(OUT / "star_data.json"))

    model, tok = load_model(dtype=torch.bfloat16)
    for p in model.parameters():
        p.requires_grad_(False)
    looper = BandLooper(model)
    adapter = MergeAdapter().cuda()
    opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)

    train = [it for it in data if it["split"] == "train" and it["label"] != "drop"]
    # length filter
    keep = []
    for it in train:
        n = len(tok(it["question"])["input_ids"])
        if n + 40 <= MAX_TOK:
            keep.append(it)
    train = keep
    pool = {"easy": [it for it in train if it["label"] == "easy"],
            "hard": [it for it in train if it["label"] == "hard"]}
    val = {lbl: rng.sample(lst, min(32, len(lst))) for lbl, lst in pool.items()}
    for lbl in pool:
        pool[lbl] = [it for it in pool[lbl] if it not in val[lbl]]
    print(f"train pool: easy={len(pool['easy'])} hard={len(pool['hard'])}", flush=True)

    log = []
    t0 = time.time()
    for step in range(STEPS):
        k, labels = K_BUCKETS[step % len(K_BUCKETS)]
        cand = [it for lbl in labels for it in pool[lbl]]
        batch = rng.sample(cand, BATCH)
        ids, msk, lab = build_train_batch(tok, batch)

        for g in opt.param_groups:
            g["lr"] = lr_at(step)
        logits = looper.loop_logits(adapter, ids, k, attention_mask=msk,
                                    use_checkpoint=True)
        loss = F.cross_entropy(
            logits[:, :-1].flatten(0, 1).float(), lab[:, 1:].flatten(),
            ignore_index=-100)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
        opt.step()

        log.append({"step": step, "k": k, "loss": loss.item()})
        if step % 10 == 0:
            print(f"step {step:4d}  k={k}  loss={loss.item():.4f}  "
                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
        if step % 100 == 99 or step == STEPS - 1:
            accs = {}
            for kk in (0, 1, 2, 4):
                accs[f"easy_k{kk}"] = val_answer_acc(looper, adapter, tok,
                                                     val["easy"], kk)
                accs[f"hard_k{kk}"] = val_answer_acc(looper, adapter, tok,
                                                     val["hard"], kk)
            print(f"  val@{step}: " +
                  "  ".join(f"{n}={v:.2f}" for n, v in accs.items()), flush=True)
            log.append({"step": step, "val": accs})
            torch.save(adapter.state_dict(), OUT / "adapter.pt")
            json.dump(log, open(OUT / "train_log.json", "w"), indent=1)

    torch.save(adapter.state_dict(), OUT / "adapter.pt")
    json.dump(log, open(OUT / "train_log.json", "w"), indent=1)
    print("done; adapter ->", OUT / "adapter.pt")


if __name__ == "__main__":
    main()