#!/usr/bin/env python3
"""Knowledge density benchmark for vinland dual-MI50.

Runs each model through llama-server one at a time, asks 40 factual
questions across 10 subjects, writes one logfile per model to results/.
Resumable: models with a completed logfile are skipped on rerun.

Usage:
    cd ~/workarea/ash/llm_knowledge_test
    python3 knowledge_bench.py

Judge afterwards by handing the results/*.log files to Claude.
"""

import glob
import json
import os
import signal
import subprocess
import sys
import time
import urllib.error
import urllib.request

# ---------------------------------------------------------------- config

SERVER_BIN = "/opt/llama.cpp-rocm/llama-server"
MODEL_DIR = os.path.expanduser("~/.cache/llama.cpp")
RESULTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
PORT = 8090
HOST = "127.0.0.1"
BASE_URL = f"http://{HOST}:{PORT}"

LOAD_TIMEOUT_S = 900        # max wait for a model to load and report healthy
QUESTION_TIMEOUT_S = 300    # max wait for one answer
MAX_TOKENS = 1200           # headroom for models that think despite the kwarg
CTX = 8192

# (short_name, path or glob relative to MODEL_DIR)
# Latest/best quant per family. Multi-file models point at shard 1 via glob.
MODELS = [
    ("gpt-oss-120b",            "openai_gpt-oss-120b-Q4_K_M/*00001-of-*.gguf"),
    ("nemotron3-super-120b",    "nvidia_Nemotron-3-Super-120B-A12B-Q4_K_M/*00001-of-*.gguf"),
    ("devstral2-123b-iq2",      "Devstral-2-123B-Instruct-2512-UD-IQ2_M.gguf"),
    ("qwen3next-80b-a3b",       "Qwen3-Next-80B-A3B-Instruct-Q4_K_S.gguf"),
    ("qwen3.6-35b-a3b",         "Qwen_Qwen3.6-35B-A3B-Q4_K_M.gguf"),
    ("qwen3.6-27b",             "Qwen_Qwen3.6-27B-Q4_K_M.gguf"),
    ("gemma4-31b",              "google_gemma-4-31B-it-Q4_K_M.gguf"),
    ("glm4.7-flash",            "zai-org_GLM-4.7-Flash-Q5_K_M.gguf"),
]

SYSTEM_PROMPT = (
    "You are being tested on factual knowledge. Answer directly and concisely, "
    "at most four sentences. Do not pad. If you are not sure, say 'not sure' "
    "and give your best guess."
)

# (subject, question)
QUESTIONS = [
    ("physics", "What is the Chandrasekhar limit in solar masses, and what does it represent?"),
    ("physics", "Which bosons mediate the weak nuclear force?"),
    ("physics", "What is the approximate value of the fine-structure constant?"),
    ("physics", "What did the Davisson-Germer experiment confirm?"),
    ("chemistry", "What is the oxidation state of manganese in potassium permanganate?"),
    ("chemistry", "Which element has the highest electronegativity after fluorine?"),
    ("chemistry", "What industrial process converts nitrogen gas to ammonia, and what catalyst does it typically use?"),
    ("chemistry", "What is a racemic mixture?"),
    ("biology", "Which nerve innervates the diaphragm, and from which spinal roots does it arise?"),
    ("biology", "What is the primary function of the loop of Henle?"),
    ("biology", "Deficiency of which vitamin causes pellagra, and what are its three classic symptoms?"),
    ("biology", "What organism causes Chagas disease and what vector transmits it?"),
    ("history", "In what year did the First Council of Nicaea convene?"),
    ("history", "Who is generally considered the last emperor of the Western Roman Empire?"),
    ("history", "What peace settlement ended the Thirty Years' War, and in what year?"),
    ("history", "Which Chinese dynasty ruled during the voyages of Zheng He?"),
    ("geography", "What is the deepest lake in the world and approximately how deep is it?"),
    ("geography", "Which two countries share the longest international land border?"),
    ("geography", "What strait separates Asia from North America?"),
    ("geography", "What is the capital of Burkina Faso?"),
    ("literature", "Who wrote The Master and Margarita?"),
    ("literature", "In Dante's Divine Comedy, who guides Dante through Paradise?"),
    ("literature", "What is the famous opening line of Anna Karenina?"),
    ("literature", "Who wrote One Hundred Years of Solitude and in what year was it first published?"),
    ("music", "What was Iron Maiden's debut studio album and in what year was it released?"),
    ("music", "Which composer wrote the Goldberg Variations?"),
    ("music", "Who produced System of a Down's album Toxicity?"),
    ("music", "The natural minor scale corresponds to which mode of the major scale?"),
    ("rf_engineering", "What is the approximate free-space path loss in dB at 1 km for a 2.4 GHz signal?"),
    ("rf_engineering", "What quantity does a Smith chart plot, and what is it used for?"),
    ("rf_engineering", "State the Shannon capacity formula and define its variables."),
    ("rf_engineering", "What modulation and spreading scheme does the GPS L1 C/A signal use?"),
    ("math", "What is the exact value of the Riemann zeta function at s=2?"),
    ("math", "State the Central Limit Theorem in one sentence."),
    ("math", "Write out Euler's identity."),
    ("math", "How many groups of order 8 exist up to isomorphism?"),
    ("econ_law", "What is the difference between nominal and real GDP?"),
    ("econ_law", "What legal principle does habeas corpus protect?"),
    ("econ_law", "Who won the Nobel Memorial Prize in Economics for prospect theory?"),
    ("econ_law", "What is Gresham's law?"),
]

COMPLETE_MARKER = "=== RUN COMPLETE ==="

# ---------------------------------------------------------------- helpers


def resolve_model_path(pattern: str) -> str | None:
    full = os.path.join(MODEL_DIR, pattern)
    if "*" in full:
        matches = sorted(glob.glob(full))
        return matches[0] if matches else None
    return full if os.path.exists(full) else None


def http_json(url: str, payload: dict | None, timeout: float):
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode())


def wait_healthy(proc: subprocess.Popen, deadline_s: float) -> bool:
    start = time.time()
    while time.time() - start < deadline_s:
        if proc.poll() is not None:
            return False  # server died during load
        try:
            r = http_json(f"{BASE_URL}/health", None, 5)
            if r.get("status") == "ok":
                return True
        except Exception:
            pass
        time.sleep(5)
    return False


def ask(question: str) -> dict:
    payload = {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": question},
        ],
        "temperature": 0,
        "max_tokens": MAX_TOKENS,
        "chat_template_kwargs": {"enable_thinking": False},
    }
    t0 = time.time()
    resp = http_json(f"{BASE_URL}/v1/chat/completions", payload, QUESTION_TIMEOUT_S)
    elapsed = time.time() - t0
    choice = resp.get("choices", [{}])[0]
    msg = choice.get("message", {})
    content = (msg.get("content") or "").strip()
    reasoning = (msg.get("reasoning_content") or "").strip()
    if not content and reasoning:
        # model spent everything thinking; keep the tail of the reasoning
        content = "[no final answer, reasoning tail follows]\n" + reasoning[-800:]
    timings = resp.get("timings", {})
    return {
        "content": content,
        "finish": choice.get("finish_reason", "?"),
        "elapsed": elapsed,
        "tps": timings.get("predicted_per_second"),
        "tokens": resp.get("usage", {}).get("completion_tokens"),
    }


def stop_server(proc: subprocess.Popen):
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=30)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
    time.sleep(8)  # let VRAM release before the next model


# ---------------------------------------------------------------- main


def run_model(name: str, model_path: str) -> None:
    logfile = os.path.join(RESULTS_DIR, f"{name}.log")
    serverlog = os.path.join(RESULTS_DIR, f"{name}.server.log")

    if os.path.exists(logfile):
        with open(logfile) as f:
            if COMPLETE_MARKER in f.read():
                print(f"[skip] {name}: already complete")
                return

    print(f"[load] {name}: {model_path}")
    env = dict(os.environ, HIP_VISIBLE_DEVICES="0,1")
    cmd = [
        SERVER_BIN, "-m", model_path,
        "-ngl", "99", "-sm", "layer",
        "-c", str(CTX), "-t", "12",
        "--host", HOST, "--port", str(PORT),
        "--jinja", "--reasoning-format", "auto",
    ]
    slog = open(serverlog, "w")
    proc = subprocess.Popen(cmd, stdout=slog, stderr=subprocess.STDOUT, env=env)
    try:
        if not wait_healthy(proc, LOAD_TIMEOUT_S):
            print(f"[FAIL] {name}: server did not become healthy, see {serverlog}")
            stop_server(proc)
            with open(logfile, "w") as f:
                f.write(f"MODEL: {name}\nPATH: {model_path}\nFAILED TO LOAD\n")
            return

        t_model = time.time()
        with open(logfile, "w") as f:
            f.write(f"MODEL: {name}\nPATH: {model_path}\n")
            f.write(f"STARTED: {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            f.flush()
            for i, (subject, q) in enumerate(QUESTIONS, 1):
                print(f"  [{name}] Q{i:02d}/{len(QUESTIONS)} ({subject})", flush=True)
                try:
                    r = ask(q)
                except Exception as e:
                    r = {"content": f"[ERROR: {e}]", "finish": "error",
                         "elapsed": 0, "tps": None, "tokens": None}
                f.write("-" * 70 + "\n")
                f.write(f"Q{i:02d} [{subject}]\n{q}\n\n")
                f.write(f"{r['content']}\n\n")
                tps = f"{r['tps']:.1f}" if r["tps"] else "?"
                f.write(f"(finish={r['finish']} tokens={r['tokens']} "
                        f"time={r['elapsed']:.1f}s tps={tps})\n\n")
                f.flush()
            f.write("-" * 70 + "\n")
            f.write(f"TOTAL: {time.time() - t_model:.0f}s for {len(QUESTIONS)} questions\n")
            f.write(COMPLETE_MARKER + "\n")
        print(f"[done] {name} -> {logfile}")
    finally:
        stop_server(proc)


def main():
    os.makedirs(RESULTS_DIR, exist_ok=True)
    if not os.path.exists(SERVER_BIN):
        sys.exit(f"llama-server not found at {SERVER_BIN}")

    todo = []
    for name, pattern in MODELS:
        path = resolve_model_path(pattern)
        if path is None:
            print(f"[warn] {name}: no file matching {pattern}, skipping")
        else:
            todo.append((name, path))

    print(f"Testing {len(todo)} models, {len(QUESTIONS)} questions each.")
    print(f"Results: {RESULTS_DIR}\n")

    t0 = time.time()
    interrupted = False

    def on_sigint(sig, frame):
        nonlocal interrupted
        interrupted = True
        raise KeyboardInterrupt

    signal.signal(signal.SIGINT, on_sigint)

    for name, path in todo:
        try:
            run_model(name, path)
        except KeyboardInterrupt:
            print("\nInterrupted. Rerun to resume; completed models are skipped.")
            break

    print(f"\nElapsed: {(time.time() - t0) / 60:.0f} min")
    if not interrupted:
        print("All done. Hand the results/*.log files to Claude for scoring.")


if __name__ == "__main__":
    main()