"""Train the merge adapter on MBPP with prompt-only ("latent planning") looping.

Same recipe as train_merge.py, with two changes:
  - loop_mask: the merge applies only to prompt positions; the code tokens are
    teacher-forced through the plain band (they still attend to looped prompt
    states each iteration) — no exposure-bias gap by construction.
  - supervision: the model's OWN passing code from prep_mbpp.py (easy: direct
    pass; hard: code extracted from the plan-first pass), CE on code tokens.
"""

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 loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt

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
WARMUP = 20
MAX_TOK = 512
K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
# --deepk 16 rescales to [(2, easy), (8, mixed), (16, hard)]

ap = argparse.ArgumentParser()
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--pause", type=int, default=0,
                help="pause-token control: p inert tokens after prompt, "
                     "feedforward adapter, no recurrence")
ap.add_argument("--feedforward", action="store_true",
                help="apply adapter once, no recurrence (pause-FF control)")
ap.add_argument("--alpha", type=float, default=0.3,
                help="merge weight (2B-tuned default 0.3; try 0.1-0.15 at 12B)")
ap.add_argument("--adaptive", action="store_true",
                help="state-dependent alpha (AdaptiveMergeAdapter)")
ap.add_argument("--deepk", type=int, default=0,
                help="scale curriculum depths by deepk/4 (e.g. 16 -> 2/8/16)")
ap.add_argument("--bptt", type=int, default=0,
                help="truncated BPTT: grads only through last N iterations")
ap.add_argument("--lr", type=float, default=1e-3)
ARGS, _ = ap.parse_known_args()
SEED = ARGS.seed
LR = ARGS.lr
SUFFIX = ((f"_s{SEED}" if SEED else "")
          + (f"_p{ARGS.pause}" if ARGS.pause else "")
          + (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else "")
          + ("_ad" if ARGS.adaptive else "")
          + (f"_dk{ARGS.deepk}" if ARGS.deepk else "")
          + (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else ""))
PAUSE_ID = 6  # <unused0>


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_code_batch(tok, items, device="cuda"):
    """Right-padded batch; labels on code tokens; loop_mask on prompt span."""
    seqs, labs = [], []
    for it in items:
        p = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX),
                add_special_tokens=False)["input_ids"]
        p = p + [PAUSE_ID] * ARGS.pause
        a = tok("```python\n" + it["sol_code"] + "\n```<end_of_turn>",
                add_special_tokens=False)["input_ids"]
        seqs.append(p + a)
        labs.append([-100] * len(p) + a)
    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
    lmask = (lab == -100) & (msk == 1)
    return (ids.to(device), msk.to(device), lab.to(device), lmask.to(device))


@torch.no_grad()
def val_loss(looper, adapter, tok, items, k):
    tot, n = 0.0, 0
    for i in range(0, len(items), BATCH):
        ids, msk, lab, lmask = build_code_batch(tok, items[i : i + BATCH])
        logits = looper.loop_logits(adapter, ids, k, attention_mask=msk,
                                    loop_mask=lmask, feedforward=ARGS.feedforward)
        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(SEED)
    torch.manual_seed(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)
    d = model.config.get_text_config().hidden_size
    if ARGS.adaptive:
        adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
    else:
        adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda()
    global K_BUCKETS
    if ARGS.deepk:
        f = ARGS.deepk / 4
        K_BUCKETS = [(max(1, int(k * f)), lbls) for k, lbls in K_BUCKETS]
        print("K_BUCKETS ->", [(k, l) for k, l in K_BUCKETS], flush=True)
    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" 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]
    pool = {"easy": [it for it in train if it["label"] == "easy"],
            "hard": [it for it in train if it["label"] == "hard"]}
    val = {lbl: lst[:16] for lbl, lst in pool.items()}
    for lbl in pool:
        pool[lbl] = pool[lbl][16:]
    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, min(BATCH, len(cand)))
        ids, msk, lab, lmask = build_code_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, loop_mask=lmask,
                                    feedforward=ARGS.feedforward,
                                    bptt=ARGS.bptt or None)
        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:
            vals = {}
            for kk in (0, 1, 2, 4):
                vals[f"easy_k{kk}"] = val_loss(looper, adapter, tok,
                                               val["easy"], kk)
                vals[f"hard_k{kk}"] = val_loss(looper, adapter, tok,
                                               val["hard"], kk)
            print(f"  val@{step}: " +
                  "  ".join(f"{n}={v:.3f}" for n, v in vals.items()), flush=True)
            log.append({"step": step, "val": vals})
            torch.save(adapter.state_dict(),
                       OUT / f"adapter_code{SUFFIX}_e{step+1}.pt")
            json.dump(log, open(OUT / f"train_code_log{SUFFIX}.json", "w"), indent=1)

    torch.save(adapter.state_dict(), OUT / f"adapter_code{SUFFIX}.pt")
    json.dump(log, open(OUT / f"train_code_log{SUFFIX}.json", "w"), indent=1)
    print("done; adapter ->", OUT / f"adapter_code{SUFFIX}.pt")


if __name__ == "__main__":
    main()