"""
Eval script for GSM8K exact-match accuracy.
Greedy decoding (do_sample=False).
Extracts final answer after #### and compares.
Usage:
    python eval_gsm8k.py --model_id MODEL_ID --batch_size 16
"""
import argparse
import re
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM


def extract_answer(text: str) -> str:
    """Extract the number after ####."""
    m = re.search(r"####\s*(-?\d+(?:,\d+)*(?:\.\d+)?)", text)
    if m:
        return m.group(1).replace(",", "")
    return ""


def build_prompt(question: str) -> str:
    return f"Question: {question}\nAnswer:"


def evaluate(model_id: str, device: str = "cuda", max_new_tokens: int = 256, batch_size: int = 16):
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    tokenizer.padding_side = "left"
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        trust_remote_code=True,
        dtype=torch.bfloat16,
        device_map=device,
    )
    model.eval()

    ds = load_dataset("openai/gsm8k", "main", split="test")
    correct = 0
    total = len(ds)

    for i in range(0, total, batch_size):
        batch = ds[i:i+batch_size]
        prompts = [build_prompt(q) for q in batch["question"]]
        inputs = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                temperature=None,
                top_p=None,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        generated = outputs[:, inputs["input_ids"].shape[1]:]
        decoded = tokenizer.batch_decode(generated, skip_special_tokens=True)

        for pred_text, ref_text in zip(decoded, batch["answer"]):
            pred = extract_answer(pred_text)
            ref = extract_answer(ref_text)
            if pred == ref:
                correct += 1

        if (i // batch_size) % 10 == 0:
            print(f"Processed {i}/{total} — running accuracy {correct / max(i, 1):.4f}")

    accuracy = correct / total
    print(f"\nFinal accuracy: {accuracy:.4f} ({correct}/{total})")
    return accuracy


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_id", type=str, required=True)
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--max_new_tokens", type=int, default=256)
    parser.add_argument("--batch_size", type=int, default=16)
    args = parser.parse_args()
    evaluate(args.model_id, args.device, args.max_new_tokens, args.batch_size)