"""GPQA Diamond client used for the turbo-attn KV cache reports (MiniCPM5-2B).

    python gpqa_client.py URL CSV OUT.json SEEDS [CONCURRENCY] [--repetition-penalty X]

URL   an OpenAI-compatible chat endpoint, e.g. http://localhost:8000/v1/chat/completions (served model name "m")
CSV   gpqa_diamond.csv from the gated dataset Idavidrein/gpqa (198 rows; sha256 41d1213cd7a49986...)
SEEDS comma-separated sampling seeds, e.g. 21,22,23 (every question is asked once per seed)

Prompt: the simple-evals GPQA template. The four options are shuffled once per question, seeded from its
Record ID, so the order is identical in every run and every seed. No max_tokens, no system prompt; all
sampling comes from the model's own generation config (vLLM --generation-config auto), plus the seed.
The answer is the last "Answer: X" after </think>; no answer line counts as wrong.
Writes ids, letters and token counts only (no question text).
"""
import asyncio, csv, hashlib, json, random, re, sys, time
import httpx

args = [a for a in sys.argv[1:] if not a.startswith("--")]
RP = None
if "--repetition-penalty" in sys.argv:
    RP = float(sys.argv[sys.argv.index("--repetition-penalty") + 1]); args.remove(str(sys.argv[sys.argv.index("--repetition-penalty") + 1]))
URL, CSV, OUT, SEEDS = args[0], args[1], args[2], [int(s) for s in args[3].split(",")]
CONC = int(args[4]) if len(args) > 4 else 256
TEMPLATE = ("Answer the following multiple choice question. The last line of your response should be of the following format: "
            "'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering.\n\n"
            "{Question}\n\nA) {A}\nB) {B}\nC) {C}\nD) {D}")
PAT = [r"(?i)Answer\s*:\s*\$?\(?\**\s*([A-D])\b", r"(?i)answer is\s*\(?\**([A-D])\b", r"\\boxed\{\s*\(?([A-D])\)?\s*\}"]

def item(r):
    rid = r["Record ID"]
    opts = [r["Correct Answer"], r["Incorrect Answer 1"], r["Incorrect Answer 2"], r["Incorrect Answer 3"]]
    perm = list(range(4)); random.Random(int(hashlib.sha256(rid.encode()).hexdigest(), 16) % 2**32).shuffle(perm)
    ch = [opts[i].strip() for i in perm]
    return rid, TEMPLATE.format(Question=r["Question"].strip(), A=ch[0], B=ch[1], C=ch[2], D=ch[3]), "ABCD"[perm.index(0)]

def extract(raw):
    t = raw.split("</think>")[-1]
    for p in PAT:
        m = re.findall(p, t)
        if m: return m[-1].upper()
    return None

async def one(c, sem, rid, prompt, gold, seed):
    async with sem:
        body = {"model": "m", "seed": seed, "messages": [{"role": "user", "content": prompt}]}
        if RP is not None: body["repetition_penalty"] = RP
        r = (await c.post(URL, json=body, timeout=14400)).json()
        ch = r["choices"][0]; raw = ch["message"]["content"] or ""
        pred = extract(raw)
        return {"id": rid, "seed": seed, "gold": gold, "pred": pred, "ok": pred == gold, "out_toks": r["usage"]["completion_tokens"],
                "in_toks": r["usage"]["prompt_tokens"], "finish": ch["finish_reason"], "closed_think": "</think>" in raw}

async def main():
    items = [item(r) for r in csv.DictReader(open(CSV))]
    sem, t0 = asyncio.Semaphore(CONC), time.time()
    async with httpx.AsyncClient(limits=httpx.Limits(max_connections=CONC + 4)) as c:
        res = await asyncio.gather(*[one(c, sem, rid, p, g, s) for s in SEEDS for rid, p, g in items])
    json.dump({"seeds": SEEDS, "n_q": len(items), "repetition_penalty": RP, "wall_s": time.time() - t0, "results": res}, open(OUT, "w"))
    print(f"accuracy {100 * sum(r['ok'] for r in res) / len(res):.1f}% over {len(res)} answers")

asyncio.run(main())