import json
import os
import re
import sys
import time

os.environ["CUDA_HOME"] = "/usr/local/cuda-13.0"
os.environ["PATH"] = f"/usr/local/cuda-13.0/bin:{os.environ.get('PATH', '')}"

MODEL_DIR = "/home/user/models/DeepSeek-V4-Flash"
sys.path.insert(0, os.path.join(MODEL_DIR, "encoding"))
from encoding_dsv4 import encode_messages

from datasets import load_dataset
from vllm import LLM, SamplingParams

OUTPUT_FILE = sys.argv[1] if len(sys.argv) > 1 else "/home/user/mmlupro_results_deepseek_v4.jsonl"

OPTION_LETTERS = "ABCDEFGHIJ"

def format_prompt(question: str, options: list[str]) -> str:
    lines = [f"Question: {question}", "", "Options:"]
    for i, opt in enumerate(options):
        letter = OPTION_LETTERS[i]
        lines.append(f"{letter}) {opt}")
    lines.extend(["", "Answer with the correct option letter only."])
    return "\n".join(lines)


def extract_answer(text: str) -> str | None:
    text = text.strip()
    m = re.search(r'\b([A-J])\b', text)
    if m:
        return m.group(1)
    return None


def main():
    print("Loading MMLU-Pro dataset...")
    ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
    print(f"Loaded {len(ds)} questions")

    print("Loading model with vLLM...")
    llm = LLM(
        model=MODEL_DIR,
        tensor_parallel_size=1,
        dtype="auto",
        kv_cache_dtype="fp8",
        max_model_len=32768,
        trust_remote_code=True,
    )

    sampling_params = SamplingParams(
        temperature=0.0,
        top_p=0.95,
        max_tokens=10,
        stop=["<｜end▁of▁sentence｜>"],
    )

    formatted_prompts = []
    metadata = []
    for example in ds:
        prompt_text = format_prompt(example["question"], example["options"])
        messages = [{"role": "user", "content": prompt_text}]
        formatted = encode_messages(messages, thinking_mode="chat")
        formatted_prompts.append(formatted)
        metadata.append({
            "question_id": example["question_id"],
            "question": example["question"],
            "options": example["options"],
            "answer": example["answer"],
            "answer_index": example["answer_index"],
            "category": example["category"],
        })

    print(f"Sample prompt: {formatted_prompts[0][:200]}...")
    print(f"Generating answers for {len(formatted_prompts)} questions...")
    start = time.time()

    outputs = llm.generate(formatted_prompts, sampling_params)

    elapsed = time.time() - start
    print(f"Generation completed in {elapsed:.2f}s")

    results = []
    correct = 0
    for out, meta in zip(outputs, metadata):
        raw = out.outputs[0].text.strip()
        predicted = extract_answer(raw)
        expected_letter = OPTION_LETTERS[meta["answer_index"]]
        is_correct = predicted == expected_letter
        if is_correct:
            correct += 1
        results.append({
            "question_id": meta["question_id"],
            "category": meta["category"],
            "question": meta["question"],
            "options": meta["options"],
            "expected": expected_letter,
            "predicted": predicted,
            "raw_output": raw,
            "correct": is_correct,
        })

    with open(OUTPUT_FILE, "w") as f:
        for r in results:
            f.write(json.dumps(r) + "\n")
    print(f"Results saved to {OUTPUT_FILE}")

    total = len(results)
    print(f"\nAccuracy: {correct}/{total} = {correct / total * 100:.2f}%")

    cats = {}
    for r in results:
        c = r["category"]
        if c not in cats:
            cats[c] = {"correct": 0, "total": 0}
        cats[c]["total"] += 1
        if r["correct"]:
            cats[c]["correct"] += 1
    print("\nPer-category:")
    for c in sorted(cats):
        v = cats[c]
        print(f"  {c}: {v['correct']}/{v['total']} = {v['correct']/v['total']*100:.1f}%")


if __name__ == "__main__":
    main()