"""E1b (item 19): label-supervised halting head on a FROZEN trained merge.

The merge adapter (curriculum ★, adapter_code.pt) is loaded and frozen;
only the halting head trains. Targets from STaR labels: easy -> halt at
iteration 1, hard -> halt at iteration 4 (BCE at every iteration: 0 below
the target depth, 1 at/above). Mixed batches, no CE, no penalty — the
head is a depth-aware difficulty classifier on the loop trajectory.
"""

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

import torch
import torch.nn.functional as F

from halting_common import HaltingMergeAdapter
from loop_common import BandLooper
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
from train_merge_code import build_code_batch, MAX_TOK

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"))

ap = argparse.ArgumentParser()
ap.add_argument("--merge", default=None,
                help="frozen merge checkpoint (default results-loop/adapter_code.pt)")
ap.add_argument("--steps", type=int, default=300)
ap.add_argument("--kmax", type=int, default=4)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--lr", type=float, default=3e-3)
ARGS = ap.parse_args()
BATCH = 4
TARGET_K = {"easy": 1, "hard": ARGS.kmax}


def main():
    rng = random.Random(ARGS.seed)
    torch.manual_seed(ARGS.seed)
    data = json.load(open(OUT / "mbpp_data.json"))
    model, tok = load_model(dtype=torch.bfloat16)
    for p in model.parameters():
        p.requires_grad_(False)
    looper = BandLooper(model)
    adapter = HaltingMergeAdapter(
        d=model.config.get_text_config().hidden_size).cuda()
    merge_ckpt = ARGS.merge or (OUT / "adapter_code.pt")
    sd = torch.load(merge_ckpt, map_location="cuda")
    adapter.load_state_dict(sd, strict=False)  # merge weights; head stays init
    for n, p in adapter.named_parameters():
        p.requires_grad_(n.startswith("halt"))
    opt = torch.optim.AdamW([p for p in adapter.parameters()
                             if p.requires_grad], lr=ARGS.lr)
    print("trainable:", sum(p.numel() for p in adapter.parameters()
                            if p.requires_grad), flush=True)

    train = [it for it in data if it["split"] == "train"
             and it["label"] != "drop" and it["sol_code"]]
    train = [it for it in train
             if len(tok(mbpp_prompt(tok, it, DIRECT_SUFFIX))["input_ids"])
             + len(tok(it["sol_code"])["input_ids"]) + 12 <= MAX_TOK]
    print(f"pool={len(train)}", flush=True)

    log = []
    t0 = time.time()
    for step in range(ARGS.steps):
        batch = rng.sample(train, BATCH)
        ids, msk, lab, lmask = build_code_batch(tok, batch)
        pl = (lmask.long().cumsum(-1).argmax(-1))
        bidx = torch.arange(len(batch), device="cuda")
        tk = torch.tensor([TARGET_K[it["label"]] for it in batch],
                          device="cuda")
        with torch.no_grad():
            calls, _ = looper.capture(ids, msk, logits_to_keep=1)
            e = looper._hin[looper.l0].detach()
            s = looper.band(e, calls)
        loss = 0.0
        for i in range(ARGS.kmax):
            with torch.no_grad():
                x = adapter(e, s)
                x = torch.where(lmask[..., None], x, e)
                s = looper.band(x, calls)
            p = adapter.halt_prob(e[bidx, pl], s[bidx, pl])
            tgt = ((i + 1) >= tk).float()
            loss = loss + F.binary_cross_entropy(p.float(), tgt)
        loss = loss / ARGS.kmax
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        log.append({"step": step, "bce": loss.item()})
        if step % 20 == 0:
            print(f"step {step:4d}  bce={loss.item():.4f}  "
                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
    torch.save(adapter.state_dict(), OUT / f"adapter_gatehead_e{ARGS.steps}.pt")
    json.dump(log, open(OUT / "train_gatehead_log.json", "w"))
    print("done", flush=True)


if __name__ == "__main__":
    main()