"""deg-MoE trainer: the W (writer) experts of a task-routed MoE (results/deg/moe-v1).

One model, one context prefill: every decoder layer of a full fine-tune deg checkpoint
(deg2b-1) gets a second FFN expert ``mlp_w`` initialised from the base model's MLP, plus
an optional writer LoRA on the attention / linear-attention projections
(:func:`deg.model.add_moe`). Routing is by token role: the context and every
decide/extract/count/anchor fork run the reader (R = the checkpoint as it is); the KNOWN
VALUES block, the generate header and the generated tokens run W and read the R trunk's
K/V and recurrent state (:class:`deg.model.WriteRoute`).

Training freezes everything R had (embeddings, attention, norms, LM head, R experts,
heads), so decide/extract outputs are unchanged by construction, and trains only the W
experts (+ writer LoRA) with the x-generate LM loss in the base layout
(:func:`deg.infer.base_generate_header`; the same items as :mod:`deg.gentrain`:
gen-v1 judge-kept fields, open-weight writer and judge only, suite guard on) plus
``--kl-base W`` * KL(base || MoE) on the same target positions, where the base Qwen
re-reads the whole sequence (the base copy exists in training only).

Each micro-batch reads every distinct context ONCE through R without autograd (R is
frozen and no W parameter touches context tokens); the W rows (KNOWN VALUES + header +
target) continue that trunk state, gradient-checkpointed. The fp32 master weights and
the AdamW moments live on ``--optim-device`` (default cpu: 12 bytes per W parameter
would not fit beside the bodies on one 24 GB L4); bf16 gradients are copied to pinned
host buffers once per optimizer step. The best ``val`` NLL checkpoint is saved as one
deg-MoE checkpoint directory (backbone.safetensors = R + shared, moe.safetensors = W
experts + writer LoRA, heads, deg.json ``moe``), loadable by
:func:`deg.model.load_model`, the server and the suite (default ``generate`` = moe).

    CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=<GPU UUID, never GPU 0> \\
    python -m deg.moetrain --ckpt /data/deg/runs/deg2b-1/step-0004000 \\
        --out /data/deg/runs/deg2b-moe-1 --kl-base 0.5
"""

from __future__ import annotations

import argparse
import json
import math
import os
import random
import time
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F

from deg.gentrain import GEN_RECORDS, GenItem, gen_items, lr_at
from deg.model import (
    DegModel,
    add_moe,
    attach_base,
    load_base_lm,
    load_model,
    moe_params,
)
from deg.train import read_jsonl

USER_GPU_UUID = "GPU-019676ed-23c9-ad9a-20cb-7bdc13ac61ca"  # GPU index 0: the user's


def check_device(dev: torch.device) -> None:
    """CUDA runs must name their GPU (CUDA_VISIBLE_DEVICES = one UUID), never GPU 0."""
    if dev.type != "cuda":
        return
    cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
    if not cvd or cvd == "0" or "," in cvd or USER_GPU_UUID.startswith(cvd):
        raise SystemExit("set CUDA_VISIBLE_DEVICES=<one GPU UUID> (GPU 0 is the user's; "
                         f"got {cvd!r})")


def micro_batches(items: list[GenItem], batch_tokens: int, rng: random.Random | None
                  ) -> list[list[GenItem]]:
    """Groups of whole records (a record's fields share its context read), length-sorted,
    at most ``batch_tokens`` padded tokens counted as contexts x longest context + rows x
    longest W part (a record over budget is a group alone); shuffled order with ``rng``."""
    by: dict[str, list[GenItem]] = {}
    for it in items:
        by.setdefault(it.uid, []).append(it)
    recs = sorted(by.values(), key=lambda g: (g[0].ctx_len, max(len(x.ids) for x in g)))
    groups: list[list[GenItem]] = []
    cur: list[list[GenItem]] = []

    def cost(rs: list[list[GenItem]]) -> int:
        ctx = max(g[0].ctx_len for g in rs)
        rows = [x for g in rs for x in g]
        return len(rs) * ctx + len(rows) * max(len(x.ids) - x.ctx_len for x in rows)

    for g in recs:
        if cur and cost(cur + [g]) > batch_tokens:
            groups.append([x for r in cur for x in r])
            cur = []
        cur.append(g)
    if cur:
        groups.append([x for r in cur for x in r])
    if rng is not None:
        rng.shuffle(groups)
    return groups


def _pack(group: list[GenItem], pad: int, dev: torch.device):
    """(ctx ids [n, Lc], ctx lengths, W ids [rows, Lw], W lengths, row -> ctx row,
    target row / column / token): column c of row r predicts W token c + 1."""
    uids: list[str] = []
    for x in group:
        if x.uid not in uids:
            uids.append(x.uid)
    first = {x.uid: x for x in reversed(group)}
    ctxs = [first[u].ids[:first[u].ctx_len] for u in uids]
    for x in group:
        assert x.ids[:x.ctx_len] == ctxs[uids.index(x.uid)], f"{x.uid}: contexts differ"
    Lc = max(len(c) for c in ctxs)
    cids = torch.full((len(ctxs), Lc), pad, dtype=torch.long)
    for i, c in enumerate(ctxs):
        cids[i, :len(c)] = torch.tensor(c)
    ws = [x.ids[x.ctx_len:] for x in group]
    Lw = max(len(w) for w in ws)
    wids = torch.full((len(ws), Lw), pad, dtype=torch.long)
    rows, cols, tgts = [], [], []
    for r, (x, w) in enumerate(zip(group, ws)):
        wids[r, :len(w)] = torch.tensor(w)
        for p in range(x.prompt_len - x.ctx_len, len(w)):
            rows.append(r)
            cols.append(p - 1)
            tgts.append(w[p])
    t = lambda v: torch.tensor(v, dtype=torch.long, device=dev)  # noqa: E731
    return (cids.to(dev), t([len(c) for c in ctxs]), wids.to(dev), t([len(w) for w in ws]),
            t([uids.index(x.uid) for x in group]), t(rows), t(cols), t(tgts))


def _target_logits(bb, cids, clens, wids, wlens, row_of, ri, ci, sel):
    """Logits at the target positions: the contexts read once without autograd (R; the
    base copy for the KL anchor), then the W rows continue them with route ``sel``."""
    with torch.no_grad():
        _, state = bb(cids, clens)
    h, _ = bb(wids, wlens, state, rows=row_of, sel=sel)
    del state
    return bb.lm_logits(h[ri, ci]).float()


def batch_losses(model: DegModel, group: list[GenItem], kl_base: float = 0.0
                 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, int]:
    """(summed target NLL, target logits, summed KL(base || MoE) or None, target tokens)
    for one micro-batch of W rows."""
    dev = model.device
    cids, clens, wids, wlens, row_of, ri, ci, tgt = _pack(group, model.renderer.pad_id, dev)
    logits = _target_logits(model.backbone, cids, clens, wids, wlens, row_of, ri, ci,
                            model.gen_route(len(group)))
    nll = F.cross_entropy(logits, tgt, reduction="sum")
    kl = None
    if kl_base > 0:
        if model.base is None:
            raise RuntimeError("--kl-base needs the base copy: attach_base(model) first")
        with torch.no_grad():
            base = _target_logits(model.base, cids, clens, wids, wlens, row_of, ri, ci,
                                  None).log_softmax(-1)
        kl = F.kl_div(logits.log_softmax(-1), base, log_target=True, reduction="sum")
    return nll, logits, kl, int(tgt.shape[0])


@torch.no_grad()
def evaluate(model: DegModel, items: list[GenItem], batch_tokens: int) -> float:
    tot, n = 0.0, 0
    for g in micro_batches(items, batch_tokens, None):
        nll, _, _, k = batch_losses(model, g)
        tot, n = tot + float(nll), n + k
    return tot / max(n, 1)


class OffloadAdamW:
    """AdamW over fp32 master copies of ``params`` on ``device`` (cpu: pinned gradient
    buffers, fused CPU AdamW); bf16 compute weights stay where they are. Two groups:
    W experts (``lr``) and writer LoRA (``writer_lr``)."""

    def __init__(self, params: list[torch.nn.Parameter], lora: list[bool], lr: float,
                 writer_lr: float, device: torch.device) -> None:
        self.params = params
        pin = device.type == "cpu" and torch.cuda.is_available() and params[0].is_cuda
        self.master = [p.detach().to(device, torch.float32, copy=True).requires_grad_(True)
                       for p in params]
        self.grads = [torch.zeros(m.shape, dtype=torch.float32, device=device,
                                  pin_memory=pin) for m in self.master]
        groups = [{"params": [m for m, lo in zip(self.master, lora) if not lo], "lr": lr,
                   "base_lr": lr},
                  {"params": [m for m, lo in zip(self.master, lora) if lo], "lr": writer_lr,
                   "base_lr": writer_lr}]
        groups = [g for g in groups if g["params"]]
        try:
            self.opt = torch.optim.AdamW(groups, weight_decay=0.0, fused=True)
        except (RuntimeError, ValueError):
            self.opt = torch.optim.AdamW(groups, weight_decay=0.0)

    def set_lr(self, factor: float) -> None:
        for g in self.opt.param_groups:
            g["lr"] = g["base_lr"] * factor

    def step(self, clip: float = 1.0) -> float:
        for p, m, g in zip(self.params, self.master, self.grads):
            if p.grad is None:
                g.zero_()
            else:
                g.copy_(p.grad, non_blocking=True)
            p.grad = None
            m.grad = g
        if self.params[0].is_cuda:
            torch.cuda.synchronize()
        norm = float(torch.nn.utils.clip_grad_norm_(self.master, clip))
        self.opt.step()
        with torch.no_grad():
            for p, m in zip(self.params, self.master):
                p.copy_(m, non_blocking=True)
        return norm


def build(ckpt: Path | str, device: torch.device, writer_rank: int,
          writer_alpha: float | None, kl_base: float, seed: int = 0,
          base_lm: torch.nn.Module | None = None, lm: torch.nn.Module | None = None,
          tokenizer: Any = None) -> tuple[DegModel, list[torch.nn.Parameter], list[bool]]:
    """deg checkpoint -> deg-MoE with W = the base MLPs, R frozen, W trainable; the base
    copy attached when ``kl_base`` > 0. ``base_lm``/``lm``/``tokenizer``: tests."""
    model = load_model(ckpt, device=device, generate="self", lm=lm, tokenizer=tokenizer)
    if base_lm is None:
        base_lm = load_base_lm(model.meta.get("backbone"), torch.bfloat16)
    model.backbone.lm.requires_grad_(False)
    model.heads.requires_grad_(False)
    params = add_moe(model, base_lm, writer_rank, writer_alpha, seed)
    for p in params:
        p.requires_grad_(True)
    lora = [p.dim() == 2 and any(p is getattr(m, "lora_A", None) or p is getattr(m, "lora_B", None)
                                 for m in model.backbone.lm.modules()) for p in params]
    if kl_base > 0:
        attach_base(model, base_lm)
    model.backbone.grad_ckpt = True
    return model, params, lora


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description="train the W experts of a deg-MoE")
    ap.add_argument("--ckpt", type=Path, required=True, help="full fine-tune deg checkpoint (R)")
    ap.add_argument("--out", type=Path, required=True, help="deg-MoE checkpoint dir (best val)")
    ap.add_argument("--records", type=Path, default=GEN_RECORDS)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--optim-device", default="cpu", help="fp32 master + AdamW moments")
    ap.add_argument("--writer-rank", type=int, default=32, help="0 = no writer attention LoRA")
    ap.add_argument("--writer-alpha", type=float, default=None, help="default 2 * rank")
    ap.add_argument("--lr", type=float, default=1e-5, help="W experts")
    ap.add_argument("--writer-lr", type=float, default=2e-4, help="writer LoRA")
    ap.add_argument("--epochs", type=float, default=2.0)
    ap.add_argument("--max-steps", type=int, default=None)
    ap.add_argument("--batch-tokens", type=int, default=4096, help="padded tokens per micro")
    ap.add_argument("--accum", type=int, default=4, help="micro-batches per optimizer step")
    ap.add_argument("--warmup", type=int, default=20)
    ap.add_argument("--eval-every", type=int, default=50)
    ap.add_argument("--patience", type=int, default=4, help="evals without val improvement")
    ap.add_argument("--kl-base", type=float, default=0.0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--limit", type=int, default=None, help="train records (smoke)")
    ap.add_argument("--val-limit", type=int, default=None)
    ap.add_argument("--no-save", action="store_true", help="smoke: never write --out")
    args = ap.parse_args(argv)

    torch.manual_seed(args.seed)
    dev = torch.device(args.device)
    check_device(dev)
    if dev.type == "cuda" and dev.index is None:
        dev = torch.device("cuda", torch.cuda.current_device())
    model, params, lora = build(args.ckpt, dev, args.writer_rank, args.writer_alpha,
                                args.kl_base, args.seed)
    opt = OffloadAdamW(params, lora, args.lr, args.writer_lr, torch.device(args.optim_device))
    counts = moe_params(model)
    def mem() -> float:
        return (round(torch.cuda.max_memory_allocated(dev) / 2**30, 2)
                if dev.type == "cuda" else 0.0)
    print(json.dumps({"event": "model", **counts, "trainable": sum(p.numel() for p in params),
                      "base_copy": model.base is not None, "mem_gb": mem()}), flush=True)

    stats = {"judge_drop": 0, "budget_drop": 0}
    train = gen_items(model, read_jsonl(args.records, "train", args.limit), stats)
    vstats = {"judge_drop": 0, "budget_drop": 0}
    val = gen_items(model, read_jsonl(args.records, "val", args.val_limit), vstats)
    print(json.dumps({"event": "data", "train": len(train), "val": len(val),
                      "train_stats": stats, "val_stats": vstats,
                      "train_target_tokens": sum(x.target_len for x in train)}), flush=True)

    rng = random.Random(args.seed)
    per_epoch = math.ceil(len(micro_batches(train, args.batch_tokens, None)) / args.accum)
    steps = max(1, int(per_epoch * args.epochs))
    if args.max_steps:
        steps = min(steps, args.max_steps)
    t0 = time.time()
    best = evaluate(model, val, args.batch_tokens)
    print(json.dumps({"event": "val", "step": 0, "nll": round(best, 4),
                      "s": round(time.time() - t0, 1), "mem_gb": mem()}), flush=True)
    meta = {"ckpt": str(args.ckpt), "records": str(args.records), "train_items": len(train),
            "val_items": len(val), "lr": args.lr, "writer_lr": args.writer_lr,
            "kl_base": args.kl_base, "seed": args.seed, "batch_tokens": args.batch_tokens,
            "accum": args.accum, "layout": "base header (deg.infer.base_generate_header)",
            "val_nll_init": round(best, 4), "params": counts}

    def save(step: int, v: float) -> None:
        if not args.no_save:
            model.save(args.out, {"moe_train": {**meta, "step": step, "val_nll": round(v, 4)}})

    best_step, bad, step, stop, tok = 0, 0, 0, False, 0
    save(0, best)
    while not stop:
        groups = micro_batches(train, args.batch_tokens, rng)
        for i in range(0, len(groups), args.accum):
            chunk = groups[i:i + args.accum]
            ntok = sum(x.target_len for g in chunk for x in g)
            opt.set_lr(lr_at(step, steps, args.warmup))
            tot_nll = tot_kl = 0.0
            for g in chunk:
                nll, _, kl, _ = batch_losses(model, g, args.kl_base)
                loss = nll + (args.kl_base * kl if kl is not None else 0.0)
                (loss / ntok).backward()
                tot_nll += float(nll)
                tot_kl += float(kl) if kl is not None else 0.0
                tok += sum(len(x.ids) for x in g)
            gnorm = opt.step()
            step += 1
            if step % 10 == 0 or step == 1:
                s = time.time() - t0
                print(json.dumps({"event": "step", "step": step, "of": steps,
                                  "nll": round(tot_nll / ntok, 4), "kl": round(tot_kl / ntok, 5),
                                  "gnorm": round(gnorm, 3),
                                  "lr": opt.opt.param_groups[0]["lr"], "s": round(s, 1),
                                  "tok_s": round(tok / max(s, 1e-9), 1), "mem_gb": mem()}),
                      flush=True)
            if step % args.eval_every == 0 or step >= steps:
                v = evaluate(model, val, args.batch_tokens)
                print(json.dumps({"event": "val", "step": step, "nll": round(v, 4)}), flush=True)
                if v < best:
                    best, bad, best_step = v, 0, step
                    save(step, v)
                else:
                    bad += 1
                if bad >= args.patience or step >= steps:
                    stop = True
                    break
    print(json.dumps({"event": "done", "best_step": best_step, "best_val_nll": round(best, 4),
                      "out": str(args.out), "s": round(time.time() - t0, 1), "mem_gb": mem()}),
          flush=True)


if __name__ == "__main__":
    main()