#!/usr/bin/env python3
"""Quick KL check of packed students against the teacher on one GPU (PyTorch), for deciding on a pod whether a run
paid off before downloading it. Mirrors the llama.cpp harness: the first N held-out eval sequences of 4,096 tokens,
plain causal attention, KL(teacher || student) over the full vocabulary, scored on the second half of each sequence.

    python scripts/pod_kl.py --packed runs/A/packed/tok_N --packed runs/B/packed/tok_N [--n 40]

Absolute values differ slightly from llama.cpp's (tokenisation of the decoded text, kernels); compare runs measured
here with each other.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import torch

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.model.student import load_packed_student  # noqa: E402

TEACHER = "/path/to/data/models/Marco-Nano-Instruct"
EVAL = "/path/to/data/shards/eval/tokens.u32"
CTX = 4096


@torch.no_grad()
def logp(model, ids: torch.Tensor) -> torch.Tensor:
    return torch.log_softmax(model(ids).logits[0, CTX // 2:].float(), -1)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--packed", action="append", required=True)
    p.add_argument("--n", type=int, default=40)
    a = p.parse_args()
    from transformers import AutoModelForCausalLM
    toks = np.memmap(EVAL, dtype=np.uint32, mode="r")
    seqs = [torch.from_numpy(toks[i * CTX:(i + 1) * CTX].astype(np.int64))[None].cuda() for i in range(a.n)]
    teacher = AutoModelForCausalLM.from_pretrained(TEACHER, dtype=torch.bfloat16).cuda().eval()
    ref = [logp(teacher, s).cpu() for s in seqs]  # ~0.6 GB per sequence on the CPU side
    del teacher
    torch.cuda.empty_cache()
    out = {}
    for d in a.packed:
        m = load_packed_student(d, device="cuda")
        kl = n = 0.0
        same = 0
        for s, lt in zip(seqs, ref):
            ls = logp(m, s)
            lt = lt.cuda()
            kl += (lt.exp() * (lt - ls)).sum(-1).sum().item()
            same += (lt.argmax(-1) == ls.argmax(-1)).sum().item()
            n += lt.shape[0]
        out[d] = {"kl": round(kl / n, 5), "same_top1_pct": round(100 * same / n, 2), "tokens": int(n)}
        print(json.dumps({d: out[d]}), flush=True)
        del m
        torch.cuda.empty_cache()


if __name__ == "__main__":
    main()