#!/usr/bin/env python3
"""Full-vocabulary KL divergence of an MLX quant vs the BF16 MLX model, on the same text as the GGUF KLD.

Mirrors llama-perplexity --kl-divergence: the text quantdiff-kld-v1.txt is tokenized once with the model's own
tokenizer, cut into 100 windows of 512 tokens, and only the second half of each window is scored (positions 256..510
predict tokens 257..511: 255 positions x 100 = 25,500), so every scored token has at least 256 tokens of context.

  ref:  python kld_mlx.py ref --model <bf16 dir> --corpus quantdiff-kld-v1.txt --out <ref dir>
        -> tokens.npy + logp_<chunk>.npy (float16 log-softmax over the full vocabulary, 255 x V per window)
  cmp:  python kld_mlx.py cmp --model <quant dir> --ref <ref dir> --out kld.json
"""
import argparse, json, math, time
from pathlib import Path
import numpy as np
import mlx.core as mx
from mlx_lm import load

CTX, CHUNKS = 512, 100
FIRST = CTX // 2


def windows(tok, text):
    ids = tok.encode(text)
    if tok.bos_token_id is not None and (not ids or ids[0] != tok.bos_token_id) and getattr(tok, "add_bos_token", False):
        ids = [tok.bos_token_id] + ids
    need = CTX * CHUNKS
    if len(ids) < need:
        raise SystemExit(f"text has only {len(ids)} tokens, need {need}")
    return np.array(ids[:need], dtype=np.int32).reshape(CHUNKS, CTX)


def logprobs(model, w):
    logits = model(mx.array(w)[None])[0, FIRST:CTX - 1].astype(mx.float32)
    return logits - mx.logsumexp(logits, axis=-1, keepdims=True)


def cmd_ref(a):
    model, tok = load(a.model)
    out = Path(a.out); out.mkdir(parents=True, exist_ok=True)
    W = windows(tok, Path(a.corpus).read_text(encoding="utf-8"))
    np.save(out / "tokens.npy", W)
    nll = 0.0; t0 = time.time()
    for i, w in enumerate(W):
        lp = logprobs(model, w)
        tgt = mx.array(w[FIRST + 1:CTX])
        nll += -mx.take_along_axis(lp, tgt[:, None], axis=-1).sum().item()
        np.save(out / f"logp_{i:03d}.npy", np.array(lp.astype(mx.float16)))
    meta = {"model": a.model, "ctx": CTX, "chunks": CHUNKS, "first": FIRST, "positions": CHUNKS * (CTX - 1 - FIRST),
            "vocab": int(lp.shape[-1]), "ppl": math.exp(nll / (CHUNKS * (CTX - 1 - FIRST))), "seconds": round(time.time() - t0, 1)}
    (out / "meta.json").write_text(json.dumps(meta, indent=1))
    print(json.dumps(meta))


def cmd_cmp(a):
    ref = Path(a.ref); meta = json.loads((ref / "meta.json").read_text())
    model, tok = load(a.model)
    W = np.load(ref / "tokens.npy")
    Wq = windows(tok, Path(a.corpus).read_text(encoding="utf-8")) if a.corpus else None
    if Wq is not None and not np.array_equal(W, Wq):
        raise SystemExit("tokenizer mismatch between quant and reference: KLD would not be sound")
    kls, same, nll_q = [], 0, 0.0; t0 = time.time()
    fwd_tokens = fwd_passes = 0  # counted for TokenStation: tokens actually run through the model (prefill only)
    for i, w in enumerate(W):
        lq = logprobs(model, w)
        fwd_tokens += int(len(w)); fwd_passes += 1
        lr = mx.array(np.load(ref / f"logp_{i:03d}.npy")).astype(mx.float32)
        if lr.shape != lq.shape:
            raise SystemExit(f"vocab/shape mismatch {lr.shape} vs {lq.shape}")
        pr = mx.exp(lr)
        kl = (pr * (lr - lq)).sum(axis=-1)
        kls.append(np.array(kl))
        same += int((mx.argmax(lr, axis=-1) == mx.argmax(lq, axis=-1)).sum().item())
        tgt = mx.array(w[FIRST + 1:CTX])
        nll_q += -mx.take_along_axis(lq, tgt[:, None], axis=-1).sum().item()
    k = np.concatenate(kls); n = k.size
    res = {"method": "mlx full-vocab KLD vs BF16 MLX (reference log-probs stored as float16)", "positions": int(n),
           "mean_kld": float(k.mean()), "median_kld": float(np.median(k)), "p99_kld": float(np.percentile(k, 99)),
           "max_kld": float(k.max()), "same_top_pct": 100.0 * same / n, "ppl_quant": math.exp(nll_q / n), "ppl_ref": meta["ppl"],
           "ctx": CTX, "chunks": CHUNKS, "vocab": meta["vocab"], "seconds": round(time.time() - t0, 1),
           "tokens_forwarded": fwd_tokens, "forward_passes": fwd_passes}
    Path(a.out).write_text(json.dumps(res, indent=1))
    print(json.dumps(res))


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sp = ap.add_subparsers(dest="cmd", required=True)
    r = sp.add_parser("ref"); r.add_argument("--model", required=True); r.add_argument("--corpus", required=True); r.add_argument("--out", required=True)
    c = sp.add_parser("cmp"); c.add_argument("--model", required=True); c.add_argument("--ref", required=True); c.add_argument("--corpus"); c.add_argument("--out", required=True)
    a = ap.parse_args()
    {"ref": cmd_ref, "cmp": cmd_cmp}[a.cmd](a)