#!/usr/bin/env python3
"""Role-fitness benchmark for vinland dual-MI50 (overnight run).

Prior rounds measured knowledge. This measures FITNESS FOR ROLE. Each
model runs three sections, and is scored per-role, not as one total:

  [main]        default daily driver: breadth, calibration, instruction-
                following, termination, concise-when-asked, honest "I
                don't know". Includes a couple of warm-sampled writing
                items.
  [escalation]  the reasoner/reviewer: multi-step problems, subtle-bug
                spotting in supplied code, false-premise resistance,
                catching an error in a plausible-looking claim.
  [worker]      bounded implementer: small self-contained coding tasks
                that are ACTUALLY EXECUTED against assertions, plus
                strict output-format compliance and termination.

Worker code tasks run in a sandboxed subprocess with a hard timeout so
no single generation can hang the overnight batch. A failed/blocked/
timed-out execution is recorded as a failure and the run continues.

Sampling is role-tuned: deterministic (temp 0) for escalation and
worker, mildly warm (temp 0.7) for the writing items in main. Thinking
is enabled; token budgets are generous but every call has a timeout.

Writes one logfile per model to results_roles/. Resumable: completed
models are skipped. Judge afterward with the companion prompt.

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

import glob
import json
import os
import signal
import subprocess
import sys
import tempfile
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_roles")
PORT = 8090
HOST = "127.0.0.1"
BASE_URL = f"http://{HOST}:{PORT}"

LOAD_TIMEOUT_S = 1200
QUESTION_TIMEOUT_S = 900     # overnight: give slow IQ2 dense room to think
MAX_TOKENS = 6144
CTX = 24576
CODE_EXEC_TIMEOUT_S = 20     # hard kill for any executed model-written code

# Curated field. Dropped: aya (weak), granite (weak for role work),
# qwen3-next-80b + qwen3.6-35b-a3b (redundant MoE), glm4.7-flash (weak
# RF signal), nemotron3-nano (no distinct role), gemma4-26b-a4b
# (redundant with 31b). Kept the role contenders + added small worker
# candidates (qwen3.5-9b, ministral3-14b). minimax gets a real retry
# with a big token budget and temp>0 to break the empty-final loop.
# (name, path-or-glob, fit_mode, minoptokens_override_or_None)
MODELS = [
    ("gpt-oss-120b",         "openai_gpt-oss-120b-Q4_K_M/*00001-of-*.gguf",                "auto", None),
    ("nemotron3-super-120b", "nvidia_Nemotron-3-Super-120B-A12B-Q4_K_M/*00001-of-*.gguf", "auto", None),
    ("minimax-m2-iq2",       "MiniMaxAI_MiniMax-M2-IQ2_S",                                 "auto", 8192),
    ("devstral2-123b-iq2",   "Devstral-2-123B-Instruct-2512-UD-IQ2_M.gguf",               "auto", None),
    ("gemma4-31b",           "google_gemma-4-31B-it-Q4_K_M.gguf",                         "full", None),
    ("qwen3.6-27b",          "Qwen_Qwen3.6-27B-Q4_K_M.gguf",                              "full", None),
    ("nemotron-cascade-30b", "nvidia_Nemotron-Cascade-2-30B-A3B-Q4_K_M.gguf",             "full", None),
    ("north-mini-code",      "North-Mini-Code-1.0-UD-Q4_K_XL.gguf",                       "full", None),
    ("gpt-oss-20b",          "openai_gpt-oss-20b-Q5_K_M.gguf",                            "full", None),
    ("devstral-small2-24b",  "mistralai_Devstral-Small-2-24B-Instruct-2512-Q4_K_M.gguf",  "full", None),
    ("ministral3-14b",       "Ministral-3-14B-Reasoning-2512-Q5_K_M.gguf",                "full", None),
    ("qwen3.5-9b",           "Qwen_Qwen3.5-9B-Q4_K_M.gguf",                               "full", None),
]

MAIN_SYSTEM = (
    "You are a helpful assistant. Answer accurately and concisely. Follow "
    "formatting instructions exactly. If you do not know something, say so "
    "plainly rather than guessing. Do not pad."
)
ESCALATION_SYSTEM = (
    "You are a senior technical reviewer. Reason carefully step by step. If a "
    "claim, premise, or piece of code is wrong, identify the error precisely "
    "and give the correct version. Do not agree with something just because it "
    "is presented confidently. End with a clear verdict."
)
WORKER_SYSTEM = (
    "You are a precise implementer. Produce exactly what is asked, in the exact "
    "format requested. When asked for code, output a single self-contained code "
    "block that runs as-is with no commentary outside the block. Always finish "
    "with a complete answer."
)

# section, id, sampling(temp), question, expected(judge)/exec-spec
# For worker code tasks, expected is a dict: {"fn": name, "tests": [...],
# "lang": "python"} and the judge/harness executes.
QUESTIONS = [

    # ============================= MAIN =============================
    ("main", "m01", 0.0,
     "In one sentence each, define latency and throughput, and state how they differ.",
     "Latency = time for one operation; throughput = operations per unit time; low latency != high throughput necessarily."),
    ("main", "m02", 0.0,
     "Answer with ONLY a single number and no other text: how many bits are in a byte?",
     "Format test: must output exactly '8' with nothing else."),
    ("main", "m03", 0.0,
     "What is the airspeed velocity of an unladen European swallow? If this is a joke reference, say so and give the joke's answer; if it has a real estimate, give it.",
     "Monty Python reference; real rough estimate ~11 m/s (24 mph). Good answer notes the reference AND gives the estimate."),
    ("main", "m04", 0.0,
     "I think Python's GIL means Python can't do any parallelism at all. Is that right?",
     "Calibration/correction: GIL limits CPU-bound threading, but multiprocessing, async I/O, C extensions, and I/O-bound threads still parallelize. Good answer corrects the overstatement without being pedantic."),
    ("main", "m05", 0.0,
     "List exactly three lossless audio codecs, as a numbered list, nothing else.",
     "Format + fact: FLAC, ALAC, WavPack, Monkey's Audio, TTA, etc. Must be numbered, exactly 3, all actually lossless (MP3/AAC/Opus are wrong)."),
    ("main", "m06", 0.7,
     "Write a two-sentence product blurb for a rugged handheld software-defined radio aimed at field engineers. No hype words like 'revolutionary' or 'seamless'.",
     "Writing/instruction-following: 2 sentences, on-topic, avoids banned hype words. Judge on fluency + constraint compliance."),
    ("main", "m07", 0.7,
     "Give me three distinct, non-obvious project name ideas for an internal tool that schedules RF spectrum measurements. One line each, name plus a 4-word rationale.",
     "Brainstorm: 3 distinct creative names, format held (name + short rationale). Judge on creativity + format."),
    ("main", "m08", 0.0,
     "What is the capital of the country whose flag is a plain green rectangle with no other features?",
     "Trick/recall: Libya used a plain green flag 1977-2011 (capital Tripoli); no country currently uses one. Good answer notes the historical Libya flag and Tripoli, or flags that no current country has this."),

    # ========================== ESCALATION ==========================
    ("escalation", "e01", 0.0,
     "Review this claim and give a verdict: 'Increasing a receiver's bandwidth always improves its sensitivity because it captures more signal energy.' Is it correct?",
     "FALSE: wider bandwidth admits more NOISE (noise power = kTB), degrading sensitivity/SNR for a fixed signal. Good answer rejects and explains the noise-bandwidth tradeoff."),
    ("escalation", "e02", 0.0,
     "This Python function is meant to return True if a list has any duplicates. Find the bug:\n\ndef has_dup(xs):\n    seen = set()\n    for x in xs:\n        if x in seen:\n            return False\n        seen.add(x)\n    return True",
     "BUG: return values inverted. Returns False on finding a dup and True at end. Should return True on dup, False at end. Good answer identifies the inversion."),
    ("escalation", "e03", 0.0,
     "A colleague says: 'We can just use a 12-bit ADC at 10 Msps to perfectly reconstruct a 6 MHz bandwidth signal, since 10 Msps is above the 6 MHz signal frequency.' What's wrong with this reasoning?",
     "Nyquist error: to reconstruct a 6 MHz BANDWIDTH baseband signal you need >12 Msps (2x highest frequency), not just above 6 MHz. 10 Msps undersamples -> aliasing. (Bandpass sampling is a separate special case.) Good answer catches the Nyquist violation."),
    ("escalation", "e04", 0.0,
     "Compute and show steps: a matched filter receiver has a noise figure of 3 dB and operates at 290 K reference. If the input noise floor is -174 dBm/Hz, what is the output noise floor in dBm/Hz, and what is the equivalent noise temperature?",
     "Output noise floor = -174 + 3 = -171 dBm/Hz. Te = T0*(10^(3/10)-1) = 290*(1.995-1) = ~288.6 K. Both parts needed."),
    ("escalation", "e05", 0.0,
     "Is this correct? 'Since HTTPS encrypts traffic, an attacker on the same WiFi network cannot see which websites (domains) you visit.' Give a verdict.",
     "MOSTLY FALSE: HTTPS encrypts content, but SNI (unless ECH), DNS lookups (unless DoH/DoT), and IP addresses can reveal the domain. Good answer rejects the absolute claim and cites SNI/DNS/IP leakage."),
    ("escalation", "e06", 0.0,
     "Find the subtle error in this reasoning: 'The series 1 + 1/2 + 1/4 + 1/8 + ... diverges because it has infinitely many terms, so adding more always increases the sum without bound.'",
     "FALSE: this geometric series CONVERGES to 2. Infinitely many positive terms can sum to a finite limit when they shrink fast enough. Good answer identifies convergence to 2."),

    # ============================ WORKER ============================
    # These are executed. expected carries the exec spec.
    ("worker", "w01", 0.0,
     "Write a Python function `fizzbuzz(n)` that returns a list of strings for 1..n: 'Fizz' if divisible by 3, 'Buzz' if by 5, 'FizzBuzz' if both, else the number as a string. Output only a single python code block defining the function.",
     {"lang": "python", "fn": "fizzbuzz",
      "tests": ["assert fizzbuzz(5) == ['1','2','Fizz','4','Buzz']",
                "assert fizzbuzz(15)[-1] == 'FizzBuzz'",
                "assert fizzbuzz(3)[2] == 'Fizz'",
                "assert fizzbuzz(1) == ['1']"]}),
    ("worker", "w02", 0.0,
     "Write a Python function `dbm_to_watts(dbm)` converting dBm to watts. Output only a single python code block.",
     {"lang": "python", "fn": "dbm_to_watts",
      "tests": ["assert abs(dbm_to_watts(0) - 0.001) < 1e-9",
                "assert abs(dbm_to_watts(30) - 1.0) < 1e-6",
                "assert abs(dbm_to_watts(-30) - 1e-6) < 1e-12"]}),
    ("worker", "w03", 0.0,
     "Write a Python function `is_palindrome(s)` that returns True if s is a palindrome ignoring case, spaces, and punctuation. Output only a single python code block.",
     {"lang": "python", "fn": "is_palindrome",
      "tests": ["assert is_palindrome('A man, a plan, a canal: Panama') == True",
                "assert is_palindrome('race a car') == False",
                "assert is_palindrome('') == True",
                "assert is_palindrome('No lemon, no melon') == True"]}),
    ("worker", "w04", 0.0,
     "Write a Python function `parse_freq(s)` that parses strings like '2.4GHz', '900MHz', '1500 kHz' and returns the frequency in Hz as a float. Handle kHz, MHz, GHz (case-insensitive) and optional space. Output only a single python code block.",
     {"lang": "python", "fn": "parse_freq",
      "tests": ["assert abs(parse_freq('2.4GHz') - 2.4e9) < 1",
                "assert abs(parse_freq('900MHz') - 900e6) < 1",
                "assert abs(parse_freq('1500 kHz') - 1500e3) < 1",
                "assert abs(parse_freq('5 GHz') - 5e9) < 1"]}),
    ("worker", "w05", 0.0,
     "Write a Python function `moving_average(xs, k)` returning the list of k-window simple moving averages (length len(xs)-k+1). Output only a single python code block.",
     {"lang": "python", "fn": "moving_average",
      "tests": ["assert moving_average([1,2,3,4], 2) == [1.5,2.5,3.5]",
                "assert moving_average([2,4,6], 3) == [4.0]",
                "assert len(moving_average(list(range(10)), 3)) == 8"]}),
    ("worker", "w06", 0.0,
     "Respond with ONLY valid JSON, no code fence, no prose: an object with keys 'name' (string 'test'), 'count' (integer 3), 'items' (array of the strings 'a','b','c').",
     {"lang": "json_exact",
      "expect": {"name": "test", "count": 3, "items": ["a", "b", "c"]}}),
    ("worker", "w07", 0.0,
     "Write a Python function `roman_to_int(s)` converting a Roman numeral string to an integer (valid input up to 3999). Output only a single python code block.",
     {"lang": "python", "fn": "roman_to_int",
      "tests": ["assert roman_to_int('IV') == 4",
                "assert roman_to_int('XLII') == 42",
                "assert roman_to_int('MCMXciv'.upper()) == 1994",
                "assert roman_to_int('III') == 3"]}),
    ("worker", "w08", 0.0,
     "Write a Python function `bit_reverse(x, width)` that reverses the lowest `width` bits of integer x and returns the result. Output only a single python code block.",
     {"lang": "python", "fn": "bit_reverse",
      "tests": ["assert bit_reverse(1, 4) == 8",
                "assert bit_reverse(0b1011, 4) == 0b1101",
                "assert bit_reverse(0, 8) == 0",
                "assert bit_reverse(0xFF, 8) == 0xFF"]}),
]

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

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


def resolve_model_path(pattern):
    full = os.path.join(MODEL_DIR, pattern)
    if "*" in full:
        m = sorted(glob.glob(full))
        return m[0] if m else None
    if os.path.isdir(full):
        shards = sorted(glob.glob(os.path.join(full, "*00001-of-*.gguf")))
        if shards:
            return shards[0]
        g = sorted(glob.glob(os.path.join(full, "*.gguf")))
        return g[0] if g 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 section_system(section):
    return {"main": MAIN_SYSTEM, "escalation": ESCALATION_SYSTEM,
            "worker": WORKER_SYSTEM}[section]


def ask(section, question, temp, max_tokens):
    payload = {
        "messages": [
            {"role": "system", "content": section_system(section)},
            {"role": "user", "content": question},
        ],
        "temperature": temp,
        "top_p": 0.95 if temp > 0 else 1.0,
        "max_tokens": max_tokens,
    }
    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()
    out = content
    if reasoning:
        out = f"[thinking]\n{reasoning[-1200:]}\n[/thinking]\n\n{content}".strip()
    if not content and not reasoning:
        out = "[empty response]"
    timings = resp.get("timings", {})
    return {
        "content": out,
        "final": content,          # the answer without thinking, for exec
        "finish": choice.get("finish_reason", "?"),
        "elapsed": elapsed,
        "tps": timings.get("predicted_per_second"),
        "tokens": resp.get("usage", {}).get("completion_tokens"),
    }


def extract_code_block(text):
    """Pull the first ```...``` block; fall back to whole text."""
    if "```" in text:
        parts = text.split("```")
        # parts[1] is inside the first fence; strip a language tag line
        block = parts[1] if len(parts) > 1 else text
        lines = block.splitlines()
        if lines and lines[0].strip().lower() in ("python", "py", "json"):
            lines = lines[1:]
        return "\n".join(lines).strip()
    return text.strip()


def run_python_task(final_answer, fn_name, tests):
    """Execute model code + asserts in a sandboxed subprocess. Returns
    (passed:int, total:int, detail:str)."""
    code = extract_code_block(final_answer)
    harness = code + "\n\n" + "\n".join(tests) + "\nprint('ALLPASS')\n"
    # per-test tallying: run tests one at a time for partial credit
    passed = 0
    detail = []
    for i, t in enumerate(tests):
        prog = code + "\n\n" + t + "\nprint('OK')\n"
        ok, msg = _exec_python(prog)
        if ok:
            passed += 1
        detail.append(f"t{i}:{'PASS' if ok else 'FAIL'}" + (f"({msg})" if msg else ""))
    return passed, len(tests), " ".join(detail)


def _exec_python(prog):
    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
        f.write(prog)
        path = f.name
    try:
        r = subprocess.run(
            [sys.executable, path],
            capture_output=True, text=True, timeout=CODE_EXEC_TIMEOUT_S,
        )
        if r.returncode == 0 and "OK" in r.stdout:
            return True, ""
        err = (r.stderr.strip().splitlines() or [""])[-1]
        return False, err[:120]
    except subprocess.TimeoutExpired:
        return False, "timeout"
    except Exception as e:
        return False, str(e)[:120]
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


def check_json_exact(final_answer, expect):
    txt = extract_code_block(final_answer)
    try:
        got = json.loads(txt)
    except Exception:
        # try to find a JSON object substring
        s, e = txt.find("{"), txt.rfind("}")
        if s >= 0 and e > s:
            try:
                got = json.loads(txt[s:e + 1])
            except Exception:
                return 0, 1, "invalid JSON"
        else:
            return 0, 1, "no JSON found"
    return (1, 1, "match") if got == expect else (0, 1, f"got {got}")


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)


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"]
    return cmd


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


def run_model(name, model_path, fit_mode, tok_override):
    logfile = os.path.join(RESULTS_DIR, f"{name}.log")
    serverlog = os.path.join(RESULTS_DIR, f"{name}.server.log")
    max_tokens = tok_override or MAX_TOKENS

    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 section, qid, temp, q, spec in QUESTIONS:
                print(f"  [{name}] {qid} ({section})", flush=True)
                try:
                    r = ask(section, q, temp, max_tokens)
                except Exception as e:
                    r = {"content": f"[ERROR: {e}]", "final": "", "finish": "error",
                         "elapsed": 0, "tps": None, "tokens": None}

                f.write("=" * 70 + "\n")
                f.write(f"{qid} [{section}] (temp={temp})\n{q}\n\n")
                f.write(f"ANSWER:\n{r['content']}\n\n")

                # worker code/json tasks get executed inline
                if section == "worker" and isinstance(spec, dict):
                    lang = spec.get("lang")
                    if lang == "python":
                        p, tot, det = run_python_task(
                            r["final"], spec["fn"], spec["tests"])
                        f.write(f"EXEC: {p}/{tot} tests passed | {det}\n")
                    elif lang == "json_exact":
                        p, tot, det = check_json_exact(r["final"], spec["expect"])
                        f.write(f"EXEC: {p}/{tot} | {det}\n")
                else:
                    f.write(f"EXPECTED (judge only): {spec}\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\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, fit_mode, tok 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, fit_mode, tok))

    n_main = sum(1 for s, *_ in QUESTIONS if s == "main")
    n_esc = sum(1 for s, *_ in QUESTIONS if s == "escalation")
    n_wrk = sum(1 for s, *_ in QUESTIONS if s == "worker")
    print(f"Role-fitness run: {len(todo)} models.")
    print(f"Sections: main={n_main} escalation={n_esc} worker={n_wrk} (executed)")
    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, tok in todo:
        try:
            run_model(name, path, fit_mode, tok)
        except KeyboardInterrupt:
            print("\nInterrupted. Rerun to resume; completed models skipped.")
            break

    print(f"\nElapsed: {(time.time() - t0) / 60:.0f} min")
    print("Hand results_roles/*.log to the judge.")


if __name__ == "__main__":
    main()