"""Frozen GSM8K evaluator: recipe prompt (role_colon + standard fewshot + boxed suffix), greedy, 512 tokens, numeric exact match."""
import sys, os, json, re, math, asyncio, hashlib, time
from common import *
import tinker
from datasets import load_dataset
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

PARSER = "boxed=tinker_cookbook.math_grading.extract_boxed(renderers.get_text_content(parse_response(tokens)[0])); x=boxed.replace(',','').replace('$','').strip().rstrip('.'); value=float(x) (non-finite or unparsable -> wrong); correct iff value==float(reference) where reference=extract_gsm8k_final_answer(answer); no boxed -> wrong"
DECODING = {"temperature": 0, "max_tokens": 512, "stop": "renderer.get_stop_sequences() of role_colon", "top_p": 1.0}


def number(s):
    try:
        v = float(s.replace(",", "").replace("$", "").strip().rstrip("."))
        return v if math.isfinite(v) else None
    except Exception:
        return None


def evaluator(lo, hi):
    code = hashlib.sha256(open(__file__, "rb").read() + open("common.py", "rb").read() +
                          f"tinker={tinker.__version__};cookbook=f79fdb2ba4cd967ef2807ba43097c52985f9f17f".encode()).hexdigest()
    prov = {"code_sha256": code,
            "prompt_template": "renderer=role_colon (tinker_cookbook f79fdb2b) tokenizer=Qwen/Qwen2.5-3B; convo=MathEnv.standard_fewshot_prefix() + [user: question + ' Provide a numerical answer without units, written inside \\\\boxed{}.']; build_generation_prompt",
            "decoding": DECODING, "parser": PARSER,
            "items": {"ids": [f"gsm8k/test/{i}" for i in range(lo, hi)]}}
    return prov, sha(prov)


async def main(arm, model_path, lo, hi, out):
    prov, eid = evaluator(lo, hi)
    tok = get_tokenizer(MODEL)
    r = renderers.get_renderer(RENDERER, tokenizer=tok)
    ds = load_dataset("openai/gsm8k", "main", split="test")
    sc = service_client()
    for attempt in range(10):
        try:
            samp = sc.create_sampling_client(model_path=model_path) if model_path != "base" else sc.create_sampling_client(base_model=MODEL)
            break
        except Exception as e:
            print("sampler create fail", repr(e)[:200]); time.sleep(30)
    params = tinker.SamplingParams(max_tokens=512, temperature=0.0, stop=r.get_stop_sequences())
    sem = asyncio.Semaphore(128)

    async def one(i):
        row = ds[i]
        convo = MathEnv.standard_fewshot_prefix() + [{"role": "user", "content": row["question"] + Gsm8kDataset.question_suffix()}]
        prompt = r.build_generation_prompt(convo)
        ref = extract_gsm8k_final_answer(row["answer"])
        async with sem:
            for attempt in range(6):
                try:
                    res = await samp.sample_async(prompt=prompt, num_samples=1, sampling_params=params)
                    break
                except Exception as e:
                    if attempt == 5: raise
                    await asyncio.sleep(10 * (attempt + 1))
        toks = res.sequences[0].tokens
        msg, _ = r.parse_response(toks)
        text = renderers.get_text_content(msg)
        try:
            boxed = extract_boxed(text)
        except Exception:
            boxed = None
        v = number(boxed) if boxed is not None else None
        correct = int(v is not None and v == float(ref))
        return {"item_id": f"gsm8k/test/{i}", "arm": arm, "seed": 0, "metric": "accuracy", "evaluator_id": eid,
                "checkpoint": model_path, "prediction": text, "parsed": boxed, "reference": ref, "value": correct}
    t0 = time.time()
    rows = await asyncio.gather(*[one(i) for i in range(lo, hi)])
    with open(out, "w") as fh:
        for row in rows: fh.write(json.dumps(row) + "\n")
    acc = sum(x["value"] for x in rows) / len(rows)
    print(json.dumps({"arm": arm, "lo": lo, "hi": hi, "acc": acc, "n": len(rows), "seconds": time.time() - t0, "evaluator_id": eid}), flush=True)

if __name__ == "__main__":
    arm, path, lo, hi, out = sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4]), sys.argv[5]
    asyncio.run(main(arm, path, lo, hi, out))