"""Frozen GSM8K test evaluator: greedy, 512 tokens, exact match on final numeric answer.
usage: evaluate.py <label> [<tinker sampler path>]  (no path = base model)"""
import asyncio, json, re, sys, time
from fractions import Fraction
from datasets import load_dataset
from tinker import types
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 common import *

DECODING = {"temperature": 0.0, "max_tokens": MAX_TOKENS, "num_samples": 1, "stop": "renderer.get_stop_sequences()"}
PROMPT = {"renderer": RENDERER, "tokenizer": MODEL, "messages": "MathEnv.standard_fewshot_prefix() + [user: question + Gsm8kDataset.question_suffix()]",
          "suffix": " Provide a numerical answer without units, written inside \\boxed{}."}
PARSER = "last \\boxed{} via tinker_cookbook extract_boxed; strip ',', '$', spaces, trailing '.'; exact Fraction equality with '####' gold; unparsable/missing -> 0"


def to_num(s):
    s = s.replace(",", "").replace("$", "").replace(" ", "").rstrip(".")
    s = re.sub(r"^\\text\{(.*)\}$", r"\1", s)
    return Fraction(s)


def score(text, gold):
    try:
        pred = extract_boxed(text)
    except Exception:
        return None, 0.0
    try:
        return pred, float(to_num(pred) == to_num(gold))
    except Exception:
        return pred, 0.0


async def main(label, path=None):
    renderer, tok = renderer_and_tokenizer()
    test = load_dataset("openai/gsm8k", "main", split="test")
    svc = service_client()
    sc = svc.create_sampling_client(model_path=path) if path else svc.create_sampling_client(base_model=MODEL)
    params = types.SamplingParams(max_tokens=MAX_TOKENS, temperature=0.0, stop=renderer.get_stop_sequences())
    sem = asyncio.Semaphore(128)
    prefix = MathEnv.standard_fewshot_prefix()

    async def one(i, row):
        convo = prefix + [{"role": "user", "content": row["question"] + PROMPT["suffix"]}]
        mi = renderer.build_generation_prompt(convo)
        async with sem:
            for attempt in range(5):
                try:
                    r = await sc.sample_async(prompt=mi, num_samples=1, sampling_params=params)
                    break
                except Exception as e:
                    if attempt == 4:
                        return {"item_id": f"gsm8k-test-{i:04d}", "error": str(e)[:300], "value": 0.0}
                    await asyncio.sleep(10 * (attempt + 1))
        text = tok.decode(r.sequences[0].tokens)
        gold = extract_gsm8k_final_answer(row["answer"])
        pred, v = score(text, gold)
        return {"item_id": f"gsm8k-test-{i:04d}", "gold": gold, "pred": pred, "value": v, "output": text,
                "stop_reason": str(getattr(r.sequences[0], "stop_reason", None)), "n_tokens": len(r.sequences[0].tokens)}

    t0 = time.time()
    rows = await asyncio.gather(*[one(i, row) for i, row in enumerate(test)])
    acc = sum(r["value"] for r in rows) / len(rows)
    with open(f"{OUT}/predictions_{label}.jsonl", "w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    summ = {"label": label, "model_path": path or f"base:{MODEL}", "accuracy": acc, "n": len(rows),
            "n_errors": sum("error" in r for r in rows), "n_unparsed": sum(r.get("pred") is None for r in rows), "seconds": time.time() - t0}
    dump(f"eval_{label}.json", summ)
    print(json.dumps(summ))


if __name__ == "__main__":
    asyncio.run(main(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else None))