"""Frozen GSM8K test evaluator: stock math_rl prompt (few-shot + boxed suffix), greedy, max_tokens=512, numeric exact match on \\boxed{}."""
import argparse, asyncio, hashlib, json, re, time
import common, tinker
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer
from tinker_cookbook.recipes.math_rl.math_env import MathEnv, Gsm8kDataset, extract_gsm8k_final_answer
from tinker_cookbook.recipes.math_rl.math_grading import extract_boxed
from datasets import load_dataset

MODEL, RENDERER, MAX_TOKENS = "Qwen/Qwen3.5-2B", "qwen3_5_disable_thinking", 512


def norm(s):
    s = s.strip().replace(",", "").replace("$", "").replace("\\%", "").replace("%", "").strip()
    s = re.sub(r"^\\text\{(.*)\}$", r"\1", s).strip()
    try:
        v = float(s)
        return str(int(v)) if v == int(v) else str(v)
    except Exception:
        return None


def score(text, ref):
    try:
        b = extract_boxed(text)
    except ValueError:
        return None, 0
    return b, int(norm(b) is not None and norm(b) == norm(ref))


async def main(a):
    t0 = time.time()
    sc = common.service()
    samp = sc.create_sampling_client(base_model=MODEL) if a.sampler_path == "base" else sc.create_sampling_client(model_path=a.sampler_path)
    r = renderers.get_renderer(RENDERER, tokenizer=get_tokenizer(MODEL))
    ds = load_dataset("openai/gsm8k", name="main", split="test")
    n = a.limit or len(ds)
    sem = asyncio.Semaphore(a.concurrency)
    stop = r.get_stop_sequences()

    async def one(i):
        row = ds[i]
        ref = extract_gsm8k_final_answer(row["answer"])
        convo = MathEnv.standard_fewshot_prefix() + [{"role": "user", "content": row["question"] + Gsm8kDataset.question_suffix()}]
        mi = r.build_generation_prompt(convo)
        async with sem:
            for k in range(8):
                try:
                    out = await samp.sample_async(prompt=mi, num_samples=1, sampling_params=tinker.SamplingParams(max_tokens=MAX_TOKENS, temperature=0.0, stop=stop))
                    break
                except Exception as e:
                    if k == 7: raise
                    await asyncio.sleep(5 * 2 ** k)
        msg, _ = r.parse_response(out.sequences[0].tokens)
        text = renderers.get_text_content(msg)
        pred, ok = score(text, ref)
        return {"item_id": f"gsm8k_test_{i}", "arm": a.arm, "seed": 0, "reference": ref, "extracted": pred, "correct": ok, "completion": text}

    res = await asyncio.gather(*[one(i) for i in range(n)])
    with open(a.out, "w") as f:
        for x in res: f.write(json.dumps(x) + "\n")
    acc = sum(x["correct"] for x in res) / n
    print(json.dumps({"arm": a.arm, "n": n, "accuracy": acc, "seconds": time.time() - t0, "code_sha256": hashlib.sha256(open(__file__, "rb").read()).hexdigest()}))


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--arm", required=True); p.add_argument("--sampler_path", default="base")
    p.add_argument("--limit", type=int, default=0); p.add_argument("--concurrency", type=int, default=256)
    p.add_argument("--out", required=True)
    asyncio.run(main(p.parse_args()))