#!/usr/bin/env python3
"""Leave-one-out nearest-neighbour benchmark for an embedding model.

Embeds a labelled corpus of Paperless documents and, for each document, predicts
its correspondent and document type from the label of its nearest neighbour
(cosine similarity, the document itself excluded). Accuracy is the share of
documents whose predicted label equals the filed one.

Corpus: JSON lines {"id", "correspondent", "doctype", "text"}, exported
read-only from the paperless database, for example on the CNPG primary:

    select json_build_object('id', d.id, 'correspondent', d.correspondent_id,
        'doctype', d.document_type_id,
        'text', left(regexp_replace(d.content, '\\s+', ' ', 'g'), 3000))
    from documents_document d
    where d.correspondent_id is not null and d.document_type_id is not null
      and length(d.content) > 200
    order by md5(d.id::text || '<seed>') limit 400

Ceiling is the share of documents whose label occurs on at least one other
document; no nearest-neighbour method can exceed it. Floor is the share of the
most common label.

Usage:
    knn-bench.py --corpus corpus.jsonl -e https://api.mistral.ai/v1/embeddings \\
        -m mistral-embed --key-env MISTRAL_API_KEY --batch 8
    knn-bench.py --corpus corpus.jsonl -m bge-m3 \\
        -e http://llama-embed.ai.svc.cluster.local:8080/v1/embeddings \\
        --via-pod "deploy/llama-server -n ai"

Standard library only. Vectors are cached per model next to the corpus.
"""
import argparse
import json
import math
import os
import re
import shlex
import subprocess
import sys
import time
import urllib.error
import urllib.request
from collections import Counter


def post(endpoint, payload, key, via_pod):
    body = json.dumps(payload).encode()
    if via_pod:
        cmd = ["kubectl", "exec", "-i", *shlex.split(via_pod), "--", "curl", "-sS", "-m", "300",
               "-H", "Content-Type: application/json", "-d", "@-", endpoint]
        proc = subprocess.run(cmd, input=body, capture_output=True, check=False)
        if proc.returncode != 0:
            raise RuntimeError(f"kubectl exec failed: {proc.stderr.decode()}")
        return json.loads(proc.stdout)
    headers = {"Content-Type": "application/json"}
    if key:
        headers["Authorization"] = f"Bearer {key}"
    req = urllib.request.Request(endpoint, data=body, headers=headers)
    for attempt in range(6):
        try:
            with urllib.request.urlopen(req, timeout=300) as resp:
                return json.load(resp)
        except urllib.error.HTTPError as err:
            if err.code == 429 and attempt < 5:
                time.sleep(2 ** attempt)
                continue
            raise RuntimeError(f"HTTP {err.code}: {err.read().decode()}") from err
    raise RuntimeError("retries exhausted")


def embed_all(docs, args, key):
    cache = f"{args.corpus}.{args.model}.{args.max_chars}.vecs.json"
    if os.path.exists(cache):
        with open(cache) as fh:
            vecs = json.load(fh)
        if len(vecs) == len(docs):
            print(f"vectors loaded from {cache}")
            return vecs
    vecs = []
    shrunk = 0
    start = time.time()
    for i in range(0, len(docs), args.batch):
        chunk = [d["text"] for d in docs[i:i + args.batch]]
        for _ in range(6):
            try:
                out = post(args.endpoint, {"model": args.model, "input": chunk}, key, args.via_pod)
                break
            except RuntimeError as err:
                # Mistral reports "Input id N has T tokens, exceeding max M tokens": cut that text to fit.
                hit = re.search(r"Input id (\d+) has (\d+) tokens, exceeding max (\d+)", str(err))
                if not hit:
                    raise
                idx, have, limit = (int(g) for g in hit.groups())
                chunk[idx] = chunk[idx][:int(len(chunk[idx]) * limit / have * 0.95)]
                shrunk += 1
        else:
            raise RuntimeError(f"documents {i}-{i + len(chunk) - 1} still over the token limit")
        if "data" not in out:
            raise RuntimeError(f"no embeddings in response for documents {i}-{i + len(chunk) - 1}: {json.dumps(out)[:500]}")
        data = sorted(out["data"], key=lambda item: item["index"])
        if len(data) != len(chunk):
            raise RuntimeError(f"asked for {len(chunk)} vectors, got {len(data)}")
        vecs.extend(item["embedding"] for item in data)
        done = len(vecs)
        if done % 50 < args.batch:
            print(f"  embedded {done}/{len(docs)} ({time.time() - start:.0f}s)", flush=True)
    if shrunk:
        print(f"  {shrunk} over-limit texts were shortened to fit the model's token limit")
    with open(cache, "w") as fh:
        json.dump(vecs, fh)
    return vecs


def normalise(vec):
    norm = math.sqrt(sum(x * x for x in vec)) or 1.0
    return [x / norm for x in vec]


def nearest_neighbours(vecs):
    n = len(vecs)
    best = [(-2.0, -1)] * n
    for i in range(n):
        vi = vecs[i]
        for j in range(i + 1, n):
            s = sum(a * b for a, b in zip(vi, vecs[j]))
            if s > best[i][0]:
                best[i] = (s, j)
            if s > best[j][0]:
                best[j] = (s, i)
    return [j for _, j in best]


def score(docs, neighbours, field):
    labels = [d[field] for d in docs]
    counts = Counter(labels)
    hits = sum(labels[i] == labels[neighbours[i]] for i in range(len(docs)))
    ceiling = sum(counts[label] > 1 for label in labels)
    floor = max(counts.values())
    n = len(docs)
    return 100 * hits / n, 100 * ceiling / n, 100 * floor / n, len(counts)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--corpus", required=True)
    ap.add_argument("-e", "--endpoint", required=True)
    ap.add_argument("-m", "--model", required=True)
    ap.add_argument("--key-env", default="", help="environment variable holding the API key")
    ap.add_argument("--via-pod", default="", help='run curl in this pod, e.g. "deploy/llama-server -n ai"')
    ap.add_argument("--batch", type=int, default=1)
    ap.add_argument("--max-chars", type=int, default=0, help="truncate each text to this many characters (0 = as exported)")
    args = ap.parse_args()

    key = os.environ.get(args.key_env, "") if args.key_env else ""
    if args.key_env and not key:
        sys.exit(f"{args.key_env} is not set")

    with open(args.corpus) as fh:
        docs = [json.loads(line) for line in fh if line.strip()]
    if args.max_chars:
        for d in docs:
            d["text"] = d["text"][:args.max_chars]
    print(f"model={args.model} documents={len(docs)} max_chars={args.max_chars} endpoint={args.endpoint}")

    start = time.time()
    vecs = [normalise(v) for v in embed_all(docs, args, key)]
    embed_s = time.time() - start
    neighbours = nearest_neighbours(vecs)

    print(f"dimensions={len(vecs[0])} embed+load={embed_s:.0f}s")
    for field in ("correspondent", "doctype"):
        acc, ceiling, floor, classes = score(docs, neighbours, field)
        print(f"{field:<14} accuracy={acc:5.1f}%  ceiling={ceiling:5.1f}%  floor={floor:5.1f}%  labels={classes}")


if __name__ == "__main__":
    main()