#!/usr/bin/env python3
"""vLLM GSM8K eval for IAPO released checkpoints.

Runs one model per process/job. This avoids vLLM engine teardown issues and
gives a cleaner wall-clock generation measurement than the Transformers script.
"""

from __future__ import annotations

import argparse
import csv
import gc
import json
import math
import os
import random
import re
import time
from pathlib import Path
from typing import Any

from datasets import load_dataset
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams


SYSTEM_PROMPT = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it.
    The assistant first thinks about the reasoning process in the mind and then provides the user
    with the answer. The reasoning process and answer are enclosed within <think> and
    <answer> tags, respectively, i.e., <think> reasoning process here </think>
    <answer> answer here </answer>.
The answer must be a single integer."""

FEWSHOT_USER = "What is 2+2?"
FEWSHOT_ASSISTANT = "<think>To calculate 2+2, we simply add the numbers together: 2 + 2 = 4.</think>\n<answer>4</answer>"

MODEL_SPECS = {
    "base-7b": {"repo": "Qwen/Qwen2.5-7B-Instruct", "subfolder": None},
    "iapo-7b-gsm8k": {"repo": "jonathanhe123/iapo", "subfolder": "Qwen2.5-7B-Instruct_GSM8K"},
    "base-0.5b": {"repo": "Qwen/Qwen2.5-0.5B-Instruct", "subfolder": None},
    "iapo-0.5b-gsm8k": {"repo": "jonathanhe123/iapo", "subfolder": "Qwen2.5-0.5B-Instruct_GSM8K"},
}


def normalize_answer(s: str | None) -> str:
    if s is None:
        return ""
    s = str(s).strip().split("=")[-1]
    s = s.replace(",", "").replace("$", "").strip()
    s = re.sub(r"\\boxed\{([^{}]+)\}", r"\1", s)
    return s.strip().rstrip(".")


def extract_gold(answer: str) -> str:
    return normalize_answer(answer.split("####")[-1])


def extract_pred(text: str) -> str:
    xml = re.search(r"<answer>(.*?)</answer>", text, flags=re.DOTALL | re.IGNORECASE)
    if xml:
        return normalize_answer(xml.group(1))
    answer_line = re.findall(r"Answer:\s*([-+]?\d[\d,]*(?:\.\d+)?)", text, flags=re.IGNORECASE)
    if answer_line:
        return normalize_answer(answer_line[-1])
    nums = re.findall(r"[-+]?\d[\d,]*(?:\.\d+)?", text)
    return normalize_answer(nums[-1]) if nums else ""


def local_model_path(repo: str, subfolder: str | None, cache_dir: str | None) -> str:
    if not subfolder:
        return repo
    patterns = [f"{subfolder}/*"]
    root = snapshot_download(repo_id=repo, allow_patterns=patterns, cache_dir=cache_dir)
    return str(Path(root) / subfolder)


def make_prompt(tokenizer: Any, question: str) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": FEWSHOT_USER},
        {"role": "assistant", "content": FEWSHOT_ASSISTANT},
        {"role": "user", "content": question},
    ]
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-key", required=True, choices=sorted(MODEL_SPECS))
    parser.add_argument("--sample-size", type=int, default=128)
    parser.add_argument("--num-return-sequences", type=int, default=8)
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-p", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.90)
    parser.add_argument("--max-model-len", type=int, default=1024)
    parser.add_argument("--output-dir", default="outputs/hf_gsm8k_vllm")
    parser.add_argument("--cache-dir", default=os.environ.get("HF_HOME"))
    args = parser.parse_args()

    random.seed(args.seed)
    spec = MODEL_SPECS[args.model_key]
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    model_path = local_model_path(spec["repo"], spec["subfolder"], args.cache_dir)
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    tokenizer.padding_side = "left"

    dataset = load_dataset("openai/gsm8k", "main", split="test")
    indices = list(range(len(dataset)))
    random.Random(args.seed).shuffle(indices)
    indices = sorted(indices[: args.sample_size])
    examples = [dataset[i] for i in indices]
    prompts = [make_prompt(tokenizer, ex["question"]) for ex in examples]
    golds = [extract_gold(ex["answer"]) for ex in examples]

    env = {
        "event": "env",
        "model_key": args.model_key,
        "repo": spec["repo"],
        "subfolder": spec["subfolder"],
        "model_path": model_path,
        "sample_size": len(examples),
        "num_return_sequences": args.num_return_sequences,
        "max_new_tokens": args.max_new_tokens,
    }
    print(json.dumps(env), flush=True)

    llm = LLM(
        model=model_path,
        tokenizer=model_path,
        trust_remote_code=True,
        dtype=args.dtype,
        gpu_memory_utilization=args.gpu_memory_utilization,
        max_model_len=args.max_model_len,
        seed=args.seed,
    )
    sampling = SamplingParams(
        n=args.num_return_sequences,
        temperature=args.temperature,
        top_p=args.top_p,
        max_tokens=args.max_new_tokens,
        seed=args.seed,
    )

    t0 = time.perf_counter()
    outputs = llm.generate(prompts, sampling, use_tqdm=True)
    gen_seconds = time.perf_counter() - t0

    per_question_correct = []
    per_question_lengths = []
    details_path = out_dir / f"{args.model_key}_generations.jsonl"
    with details_path.open("w") as f:
        for ex_idx, (dataset_idx, gold, request_output) in enumerate(zip(indices, golds, outputs)):
            correct_row = []
            length_row = []
            for sample_idx, completion in enumerate(request_output.outputs):
                text = completion.text
                token_ids = completion.token_ids or []
                pred = extract_pred(text)
                ok = int(pred == gold)
                correct_row.append(ok)
                length_row.append(len(token_ids))
                f.write(json.dumps({
                    "model": args.model_key,
                    "dataset_index": dataset_idx,
                    "sample_in_group": sample_idx,
                    "gold": gold,
                    "prediction": pred,
                    "correct": ok,
                    "completion_tokens": len(token_ids),
                    "completion": text,
                }) + "\n")
            per_question_correct.append(correct_row)
            per_question_lengths.append(length_row)

    k_values = [k for k in [1, 2, 4, 8, 16, 32] if k <= args.num_return_sequences]
    generated_count = len(examples) * args.num_return_sequences
    summary = {
        "event": "summary",
        "model": args.model_key,
        "repo": spec["repo"],
        "subfolder": spec["subfolder"],
        "sample_size": len(examples),
        "num_return_sequences": args.num_return_sequences,
        "max_new_tokens": args.max_new_tokens,
        "generated_count": generated_count,
        "generation_seconds": gen_seconds,
        "avg_seconds_per_completion": gen_seconds / max(1, generated_count),
    }
    for k in k_values:
        pass_at_k = sum(int(any(row[:k])) for row in per_question_correct) / len(per_question_correct)
        length_at_k = sum(sum(row[:k]) for row in per_question_lengths) / (len(per_question_lengths) * k)
        summary[f"pass@{k}"] = pass_at_k
        summary[f"length@{k}"] = length_at_k
        summary[f"ratio@{k}"] = pass_at_k / length_at_k if length_at_k else math.nan

    summary_path = out_dir / f"{args.model_key}_summary.json"
    csv_path = out_dir / f"{args.model_key}_summary.csv"
    summary_path.write_text(json.dumps(summary, indent=2))
    with csv_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=sorted(summary))
        writer.writeheader()
        writer.writerow(summary)

    print(json.dumps(summary), flush=True)
    print(json.dumps({
        "event": "wrote_outputs",
        "summary": str(summary_path),
        "csv": str(csv_path),
        "details": str(details_path),
    }), flush=True)

    del llm
    gc.collect()


if __name__ == "__main__":
    main()