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

Adds the models untested in round 1 and fixes the OOM that killed
Nemotron-3-Super-120B: instead of forcing -ngl 99 (which pinned all
layers and blew up device 0), this version lets llama.cpp auto-fit
layers to the two cards, and only falls back to a fixed high ngl if
auto-fit is unavailable.

Same output format as v1 (results/<name>.log), same COMPLETE_MARKER,
same question set, so v1 and v2 logs are graded identically. Models
already completed in results/ are skipped, so this can run alongside
the v1 outputs without redoing them.

Usage:
    cd ~/workarea/ash/llm_knowledge_test
    python3 knowledge_bench2.py
"""

import glob
import json
import os
import signal
import subprocess
import sys
import time
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 = 1200       # dense IQ2 123B + big MoE can be slow to load
QUESTION_TIMEOUT_S = 400    # IQ2 dense is slow per token; give room
MAX_TOKENS = 1200
CTX = 8192

# Models NOT already run in v1. Set includes the round-1 failure
# (nemotron-super-120b) so it gets a second chance with auto-fit.
# (short_name, path-or-glob relative to MODEL_DIR, fit_mode)
#   fit_mode "auto"  -> no -ngl, llama.cpp splits across both cards
#   fit_mode "full"  -> -ngl 99 (fits comfortably in 64GB)
MODELS = [
    ("nemotron3-super-120b",   "nvidia_Nemotron-3-Super-120B-A12B-Q4_K_M/*00001-of-*.gguf", "auto"),
    ("minimax-m2-iq2",         "MiniMaxAI_MiniMax-M2-IQ2_S/*00001-of-*.gguf",               "auto"),
    ("minimax-m2-iq2",         "MiniMaxAI_MiniMax-M2-IQ2_S",                                 "auto"),  # dir fallback
    ("devstral2-123b-iq2",     "Devstral-2-123B-Instruct-2512-UD-IQ2_M.gguf",               "auto"),
    ("nemotron-cascade-30b",   "nvidia_Nemotron-Cascade-2-30B-A3B-Q4_K_M.gguf",             "full"),
    ("nemotron3-nano-30b",     "Nemotron-3-Nano-30B-A3B-Q5_K_M.gguf",                       "full"),
    ("devstral-small2-24b",    "mistralai_Devstral-Small-2-24B-Instruct-2512-Q4_K_M.gguf",  "full"),
    ("gemma4-26b-a4b",         "google_gemma-4-26B-A4B-it-Q4_K_M.gguf",                     "full"),
    ("ministral3-14b",         "Ministral-3-14B-Reasoning-2512-Q5_K_M.gguf",                "full"),
    ("north-mini-code",        "North-Mini-Code-1.0-UD-Q4_K_XL.gguf",                       "full"),
    ("aya-23-8b",              "aya-23-8B-Q5_K_M.gguf",                                     "full"),
    ("granite4-micro",         "granite-4.0-micro-Q4_K_M.gguf",                             "full"),
    ("gpt-oss-20b",            "openai_gpt-oss-20b-Q5_K_M.gguf",                            "full"),
    ("glm4.6v-iq2",            "GLM-4.6V-UD-IQ2_M.gguf",                                     "auto"),
]

# GLM-4.6V is a vision model; it's included for completeness but its
# text-only knowledge score is not directly comparable to the dense LLMs.

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."
)

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):
    full = os.path.join(MODEL_DIR, pattern)
    if "*" in full:
        matches = sorted(glob.glob(full))
        return matches[0] if matches else None
    if os.path.isdir(full):
        shards = sorted(glob.glob(os.path.join(full, "*00001-of-*.gguf")))
        if shards:
            return shards[0]
        ggufs = sorted(glob.glob(os.path.join(full, "*.gguf")))
        return ggufs[0] if ggufs else None
    return full if os.path.exists(full) else None


def http_json(url, payload, timeout):
    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, deadline_s):
    start = time.time()
    while time.time() - start < deadline_s:
        if proc.poll() is not None:
            return False
        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):
    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:
        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):
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=30)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
    time.sleep(10)  # release VRAM before next model


def build_cmd(model_path, fit_mode):
    cmd = [
        SERVER_BIN, "-m", model_path,
        "-sm", "layer", "-c", str(CTX), "-t", "12",
        "--host", HOST, "--port", str(PORT),
        "--jinja", "--reasoning-format", "auto",
    ]
    if fit_mode == "full":
        cmd += ["-ngl", "99"]
    # auto: omit -ngl entirely so llama.cpp fits layers across both GPUs
    return cmd


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


def run_model(name, model_path, fit_mode):
    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} ({fit_mode}): {model_path}")
    env = dict(os.environ, HIP_VISIBLE_DEVICES="0,1")
    cmd = build_cmd(model_path, fit_mode)
    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}: 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 = []
    seen = set()
    for name, pattern, fit_mode in MODELS:
        if name in seen:
            continue  # already resolved via an earlier glob variant
        path = resolve_model_path(pattern)
        if path is None:
            continue
        seen.add(name)
        todo.append((name, path, fit_mode))

    # report anything that resolved to nothing at all
    resolved = {n for n, _, _ in todo}
    for name, pattern, _ in MODELS:
        if name not in resolved and name not in seen:
            print(f"[warn] {name}: no file matching {pattern}, skipping")

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

    t0 = time.time()

    def on_sigint(sig, frame):
        raise KeyboardInterrupt

    signal.signal(signal.SIGINT, on_sigint)

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

    print(f"\nElapsed: {(time.time() - t0) / 60:.0f} min")
    print("Hand results/*.log to Claude for scoring.")


if __name__ == "__main__":
    main()
