"""AIME25 avg@k 评测。

针对 non-thinking 模型: prompt 必须用 enable_thinking=False 渲染
(尾部会带 <think>\n\n</think>\n\n)。见 MEMORY: nothink-flag-for-1p7b-sft。

答案抽取按优先级: \boxed{} -> "Answer: X" -> 文末最后一个整数。
AIME 答案恒为 0-999 的整数, 比较时按 int 比。
"""
import argparse
import json
import os
import re
import sys
import time
from collections import Counter

p = argparse.ArgumentParser()
p.add_argument("--model", required=True)
p.add_argument("--data", default="/home/haojinw2/datasets/aime25/test.jsonl")
p.add_argument("--out", required=True)
p.add_argument("--n", type=int, default=16)
p.add_argument("--max-new", type=int, default=8192)
p.add_argument("--ctx", type=int, default=10240)
p.add_argument("--tp", type=int, default=1)
p.add_argument("--temp", type=float, default=0.7)
p.add_argument("--top-p", type=float, default=0.8)
p.add_argument("--top-k", type=int, default=20)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--enable-thinking", type=lambda s: s.lower() == "true", default=False,
               help="False = 关闭 think(non-thinking 模型必须用 False)")
p.add_argument("--gpu-util", type=float, default=0.85)
a = p.parse_args()


def log(*x):
    print(f"[{time.strftime('%H:%M:%S')}]", *x, flush=True)


INSTR = "Please reason step by step, and put your final answer within \\boxed{}."

BOXED = re.compile(r"\\boxed\s*\{([^{}]*)\}")
ANSW = re.compile(r"[Aa]nswer\s*[:：]\s*\$?(-?\d+)")
INTS = re.compile(r"-?\d+")


def _safe_int(s):
    """AIME 答案恒为 0-999。超长数字是模型退化的产物(见过 28000 位),
    直接 int() 会撞 Python 3.11 的 int_max_str_digits=4300 抛 ValueError。"""
    if s is None:
        return None
    t = s.lstrip("-")
    if len(t) > 10 or not t:
        return None
    try:
        return int(s)
    except ValueError:
        return None


def extract(text):
    m = BOXED.findall(text)
    if m:
        digits = INTS.findall(m[-1])
        if digits:
            v = _safe_int(digits[-1])
            if v is not None:
                return v
    m = ANSW.findall(text)
    if m:
        v = _safe_int(m[-1])
        if v is not None:
            return v
    m = INTS.findall(text)
    if m:
        return _safe_int(m[-1])
    return None


def _kill_descendants():
    """递归杀掉本进程的所有子孙(vLLM worker 是独立进程, 不杀会变孤儿占满卡)。"""
    import signal
    import subprocess
    me = os.getpid()
    try:
        out = subprocess.run(["ps", "-eo", "pid,ppid", "--no-headers"],
                             capture_output=True, text=True, timeout=20).stdout
    except Exception:
        return
    children = {}
    for ln in out.splitlines():
        parts = ln.split()
        if len(parts) != 2:
            continue
        try:
            pid, ppid = int(parts[0]), int(parts[1])
        except ValueError:
            continue
        children.setdefault(ppid, []).append(pid)
    doomed, stack = [], [me]
    while stack:
        for c in children.get(stack.pop(), []):
            doomed.append(c)
            stack.append(c)
    for pid in reversed(doomed):           # 先叶子后根
        try:
            os.kill(pid, signal.SIGKILL)
        except OSError:
            pass


def _hard_exit(code=0):
    sys.stdout.flush()
    sys.stderr.flush()
    _kill_descendants()
    os._exit(code)


def _die(exc_type, exc, tb):
    import traceback as _tb
    _tb.print_exception(exc_type, exc, tb)
    _hard_exit(1)        # 异常路径同样要把 worker 一起带走


sys.excepthook = _die

rows = [json.loads(l) for l in open(a.data) if l.strip()]
log(f"{len(rows)} 题 x n={a.n}, enable_thinking={a.enable_thinking}")

from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)
prompts = []
for r in rows:
    msgs = [{"role": "user", "content": r["problem"].strip() + "\n\n" + INSTR}]
    prompts.append(tok.apply_chat_template(
        msgs, add_generation_prompt=True, tokenize=False,
        enable_thinking=a.enable_thinking))
log("prompt 尾部抽检: " + repr(prompts[0][-60:]))

def _fit_tp(model_path, tp):
    """attn heads 必须能被 TP 整除, 否则 vLLM 起不来(R1-Distill-1.5B 是 12 heads)。"""
    try:
        cfg = json.load(open(os.path.join(model_path, "config.json")))
        h = cfg.get("num_attention_heads")
    except Exception:
        return tp
    if not h or h % tp == 0:
        return tp
    for cand in range(tp, 0, -1):
        if h % cand == 0:
            log(f"注意: {h} 个 attention head 不能被 TP={tp} 整除, 自动降为 TP={cand}")
            return cand
    return 1


a.tp = _fit_tp(a.model, a.tp)
llm = LLM(model=a.model, tensor_parallel_size=a.tp,
          gpu_memory_utilization=a.gpu_util, max_model_len=a.ctx,
          trust_remote_code=True, seed=a.seed)
# 不给 SamplingParams 传 seed: per-request seed 会让 n>1 退出批量采样快路径,
# 480 个 seeded 样本会逐个采样、慢到形同挂死。可复现性由 LLM(seed=) 保证。
sp = SamplingParams(n=a.n, temperature=a.temp, top_p=a.top_p, top_k=a.top_k,
                    max_tokens=a.max_new)

t0 = time.time()
outs = llm.generate(prompts, sp)
log(f"生成 {time.time()-t0:.0f}s")

recs, n_correct, n_tot, n_trunc, lens = [], 0, 0, 0, []
per_problem = []
for r, out in zip(rows, outs):
    gt = int(str(r["answer"]).strip())
    hits, samples = 0, []
    for o in out.outputs:
        pred = extract(o.text)
        ok = (pred is not None and pred == gt)
        hits += ok
        n_correct += ok
        n_tot += 1
        n_trunc += (o.finish_reason == "length")
        lens.append(len(o.token_ids))
        samples.append({"pred": pred, "ok": bool(ok),
                        "finish": o.finish_reason, "len": len(o.token_ids),
                        "text": o.text})
    per_problem.append({"id": r.get("id"), "gt": gt, "hits": hits, "n": len(out.outputs)})
    recs.append({"id": r.get("id"), "problem": r["problem"], "gt": gt,
                 "hits": hits, "samples": samples})

os.makedirs(os.path.dirname(os.path.abspath(a.out)) or ".", exist_ok=True)
with open(a.out, "w") as f:
    for rec in recs:
        f.write(json.dumps(rec, ensure_ascii=False) + "\n")

avg = n_correct / n_tot
solved_any = sum(1 for x in per_problem if x["hits"] > 0)
solved_all = sum(1 for x in per_problem if x["hits"] == x["n"])
summary = {
    "model": a.model, "data": a.data, "n": a.n,
    "enable_thinking": a.enable_thinking,
    "temperature": a.temp, "top_p": a.top_p, "top_k": a.top_k,
    "max_new": a.max_new, "seed": a.seed,
    "problems": len(rows), "samples": n_tot,
    f"avg@{a.n}": avg,
    "solved_any": solved_any, "solved_all": solved_all,
    "truncated_rate": n_trunc / n_tot,
    "resp_len_mean": sum(lens) / len(lens),
    "resp_len_max": max(lens),
    "hits_hist": dict(Counter(x["hits"] for x in per_problem)),
}
with open(a.out.replace(".jsonl", "_summary.json"), "w") as f:
    json.dump(summary, f, indent=2, ensure_ascii=False)

print()
print(json.dumps(summary, indent=2, ensure_ascii=False), flush=True)
print()
print(f"AIME25  avg@{a.n} = {avg:.4f}   ({n_correct}/{n_tot})", flush=True)
print(f"        解出>=1次 {solved_any}/{len(rows)},  全对 {solved_all}/{len(rows)}", flush=True)
print(f"        截断率 {n_trunc/n_tot:.3f},  平均长度 {sum(lens)/len(lens):.0f}", flush=True)

# vLLM TP>1 的 worker 常在析构时挂尸、占着整卡不放, 后续任务就起不来。
# 结果已全部写盘, 直接跳过析构强制退出。
_hard_exit(0)