#!/usr/bin/env python3
"""Run reasoning eval against an OpenAI-compatible /v1/chat/completions endpoint.

Usage:
  python3 run_eval.py --base http://LLAMA_SWAP_HOST:PORT --model <id> --tag distilled
  python3 run_eval.py --base ... --model <id> --tag stock --key sk-...
Then compare:
  python3 run_eval.py --compare results/distilled.jsonl results/stock.jsonl
"""
import json, os, re, time, argparse, urllib.request, glob
from concurrent.futures import ThreadPoolExecutor

HERE = os.path.dirname(__file__)
DATA = os.path.join(HERE, "data")
RES = os.path.join(HERE, "results"); os.makedirs(RES, exist_ok=True)

# Identical for both models — measures reasoning ABILITY, not default verbosity.
SYS = ("You are a careful problem solver. Think step by step, then end your reply "
       "with a single final line in exactly this format:\nANSWER: <your final answer>\n"
       "For multiple-choice, the final answer is the single letter of the correct option.")

def load_items():
    items = []
    for f in sorted(glob.glob(os.path.join(DATA, "*.jsonl"))):
        with open(f) as fh:
            items += [json.loads(l) for l in fh if l.strip()]
    return items

def call(base, model, key, question, max_tokens, temp, timeout, top_p=0.95, think=False):
    payload = {"model": model, "temperature": temp, "top_p": top_p,
               "max_tokens": max_tokens,
               "messages": [{"role": "system", "content": SYS},
                            {"role": "user", "content": question}]}
    if think:
        payload["chat_template_kwargs"] = {"enable_thinking": True}
    body = json.dumps(payload).encode()
    req = urllib.request.Request(base.rstrip("/") + "/v1/chat/completions", data=body,
                                 headers={"Content-Type": "application/json",
                                          **({"Authorization": f"Bearer {key}"} if key else {})})
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout) as r:
        d = json.load(r)
    dt = time.time() - t0
    msg = d["choices"][0]["message"]
    text = msg.get("content") or msg.get("reasoning_content") or ""
    usage = d.get("usage", {})
    return text, usage.get("completion_tokens"), dt, d.get("model")

def extract(text):
    m = list(re.finditer(r"ANSWER:\s*(.+)", text, re.I))
    if m: return m[-1].group(1).strip()
    b = list(re.finditer(r"\\boxed\{([^}]*)\}", text))
    if b: return b[-1].group(1).strip()
    return text.strip().splitlines()[-1].strip() if text.strip() else ""

def norm_num(s):
    s = s.replace(",", "").replace("$", "").replace("\\", "").strip().rstrip(".")
    m = re.search(r"-?\d+\.?\d*", s)
    return m.group(0) if m else None

def norm_text(s):
    return re.sub(r"[\s${}\\]", "", s).lower()

def score(cat, pred, gold):
    p = extract(pred) if False else pred  # pred already extracted by caller
    if cat in ("bbh", "mmlu_pro"):
        up = p.upper()
        m = re.search(r"\b([A-J])\b", up) or re.search(r"([A-J])\)", up) or re.search(r"[A-J]", up)
        return bool(m) and m.group(1 if m.re.groups else 0) == gold.strip().upper()
    if cat == "gsm8k":
        a, b = norm_num(p), norm_num(gold)
        if a is None or b is None: return False
        try: return abs(float(a) - float(b)) < 1e-6
        except: return a == b
    # math: numeric if possible, else normalized string
    a, b = norm_num(p), norm_num(gold)
    if a is not None and b is not None:
        try:
            if abs(float(a) - float(b)) < 1e-6: return True
        except: pass
    return norm_text(p) == norm_text(gold)

def run(args):
    items = load_items()
    out = os.path.join(RES, f"{args.tag}.jsonl")
    print(f"[{args.tag}] {len(items)} items -> {args.model} @ {args.base}")

    def work(it):
        try:
            text, ctoks, dt, served = call(args.base, args.model, args.key, it["question"],
                                           args.max_tokens, args.temp, args.timeout, args.top_p,
                                           args.think)
            pred = extract(text)
            ok = bool(score(it["category"], pred, it["answer"]))
            return {**{k: it[k] for k in ("id", "category", "answer")},
                    "pred": pred, "correct": ok, "completion_tokens": ctoks,
                    "latency_s": round(dt, 2), "served_model": served, "raw": text}
        except Exception as e:
            return {**{k: it[k] for k in ("id", "category", "answer")},
                    "pred": None, "correct": False, "error": str(e)}

    rows = []
    with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
        for i, r in enumerate(ex.map(work, items), 1):
            rows.append(r)
            flag = "ok " if r["correct"] else ("ERR" if r.get("error") else "x  ")
            print(f"  [{i:>3}/{len(items)}] {flag} {r['id']:<14} "
                  f"gold={str(r['answer'])[:18]:<18} pred={str(r['pred'])[:24]}")
    with open(out, "w") as f:
        for r in rows: f.write(json.dumps(r) + "\n")
    served = sorted({r.get("served_model") for r in rows if r.get("served_model")})
    print(f"\nSERVED MODEL(S) this run: {served}")
    summarize(rows, args.tag)
    print(f"saved -> {out}")

def summarize(rows, tag):
    cats = {}
    for r in rows:
        c = cats.setdefault(r["category"], {"n": 0, "ok": 0, "tok": 0, "tn": 0})
        c["n"] += 1; c["ok"] += int(r["correct"])
        if r.get("completion_tokens"): c["tok"] += r["completion_tokens"]; c["tn"] += 1
    print(f"\n=== {tag} ===")
    print(f"{'category':<12}{'acc':>10}{'avg_tok':>10}")
    tot_n = tot_ok = 0
    for c, v in sorted(cats.items()):
        tot_n += v["n"]; tot_ok += v["ok"]
        at = (v["tok"] / v["tn"]) if v["tn"] else 0
        print(f"{c:<12}{v['ok']:>3}/{v['n']:<3}{100*v['ok']/v['n']:>5.0f}%{at:>10.0f}")
    print(f"{'OVERALL':<12}{tot_ok:>3}/{tot_n:<3}{100*tot_ok/tot_n:>5.0f}%")

def compare(a, b):
    def load(p): return {r["id"]: r for r in (json.loads(l) for l in open(p))}
    A, B = load(a), load(b)
    ta, tb = os.path.basename(a).split(".")[0], os.path.basename(b).split(".")[0]
    cats = {}
    for i in A:
        if i not in B: continue
        c = cats.setdefault(A[i]["category"], [0, 0, 0])
        c[0] += 1; c[1] += int(A[i]["correct"]); c[2] += int(B[i]["correct"])
    print(f"{'category':<12}{ta:>10}{tb:>10}{'Δ':>8}")
    TA = TB = TN = 0
    for c, (n, x, y) in sorted(cats.items()):
        TN += n; TA += x; TB += y
        print(f"{c:<12}{100*x/n:>9.0f}%{100*y/n:>9.0f}%{100*(x-y)/n:>+7.0f}")
    print(f"{'OVERALL':<12}{100*TA/TN:>9.0f}%{100*TB/TN:>9.0f}%{100*(TA-TB)/TN:>+7.0f}")
    print(f"\nN={TN}. Note: ~±{int(100*(0.5/ (TN**0.5))) }pt rough 1-sigma per model; "
          "treat small gaps as noise.")

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--base"); ap.add_argument("--model"); ap.add_argument("--key")
    ap.add_argument("--tag", default="run")
    ap.add_argument("--max-tokens", type=int, default=8192, dest="max_tokens")
    ap.add_argument("--temp", type=float, default=0.6)
    ap.add_argument("--top-p", type=float, default=0.95, dest="top_p")
    ap.add_argument("--timeout", type=int, default=300)
    ap.add_argument("--think", action="store_true", help="send chat_template_kwargs enable_thinking=true")
    ap.add_argument("--concurrency", type=int, default=4)
    ap.add_argument("--compare", nargs=2)
    a = ap.parse_args()
    if a.compare: compare(*a.compare)
    else: run(a)