# E15 matched-compute baselines (standard methods: self-consistency Wang et al. 2022, best-of-N by mean token logprob; no novelty claim).
# Hand-written by the AGI loop 10:58 AM ET Oct 27 2026: Cursor CLI refused non-Grok models (Codex 5.3 not in weekly usage), Grok models excluded by owner.
# Prompts/extraction/test execution identical to bench\server_eval.py (GSM8K) and bench\humaneval_eval.py.
import argparse, json, re, subprocess, sys, tempfile, time, os, requests
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from datasets import load_dataset
ap = argparse.ArgumentParser()
ap.add_argument("--url", default="http://127.0.0.1:8090"); ap.add_argument("--task", required=True, choices=["gsm8k", "humaneval"])
ap.add_argument("--tag", required=True); ap.add_argument("--n_items", type=int, default=0); ap.add_argument("--n_samples", type=int, default=0)
ap.add_argument("--temp", type=float, default=0.7); ap.add_argument("--max_tok", type=int, default=0); ap.add_argument("--seed_base", type=int, default=1000)
ap.add_argument("--workers", type=int, default=8); ap.add_argument("--resume", action="store_true")
a = ap.parse_args(); G = a.task == "gsm8k"
NS = a.n_samples or (8 if G else 4); MT = a.max_tok or (320 if G else 512)
os.makedirs("results/quality", exist_ok=True)
OUT = f"results/quality/e15_{a.task}_{a.tag}.jsonl"
def comp(prompt, greedy, seed):
    body = {"prompt": prompt, "n_predict": MT, "n_probs": 1, "cache_prompt": False, "seed": seed}
    body.update({"temperature": 0, "top_k": 1} if greedy else {"temperature": a.temp, "top_k": 40, "top_p": 0.95, "min_p": 0})
    t = time.time(); r = requests.post(a.url + "/completion", json=body, timeout=3600); r.raise_for_status(); j = r.json()
    lps = [c.get("logprob", 0.0) for c in (j.get("completion_probabilities") or [])]
    n = j.get("timings", {}).get("predicted_n", len(lps))
    return {"text": j.get("content", ""), "tokens": n, "sum_logprob": sum(lps), "mean_logprob": sum(lps) / max(len(lps), 1),
            "wall_sec": round(time.time() - t, 2)}
def gsm_pred(txt):
    m = re.findall(r"-?\d[\d,]*\.?\d*", txt.split("answer is")[-1] if "answer is" in txt else txt)
    return m[-1].replace(",", "").rstrip(".") if m else ""
def gsm_ok(pred, gold):
    try: return abs(float(pred) - float(gold)) < 1e-6
    except: return False
def norm(p):
    try: return str(float(p))
    except: return p
def he_extract(txt, entry):
    blocks = re.findall(r"```(?:python|py)?\s*\n(.*?)```", txt, re.S)
    if not blocks: blocks = [txt.split("```")[0]]
    for b in blocks:
        if f"def {entry}" in b: return b
    return blocks[0]
def he_pass(ex, txt):
    prog = ex["prompt"] + "\n    pass\n\n" + he_extract(txt, ex["entry_point"]) + "\n\n" + ex["test"] + f"\ncheck({ex['entry_point']})\n"
    fd, fn = tempfile.mkstemp(suffix=".py"); os.write(fd, prog.encode("utf-8")); os.close(fd)
    try: ok = subprocess.run([sys.executable, fn], capture_output=True, timeout=15).returncode == 0
    except subprocess.TimeoutExpired: ok = False
    os.remove(fn); return bool(ok)
if G:
    ds = load_dataset("openai/gsm8k", "main", split="test"); N = a.n_items or 128
else:
    ds = load_dataset("openai/openai_humaneval", split="test"); N = a.n_items or 164
done = set()
if a.resume and os.path.exists(OUT):
    for l in open(OUT, encoding="utf-8"):
        try: done.add(json.loads(l)["item"])
        except: pass
pool = ThreadPoolExecutor(a.workers); tw = time.time()
for k in range(min(N, len(ds))):
    if k in done: continue
    ex = ds[k]; t0 = time.time()
    if G:
        gold = ex["answer"].split("####")[-1].strip().replace(",", "")
        p = ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n" + ex["question"] +
             "\nSolve step by step, then finish with a line 'The answer is N.'<|im_end|>\n<|im_start|>assistant\n")
    else:
        p = ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nComplete the following Python function. "
             "Return the complete function (with any imports it needs) in a single ```python code block.\n\n```python\n" + ex["prompt"] +
             "```<|im_end|>\n<|im_start|>assistant\n")
    futs = [pool.submit(comp, p, True, 0)] + [pool.submit(comp, p, False, a.seed_base + j) for j in range(NS)]
    outs = [f.result() for f in futs]
    for o in outs:
        if G: o["pred"] = gsm_pred(o["text"]); o["correct"] = gsm_ok(o["pred"], gold)
        else: o["correct"] = he_pass(ex, o["text"])
    rec = {"item": k, "id": ex.get("task_id", k), "gold": gold if G else None, "greedy": outs[0], "samples": outs[1:],
           "item_wall_sec": round(time.time() - t0, 2)}
    with open(OUT, "a", encoding="utf-8") as f: f.write(json.dumps(rec) + "\n")
    print(f"{a.task} {k} greedy={int(outs[0]['correct'])} samp_ok={sum(o['correct'] for o in outs[1:])}/{NS} wall={rec['item_wall_sec']}s", flush=True)
recs = sorted((json.loads(l) for l in open(OUT, encoding="utf-8")), key=lambda r: r["item"])
recs = [r for r in recs if r["item"] < N]; n = len(recs)
def bon(ss): return max(ss, key=lambda s: s["mean_logprob"])
def sc(ss):
    c = Counter(norm(s["pred"]) for s in ss if s["pred"] != "")
    if not c: return ss[0]
    top = max(c.values()); tied = [s for s in ss if s["pred"] != "" and c[norm(s["pred"])] == top]
    return max(tied, key=lambda s: s["mean_logprob"])
S = {"tag": a.tag, "task": a.task, "n": n, "n_samples": NS, "temp": a.temp, "wall_sec_total": round(sum(r["item_wall_sec"] for r in recs), 1)}
S["greedy"] = {"acc": sum(r["greedy"]["correct"] for r in recs) / n, "tok_per_problem": sum(r["greedy"]["tokens"] for r in recs) / n}
Ns = [4, 8] if G and NS >= 8 else [min(4, NS)]
for m in Ns:
    tok = sum(sum(s["tokens"] for s in r["samples"][:m]) for r in recs) / n
    S[f"bon_meanlogprob_N{m}"] = {"acc": sum(bon(r["samples"][:m])["correct"] for r in recs) / n, "tok_per_problem": tok}
    S[f"oracle_N{m}"] = {"acc": sum(any(s["correct"] for s in r["samples"][:m]) for r in recs) / n, "tok_per_problem": tok, "note": "reference only, uses gold"}
    if G: S[f"selfconsistency_N{m}"] = {"acc": sum(sc(r["samples"][:m])["correct"] for r in recs) / n, "tok_per_problem": tok}
S["sample_mean_acc"] = sum(sum(s["correct"] for s in r["samples"]) / len(r["samples"]) for r in recs) / n
print("E15SUMMARY", json.dumps(S), flush=True)
json.dump(S, open(f"results/quality/e15_{a.task}_{a.tag}_summary.json", "w"), indent=1)