"""Quality comparison against other small models (the human, 2026-10-27): the RS-fast items asked through each model's own
chat template, so models with different tokenizers and templates see the same questions.

Multiple choice (MMLU-Redux, Global-MMLU-Lite x5, Belebele x5): the RS-fast item text with an answer-with-the-letter
instruction, rendered by the model's own template up to the start of its reply (/apply-template, thinking off), then
"The answer is" appended as plain text and the next token held to A-D by a grammar (/completion), so greedy decoding
returns the model's most likely letter at the answer position. (Free-text parsing failed: Marco explains before
answering and Hungarian's article "A" reads as an answer; a grammar on the reply's first token alone gave mostly "A";
a pre-filled assistant message let Gemma 4's template insert an empty thought channel before it, which turned 54% into
40% with 67% "D" on 120 MMLU items.)
GSM8K-250: question + MATH_SUFFIX, greedy, scored as in RS-fast. Thinking is off everywhere.

--suite full (the human, 22:10: "do more benchmarks") adds the RS-full sets: MMLU-Pro 2,000 (direct answer, letters
A-J by the same grammar; RS-full's own MMLU-Pro is chain-of-thought), Global-MMLU 1,000 in five more languages
(de, es, fr, ja, ru), MATH-500 (math-verify against the last \\boxed{}) and IFEval 541 (Google's checks via lm-eval;
prompt-level strict accuracy, instruction-level kept per item). Scoring runs in the main thread (math-verify's
timeouts use signals).

One llama-server per model on a spare port (never the router's :8081). Per-item results go to <out>/<model>.jsonl and are
reused on a rerun, so an interrupted run resumes.

  python scripts/compare_quality.py --server .../llama-server --out reports/compare \
      --model marco-q8=/path/teacher-q8_0.gguf --model gemma4-e2b=/path/gemma.gguf ...
"""
from __future__ import annotations

import argparse
import json
import re
import signal
import subprocess
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO / "configs" / "eval" / "lm_eval_tasks"))
sys.path.insert(0, str(REPO / "src"))
import rs_utils  # noqa: E402
from lowbit.eval.gsm8k import MATH_SUFFIX, extract_answer, same_number  # noqa: E402

LANGS = ["zh", "ar", "ko", "hu", "bn"]
MC_INSTRUCTION = "Answer with the letter of the correct option (A, B, C or D) only."


def mc_items() -> list[dict]:
    out = []
    for d in rs_utils.mmlu_redux()["test"]:
        text = (f"The following are multiple choice questions (with answers) about {d['subject']}.\n\n"
                + rs_utils._mc_block(d["question"], d["choices"]))
        out.append({"task": "mmlu_redux", "id": d["item_id"], "text": text, "gold": "ABCD"[d["answer"]]})
    for lang in LANGS:
        for d in rs_utils.global_mmlu_lite(lang)["test"]:
            out.append({"task": f"gmmlu_{lang}", "id": d["item_id"], "text": rs_utils._mc_block(d["question"], d["choices"]),
                        "gold": "ABCD"[d["answer"]]})
        for d in rs_utils.belebele(lang)["test"]:
            text = f"P: {d['passage']}\nQ: {d['question'].strip()}" + "".join(
                f"\n{'ABCD'[i]}: {c}" for i, c in enumerate(d["choices"])) + "\nAnswer:"
            out.append({"task": f"belebele_{lang}", "id": d["item_id"], "text": text, "gold": "ABCD"[d["answer"]]})
    for it in out:  # the RS-fast text ends in "Answer:"; a chat model is told what to answer with instead
        it["prompt"] = it.pop("text").removesuffix("Answer:") + MC_INSTRUCTION
        it["letters"] = "ABCD"
    return out


FULL_LANGS = ["de", "es", "fr", "ja", "ru"]
GEN_TOKENS = {"gsm8k": 1536, "math500": 2048, "ifeval": 1536}


def full_items() -> list[dict]:
    out = []
    for d in rs_utils.mmlu_pro()["test"]:
        letters = "ABCDEFGHIJ"[:len(d["options"])]
        opts = "".join(f"\n{letters[i]}. {o}" for i, o in enumerate(d["options"]))
        out.append({"task": "mmlu_pro", "id": d["item_id"], "gold": d["gold"], "letters": letters,
                    "prompt": f"The following is a multiple choice question about {d['category']}.\n\n"
                              f"Question: {d['question'].strip()}\nOptions:{opts}\n"
                              f"Answer with the letter of the correct option (A to {letters[-1]}) only."})
    for lang in FULL_LANGS:
        for d in rs_utils.global_mmlu(lang)["test"]:
            out.append({"task": f"gmmlux_{lang}", "id": d["item_id"], "gold": "ABCD"[d["answer"]], "letters": "ABCD",
                        "prompt": rs_utils._mc_block(d["question"], d["choices"]).removesuffix("Answer:") + MC_INSTRUCTION})
    for d in rs_utils.math500()["test"]:
        out.append({"task": "math500", "id": d["item_id"], "gold": d["gold"], "prompt": rs_utils.math500_prompt(d)})
    for d in rs_utils.ifeval()["test"]:
        out.append({"task": "ifeval", "id": d["item_id"], "gold": None, "prompt": d["prompt"],
                    "doc": {k: d[k] for k in ("key", "prompt", "instruction_id_list", "kwargs")}})
    return out


def gsm_items() -> list[dict]:
    return [{"task": "gsm8k", "id": d["item_id"], "prompt": d["question"] + MATH_SUFFIX, "gold": d["gold"]}
            for d in rs_utils.gsm8k_250()["test"]]


THINK = re.compile(r"<think>.*?</think>", re.S)
MC_PREFILL = "The answer is"
MC_GRAMMAR = 'root ::= [ ]? [(]? [{}]'  # the item's letters


def post(port: int, path: str, body: dict) -> dict:
    req = urllib.request.Request(f"http://127.0.0.1:{port}{path}", json.dumps(body).encode(),
                                 {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=900) as r:
        return json.load(r)


def ask(port: int, it: dict) -> dict:
    """The model's raw answer; scoring happens in the main thread (score)."""
    mc = "letters" in it
    msgs = [{"role": "user", "content": it["prompt"]}]
    kw = {"chat_template_kwargs": {"enable_thinking": False}}
    try:
        if mc:
            prompt = post(port, "/apply-template", {"messages": msgs} | kw)["prompt"] + MC_PREFILL
            res = post(port, "/completion", {"prompt": prompt, "n_predict": 4, "temperature": 0.0,
                                             "grammar": MC_GRAMMAR.format(it["letters"])})
            return {"content": res["content"], "n_out": res.get("tokens_predicted")}
        res = post(port, "/v1/chat/completions", {"messages": msgs, "temperature": 0.0,
                                                  "max_tokens": GEN_TOKENS[it["task"]]} | kw)
        return {"content": res["choices"][0]["message"].get("content") or "",
                "n_out": res.get("usage", {}).get("completion_tokens")}
    except Exception as e:  # noqa: BLE001 — an item the server refuses (e.g. over the slot context) is scored wrong
        return {"error": str(e)[:300]}


def score(it: dict, a: dict) -> dict:
    row = {"task": it["task"], "id": it["id"], "gold": it["gold"]}
    if "error" in a:
        return row | {"pred": None, "how": "error", "correct": False, "error": a["error"]}
    content = a["content"]
    if "letters" in it:
        pred = content.strip().lstrip("(") or None
        how, ok = ("grammar" if pred and pred in it["letters"] else "none"), pred == it["gold"]
    elif it["task"] == "gsm8k":
        pred, how = extract_answer(THINK.sub("", content)), "gsm8k"
        ok = same_number(pred, it["gold"])
    elif it["task"] == "math500":
        pred, how = None, "math-verify"
        ok = bool(rs_utils.math500_process_results(it, [THINK.sub("", content)])["exact_match"])
    else:  # ifeval
        r = rs_utils.ifeval_process_results(it["doc"], [THINK.sub("", content)])
        pred, how, ok = None, "ifeval", bool(r["prompt_level_strict_acc"])
        row["inst_strict"] = list(map(bool, r["inst_level_strict_acc"]))
    return row | {"pred": pred, "how": how, "correct": bool(ok), "out": content[-400:], "n_out": a.get("n_out")}


def _up(port: int) -> bool:
    try:
        with urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2) as r:
            return json.load(r).get("status") == "ok"
    except Exception:  # noqa: BLE001 — nothing there, or still loading
        return False


def serve(server: str, gguf: str, port: int, ngl: str, log: Path) -> subprocess.Popen:
    """Start llama-server for one model and return once that same process answers with that same model. An orphaned
    server from an interrupted run once held the port, so a model's own server could not bind and every question was
    answered by the orphan's model (27 Oct 23:32; misstep #32): the port must be free first, and the answering server
    must be this process with this file."""
    if _up(port):
        raise RuntimeError(f"port {port} already has a llama-server; stop it first")
    cmd = [server, "-m", gguf, "--host", "127.0.0.1", "--port", str(port), "-ngl", ngl, "-c", "32768", "-np", "8",
           "-fa", "on", "--jinja", "--reasoning-budget", "0", "--no-webui", "-t", "16"]
    p = subprocess.Popen(cmd, stdout=log.open("w"), stderr=subprocess.STDOUT, start_new_session=True)
    try:
        for _ in range(600):
            if p.poll() is not None:
                raise RuntimeError(f"llama-server exited ({p.returncode}); see {log}")
            if _up(port):
                props = post_get(port, "/props")
                if p.poll() is not None or Path(props.get("model_path", "")).name != Path(gguf).name:
                    raise RuntimeError(f"port {port} answered with {props.get('model_path')}, not {gguf}")
                return p
            time.sleep(1)
        raise RuntimeError("llama-server not ready after 10 min")
    except BaseException:  # interrupted while loading (SIGTERM -> SystemExit) or failed: never leave it running
        p.kill()
        p.wait(30)
        raise


def post_get(port: int, path: str) -> dict:
    with urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=10) as r:
        return json.load(r)


GROUPS = {"MMLU-Redux": ["mmlu_redux"], "Global-MMLU (5 lang)": [f"gmmlu_{l}" for l in LANGS],
          "Belebele (5 lang)": [f"belebele_{l}" for l in LANGS], "GSM8K-250": ["gsm8k"]}
FULL_GROUPS = {"MMLU-Pro (direct)": ["mmlu_pro"], "Global-MMLU (de es fr ja ru)": [f"gmmlux_{l}" for l in FULL_LANGS],
               "MATH-500": ["math500"], "IFEval (prompt strict)": ["ifeval"]}


def summarise(name: str, rows: list[dict], groups_def: dict = GROUPS) -> dict:
    by = {}
    for r in rows:
        by.setdefault(r["task"], []).append(r["correct"])
    tasks = {t: sum(v) / len(v) for t, v in sorted(by.items())}
    groups = {g: sum(tasks[t] for t in ts) / len(ts) for g, ts in groups_def.items() if all(t in tasks for t in ts)}
    inst = [x for r in rows for x in r.get("inst_strict", [])]
    if inst:
        groups["IFEval (instruction strict)"] = sum(inst) / len(inst)
    hows = {}
    for r in rows:
        hows[r["how"]] = hows.get(r["how"], 0) + 1
    return {"model": name, "groups": groups, "tasks": tasks, "answer_parse": hows, "n": len(rows)}


def main() -> None:
    # a SIGTERM must still reach the finally below, or the model's llama-server (own session) outlives the run
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(143))
    ap = argparse.ArgumentParser()
    ap.add_argument("--server", required=True)
    ap.add_argument("--model", action="append", required=True, help="NAME=PATH.gguf")
    ap.add_argument("--out", type=Path, default=REPO / "reports" / "compare")
    ap.add_argument("--port", type=int, default=18190)
    ap.add_argument("--ngl", default="99")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--limit", type=int, default=0, help="first N items per task (smoke test)")
    ap.add_argument("--suite", choices=["fast", "full", "both"], default="fast")
    a = ap.parse_args()
    a.out.mkdir(parents=True, exist_ok=True)
    suites = {"fast": ["fast"], "full": ["full"], "both": ["fast", "full"]}[a.suite]
    items = (mc_items() + gsm_items() if "fast" in suites else []) + (full_items() if "full" in suites else [])
    if a.limit:
        count = {}
        for it in items:
            count[it["task"]] = count.get(it["task"], 0) + 1
            it["keep"] = count[it["task"]] <= a.limit
        items = [it for it in items if it.pop("keep")]
    print(f"{len(items)} items", flush=True)
    for spec in a.model:
        name, gguf = spec.split("=", 1)
        path = a.out / f"{name}.jsonl"
        done = {(r["task"], r["id"]): r for r in map(json.loads, path.open())} if path.exists() else {}
        todo = [it for it in items if (it["task"], it["id"]) not in done or done[(it["task"], it["id"])]["how"] == "error"]
        if todo:
            t0 = time.time()
            p = serve(a.server, gguf, a.port, a.ngl, a.out / f"{name}.server.log")
            try:
                (a.out / f"{name}.prompt.txt").write_text(post(a.port, "/apply-template", {
                    "messages": [{"role": "user", "content": items[0]["prompt"]}],
                    "chat_template_kwargs": {"enable_thinking": False}})["prompt"] + MC_PREFILL)
                with ThreadPoolExecutor(a.workers) as ex, path.open("a") as f:
                    for i, (it, ans) in enumerate(zip(todo, ex.map(lambda it: ask(a.port, it), todo))):
                        r = score(it, ans)
                        done[(r["task"], r["id"])] = r
                        f.write(json.dumps(r, ensure_ascii=False) + "\n")
                        if (i + 1) % 500 == 0:
                            print(f"{name}: {i + 1}/{len(todo)} ({time.time() - t0:.0f}s)", flush=True)
            finally:
                p.terminate()
                p.wait(60)
            print(f"{name}: {len(todo)} items in {time.time() - t0:.0f}s", flush=True)
        for suite, gdef, tag in (("fast", GROUPS, ""), ("full", FULL_GROUPS, ".full")):
            if suite not in suites:
                continue
            tasks = {t for ts in gdef.values() for t in ts}
            rows = [done[(it["task"], it["id"])] for it in items if it["task"] in tasks]
            s = summarise(name, rows, gdef) | {"gguf": gguf, "bytes": Path(gguf).stat().st_size}
            (a.out / f"{name}{tag}.summary.json").write_text(json.dumps(s, indent=1))
            print(json.dumps({"model": name, "suite": suite, **{k: round(v, 4) for k, v in s["groups"].items()},
                              "parse": s["answer_parse"]}), flush=True)


if __name__ == "__main__":
    main()