#!/usr/bin/env python3
# coding=utf-8
# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Minimal vLLM MMLU-Pro single-sample inference example.

Example:
  # Use embedded MMLU-Pro example sample (no dataset file needed)
  python run_text_vllm_example.py --model-path $(pwd)/../../checkpoint_folder_textonly

  # Or use a real MMLU-Pro json file
  python run_text_vllm_example.py \
      --model-path /path/to/model \
      --mmlupro-json /path/to/mmlu_pro/test.json \
      --sample-idx 0
"""

import argparse
import json
import re
from pathlib import Path
from typing import Any

from vllm import LLM, SamplingParams


SYSTEM_PROMPT = (
    "<|im_start|>system\n"
    "You are a helpful and harmless assistant.\n\n"
    "You are not allowed to use any tools."
    "<|im_end|>\n"
)

CHOICES = list("ABCDEFGHIJKLMNOP")
STOP_MARKERS = ("<|im_end|>", "<|end_of_text|>", "<|eot_id|>")
EXAMPLE_MMLUPRO_SAMPLE = {
    "question": "Which organelle is primarily responsible for ATP production in eukaryotic cells?",
    "options": [
        "Golgi apparatus",
        "Mitochondrion",
        "Lysosome",
        "Endoplasmic reticulum",
    ],
    "answer": "B",
}


def build_mmlupro_user_prompt(sample: dict[str, Any]) -> str:
    """Build the MMLU-Pro user prompt with boxed-answer instruction."""
    options = [opt for opt in sample["options"] if opt != "N/A"]
    prompt = "Question:\n" + sample["question"] + "\n\nAnswer Choices:"
    for i, opt in enumerate(options):
        prompt += f"\n({CHOICES[i]}) {opt}"
    prompt += (
        "\n\nConclude your response with the sentence "
        "`The answer is \\boxed{{X}}.`, in which X is the correct capital letter "
        "of your choice."
    )
    return prompt.strip() + "\n"


def build_chatml_prompt(user_prompt: str, think: bool = True) -> str:
    assistant_prefix = "<think>\n" if think else ""
    return (
        SYSTEM_PROMPT
        + "<|im_start|>user\n"
        + user_prompt
        + "<|im_end|>\n"
        + "<|im_start|>assistant\n"
        + assistant_prefix
    )


def clean_generation(text: str) -> str:
    """Trim common end markers used in the eval scripts."""
    cleaned = text
    for marker in STOP_MARKERS:
        idx = cleaned.find(marker)
        if idx != -1:
            cleaned = cleaned[:idx]
    return cleaned.strip()


def extract_boxed_answer(text: str) -> str | None:
    """Extract answer letter from `The answer is \\boxed{X}.`"""
    match = re.search(r"The answer is\s*\\boxed\{([A-P])\}\.?", text)
    if match:
        return match.group(1)
    match = re.search(r"\\boxed\{([A-P])\}", text)
    return match.group(1) if match else None


def load_mmlupro_sample(path: Path, sample_idx: int) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as f:
        data = json.load(f)
    if not (0 <= sample_idx < len(data)):
        raise IndexError(f"sample_idx={sample_idx} out of range [0, {len(data) - 1}]")
    return data[sample_idx]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Minimal vLLM MMLU-Pro inference with reasoning template."
    )
    parser.add_argument("--model-path", type=str, required=True, help="Model path for vLLM.")
    parser.add_argument(
        "--mmlupro-json",
        type=str,
        default=None,
        help="Optional path to MMLU-Pro test.json (list of {question, options, ...}).",
    )
    parser.add_argument("--sample-idx", type=int, default=0, help="MMLU-Pro sample index.")
    parser.add_argument("--tensor-parallel-size", type=int, default=1)
    parser.add_argument("--max-tokens", type=int, default=131072)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-p", type=float, default=0.95)
    parser.add_argument("--seed", type=int, default=100)
    parser.add_argument("--disable-thinking", action="store_true")
    parser.add_argument("--fp16", action="store_true", help="Use float16 instead of bfloat16.")
    parser.add_argument("--print-prompt", action="store_true", help="Print full prompt.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if args.mmlupro_json:
        sample = load_mmlupro_sample(Path(args.mmlupro_json), args.sample_idx)
        print(f"Loaded sample {args.sample_idx} from: {args.mmlupro_json}")
    else:
        sample = EXAMPLE_MMLUPRO_SAMPLE
        print("Using embedded MMLU-Pro example sample.")
    user_prompt = build_mmlupro_user_prompt(sample)
    prompt = build_chatml_prompt(user_prompt, think=not args.disable_thinking)

    if args.print_prompt:
        print("=== PROMPT ===")
        print(prompt)
        print("==============")

    dtype = "float16" if args.fp16 else "bfloat16"
    model = LLM(
        args.model_path,
        dtype=dtype,
        tensor_parallel_size=args.tensor_parallel_size,
        trust_remote_code=True,
        enable_prefix_caching=True,
        enforce_eager=False,
    )
    sampling_params = SamplingParams(
        temperature=args.temperature,
        top_p=args.top_p,
        max_tokens=args.max_tokens,
        seed=args.seed,
    )

    output = model.generate([prompt], sampling_params)[0].outputs[0].text
    output = clean_generation(output)
    pred = extract_boxed_answer(output)

    print("\n=== QUESTION ===")
    print(sample.get("question", ""))
    print("\n=== MODEL OUTPUT ===")
    print(output)
    print("\n=== PARSED PREDICTION ===")
    print(pred if pred is not None else "No boxed answer found")

    if "answer" in sample:
        print("\n=== REFERENCE ANSWER ===")
        print(sample["answer"])


if __name__ == "__main__":
    main()