"""Frozen GSM8K evaluator: zero-shot chat prompt (thinking disabled), greedy, max_tokens 512, '####' answer parser."""
import json, re, hashlib, sys, time, concurrent.futures as cf
from tinker import types
from transformers import AutoTokenizer

MODEL = "Qwen/Qwen3.5-2B"
INSTRUCTION = "Solve the following math problem step by step. Put the final numeric answer on the last line after '####'."
PROMPT_TEMPLATE = "{instruction}\n\n{question}"
STOP = [248046]  # <|im_end|>
_tok = None

def tokenizer():
    global _tok
    if _tok is None:
        _tok = AutoTokenizer.from_pretrained(MODEL)
    return _tok

def prompt_tokens(question):
    msgs = [{"role": "user", "content": PROMPT_TEMPLATE.format(instruction=INSTRUCTION, question=question)}]
    return tokenizer().apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, enable_thinking=False, return_dict=False)

_NUM = re.compile(r"-?\d[\d,]*\.?\d*")

def parse_answer(text):
    """Number after the last '####'; else the last number in the text; None if no number."""
    seg = text.rsplit("####", 1)[1] if "####" in text else text
    nums = _NUM.findall(seg)
    if not nums:
        return None
    s = (nums[0] if "####" in text else nums[-1]).replace(",", "").rstrip(".")
    return s

def reference_answer(answer_field):
    return answer_field.rsplit("####", 1)[1].strip().replace(",", "")

def is_correct(pred, ref):
    if pred is None:
        return 0
    try:
        a, b = float(pred), float(ref)
    except (ValueError, OverflowError):
        return 0
    if a != a or abs(a) == float("inf"):
        return 0
    return int(abs(a - b) < 1e-6)

def load(path):
    return [json.loads(l) for l in open(path)]

def evaluator_provenance():
    import tinker, transformers
    code = open(__file__, "rb").read()
    items = load("gsm8k_test500.jsonl")
    return {
        "code_sha256": hashlib.sha256(code + f"tinker=={tinker.__version__};transformers=={transformers.__version__}".encode()).hexdigest(),
        "prompt_template": {"user": PROMPT_TEMPLATE, "instruction": INSTRUCTION, "chat_template": f"{MODEL} tokenizer apply_chat_template(add_generation_prompt=True, enable_thinking=False)"},
        "decoding": {"temperature": 0, "max_tokens": 512, "stop": STOP},
        "parser": "gsm8k_eval.parse_answer/is_correct (number after last '####' else last number; float equality 1e-6; reference = text after '####' with commas removed)",
        "items": {"ids": [it["id"] for it in items]},
    }

def evaluator_id():
    return hashlib.sha256(json.dumps(evaluator_provenance(), sort_keys=True, separators=(",", ":"), ensure_ascii=True, allow_nan=False).encode()).hexdigest()

def evaluate(sampling_client, arm, seed, checkpoint, out_path, workers=64):
    items = load("gsm8k_test500.jsonl")
    eid = evaluator_id()
    params = types.SamplingParams(max_tokens=512, temperature=0.0, stop=STOP)
    def one(it):
        for attempt in range(6):
            try:
                r = sampling_client.sample(types.ModelInput.from_ints(prompt_tokens(it["question"])), 1, params).result()
                return tokenizer().decode(r.sequences[0].tokens, skip_special_tokens=True)
            except Exception as e:
                err = repr(e); time.sleep(2 ** attempt)
        return "<<SAMPLING_ERROR " + err + ">>"
    t0 = time.time()
    with cf.ThreadPoolExecutor(workers) as ex:
        texts = list(ex.map(one, items))
    rows = []
    for it, txt in zip(items, texts):
        ref = reference_answer(it["answer"]); pred = parse_answer(txt)
        rows.append({"item_id": it["id"], "arm": arm, "seed": seed, "checkpoint": checkpoint, "metric": "accuracy",
                     "evaluator_id": eid, "prediction": txt, "parsed": pred, "reference": ref, "value": is_correct(pred, ref)})
    with open(out_path, "w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    acc = sum(r["value"] for r in rows) / len(rows)
    return acc, time.time() - t0

if __name__ == "__main__":
    from common import service_client
    sc = service_client()
    t0 = time.time()
    samp = sc.create_sampling_client(base_model=MODEL)
    acc, secs = evaluate(samp, "base", None, f"base:{MODEL}", "predictions_base.jsonl")
    print(json.dumps({"base_accuracy": acc, "eval_seconds": secs, "total_seconds": time.time() - t0, "evaluator_id": evaluator_id()}))