"""Phase C: fine-tune the nemotron_h MTP head on dumped triples.

Trains the BF16 `mtp.*` head from the merged3 checkpoint against the trunk's
own conditionals (see dump_triples.py): for each completion position t the head
sees (hidden[t], token[t+1]) and must predict token[t+2], with the trunk's
top-64 logprobs at t+1 as a soft KL target.  Loss = CE + lambda * KL.

Faithful to serving (mtp_probe recipe + MTPModule.__call__):
  - fuse enorm(embed(t+1)) || hnorm(hidden[t]) -> eh_proj
  - + attention block (causal over the completion region -- the serving chain's
    committed KV is exactly the prior completion positions; NoPE, no rope)
  - + MoE block, final_layernorm, shared lm_head

Phase 1 trainables (~120M): eh_proj, enorm/hnorm/norms, attention, gate.weight,
latent projections, shared experts.  Frozen: 512 routed experts (switch_mlp),
gate.e_score_correction_bias (selection-only, gradient-free), embed, lm_head.
Trainables are trained as fp32 masters and written back bf16.

Usage:
  python3 train_mtp_head.py --eval-only            # pipeline check: ~52% top-1
  python3 train_mtp_head.py [--epochs 2] [--batch-positions 4096]
                            [--lr 1e-4] [--kl-lambda 1.0] [--limit-shards N]

Output: mtp_trained.safetensors (disk-named mtp.* tensors, splice-ready) +
training log lines on stdout.
"""

import argparse
import glob as globmod
import json
import math
import os
import random
import subprocess
import time
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np
from mlx.utils import tree_flatten, tree_unflatten

ROOT = Path(__file__).parent
BF16_DIR = Path("/Users/david/AI/NVIDIA-Nemotron-3-Super-120B-merged3")


# --------------------------------------------------------------------------- #
# Model
# --------------------------------------------------------------------------- #
class LoRASwitchLinear(nn.Module):
    """Per-expert LoRA over a frozen SwitchLinear.

    Phase 1 showed train `agree` == eval `agree` (0.6087 vs 0.6086), i.e. the
    head could not raise argmax accuracy even on data it was training on --
    a capacity limit, not underfitting. The 512 routed experts hold ~2.8B of
    the head's ~3B params and were frozen, so this adds low-rank adapters to
    them through the same gather_mm dispatch the base layer uses.

    B is zero-initialised, so the wrapped layer starts numerically identical
    to the frozen base -- training resumes exactly from the Phase 1 optimum
    rather than perturbing it.
    """

    def __init__(self, base, rank: int, scale: float = 2.0):
        super().__init__()
        self.base = base
        n_experts, out_dims, in_dims = base.weight.shape
        bound = 1.0 / math.sqrt(in_dims)
        self.lora_a = mx.random.uniform(
            low=-bound, high=bound, shape=(n_experts, rank, in_dims)
        )
        self.lora_b = mx.zeros((n_experts, out_dims, rank))
        self.scale = scale

    def __call__(self, x, indices, sorted_indices=False):
        y = self.base(x, indices, sorted_indices=sorted_indices)
        z = mx.gather_mm(
            x,
            self["lora_a"].swapaxes(-1, -2),
            rhs_indices=indices,
            sorted_indices=sorted_indices,
        )
        z = mx.gather_mm(
            z,
            self["lora_b"].swapaxes(-1, -2),
            rhs_indices=indices,
            sorted_indices=sorted_indices,
        )
        return y + self.scale * z


def build_head(expert_lora_rank: int = 0, expert_lora_scale: float = 2.0):
    """Construct MTPModule with bf16 weights from the merged3 checkpoint,
    plus frozen embedding + lm_head tables."""
    from omlx.patches.mlx_lm_mtp import nemotron_h_model as nhm

    nhm.apply()
    nhm.set_mtp_active(True)
    from mlx_lm.models import nemotron_h as nh

    config = json.load(open(BF16_DIR / "config.json"))
    args = nh.ModelArgs.from_dict(config)
    head = nh.MTPModule(args)

    index = json.load(open(BF16_DIR / "model.safetensors.index.json"))["weight_map"]
    need_files = {index[k] for k in index if k.startswith("mtp.")}
    need_files.add(index["backbone.embeddings.weight"])
    need_files.add(index["lm_head.weight"])

    mtp_w, emb_w, lm_w = {}, None, None
    for fname in sorted(need_files):
        shard = mx.load(str(BF16_DIR / fname))
        for k, v in shard.items():
            if k.startswith("mtp."):
                mtp_w[k[len("mtp."):]] = v
            elif k == "backbone.embeddings.weight":
                emb_w = v
            elif k == "lm_head.weight":
                lm_w = v
    assert emb_w is not None and lm_w is not None

    # Stack routed experts exactly like Model.sanitize does.
    E = config["n_routed_experts"]
    ep = "layers.1.mixer.experts"
    stacked = {
        "layers.1.mixer.switch_mlp.fc1.weight": mx.stack(
            [mtp_w.pop(f"{ep}.{e}.up_proj.weight") for e in range(E)]
        ),
        "layers.1.mixer.switch_mlp.fc2.weight": mx.stack(
            [mtp_w.pop(f"{ep}.{e}.down_proj.weight") for e in range(E)]
        ),
    }
    mtp_w.update(stacked)
    head.load_weights(list(mtp_w.items()), strict=True)

    # Freeze routed experts + the selection-only gate bias (no gradient flows
    # through argtopk; an optimizer step would only decay/perturb it).
    head.layers[1].mixer.switch_mlp.freeze()
    head.layers[1].mixer.gate.freeze(keys=["e_score_correction_bias"])

    # Phase 2: low-rank adapters on the routed experts. Wrap AFTER the freeze
    # above so the base SwitchLinears stay frozen and only lora_a/lora_b pick
    # up gradients; re-freeze the bases explicitly since the wrapper is new.
    if expert_lora_rank > 0:
        switch_mlp = head.layers[1].mixer.switch_mlp
        switch_mlp.fc1 = LoRASwitchLinear(
            switch_mlp.fc1, expert_lora_rank, expert_lora_scale)
        switch_mlp.fc2 = LoRASwitchLinear(
            switch_mlp.fc2, expert_lora_rank, expert_lora_scale)
        switch_mlp.fc1.base.freeze()
        switch_mlp.fc2.base.freeze()

    # fp32 master weights for everything trainable.
    trainable = tree_flatten(head.trainable_parameters())
    head.update(tree_unflatten([(k, v.astype(mx.float32)) for k, v in trainable]))

    n_train = sum(v.size for _, v in trainable)
    print(f"head built: {n_train/1e6:.1f}M trainable params (fp32 masters)")
    return head, emb_w, lm_w


def head_forward(head, emb_w, hidden, next_ids):
    """Batched training forward, mirroring MTPModule.__call__ with a causal
    mask over the whole (padded) window.  hidden (B,S,H) fp32, next_ids (B,S)."""
    l0, l1 = head.layers
    e = l0.enorm(emb_w[next_ids].astype(mx.float32))
    h = l0.hnorm(hidden)
    fused = l0.eh_proj(mx.concatenate([e, h], axis=-1))
    x = fused + l0.mixer(l0.norm(fused), mask="causal", cache=None)
    x = x + l1.mixer(l1.norm(x))
    return l1.final_layernorm(x)


# --------------------------------------------------------------------------- #
# Data
# --------------------------------------------------------------------------- #
class Triples:
    """Doc-granular access over triples-*.npz shards (kept in RAM as numpy)."""

    def __init__(self, shard_glob, limit_shards=None):
        self.docs = []  # (shard_i, start, end, tok_off)
        self.shards = []
        files = sorted(globmod.glob(shard_glob))
        if limit_shards:
            files = files[:limit_shards]
        for si, f in enumerate(files):
            z = np.load(f)
            sh = {k: z[k] for k in z.files}
            self.shards.append(sh)
            for di, (s, e) in enumerate(sh["doc_bounds"]):
                self.docs.append((si, int(s), int(e), int(s) + 2 * di))
        n_pos = sum(e - s for _, s, e, _ in self.docs)
        print(f"{len(files)} shards, {len(self.docs)} docs, {n_pos:,} positions")

    def fetch(self, doc):
        si, s, e, toff = doc
        sh = self.shards[si]
        n = e - s
        hid_u16 = sh["hiddens"][s:e]
        toks = sh["tokens"][toff : toff + n + 2].astype(np.int64)
        return (
            hid_u16,                       # (n, H) uint16 bf16-bits
            toks[1 : n + 1],               # input token t+1
            toks[2 : n + 2],               # hard target t+2
            sh["topk_ids"][s:e].astype(np.int64),
            sh["topk_lp"][s:e].astype(np.float32),
        )


def make_batches(dataset, doc_ids, batch_positions, seed):
    """Length-bucketed padded batches: list of lists of doc indices."""
    order = sorted(doc_ids, key=lambda i: dataset.docs[i][2] - dataset.docs[i][1])
    batches, cur, cur_max = [], [], 0
    for i in order:
        n = dataset.docs[i][2] - dataset.docs[i][1]
        m = max(cur_max, n)
        if cur and m * (len(cur) + 1) > batch_positions:
            batches.append(cur)
            cur, cur_max = [], 0
            m = n
        cur.append(i)
        cur_max = m
    if cur:
        batches.append(cur)
    random.Random(seed).shuffle(batches)
    return batches


def collate(dataset, batch):
    docs = [dataset.fetch(dataset.docs[i]) for i in batch]
    B = len(docs)
    S = max(d[0].shape[0] for d in docs)
    H = docs[0][0].shape[1]
    K = docs[0][4].shape[1]
    hid = np.zeros((B, S, H), np.uint16)
    nxt = np.zeros((B, S), np.int64)
    tgt = np.zeros((B, S), np.int64)
    kid = np.zeros((B, S, K), np.int64)
    klp = np.full((B, S, K), -1e9, np.float32)
    msk = np.zeros((B, S), np.float32)
    for b, (h, nx, tg, ki, kl) in enumerate(docs):
        n = h.shape[0]
        hid[b, :n], nxt[b, :n], tgt[b, :n] = h, nx, tg
        kid[b, :n], klp[b, :n], msk[b, :n] = ki, kl, 1.0
    hidden = mx.array(hid).view(mx.bfloat16).astype(mx.float32)
    return (hidden, mx.array(nxt), mx.array(tgt), mx.array(kid),
            mx.array(klp), mx.array(msk))


# --------------------------------------------------------------------------- #
# Loss / metrics
# --------------------------------------------------------------------------- #
def batch_stats(head, emb_w, lm_w, batch, kl_lambda, ce_lambda=1.0):
    hidden, nxt, tgt, kid, klp, msk = batch
    out = head_forward(head, emb_w, hidden, nxt)          # (B,S,H) fp32
    logits = out @ lm_w.T.astype(mx.float32)              # (B,S,V)
    lse = mx.logsumexp(logits, axis=-1)                   # (B,S)
    tgt_logit = mx.take_along_axis(logits, tgt[..., None], axis=-1)[..., 0]
    ce = lse - tgt_logit
    head_klp = mx.take_along_axis(logits, kid, axis=-1) - lse[..., None]
    p = mx.exp(klp)                                       # trunk top-64 probs
    kl = (p * (klp - head_klp)).sum(axis=-1)
    denom = msk.sum()
    # KL is the real objective, not a regularizer: `agree` below measures
    # agreement with the TRUNK's top-1 (agree_ref comes from klp), which is what
    # speculative acceptance actually is. CE pulls toward corpus tokens instead.
    # Measured 2026-07-18: kl_lambda 0.3 -> agree 0.5726, 1.0 -> 0.6086,
    # 3.0 -> 0.6355. ce_lambda exists to test pushing CE's weight toward 0.
    loss = ((ce_lambda * ce + kl_lambda * kl) * msk).sum() / denom
    agree_ref = mx.take_along_axis(kid, mx.argmax(klp, axis=-1)[..., None], axis=-1)[..., 0]
    agree = ((mx.argmax(logits, axis=-1) == agree_ref) * msk).sum() / denom
    ce_m = (ce * msk).sum() / denom
    kl_m = (kl * msk).sum() / denom
    return loss, (ce_m, kl_m, agree, denom)


def run_eval(head, emb_w, lm_w, dataset, batches, kl_lambda):
    tot = {"ce": 0.0, "kl": 0.0, "agree": 0.0, "n": 0.0}
    for b in batches:
        _, (ce, kl, ag, n) = batch_stats(
            head, emb_w, lm_w, collate(dataset, b), kl_lambda)
        mx.eval(ce, kl, ag, n)
        n = n.item()
        tot["ce"] += ce.item() * n
        tot["kl"] += kl.item() * n
        tot["agree"] += ag.item() * n
        tot["n"] += n
    n = max(tot["n"], 1)
    return tot["ce"] / n, tot["kl"] / n, tot["agree"] / n


# --------------------------------------------------------------------------- #
# Save
# --------------------------------------------------------------------------- #
def save_trained(head, path):
    """Trained (non-expert) tensors, bf16, with on-disk mtp.* names."""
    out = {}
    for k, v in tree_flatten(head.trainable_parameters()):
        out["mtp." + k] = v.astype(mx.bfloat16)
    mx.save_safetensors(str(path), out)
    print(f"saved {len(out)} tensors -> {path}")


# --------------------------------------------------------------------------- #
def running_dumps():
    """Pids of any live dump_triples.py.

    A concurrent dump means the shards are still incomplete *and* ~85 GB is
    already committed to its trunk. Training on top of that OOM-killed the
    machine on 2026-07-18 (Jetsam took WindowServer with it), so this is a
    hard stop rather than a warning.
    """
    try:
        out = subprocess.run(
            ["pgrep", "-f", "dump_triples.py"], capture_output=True, text=True
        ).stdout.split()
    except FileNotFoundError:
        return []
    return [p for p in out if p != str(os.getpid())]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--shard-glob", default=str(ROOT / "triples" / "triples-*.npz"))
    ap.add_argument("--limit-shards", type=int, default=None)
    ap.add_argument("--eval-only", action="store_true")
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--batch-positions", type=int, default=4096)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--warmup", type=int, default=100)
    ap.add_argument("--kl-lambda", type=float, default=1.0)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--eval-docs", type=int, default=100)
    ap.add_argument("--eval-every", type=int, default=200)
    ap.add_argument("--save-every", type=int, default=500)
    ap.add_argument("--out", default=str(ROOT / "mtp_trained.safetensors"))
    ap.add_argument("--seed", type=int, default=17)
    ap.add_argument("--ce-lambda", type=float, default=1.0,
                    help="weight on the hard-target CE term (0 = pure KL)")
    ap.add_argument("--expert-lora-rank", type=int, default=0,
                    help="Phase 2: LoRA rank on the 512 routed experts (0=off)")
    ap.add_argument("--expert-lora-scale", type=float, default=2.0)
    ap.add_argument(
        "--ignore-running-dump",
        action="store_true",
        help="start even if dump_triples.py is live (only with real headroom)",
    )
    args = ap.parse_args()

    busy = running_dumps()
    if busy and not args.ignore_running_dump:
        raise SystemExit(
            f"refusing to start: dump_triples.py still running (pid {', '.join(busy)}).\n"
            "Its shards are incomplete and it holds ~85 GB; this run grows to "
            "~60 GB and the pair has OOM-killed the machine before.\n"
            "Wait for the dump to print DONE, or pass --ignore-running-dump."
        )

    dataset = Triples(args.shard_glob, args.limit_shards)
    ids = list(range(len(dataset.docs)))
    random.Random(args.seed).shuffle(ids)
    eval_ids, train_ids = ids[: args.eval_docs], ids[args.eval_docs :]
    eval_batches = make_batches(dataset, eval_ids, args.batch_positions, 0)

    head, emb_w, lm_w = build_head(
        args.expert_lora_rank, args.expert_lora_scale)

    ce, kl, ag = run_eval(head, emb_w, lm_w, dataset, eval_batches, args.kl_lambda)
    print(f"[baseline] ce={ce:.4f} kl={kl:.4f} agree={ag:.4f}")
    if args.eval_only:
        return

    steps_per_epoch = max(
        1, len(make_batches(dataset, train_ids, args.batch_positions, 0)))
    total_steps = steps_per_epoch * args.epochs
    sched = optim.join_schedules(
        [optim.linear_schedule(0.0, args.lr, args.warmup),
         optim.cosine_decay(args.lr, max(1, total_steps - args.warmup))],
        [args.warmup],
    )
    opt = optim.Adam(learning_rate=sched)

    def loss_fn(head_, batch):
        loss, aux = batch_stats(
            head_, emb_w, lm_w, batch, args.kl_lambda, args.ce_lambda)
        return loss, aux

    vg = nn.value_and_grad(head, loss_fn)

    step, t0 = 0, time.time()
    best_agree = ag
    for epoch in range(args.epochs):
        batches = make_batches(
            dataset, train_ids, args.batch_positions, args.seed + epoch)
        for b in batches:
            (loss, (ce, kl, ag_b, npos)), grads = vg(head, collate(dataset, b))
            if args.grad_clip > 0:
                grads, _ = optim.clip_grad_norm(grads, args.grad_clip)
            opt.update(head, grads)
            mx.eval(head.parameters(), opt.state, loss)
            step += 1
            if step % 20 == 0:
                dt = time.time() - t0
                print(f"[{dt/60:5.1f}m] step {step}/{total_steps} "
                      f"loss={loss.item():.4f} ce={ce.item():.4f} "
                      f"kl={kl.item():.4f} agree={ag_b.item():.4f}",
                      flush=True)
            if step % args.eval_every == 0:
                ce_e, kl_e, ag_e = run_eval(
                    head, emb_w, lm_w, dataset, eval_batches, args.kl_lambda)
                print(f"[eval @ {step}] ce={ce_e:.4f} kl={kl_e:.4f} "
                      f"agree={ag_e:.4f} (baseline {best_agree:.4f})", flush=True)
            if step % args.save_every == 0:
                save_trained(head, args.out)
    save_trained(head, args.out)
    ce, kl, ag_f = run_eval(head, emb_w, lm_w, dataset, eval_batches, args.kl_lambda)
    print(f"DONE {step} steps in {(time.time()-t0)/60:.1f}m: "
          f"ce={ce:.4f} kl={kl:.4f} agree={ag_f:.4f} (baseline {best_agree:.4f})")


if __name__ == "__main__":
    main()