"""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,
                         NoisyMergeAdapter, ParcaeAdapter, PerDepthAdapter,
                         RecurrentAdapter, TiedAlphaAdapter)
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)
ap.add_argument("--warm", default=None, help="warm-start adapter checkpoint")
ap.add_argument("--rec", action="store_true",
                help="Huginn-style regime: RecurrentAdapter (learned A/B, "
                     "noise h0) + log-uniform random depth 1..recmax, "
                     "truncated bptt (default 4)")
ap.add_argument("--recmax", type=int, default=16)
ap.add_argument("--parcae", action="store_true",
                help="rec regime with rho(A)<1 by construction "
                     "(diag-negative-exp ZOH parameterization)")
ap.add_argument("--randk", action="store_true",
                help="factorial cell: standard MergeAdapter but rec-style "
                     "log-uniform random depth (isolates the curriculum)")
ap.add_argument("--noises0", action="store_true",
                help="factorial cell: standard merge + noise s0")
ap.add_argument("--hidden", type=int, default=512,
                help="adapter MLP hidden width (capacity control)")
ap.add_argument("--tiedalpha", action="store_true",
                help="merge with learned per-dim constant alpha, B tied "
                     "to (1-a); standard curriculum")
ap.add_argument("--perdepth", type=int, default=0,
                help="PerDepthAdapter: one merge per iteration depth "
                     "(Bae-style relaxation), standard curriculum")
ARGS, _ = ap.parse_known_args()
if ARGS.parcae:
    ARGS.rec = True
if (ARGS.rec or ARGS.randk) and not ARGS.bptt:
    ARGS.bptt = 4
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 "")
          + ("_warm" if ARGS.warm else "")
          + (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax)
             if ARGS.rec else "")
          + (f"_pd{ARGS.perdepth}" if ARGS.perdepth else "")
          + ("_ta" if ARGS.tiedalpha else "")
          + (f"_rk{ARGS.recmax}" if ARGS.randk else "")
          + ("_ns" if ARGS.noises0 else "")
          + (f"_h{ARGS.hidden}" if ARGS.hidden != 512 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
    H = ARGS.hidden
    if ARGS.tiedalpha:
        adapter = TiedAlphaAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
    elif ARGS.noises0:
        adapter = NoisyMergeAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
    elif ARGS.perdepth:
        adapter = PerDepthAdapter(d=d, alpha=ARGS.alpha,
                                  n_depth=ARGS.perdepth).cuda()
    elif ARGS.parcae:
        adapter = ParcaeAdapter(d=d, alpha=ARGS.alpha).cuda()
    elif ARGS.rec:
        adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()
    elif ARGS.adaptive:
        adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
    else:
        adapter = MergeAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
    if ARGS.warm:
        adapter.load_state_dict(torch.load(ARGS.warm, map_location="cuda"))
        print("warm-started from", ARGS.warm, flush=True)
    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):
        if ARGS.rec or ARGS.randk:
            # randomized depth, log-uniform in [1, recmax], any difficulty
            k = min(ARGS.recmax,
                    max(1, int(math.exp(rng.uniform(0, math.log(ARGS.recmax))))))
            labels = ("easy", "hard")
        else:
            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:
            rho = (f"  rho(A)={adapter.rho():.3f}"
                   if hasattr(adapter, "rho") else "")
            print(f"step {step:4d}  k={k}  loss={loss.item():.4f}{rho}  "
                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
            if rho:
                log.append({"step": step, "rho": adapter.rho()})
        if step % 100 == 99 or step == STEPS - 1:
            vals = {}
            for kk in ((0, 1, 2, 4, 8, 16) if ARGS.rec else (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()
