# /// script
# dependencies = [
#     "vllm",
#     "transformers>=4.40.0",
#     "datasets>=2.14",
#     "pyarrow<19.0",
#     "torch>=2.0.0",
#     "huggingface-hub",
#     "tqdm>=4.65",
# ]
# ///

import os
import re
from datasets import load_dataset
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

hf_token = os.getenv("HF_TOKEN")
if hf_token:
    os.environ.setdefault("HF_HUB_TOKEN", hf_token)

MODEL = os.getenv("MODEL", "Qwen/Qwen3.5-4B")
MAX_SAMPLES = int(os.getenv("MAX_SAMPLES", "250"))

def extract_answer(text):
    text = text.strip()
    m = re.search(r"####\s*(-?\d+(?:\.\d+)?)", text)
    if m:
        return m.group(1).strip()
    m = re.search(r"ANSWER:\s*(-?\d+(?:\.\d+)?)", text, re.IGNORECASE)
    if m:
        return m.group(1).strip()
    nums = re.findall(r"-?\d+(?:\.\d+)?", text)
    if nums:
        return nums[-1].strip()
    return ""

def build_prompt(question):
    messages = [
        {"role": "system", "content": "You are a helpful assistant that thinks step by step."},
        {"role": "user", "content": question},
    ]
    return messages

if __name__ == '__main__':
    print(f"Loading GSM8K dataset (max {MAX_SAMPLES} samples)...")
    ds = load_dataset("gsm8k", "main", split="test")
    if MAX_SAMPLES:
        ds = ds.select(range(min(MAX_SAMPLES, len(ds))))

    print(f"Loading model {MODEL} with vLLM...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
    llm = LLM(
        model=MODEL,
        trust_remote_code=True,
        dtype="half",
        gpu_memory_utilization=0.9,
    )

    questions = []
    ground_truth = []
    for ex in ds:
        msgs = build_prompt(ex["question"])
        formatted = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        questions.append(formatted)
        ground_truth.append(extract_answer(ex["answer"]))

    # SAME sampling params as tuned eval (temperature=0.7, do_sample)
    sampling_params = SamplingParams(
        temperature=0.7,
        top_p=0.9,
        max_tokens=512,
        stop=["<|im_end|>", "<|endoftext|>"],
    )

    print(f"Generating answers for {len(questions)} questions...")
    outputs = llm.generate(questions, sampling_params)

    correct = 0
    results = []
    for i, out in enumerate(outputs):
        pred_text = out.outputs[0].text
        pred_answer = extract_answer(pred_text)
        gt = ground_truth[i]
        is_correct = pred_answer == gt
        if is_correct:
            correct += 1
        results.append({
            "question": ds[i]["question"],
            "ground_truth": gt,
            "predicted": pred_answer,
            "correct": is_correct,
            "raw_output": pred_text,
        })
        if i < 5:
            print(f"Q{i+1}: GT={gt} Pred={pred_answer} OK={is_correct}")

    accuracy = correct / len(ds) * 100
    print(f"\n=== GSM8K BASELINE (FAIR) ===")
    print(f"Correct: {correct}/{len(ds)}")
    print(f"Accuracy: {accuracy:.2f}%")

    import json
    with open("/tmp/gsm8k_baseline_fair_results.json", "w") as f:
        json.dump({
            "model": MODEL,
            "accuracy": accuracy,
            "correct": correct,
            "total": len(ds),
            "results": results,
        }, f, indent=2)
    print("Results saved to /tmp/gsm8k_baseline_fair_results.json")