#!/usr/bin/env python3
"""Measure how far a quantised model's next-token distributions have moved from the original.

Purpose: the portable version of llama.cpp's KL-divergence mode, for the cases that tool cannot
    reach: an MLX model, an AWQ checkpoint, two different engines, or any pair of models served
    behind the Part 9 gateway. It walks a fixed text, asks both servers for the next-token
    distribution at the same positions, and reports the divergence between them along with how
    often they would have chosen the same token.
Platform: all (pure Python over HTTP; the servers may be on any track or another machine)
Minimum memory: negligible for this script. The servers behind it need whatever their models
    need, and running the reference and the quantised model one after the other on a small
    machine works: pass --save-ref and --load-ref instead of two live endpoints.
Assumes: Python 3.9 or later, and OpenAI-compatible servers exposing /v1/completions with
    log-probabilities (llama-server, vLLM, mlx_lm.server and the Part 9 gateway all do). The two
    models must share a tokeniser, which is guaranteed when one is a quantisation of the other
    and is the reason this comparison is meaningful at all.

    HONEST LIMITATION: the API returns only the top few candidate tokens, so the divergence is
    computed over a truncated support with a floor for the unseen mass. It is comparable between
    quantisations measured this way against the same reference on the same text, and it is NOT
    comparable with a figure from llama-perplexity --kl-divergence, which sees the full
    distribution. The output records which method produced it.

Usage: python3 kl-divergence.py --text calibration.txt \
           --base-url-ref http://127.0.0.1:8080/v1 --model-ref reference \
           --base-url-quant http://127.0.0.1:8081/v1 --model-quant q4-k-m \
           --quant-label Q4_K_M --out kld-Q4_K_M.json --labbook labbook.md
       python3 kl-divergence.py --text calibration.txt --base-url-ref http://127.0.0.1:8080/v1 \
           --model-ref reference --save-ref ref-logprobs.json
       python3 kl-divergence.py --text calibration.txt --load-ref ref-logprobs.json \
           --base-url-quant http://127.0.0.1:8080/v1 --model-quant q4-k-m --quant-label Q4_K_M
"""

from __future__ import annotations

import argparse
import json
import math
import platform
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Optional

# Probability mass the API did not show us has to be given some value. Every token outside a
# model's returned top-k is treated as one natural log unit below the smallest logprob it did
# return, which keeps the divergence finite without pretending the tail is empty. It is an
# assumption, it is applied identically to both models, and it is recorded in the output.
UNSEEN_MARGIN = 1.0


def post_json(url: str, payload: dict, api_key: Optional[str], timeout: int) -> dict:
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


def extract_top_logprobs(body: dict) -> dict:
    """Pull {token: logprob} out of either the completions or the chat-completions shape."""
    choice = body["choices"][0]
    logprobs = choice.get("logprobs") or {}

    # Legacy completions: {"top_logprobs": [{token: logprob, ...}]}
    top = logprobs.get("top_logprobs")
    if isinstance(top, list) and top and isinstance(top[0], dict):
        first = top[0]
        if all(isinstance(v, (int, float)) for v in first.values()):
            return {str(k): float(v) for k, v in first.items()}

    # Chat completions: {"content": [{"top_logprobs": [{"token": …, "logprob": …}, …]}]}
    content = logprobs.get("content")
    if isinstance(content, list) and content:
        entries = content[0].get("top_logprobs") or []
        if entries:
            return {str(e["token"]): float(e["logprob"]) for e in entries}

    raise RuntimeError(
        "the server returned no top log-probabilities. Start llama-server without --log-disable "
        "and check that the endpoint is /v1/completions; some servers need the request to ask "
        "for logprobs explicitly, which this script does."
    )


def make_prefixes(text: str, count: int, min_words: int) -> list:
    """Cut the text into increasing prefixes, spread evenly, at whitespace boundaries.

    Positions are chosen deterministically so that two runs of this script on the same text
    measure the same places. Nothing is sampled anywhere in this script.
    """
    words = text.split()
    if len(words) < min_words + count:
        raise SystemExit(
            f"the text has only {len(words)} words; need at least {min_words + count}. "
            "Use a longer calibration text."
        )
    span = len(words) - min_words
    step = max(1, span // count)
    prefixes = []
    for i in range(count):
        end = min(len(words) - 1, min_words + i * step)
        prefixes.append(" ".join(words[:end]))
    return prefixes


def ask(base_url: str, model: str, prompt: str, top_k: int, api_key: Optional[str],
        timeout: int) -> dict:
    payload = {
        "model": model,
        "prompt": prompt,
        "max_tokens": 1,
        "temperature": 0.0,
        "logprobs": top_k,
        "top_logprobs": top_k,
        "seed": 0,
    }
    body = post_json(base_url.rstrip("/") + "/completions", payload, api_key, timeout)
    return extract_top_logprobs(body)


def to_probabilities(logprobs: dict, support: list) -> dict:
    """Renormalise a truncated distribution over a shared support."""
    floor = min(logprobs.values()) - UNSEEN_MARGIN
    raw = {token: math.exp(logprobs.get(token, floor)) for token in support}
    total = sum(raw.values()) or 1.0
    return {token: value / total for token, value in raw.items()}


def divergence(ref: dict, quant: dict) -> tuple:
    support = sorted(set(ref) | set(quant))
    p = to_probabilities(ref, support)
    q = to_probabilities(quant, support)
    kld = sum(p[t] * (math.log(p[t]) - math.log(q[t])) for t in support if p[t] > 0)
    top_ref = max(p, key=p.get)
    top_quant = max(q, key=q.get)
    return kld, top_ref == top_quant, p[top_ref] - q[top_ref]


def percentile(values: list, fraction: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, int(round(fraction * (len(ordered) - 1))))
    return ordered[index]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--text", required=True, help="plain text to walk; not your evaluation set")
    parser.add_argument("--positions", type=int, default=64,
                        help="how many positions to measure; each costs one request per model")
    parser.add_argument("--min-words", type=int, default=32,
                        help="shortest prefix, so the first positions are not measured on nothing")
    parser.add_argument("--top-logprobs", type=int, default=20,
                        help="candidates requested per position; servers cap this, often at 20")
    parser.add_argument("--base-url-ref", default=None)
    parser.add_argument("--model-ref", default=None)
    parser.add_argument("--save-ref", default=None,
                        help="write the reference distributions here and stop, so the reference "
                             "and the quantised model can be served one after the other")
    parser.add_argument("--load-ref", default=None, help="read reference distributions from a file")
    parser.add_argument("--base-url-quant", default=None)
    parser.add_argument("--model-quant", default=None)
    parser.add_argument("--quant-label", default="unlabelled",
                        help="the quantisation's name, as it will appear in your results table")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--out", default=None)
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    text = Path(args.text).read_text(encoding="utf-8", errors="replace")
    prefixes = make_prefixes(text, args.positions, args.min_words)

    # ---- reference -------------------------------------------------------
    if args.load_ref:
        stored = json.loads(Path(args.load_ref).read_text(encoding="utf-8"))
        if stored["prefix_hashes"] != [hash_prefix(p) for p in prefixes]:
            sys.exit("the stored reference was measured on different positions; regenerate it "
                     "with the same --text, --positions and --min-words")
        reference = stored["distributions"]
        ref_model = stored["model"]
    else:
        if not (args.base_url_ref and args.model_ref):
            sys.exit("give --base-url-ref and --model-ref, or --load-ref")
        print(f"==> Reference: {args.model_ref} at {args.base_url_ref}")
        reference = []
        for i, prefix in enumerate(prefixes, 1):
            reference.append(ask(args.base_url_ref, args.model_ref, prefix,
                                 args.top_logprobs, args.api_key, args.timeout))
            print(f"    {i}/{len(prefixes)}", end="\r", flush=True)
        print()
        ref_model = args.model_ref

    if args.save_ref:
        Path(args.save_ref).write_text(json.dumps({
            "model": ref_model,
            "text": args.text,
            "positions": args.positions,
            "min_words": args.min_words,
            "top_logprobs": args.top_logprobs,
            "prefix_hashes": [hash_prefix(p) for p in prefixes],
            "distributions": reference,
        }), encoding="utf-8")
        print(f"==> Saved the reference distributions to {args.save_ref}")
        if not (args.base_url_quant and args.model_quant):
            return

    # ---- the quantisation ------------------------------------------------
    if not (args.base_url_quant and args.model_quant):
        sys.exit("give --base-url-quant and --model-quant")
    print(f"==> Quantised: {args.model_quant} at {args.base_url_quant}")
    started = time.time()
    klds, agreements, deltas = [], [], []
    for i, prefix in enumerate(prefixes):
        quant = ask(args.base_url_quant, args.model_quant, prefix,
                    args.top_logprobs, args.api_key, args.timeout)
        kld, same_top, delta = divergence(reference[i], quant)
        klds.append(kld)
        agreements.append(same_top)
        deltas.append(delta)
        print(f"    {i + 1}/{len(prefixes)}", end="\r", flush=True)
    print()

    summary = {
        "lab": "part-16/lab-quantise-five-ways-and-measure/kld",
        "run_id": time.strftime("%Y%m%dT%H%M%S"),
        "method": "truncated-support KL over an OpenAI-compatible API; not comparable with "
                  "llama-perplexity --kl-divergence",
        "unseen_margin_nats": UNSEEN_MARGIN,
        "reference_model": ref_model,
        "quant_model": args.model_quant,
        "quant_label": args.quant_label,
        "text": Path(args.text).name,
        "positions": len(klds),
        "top_logprobs_requested": args.top_logprobs,
        "mean_kld": round(sum(klds) / len(klds), 6),
        "median_kld": round(percentile(klds, 0.5), 6),
        "p99_kld": round(percentile(klds, 0.99), 6),
        "max_kld": round(max(klds), 6),
        "same_top_token": round(sum(agreements) / len(agreements), 4),
        "mean_delta_p": round(sum(deltas) / len(deltas), 6),
        "rms_delta_p": round(math.sqrt(sum(d * d for d in deltas) / len(deltas)), 6),
        "max_abs_delta_p": round(max(abs(d) for d in deltas), 6),
        "seconds": round(time.time() - started, 1),
        "host": platform.platform(),
        "date": time.strftime("%Y-%m-%d"),
    }

    print(json.dumps(summary, indent=2))
    print("\nRead p99_kld and same_top_token before mean_kld. A quantisation that agrees with the "
          "\noriginal almost everywhere and diverges sharply at a few positions has a small mean "
          "\nand a large tail, and it is the tail that a reader sees.")

    if args.out:
        Path(args.out).write_text(json.dumps(
            {"summary": summary, "kld_per_position": klds,
             "same_top_per_position": agreements}, indent=2), encoding="utf-8")
        print(f"written to {args.out}")
    if args.labbook:
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(summary) + "\n")
        print(f"recorded in {args.labbook}")


def hash_prefix(prefix: str) -> str:
    """A short, stable fingerprint so a saved reference cannot be paired with different prompts."""
    import hashlib
    return hashlib.sha256(prefix.encode("utf-8")).hexdigest()[:16]


if __name__ == "__main__":
    main()