hugging/td_fuse/td_benchmark.py

"""
TD Benchmark Suite v1 — Model Report Card

Tests the model across 6 categories and gives a score per category.
This is Step 1A of the self-improvement loop.

Categories:
  1. Math (arithmetic, algebra, word problems)
  2. Code (write functions, debug, explain)
  3. Reasoning (logic, multi-step, common sense)
  4. Creativity (writing, rewriting, explaining)
  5. Knowledge (science, history, general facts)
  6. Instruction Following (do exactly what's asked)

Output: JSON report with scores per category + list of failed questions.
The weakness finder uses this to decide what to train on next.
"""

import torch
import json
import time
import re
import gc
from pathlib import Path
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass, field


@dataclass
class BenchmarkConfig:
    model_path: str = ""
    output_dir: str = "td_fuse_outputs/benchmark"
    output_path: str = ""  # If set, save results to this exact path
    max_new_tokens: int = 256  # SPEED: 512 was overkill, most answers are short
    temperature: float = 0.0  # Greedy — deterministic answers
    batch_questions: int = 1  # Questions per batch (1 = sequential)
    disable_thinking: bool = True  # SPEED: Skip <think> blocks for benchmark


# ============================================================
# TEST BANKS — questions with verifiable answers
# ============================================================

MATH_TESTS = [
    # Arithmetic
    {"q": "What is 47 * 23?", "answer": "1081", "type": "arithmetic", "difficulty": "easy"},
    {"q": "What is 156 + 289?", "answer": "445", "type": "arithmetic", "difficulty": "easy"},
    {"q": "What is 1000 - 387?", "answer": "613", "type": "arithmetic", "difficulty": "easy"},
    {"q": "What is 144 / 12?", "answer": "12", "type": "arithmetic", "difficulty": "easy"},
    {"q": "What is 17 * 19?", "answer": "323", "type": "arithmetic", "difficulty": "easy"},
    {"q": "What is 2.5 * 4.8?", "answer": "12", "type": "arithmetic", "difficulty": "medium"},
    {"q": "What is 15% of 240?", "answer": "36", "type": "arithmetic", "difficulty": "medium"},
    {"q": "What is 3/4 + 2/3? Give your answer as a fraction.", "answer": "17/12", "type": "fractions", "difficulty": "medium"},

    # Algebra
    {"q": "Solve for x: 3x + 7 = 22", "answer": "5", "type": "algebra", "difficulty": "easy"},
    {"q": "Solve for x: 2x - 5 = 13", "answer": "9", "type": "algebra", "difficulty": "easy"},
    {"q": "Solve for x: x^2 = 144", "answer": "12", "type": "algebra", "difficulty": "medium"},
    {"q": "If f(x) = 2x + 3, what is f(7)?", "answer": "17", "type": "algebra", "difficulty": "medium"},

    # Word problems
    {"q": "A store sells apples for $2 each and oranges for $3 each. If I buy 4 apples and 3 oranges, how much do I spend?", "answer": "17", "type": "word_problem", "difficulty": "easy"},
    {"q": "A train travels at 60 mph. How far does it go in 2.5 hours?", "answer": "150", "type": "word_problem", "difficulty": "easy"},
    {"q": "If 5 workers can build a wall in 10 days, how many days would it take 10 workers?", "answer": "5", "type": "word_problem", "difficulty": "medium"},
    {"q": "A rectangle has a perimeter of 30 cm and a width of 5 cm. What is its length?", "answer": "10", "type": "word_problem", "difficulty": "medium"},
    {"q": "I have 3 times as many cats as dogs. I have 2 dogs. How many animals do I have in total?", "answer": "8", "type": "word_problem", "difficulty": "easy"},
    {"q": "A car uses 8 liters of fuel per 100km. How many liters does it need for a 350km trip?", "answer": "28", "type": "word_problem", "difficulty": "medium"},

    # Harder math
    {"q": "What is the sum of the first 10 positive integers?", "answer": "55", "type": "series", "difficulty": "medium"},
    {"q": "What is 2^10?", "answer": "1024", "type": "exponents", "difficulty": "medium"},
    {"q": "What is the greatest common divisor of 48 and 36?", "answer": "12", "type": "number_theory", "difficulty": "medium"},
    {"q": "If a triangle has sides of length 3, 4, and 5, what is its area?", "answer": "6", "type": "geometry", "difficulty": "medium"},
    {"q": "What is the square root of 169?", "answer": "13", "type": "roots", "difficulty": "easy"},
    {"q": "A coin is flipped 3 times. How many possible outcomes are there?", "answer": "8", "type": "probability", "difficulty": "medium"},
    {"q": "What is 7! (7 factorial)?", "answer": "5040", "type": "factorial", "difficulty": "hard"},
]

CODE_TESTS = [
    # Write functions
    {"q": "Write a Python function called 'is_even' that returns True if a number is even, False otherwise. Just give the function, no explanation.",
     "check": "is_even", "verify": "callable", "type": "write", "difficulty": "easy"},
    {"q": "Write a Python function called 'factorial' that computes the factorial of a non-negative integer. Just give the function.",
     "check": "factorial", "verify": "callable", "type": "write", "difficulty": "medium"},
    {"q": "Write a Python function called 'reverse_string' that takes a string and returns it reversed. Just give the function.",
     "check": "reverse_string", "verify": "callable", "type": "write", "difficulty": "easy"},
    {"q": "Write a Python function called 'max_of_three' that takes three numbers and returns the largest. Just give the function.",
     "check": "max_of_three", "verify": "callable", "type": "write", "difficulty": "easy"},
    {"q": "Write a Python function called 'count_vowels' that counts the number of vowels (a,e,i,o,u) in a string. Just give the function.",
     "check": "count_vowels", "verify": "callable", "type": "write", "difficulty": "medium"},
    {"q": "Write a Python function called 'fibonacci' that returns the nth Fibonacci number (0-indexed, so fibonacci(0)=0, fibonacci(1)=1). Just give the function.",
     "check": "fibonacci", "verify": "callable", "type": "write", "difficulty": "medium"},
    {"q": "Write a Python function called 'is_palindrome' that returns True if a string reads the same forwards and backwards. Just give the function.",
     "check": "is_palindrome", "verify": "callable", "type": "write", "difficulty": "medium"},
    {"q": "Write a Python function called 'flatten' that takes a list of lists and returns a single flat list. Example: flatten([[1,2],[3,4]]) returns [1,2,3,4]. Just give the function.",
     "check": "flatten", "verify": "callable", "type": "write", "difficulty": "medium"},

    # Debug
    {"q": "This Python code has a bug. Fix it and give ONLY the corrected code:\ndef add_list(lst):\n    total = 0\n    for i in range(len(lst) + 1):\n        total += lst[i]\n    return total",
     "check": "range(len(lst))", "verify": "contains", "type": "debug", "difficulty": "easy"},
    {"q": "This Python code has a bug. Fix it and give ONLY the corrected code:\ndef greet(name):\n    return 'Hello' + name",
     "check": ["Hello ' + name", "Hello \" + name", "'Hello ' +", "'Hello, '", "f'Hello {", "f\"Hello {", "Hello \" +"], "verify": "contains_any", "type": "debug", "difficulty": "easy"},

    # Explain
    {"q": "In one sentence, what does this Python code do?\n[x**2 for x in range(10) if x % 2 == 0]",
     "check": "square", "verify": "contains", "type": "explain", "difficulty": "easy"},
    {"q": "In one sentence, what does this Python code do?\nlen(set(my_list))",
     "check": "unique", "verify": "contains", "type": "explain", "difficulty": "easy"},
]

REASONING_TESTS = [
    # Logic
    {"q": "If all roses are flowers and all flowers need water, do roses need water? Answer yes or no.",
     "answer": "yes", "type": "syllogism", "difficulty": "easy"},
    {"q": "If it's raining, the ground is wet. The ground is wet. Is it definitely raining? Answer yes or no.",
     "answer": "no", "type": "logic_fallacy", "difficulty": "medium"},
    {"q": "All cats are animals. Some animals are pets. Can we conclude all cats are pets? Answer yes or no.",
     "answer": "no", "type": "syllogism", "difficulty": "medium"},
    {"q": "If A is taller than B, and B is taller than C, who is the shortest?",
     "answer": "c", "type": "transitive", "difficulty": "easy"},
    {"q": "If no fish can fly, and all salmon are fish, can salmon fly? Answer yes or no.",
     "answer": "no", "type": "syllogism", "difficulty": "easy"},

    # Pattern recognition
    {"q": "What comes next: 2, 4, 8, 16, ?", "answer": "32", "type": "pattern", "difficulty": "easy"},
    {"q": "What comes next: 1, 1, 2, 3, 5, 8, ?", "answer": "13", "type": "pattern", "difficulty": "medium"},
    {"q": "What comes next: 3, 6, 9, 12, ?", "answer": "15", "type": "pattern", "difficulty": "easy"},
    {"q": "What comes next: 1, 4, 9, 16, 25, ?", "answer": "36", "type": "pattern", "difficulty": "medium"},

    # Multi-step reasoning
    {"q": "There are 3 boxes. Box A is heavier than Box B. Box C is lighter than Box B. Which box is the heaviest?",
     "answer": "a", "type": "multi_step", "difficulty": "easy"},
    {"q": "Tom is twice as old as Jerry. Jerry is 3 years older than Spike. Spike is 4 years old. How old is Tom?",
     "answer": "14", "type": "multi_step", "difficulty": "medium"},
    {"q": "A snail climbs 3 meters during the day but slips back 2 meters at night. How many days does it take to climb a 10 meter wall?",
     "answer": "8", "type": "multi_step", "difficulty": "hard"},
    {"q": "In a race, you overtake the person in 2nd place. What position are you now in?",
     "answer": "2", "type": "trick", "difficulty": "medium"},

    # Common sense
    {"q": "If I put a ice cube in a hot pan, what happens to it?",
     "answer": "melt", "type": "common_sense", "difficulty": "easy"},
    {"q": "Can a dead person eat lunch? Answer yes or no.",
     "answer": "no", "type": "common_sense", "difficulty": "easy"},
]

CREATIVITY_TESTS = [
    # These are scored differently — we check for length, variety, and coherence
    {"q": "Write a 4-line poem about the ocean.",
     "min_lines": 3, "min_words": 15, "type": "poem", "difficulty": "easy"},
    {"q": "Describe a sunset to someone who has never seen one, in 2-3 sentences.",
     "min_lines": 1, "min_words": 20, "type": "description", "difficulty": "easy"},
    {"q": "Come up with 3 creative names for a coffee shop run by robots.",
     "min_items": 3, "type": "brainstorm", "difficulty": "easy"},
    {"q": "Write a very short story (3-5 sentences) about a cat who learns to fly.",
     "min_lines": 2, "min_words": 25, "type": "story", "difficulty": "medium"},
    {"q": "Explain how a refrigerator works to a 5 year old, in 2-3 simple sentences.",
     "min_lines": 1, "min_words": 15, "type": "explain_simple", "difficulty": "easy"},
    {"q": "Rewrite this sentence to make it more exciting: 'The man walked to the store.'",
     "min_words": 5, "type": "rewrite", "difficulty": "easy",
     "must_not_contain": "The man walked to the store"},
    {"q": "Give me an analogy that explains what a CPU does in a computer.",
     "min_words": 10, "type": "analogy", "difficulty": "medium"},
    {"q": "Write a haiku about programming (5-7-5 syllable pattern).",
     "min_lines": 3, "min_words": 8, "type": "poem", "difficulty": "medium"},
]

KNOWLEDGE_TESTS = [
    # Science
    {"q": "What planet is closest to the Sun?", "answer": "mercury", "type": "science", "difficulty": "easy"},
    {"q": "What is the chemical symbol for water?", "answer": "h2o", "type": "science", "difficulty": "easy"},
    {"q": "What gas do plants absorb from the atmosphere?", "answer": "carbon dioxide", "type": "science", "difficulty": "easy",
     "alt_answers": ["co2"]},
    {"q": "How many chromosomes do humans have?", "answer": "46", "type": "science", "difficulty": "medium"},
    {"q": "What is the speed of light in km/s approximately?", "answer": "300000", "type": "science", "difficulty": "medium",
     "alt_answers": ["300,000", "3×10^8", "3x10^8", "299792"]},

    # History
    {"q": "In what year did World War 2 end?", "answer": "1945", "type": "history", "difficulty": "easy"},
    {"q": "Who was the first person to walk on the Moon?", "answer": "armstrong", "type": "history", "difficulty": "easy",
     "alt_answers": ["neil armstrong"]},
    {"q": "What country built the Great Wall?", "answer": "china", "type": "history", "difficulty": "easy"},

    # Geography
    {"q": "What is the largest ocean on Earth?", "answer": "pacific", "type": "geography", "difficulty": "easy"},
    {"q": "What is the capital of Japan?", "answer": "tokyo", "type": "geography", "difficulty": "easy"},
    {"q": "What is the longest river in the world?", "answer": "nile", "type": "geography", "difficulty": "medium",
     "alt_answers": ["amazon"]},  # Both are accepted — depends on measurement method
    {"q": "How many continents are there?", "answer": "7", "type": "geography", "difficulty": "easy"},

    # General
    {"q": "What programming language is known for its use in web browsers?", "answer": "javascript", "type": "tech", "difficulty": "easy",
     "alt_answers": ["js"]},
    {"q": "What does CPU stand for?", "answer": "central processing unit", "type": "tech", "difficulty": "easy"},
    {"q": "What does DNA stand for?", "answer": "deoxyribonucleic acid", "type": "science", "difficulty": "medium"},
]

INSTRUCTION_TESTS = [
    # Follow exact instructions
    {"q": "Respond with only the word 'hello'. Nothing else.",
     "check": "hello", "verify": "exact_word", "type": "exact", "difficulty": "easy"},
    {"q": "List exactly 3 colors, one per line. Nothing else.",
     "check": 3, "verify": "line_count", "type": "format", "difficulty": "easy"},
    {"q": "Write a sentence that contains exactly 5 words.",
     "check": 5, "verify": "word_count", "type": "format", "difficulty": "medium"},
    {"q": "Give me a number between 1 and 10. Just the number, nothing else.",
     "check": "number_1_10", "verify": "number_range", "type": "exact", "difficulty": "easy"},
    {"q": "Repeat the following sentence exactly: 'The quick brown fox jumps over the lazy dog.'",
     "check": "the quick brown fox jumps over the lazy dog", "verify": "contains_exact", "type": "repeat", "difficulty": "easy"},
    {"q": "Name 5 fruits. Number them 1-5.",
     "check": 5, "verify": "numbered_items", "type": "format", "difficulty": "easy"},
    {"q": "Answer this question with exactly one word: What color is the sky on a clear day?",
     "check": 1, "verify": "word_count", "type": "exact", "difficulty": "medium"},
    {"q": "Write the alphabet in reverse (Z to A) with commas between each letter.",
     "check": "z", "verify": "starts_with", "type": "exact", "difficulty": "hard"},
    {"q": "Respond with ONLY a JSON object with keys 'name' and 'age'. Use any values.",
     "check": "{", "verify": "starts_with_json", "type": "format", "difficulty": "medium"},
    {"q": "Say 'yes' if 2+2=4, otherwise say 'no'. Only one word.",
     "check": "yes", "verify": "exact_word", "type": "exact", "difficulty": "easy"},
]


# ============================================================
# ANSWER CHECKING
# ============================================================

def strip_thinking(response: str) -> str:
    """
    BUG FIX: Qwen3-VL wraps reasoning in <think>...</think> tags.
    These aren't special tokens so they stay in the decoded output.
    We need to strip them before checking answers, otherwise:
    - instruction checks see "<think>" as the first word
    - code exec tries to run thinking text as Python
    - exact match checks fail because of extra text
    """
    # Remove <think>...</think> blocks
    cleaned = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
    # Also handle unclosed think tags (model didn't close it)
    cleaned = re.sub(r'<think>.*$', '', cleaned, flags=re.DOTALL)
    return cleaned.strip()


def check_math_answer(response: str, expected: str) -> bool:
    """Check if the response contains the correct number."""
    response_clean = strip_thinking(response).lower().strip()
    # Strip commas from numbers (model might say "1,081" for 1081)
    response_no_commas = response_clean.replace(",", "")
    # Extract numbers from response
    numbers = re.findall(r'-?\d+\.?\d*', response_no_commas)

    # Also check for fraction answers
    if "/" in expected:
        if expected.lower() in response_clean:
            return True
        # Check decimal equivalent
        parts = expected.split("/")
        if len(parts) == 2:
            try:
                decimal_val = str(round(int(parts[0]) / int(parts[1]), 4))
                if decimal_val in response_no_commas:
                    return True
            except (ValueError, ZeroDivisionError):
                pass

    # BUG FIX: Also match "12" to "12.0" (correct but different format)
    for num_str in numbers:
        try:
            if float(num_str) == float(expected):
                return True
        except ValueError:
            continue

    return expected in numbers


def check_knowledge_answer(response: str, test: dict) -> bool:
    """Check if response contains the expected answer or alternatives."""
    response_lower = strip_thinking(response).lower().strip()
    if test["answer"].lower() in response_lower:
        return True
    for alt in test.get("alt_answers", []):
        if alt.lower() in response_lower:
            return True
    return False


def check_reasoning_answer(response: str, expected: str) -> bool:
    """Check reasoning answer — flexible matching."""
    response_lower = strip_thinking(response).lower().strip()
    expected_lower = expected.lower()

    # Direct match
    if expected_lower in response_lower:
        return True

    # For yes/no questions
    if expected_lower in ["yes", "no"]:
        # Check first word
        first_word = response_lower.split()[0] if response_lower.split() else ""
        first_word = first_word.strip(".,!?")
        return first_word == expected_lower

    # For number answers
    numbers = re.findall(r'-?\d+\.?\d*', response_lower)
    if expected_lower in numbers:
        return True

    return False


def check_code_answer(response: str, test: dict) -> bool:
    """Check code answers — verify function exists and/or contains expected code."""
    verify = test["verify"]
    check = test["check"]
    response = strip_thinking(response)

    if verify == "callable":
        # Try to extract and compile the function
        code_block = response
        # Extract from markdown code block if present
        md_match = re.search(r'```(?:python)?\s*\n(.*?)```', response, re.DOTALL)
        if md_match:
            code_block = md_match.group(1)

        # BUG FIX: If no code block found, try to extract just the def...
        if not md_match and "def " in code_block:
            # Find from first 'def' to end — skip any explanation text before it
            def_start = code_block.index("def ")
            code_block = code_block[def_start:]

        try:
            import signal

            def _timeout_handler(signum, frame):
                raise TimeoutError("Code execution timed out")

            old_handler = signal.signal(signal.SIGALRM, _timeout_handler)
            # BUG FIX: Set alarm BEFORE exec(), not after!
            # If model writes `while True: pass` at module level, exec()
            # hangs forever and the alarm never gets set.
            signal.alarm(10)  # 10 seconds for compile+exec+tests

            try:
                namespace = {}
                exec(code_block.strip(), namespace)
                func = namespace.get(check)
                if func is None or not callable(func):
                    return False

                # Run basic tests based on function name
                if check == "is_even":
                    result = func(4) == True and func(3) == False
                elif check == "factorial":
                    result = func(5) == 120 and func(0) == 1
                elif check == "reverse_string":
                    result = func("hello") == "olleh"
                elif check == "max_of_three":
                    result = func(1, 5, 3) == 5
                elif check == "count_vowels":
                    result = func("hello") == 2
                elif check == "fibonacci":
                    result = func(0) == 0 and func(1) == 1 and func(6) == 8
                elif check == "is_palindrome":
                    result = func("racecar") == True and func("hello") == False
                elif check == "flatten":
                    result = func([[1, 2], [3, 4]]) == [1, 2, 3, 4]
                else:
                    result = True
                return result
            finally:
                signal.alarm(0)
                signal.signal(signal.SIGALRM, old_handler)
        except (Exception, TimeoutError):
            return False

    elif verify == "contains":
        return check.lower() in response.lower()

    elif verify == "contains_any":
        return any(c.lower() in response.lower() for c in check)

    return False


def check_creativity_answer(response: str, test: dict) -> bool:
    """Check creativity — basic quality checks."""
    response = strip_thinking(response)
    lines = [l.strip() for l in response.strip().split("\n") if l.strip()]
    words = response.split()

    # Must not contain forbidden text (for rewrite tasks)
    if "must_not_contain" in test:
        if test["must_not_contain"].lower() in response.lower():
            return False

    # Check minimum lines
    if "min_lines" in test and len(lines) < test["min_lines"]:
        return False

    # Check minimum words
    if "min_words" in test and len(words) < test["min_words"]:
        return False

    # Check minimum items (for brainstorm)
    if "min_items" in test:
        # Count numbered items or bullet points
        items = [l for l in lines if re.match(r'^\d+[\.\):]|^[-*•]', l)]
        if len(items) < test["min_items"]:
            # Also count plain lines as items
            if len(lines) < test["min_items"]:
                return False

    # Basic coherence — not just random characters
    if len(response.strip()) < 10:
        return False

    return True


def check_instruction_answer(response: str, test: dict) -> bool:
    """Check instruction following — strict format checks."""
    verify = test["verify"]
    check = test["check"]
    response_clean = strip_thinking(response).strip()
    response_lower = response_clean.lower()

    if verify == "exact_word":
        # Should be just this word (maybe with punctuation)
        first_word = response_lower.split()[0] if response_lower.split() else ""
        first_word = re.sub(r'[.,!?"\']', '', first_word)
        return first_word == check.lower()

    elif verify == "line_count":
        lines = [l.strip() for l in response_clean.split("\n") if l.strip()]
        return len(lines) == check

    elif verify == "word_count":
        # Find the main sentence (skip thinking/explanation)
        # Look for a clean sentence
        lines = [l.strip() for l in response_clean.split("\n") if l.strip()]
        for line in lines:
            # Skip lines that look like meta-commentary
            if any(skip in line.lower() for skip in ["here is", "here's", "sure", "okay"]):
                continue
            words = line.split()
            if len(words) == check:
                return True
        # Also check total
        return len(response_clean.split()) == check

    elif verify == "number_range":
        numbers = re.findall(r'\d+', response_clean)
        if numbers:
            n = int(numbers[0])
            return 1 <= n <= 10
        return False

    elif verify == "contains_exact":
        return check.lower() in response_lower

    elif verify == "numbered_items":
        items = re.findall(r'^\d+[\.\):]', response_clean, re.MULTILINE)
        return len(items) >= check

    elif verify == "starts_with":
        return response_lower.startswith(check.lower())

    elif verify == "starts_with_json":
        stripped = response_clean.lstrip()
        # May be in a code block
        if "```" in stripped:
            md_match = re.search(r'```(?:json)?\s*\n(.*?)```', stripped, re.DOTALL)
            if md_match:
                stripped = md_match.group(1).strip()
        try:
            obj = json.loads(stripped)
            return "name" in obj and "age" in obj
        except (json.JSONDecodeError, TypeError):
            return stripped.startswith("{")

    return False


# ============================================================
# MODEL RUNNER
# ============================================================

def load_model(model_path: str):
    """Load model and tokenizer. Validates checkpoint integrity first."""
    from transformers import AutoModelForImageTextToText, AutoTokenizer

    # Validate checkpoint before loading — catch stale/corrupt checkpoints
    p = Path(model_path)
    if p.is_dir():
        config_file = p / "config.json"
        if not config_file.exists():
            raise FileNotFoundError(f"No config.json in {model_path} — stale or corrupt checkpoint!")
        st_files = list(p.glob("*.safetensors"))
        if not st_files and not list(p.glob("*.bin")):
            raise FileNotFoundError(f"No model weights in {model_path} — checkpoint is empty!")
        total_size = sum(f.stat().st_size for f in st_files) / 1e9 if st_files else 0
        if st_files and total_size < 0.5:
            raise ValueError(f"Model weights only {total_size:.2f} GB — likely corrupt checkpoint!")

    print(f"  Loading model from {model_path}...")
    try:
        model = AutoModelForImageTextToText.from_pretrained(
            model_path, dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True
        )
    except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise RuntimeError(f"Failed to load model (OOM or corrupt): {e}") from e

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    return model, tokenizer


def ask_model(model, tokenizer, question: str, max_tokens: int = 256,
              disable_thinking: bool = True) -> str:
    """Ask the model a question and get the response.

    SPEED: disable_thinking=True adds /no_think to skip <think> blocks.
    Qwen3-VL's thinking mode generates 100-500+ tokens of reasoning
    before giving a 1-word answer. For benchmarks this wastes ~80% of
    inference time. /no_think makes it answer directly.
    """
    # Build prompt — /no_think goes in the system message
    if disable_thinking:
        chat = (
            "<|im_start|>system\n/no_think<|im_end|>\n"
            f"<|im_start|>user\n{question}<|im_end|>\n"
            "<|im_start|>assistant\n"
        )
    else:
        chat = f"<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"

    ids = tokenizer(chat, return_tensors="pt").to(model.device)

    try:
        with torch.no_grad():
            out = model.generate(
                **ids,
                max_new_tokens=max_tokens,
                do_sample=False,
                temperature=None,
                top_p=None,
            )
    except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
        # OOM during generation — clean up and return empty
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        print(f"    OOM during generation: {e}")
        return "[OOM — no response generated]"

    response = tokenizer.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)

    # Guard against garbage output: if response is absurdly long, truncate
    if len(response) > 5000:
        response = response[:5000]

    return response.strip()


# ============================================================
# RUN BENCHMARK
# ============================================================

def run_category(model, tokenizer, category_name: str, tests: list,
                 check_fn, cfg: BenchmarkConfig) -> Dict:
    """Run all tests in a category and return results."""
    print(f"\n  === {category_name.upper()} ({len(tests)} questions) ===")

    results = []
    passed = 0
    failed_questions = []

    for i, test in enumerate(tests):
        q = test["q"]
        q_start = time.time()

        try:
            response = ask_model(model, tokenizer, q, cfg.max_new_tokens,
                                 disable_thinking=cfg.disable_thinking)
        except Exception as e:
            print(f"    [ERROR] Question {i+1} crashed: {e}")
            response = "[ERROR — question crashed]"
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        q_time = time.time() - q_start

        try:
            if check_fn == check_code_answer:
                correct = check_fn(response, test)
            elif check_fn == check_creativity_answer:
                correct = check_fn(response, test)
            elif check_fn == check_instruction_answer:
                correct = check_fn(response, test)
            elif check_fn == check_knowledge_answer:
                correct = check_fn(response, test)
            elif check_fn == check_reasoning_answer:
                correct = check_fn(response, test.get("answer", ""))
            else:
                # Math
                correct = check_fn(response, test.get("answer", ""))
        except Exception as e:
            print(f"    [ERROR] Answer checker crashed on question {i+1}: {e}")
            correct = False

        status = "PASS" if correct else "FAIL"
        if correct:
            passed += 1
        else:
            # Store the clean response (without thinking) for the failure report
            clean_response = strip_thinking(response)
            failed_questions.append({
                "question": q,
                "response": clean_response[:300],
                "expected": test.get("answer", test.get("check", "N/A")),
                "type": test.get("type", "unknown"),
                "difficulty": test.get("difficulty", "unknown"),
            })

        # Print compact result
        short_q = q[:60] + "..." if len(q) > 60 else q
        print(f"    [{status}] {short_q} ({q_time:.1f}s)")
        if not correct:
            clean_r = strip_thinking(response)
            short_r = clean_r[:80] + "..." if len(clean_r) > 80 else clean_r
            print(f"          Got: {short_r}")

        results.append({
            "question": q,
            "correct": correct,
            "type": test.get("type", "unknown"),
            "difficulty": test.get("difficulty", "unknown"),
        })

    total = len(tests)
    # CRITICAL FIX: Score as 0.0-1.0 (NOT 0-100) so td_weakness.py targets work
    score = passed / total if total > 0 else 0.0

    print(f"\n    SCORE: {passed}/{total} ({score:.0%})")

    # Per-type breakdown
    type_scores = {}
    for r in results:
        t = r["type"]
        if t not in type_scores:
            type_scores[t] = {"correct": 0, "total": 0}
        type_scores[t]["total"] += 1
        if r["correct"]:
            type_scores[t]["correct"] += 1

    for t, s in sorted(type_scores.items()):
        pct = s["correct"] / s["total"] if s["total"] > 0 else 0.0
        print(f"      {t}: {s['correct']}/{s['total']} ({pct:.0%})")

    return {
        "category": category_name,
        "score": score,  # 0.0-1.0 scale
        "passed": passed,
        "total": total,
        "type_scores": {t: {"correct": s["correct"], "total": s["total"],
                            "pct": s["correct"] / s["total"] if s["total"] > 0 else 0.0}
                        for t, s in type_scores.items()},
        "failed_questions": failed_questions,
        "results": results,
    }


def run_benchmark(cfg: BenchmarkConfig = None) -> Dict:
    """Run the full benchmark suite."""
    if cfg is None:
        cfg = BenchmarkConfig()

    # Auto-detect model
    if not cfg.model_path:
        base = Path("td_fuse_outputs/self_improve")
        if base.exists():
            for n in range(50, 0, -1):
                d = base / f"improved_cycle{n}"
                if d.exists() and list(d.glob("*.safetensors")):
                    cfg.model_path = str(d)
                    break
        if not cfg.model_path:
            healed = Path("td_fuse_outputs/reasoning_healed")
            if healed.exists():
                cfg.model_path = str(healed)
        if not cfg.model_path:
            raise FileNotFoundError("No model found!")

    start = time.time()
    print("=" * 60)
    print(f"TD BENCHMARK SUITE v1")
    print(f"Model: {cfg.model_path}")
    print(f"Started: {time.strftime('%H:%M:%S')}")
    print("=" * 60)

    model, tokenizer = load_model(cfg.model_path)

    # Run all categories
    # SPEED: Per-category max tokens. Math/knowledge/reasoning answers are
    # 1-10 tokens. Code needs 200+. Creativity needs 150+.
    # Using 256 for a "yes/no" question wastes inference time.
    CATEGORY_MAX_TOKENS = {
        "math": 64,
        "code": 300,  # Functions can be long
        "reasoning": 64,
        "creativity": 200,  # Stories/poems need room
        "knowledge": 64,
        "instruction_following": 128,  # JSON output, alphabet, etc.
    }

    categories = [
        ("math", MATH_TESTS, check_math_answer),
        ("code", CODE_TESTS, check_code_answer),
        ("reasoning", REASONING_TESTS, check_reasoning_answer),
        ("creativity", CREATIVITY_TESTS, check_creativity_answer),
        ("knowledge", KNOWLEDGE_TESTS, check_knowledge_answer),
        ("instruction_following", INSTRUCTION_TESTS, check_instruction_answer),
    ]

    all_results = {}
    total_passed = 0
    total_questions = 0

    # Save original max_new_tokens and swap per-category
    orig_max_tokens = cfg.max_new_tokens
    for cat_name, tests, check_fn in categories:
        cfg.max_new_tokens = CATEGORY_MAX_TOKENS.get(cat_name, orig_max_tokens)
        result = run_category(model, tokenizer, cat_name, tests, check_fn, cfg)
        all_results[cat_name] = result
        total_passed += result["passed"]
        total_questions += result["total"]
    cfg.max_new_tokens = orig_max_tokens  # Restore

    # Cleanup
    del model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    elapsed = (time.time() - start) / 60
    overall_score = total_passed / total_questions if total_questions > 0 else 0.0

    # Print summary
    print(f"\n{'='*60}")
    print(f"BENCHMARK COMPLETE — {elapsed:.1f} min")
    print(f"{'='*60}")
    print(f"\n  REPORT CARD:")
    print(f"  {'Category':<25} {'Score':>10} {'Grade':>8}")
    print(f"  {'-'*45}")

    for cat_name in ["math", "code", "reasoning", "creativity", "knowledge", "instruction_following"]:
        r = all_results[cat_name]
        score = r["score"]  # 0.0-1.0
        if score >= 0.90: grade = "A"
        elif score >= 0.80: grade = "B"
        elif score >= 0.70: grade = "C"
        elif score >= 0.60: grade = "D"
        else: grade = "F"
        bar = "█" * int(score * 20)
        print(f"  {cat_name:<25} {r['passed']:>3}/{r['total']:<3} ({score:>5.0%})  {grade:>4}  {bar}")

    print(f"\n  OVERALL: {total_passed}/{total_questions} ({overall_score:.0%})")
    print(f"{'='*60}")

    # Identify weaknesses (bottom 3 sub-categories)
    all_types = []
    for cat_name, result in all_results.items():
        for type_name, type_score in result["type_scores"].items():
            all_types.append({
                "category": cat_name,
                "type": type_name,
                "pct": type_score["pct"],
                "correct": type_score["correct"],
                "total": type_score["total"],
            })

    weakest = sorted(all_types, key=lambda x: x["pct"])[:5]
    print(f"\n  TOP 5 WEAKNESSES:")
    for w in weakest:
        print(f"    {w['category']}/{w['type']}: {w['correct']}/{w['total']} ({w['pct']:.0%})")

    # Save report
    out_dir = Path(cfg.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    report = {
        "model_path": cfg.model_path,
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "duration_min": elapsed,
        "overall_score": overall_score,
        "total_passed": total_passed,
        "total_questions": total_questions,
        "categories": {name: {
            "score": r["score"],
            "passed": r["passed"],
            "total": r["total"],
            "type_scores": r["type_scores"],
            "failed_questions": r["failed_questions"],
        } for name, r in all_results.items()},
        "weakest_areas": weakest,
    }

    # Support both output_path (exact path) and output_dir (auto-named)
    if cfg.output_path:
        report_path = Path(cfg.output_path)
        report_path.parent.mkdir(parents=True, exist_ok=True)
    else:
        report_path = out_dir / f"benchmark_{time.strftime('%Y%m%d_%H%M%S')}.json"

    try:
        with open(report_path, "w") as f:
            json.dump(report, f, indent=2)
        print(f"\n  Report saved: {report_path}")
    except OSError as e:
        print(f"\n  WARNING: Could not save benchmark report: {e}")
        print(f"  Results are still returned in memory — loop can continue.")

    return report


if __name__ == "__main__":
    run_benchmark()