"""Frozen greedy GSM8K evaluator for run_40338c73ff.
usage: python eval_gsm8k.py <arm> <base|tinker://...sampler path> <selection|confirmation>"""
import asyncio, hashlib, json, math, re, sys, time
from importlib import metadata
from datasets import load_dataset
import tinker
from tinker import types
import common
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer
from tinker_cookbook.recipes.math_rl.math_env import MathEnv, extract_gsm8k_final_answer
from tinker_cookbook.recipes.math_rl.math_grading import extract_boxed
from research_exchange.evaluator_provenance import evaluator_id

SPLITS = {"selection": range(0, 200), "confirmation": range(200, 1319)}
SUFFIX = " Provide a numerical answer without units, written inside \\boxed{}."
MAX_TOKENS = 512
NUM = re.compile(r"-?\d[\d,]*(?:\.\d+)?|-?\.\d+")

def to_number(s):
    s = s.replace(",", "").replace("$", "").strip()
    try:
        v = float(s)
    except (ValueError, OverflowError):
        return None
    return v if math.isfinite(v) else None

def parse(text):
    """Final number: last number inside the last \\boxed{}, else last number in the response."""
    try:
        region = extract_boxed(text)
    except ValueError:
        region = text
    nums = NUM.findall(region) or NUM.findall(text)
    return to_number(nums[-1]) if nums else None

def correct(pred, ref):
    r = to_number(ref)
    return int(pred is not None and r is not None and pred == r)

def provenance(split):
    code = b"".join(open(f, "rb").read() for f in ("eval_gsm8k.py", "common.py"))
    deps = {p: metadata.version(p) for p in ("tinker", "tinker-cookbook", "transformers", "datasets")}
    deps["tinker-cookbook-git"] = "f79fdb2ba4cd967ef2807ba43097c52985f9f17f"
    return {"code_sha256": hashlib.sha256(code + json.dumps(deps, sort_keys=True).encode()).hexdigest(),
            "prompt_template": {"renderer": common.RENDERER, "tokenizer": common.MODEL,
                                "convo": "MathEnv.standard_fewshot_prefix() + [user: question + Gsm8kDataset.question_suffix()]",
                                "question_suffix": SUFFIX, "build": "renderer.build_generation_prompt(convo)"},
            "decoding": {"temperature": 0.0, "max_tokens": MAX_TOKENS, "stop": "renderer.get_stop_sequences() (role_colon)"},
            "parser": {"code": "eval_gsm8k.parse/correct", "rule": "last number in last \\boxed{} (else last number in text), strip ',' '$', float; correct iff finite and == float(reference '####' answer)"},
            "items": {"ids": [f"gsm8k/test/{i}" for i in SPLITS[split]]}}

async def main(arm, model, split):
    prov = provenance(split); eid = evaluator_id(prov)
    tok = get_tokenizer(common.MODEL); r = renderers.get_renderer(common.RENDERER, tok)
    ds = load_dataset("openai/gsm8k", "main", split="test")
    sc = common.service()
    for attempt in range(8):
        try:
            client = (await sc.create_sampling_client_async(base_model=common.MODEL) if model == "base"
                      else await sc.create_sampling_client_async(model_path=model)); break
        except Exception as e:
            print("sampler retry", repr(e)[:300], flush=True); await asyncio.sleep(min(60 * 2 ** attempt, 240))
    params = types.SamplingParams(max_tokens=MAX_TOKENS, temperature=0.0, stop=r.get_stop_sequences())
    sem = asyncio.Semaphore(128); t0 = time.time()
    async def one(i):
        row = ds[i]; ref = extract_gsm8k_final_answer(row["answer"])
        convo = MathEnv.standard_fewshot_prefix() + [{"role": "user", "content": row["question"] + SUFFIX}]
        prompt = r.build_generation_prompt(convo)
        async with sem:
            for attempt in range(6):
                try:
                    res = await client.sample_async(prompt=prompt, num_samples=1, sampling_params=params); break
                except Exception as e:
                    if attempt == 5: raise
                    await asyncio.sleep(10 * 2 ** attempt)
        toks = res.sequences[0].tokens
        msg, _ = r.parse_response(toks); text = renderers.get_text_content(msg)
        p = parse(text)
        return {"item_id": f"gsm8k/test/{i}", "arm": arm, "seed": 0, "metric": "accuracy", "evaluator_id": eid,
                "checkpoint": model, "prediction": text, "parsed": p, "reference": ref, "value": correct(p, ref),
                "n_tokens": len(toks)}
    rows = await asyncio.gather(*(one(i) for i in SPLITS[split]))
    secs = time.time() - t0
    out = f"ev/pred_{split}_{arm}.jsonl"
    with open(out, "w") as f:
        for row in rows: f.write(json.dumps(row) + "\n")
    acc = sum(r_["value"] for r_ in rows) / len(rows)
    json.dump({"arm": arm, "split": split, "checkpoint": model, "evaluator_id": eid, "provenance": prov,
               "accuracy": acc, "n": len(rows), "seconds": secs}, open(f"ev/score_{split}_{arm}.json", "w"), indent=1)
    print(json.dumps({"arm": arm, "split": split, "accuracy": acc, "n": len(rows), "seconds": round(secs, 1), "evaluator_id": eid}), flush=True)

if __name__ == "__main__":
    asyncio.run(main(*sys.argv[1:4]))