"""RLCD fine-tuning loop for Laya System One models.

The published Laya repository ships the building blocks but no training loop
(checked 2026-09-20: no `backward`, no optimizer in any of its .py files).
This loop implements the method described in the Laya README:

  "The policy reports a distribution; exploration adds zero-mean Gaussian noise
   to the logits; the reward is a strictly proper scoring rule (log + spherical,
   plus ranked probability score for ordinal questions). Expected reward is
   maximised only by reporting honest probabilities. Updates are REINFORCE with
   a group-mean baseline (GRPO-style)."

The gradient: the policy is a Gaussian with mean = logits(theta) and fixed sigma.
For a sample a = logits.detach() + eps,

    grad_theta log pi(a) = (eps / sigma^2) * grad_theta logits,

hence the term (eps.detach() * logits).sum() / sigma^2 in the loss. The baseline
is the mean reward of the G samples of the SAME item (GRPO), not a learned critic.

Usage:
    LAYA_SNAP=/path/to/laya/snapshot python train.py \
        --data train.jsonl --gold gold.jsonl --questions questions.json \
        --n 20000 --batch 16 --group 32 --stage 400 --save out/
"""
import argparse, json, os, random, sys, time
import numpy as np, torch

LAYA_REPO = "convaiinnovations/laya"

SNAP = os.environ.get("LAYA_SNAP")
if not SNAP:
    # Fetch the base model on first use. Set LAYA_SNAP to point at a local copy
    # instead -- on an air-gapped machine, or to train from your own checkpoint.
    from huggingface_hub import snapshot_download
    print("fetching %s ..." % LAYA_REPO, flush=True)
    SNAP = snapshot_download(LAYA_REPO)
sys.path.insert(0, SNAP)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from rl_common import (QTYPES, build_model, collate_items, encode_record,  # noqa: E402
                       proper_reward, auroc, spearman, seed_all)
import data as data_mod                                                    # noqa: E402


def load_model(variant="multilingual", device="cuda"):
    from safetensors.torch import load_file
    base = os.path.join(SNAP, variant) if variant else SNAP
    cfg = json.load(open(os.path.join(base, "rl_agent_config.json")))
    m = build_model(cfg, encoder_dir=os.path.join(base, "encoder"))
    sd = load_file(os.path.join(base, "model.safetensors"))
    missing, _ = m.load_state_dict(sd, strict=False)
    if missing:
        print("  missing weights:", len(missing), missing[:3])
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(os.path.join(base, "tokenizer"))
    return m.to(device), tok, cfg


def batches(recs, tok, cfg, rng, train, size):
    groups = []
    for r in recs:
        g = encode_record(r, tok, cfg, rng, train)
        for it in g:            # encode_record only knows q_index; carry the name
            qi = it.get("q_index")
            if qi is not None and qi < len(r["qs"]):
                it["question"] = r["qs"][qi]["name"]
        if g:
            groups.append(g)
    for i in range(0, len(groups), size):
        b = collate_items(groups[i:i + size], tok.pad_token_id)
        if b:
            yield b


def to_device(b, dev):
    return {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in b.items()}


@torch.no_grad()
def evaluate(m, recs, tok, cfg, dev, size=8):
    """Score against the gold set, PER QUESTION.

    Averaging over the question type hides everything: on our news corpus the
    four ordinal questions ranged from 0.54 (urgency) to 0.74 (sentiment) even
    for the teacher model. Reporting a single "score" number tells you neither
    which question is lagging nor where the gold set needs more examples.
    """
    m.eval()
    per = {}
    for b in batches(recs, tok, cfg, None, False, size):
        bd = to_device(b, dev)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits, _ = m(bd["input_ids"], bd["attention_mask"], bd["marker_pos"],
                          bd["marker_mask"], bd["qtype"])
        pr = torch.softmax(logits.float(), -1) * bd["marker_mask"]
        pick = pr.argmax(-1).cpu().numpy()
        want = b["target"].argmax(-1).numpy()
        p1 = pr[:, 1].cpu().numpy()
        for i, meta in enumerate(b["meta"]):
            n = meta.get("question", "?")
            d = per.setdefault(n, {"t": int(b["qtype"][i]), "pick": [], "want": [], "p": []})
            d["pick"].append(int(pick[i])); d["want"].append(int(want[i]))
            d["p"].append(float(p1[i]))
    m.train()
    out = {}
    for n, d in per.items():
        w, z = np.array(d["pick"]), np.array(d["want"])
        e = {"n": len(w), "acc": round(float(np.mean(w == z)), 4)}
        if d["t"] == QTYPES["score"]:
            e["pm1"] = round(float(np.mean(np.abs(w - z) <= 1)), 4)
            e["rho"] = round(spearman(w.astype(float), z.astype(float)), 4)
        if d["t"] == QTYPES["noul"] and len(set(z.tolist())) > 1:
            e["auc"] = round(auroc(np.array(d["p"]), z.astype(float)), 4)
        out[n] = e
    return out


def train(a):
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    seed_all(a.seed); rng = random.Random(a.seed)
    print("loading model (%s) ..." % a.variant, flush=True)
    m, tok, cfg = load_model(a.variant, dev)
    if a.chars:
        cfg["max_len"] = a.chars

    questions = data_mod.load_questions(a.questions)
    recs = data_mod.records(a.data, questions, max_chars=a.max_chars)[:a.n]
    gold = data_mod.records(a.gold, questions, max_chars=a.max_chars)
    gold_ids = {r["src"] for r in gold}
    recs = [r for r in recs if r["src"] not in gold_ids]   # never train on gold
    print("training %d records | gold %d (disjoint)" % (len(recs), len(gold)), flush=True)

    # --init is NOT --resume. Resuming continues an interrupted run on its own
    # schedule; init starts a NEW run from trained weights, with a fresh optimizer
    # and a fresh learning-rate curve. Use it to add a question, or to retrain one
    # question on a rebalanced set without starting from the base model.
    #
    # This has to happen BEFORE the baseline measurement. Loading the weights
    # after it prints a "before" score for the untrained model -- which is not
    # the number you are comparing against, and the run looks like a bigger win
    # than it was.
    if a.init:
        st = torch.load(a.init, map_location=dev, weights_only=False)
        m.load_state_dict(st["model"] if "model" in st else st)
        print("  initialised from %s (fresh optimizer and schedule)" % a.init, flush=True)

    print("BEFORE:", json.dumps(evaluate(m, gold, tok, cfg, dev)), flush=True)
    if a.eval_only:
        return

    # Memory. A 12 GB consumer card cannot hold AdamW over 322M parameters:
    # params + grads + two moments is roughly 5.2 GB before activations, and a
    # full fine-tune OOMs at batch 4. --freeze is the way out on such a card;
    # on an 80 GB accelerator leave it at 0.
    m.head_checkpointing = True
    if hasattr(m.encoder, "gradient_checkpointing_enable"):
        m.encoder.gradient_checkpointing_enable()
    if a.freeze:
        # Freeze the lower encoder layers: they carry general language, the
        # task-specific behaviour sits on top. Saves gradients AND optimizer state.
        for name, par in m.encoder.named_parameters():
            part = name.split(".")
            nr = int(part[1]) if len(part) > 1 and part[0] == "layers" and part[1].isdigit() else -1
            if 0 <= nr < a.freeze or name.startswith("embeddings"):
                par.requires_grad = False
        free = sum(p.numel() for p in m.parameters() if p.requires_grad)
        print("  frozen layers 0..%d + embeddings | trainable %.0f M"
              % (a.freeze - 1, free / 1e6), flush=True)

    opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad],
                            lr=a.lr, weight_decay=0.01)
    # Learning-rate decay is not optional here. Without it REINFORCE keeps
    # jumping until the very last step: measured on a 1200-step run, one
    # question oscillated between 0.653 and 0.756 without settling, and the
    # final checkpoint scored 4.6 points below the best one.
    total = a.max_steps or (len(recs) // max(a.batch, 1)) * max(a.epochs, 1)
    plan = torch.optim.lr_scheduler.OneCycleLR(
        opt, max_lr=a.lr, total_steps=max(total, 2),
        pct_start=0.1, anneal_strategy="cos", div_factor=10, final_div_factor=100)
    print("  lr %.1e, 10%% warmup, then cosine decay over %d steps" % (a.lr, total),
          flush=True)

    step = 0; t0 = time.time(); rewards = []; best = (0.0, 0); last_push = 0.0

    # RESUMING after an interruption. Weights alone are NOT enough: AdamW carries
    # two running moments per parameter, and OneCycleLR has to know where on its
    # curve it stopped. Without both it is a warm start with a fresh warmup cycle,
    # not a continuation -- the learning rate jumps back up and shakes the weights.
    src = a.resume or os.environ.get("LAYA_RESUME")
    if src and os.path.exists(src):
        st = torch.load(src, map_location=dev, weights_only=False)
        m.load_state_dict(st["model"]); opt.load_state_dict(st["optimizer"])
        plan.load_state_dict(st["plan"])
        step = st["step"]; best = tuple(st["best"])
        print("  RESUMED at step %d (best so far %.4f at %d)" % (step, best[0], best[1]),
              flush=True)

    for _ in range(a.epochs):
        rng.shuffle(recs)
        for b in batches(recs, tok, cfg, rng, True, a.batch):
            b = to_device(b, dev)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                logits, _ = m(b["input_ids"], b["attention_mask"], b["marker_pos"],
                              b["marker_mask"], b["qtype"])
            logits = logits.float()
            mask = b["marker_mask"]
            # G samples per item; noise only on the valid options
            eps = torch.randn((a.group,) + logits.shape, device=dev) * a.sigma
            eps = eps * mask
            q = torch.softmax(logits.detach()[None] + eps, -1) * mask
            q = q / q.sum(-1, keepdim=True).clamp_min(1e-9)
            r = proper_reward(q, b["target"], b["qtype"], mask)      # [G, N]
            adv = r - r.mean(0, keepdim=True)                        # GRPO baseline
            adv = adv / adv.std(0, keepdim=True).clamp_min(1e-6)
            score = (eps.detach() * logits[None]).sum(-1) / (a.sigma ** 2)
            loss = -(adv.detach() * score).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
            opt.step()
            if step + 1 < plan.total_steps:
                plan.step()
            rewards.append(float(r.mean())); step += 1

            if a.stage and step % a.stage == 0:
                row = {"step": step, "seen": step * a.batch,
                       "reward": round(float(np.mean(rewards[-a.stage:])), 4),
                       "lr": round(plan.get_last_lr()[0], 8),
                       "gold": evaluate(m, gold, tok, cfg, dev)}
                print("  STAGE " + json.dumps(row), flush=True)
                if a.history:
                    with open(a.history, "a") as fh:
                        fh.write(json.dumps(row) + "\n")
                mean = float(np.mean([v["acc"] for v in row["gold"].values()]))

                # Two different things, and conflating them costs you a run:
                #
                # checkpoint.pt is the LATEST state, written every stage. It is
                # what --resume reads, so an interruption costs one stage at most.
                # (Writing it only on improvement looks harmless and is not: a run
                # whose best score falls early would resume hours behind.)
                #
                # model-best.pt is the BEST state. REINFORCE oscillates, so the
                # final weights are a lottery ticket -- up to 6 points between
                # adjacent stages, and one of our runs ended 4.6 points below its
                # own best. That file is the one you deploy.
                improved = mean > best[0]
                if a.save and improved:
                    best = (mean, step)
                if a.save:
                    os.makedirs(a.save, exist_ok=True)
                    torch.save({"model": m.state_dict(), "optimizer": opt.state_dict(),
                                "plan": plan.state_dict(), "step": step,
                                "best": list(best), "seen": step * a.batch},
                               os.path.join(a.save, "checkpoint.pt"))
                if a.save and improved:
                    torch.save(m.state_dict(), os.path.join(a.save, "model-best.pt"))
                    with open(os.path.join(a.save, "best.json"), "w") as fh:
                        json.dump({"step": step, "gold_mean": round(mean, 4),
                                   "gold": row["gold"]}, fh, indent=1)
                    print("    -> best so far (gold mean %.4f)" % mean, flush=True)
                    # Push intermediate results: a multi-hour run must not lose
                    # everything to a timeout. Throttled, because every upload
                    # runs on the accelerator's clock.
                    dest = os.environ.get("LAYA_PUSH")
                    if dest and time.time() - last_push < a.push_every:
                        dest = None
                    if dest:
                        last_push = time.time()
                        try:
                            from huggingface_hub import HfApi
                            api = HfApi(); tk = os.environ.get("HF_TOKEN")
                            api.create_repo(dest, private=True, exist_ok=True, token=tk)
                            api.upload_folder(folder_path=a.save, repo_id=dest, token=tk)
                            print("    -> pushed to %s (step %d)" % (dest, step), flush=True)
                        except Exception as e:
                            print("    !! push failed: %s" % e, flush=True)

            if step % a.report == 0:
                print("  step %5d  reward %+.4f  %.1f s/100"
                      % (step, float(np.mean(rewards[-a.report:])),
                         (time.time() - t0) / max(step, 1) * 100), flush=True)
            if a.max_steps and step >= a.max_steps:
                break
        if a.max_steps and step >= a.max_steps:
            break

    print("AFTER:", json.dumps(evaluate(m, gold, tok, cfg, dev)), flush=True)
    if a.save:
        os.makedirs(a.save, exist_ok=True)
        torch.save(m.state_dict(), os.path.join(a.save, "model-last.pt"))
        print("saved to %s (best at step %d, gold mean %.4f)" % (a.save, best[1], best[0]))


def cli():
    p = argparse.ArgumentParser()
    p.add_argument("--data", required=True, help="training records, one JSON object per line")
    p.add_argument("--gold", required=True, help="held-out records, same format")
    p.add_argument("--questions", required=True, help="question definitions (JSON)")
    p.add_argument("--n", type=int, default=20000, help="cap on training records")
    p.add_argument("--batch", type=int, default=16)
    p.add_argument("--group", type=int, default=32, help="G samples per item (GRPO baseline)")
    p.add_argument("--sigma", type=float, default=1.0, help="exploration noise on the logits")
    p.add_argument("--lr", type=float, default=1e-5)
    p.add_argument("--epochs", type=int, default=1)
    p.add_argument("--max-steps", type=int, default=0, dest="max_steps")
    p.add_argument("--stage", type=int, default=400, help="evaluate every N steps")
    p.add_argument("--report", type=int, default=200)
    p.add_argument("--variant", default="multilingual", help="subfolder of the Laya snapshot")
    p.add_argument("--chars", type=int, default=0, help="override the model's max_len (tokens)")
    p.add_argument("--max-chars", type=int, default=data_mod.MAX_CHARS, dest="max_chars",
                   help="cut each text to this many characters before encoding")
    p.add_argument("--freeze", type=int, default=0, help="freeze the lowest N encoder layers")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--save", default="", help="output directory")
    p.add_argument("--resume", default="", help="checkpoint.pt to continue an interrupted run")
    p.add_argument("--init", default="", help="weights to start a NEW run from "
                                              "(fresh optimizer and schedule)")
    p.add_argument("--history", default="", help="append per-stage metrics as JSONL")
    p.add_argument("--push-every", type=int, default=1800, dest="push_every",
                   help="minimum seconds between hub pushes (LAYA_PUSH)")
    p.add_argument("--eval-only", action="store_true", dest="eval_only")
    return p.parse_args()


if __name__ == "__main__":
    train(cli())