"""MATH-500 benchmark loader for competition mathematics.

MATH-500 is a subset of the MATH benchmark containing 500 diverse competition
math problems. These are harder than GSM8K and test more advanced mathematical
reasoning across algebra, geometry, number theory, etc.
"""

import re
from typing import Iterator

from ember.api.data import HuggingFaceSource
from embergraphrouter.evals.benchmarks.base import (
    Benchmark,
    BenchmarkDomain,
    BenchmarkExample,
    EvaluationScore,
    MetricType,
)
from embergraphrouter.evals.benchmarks.registry import register_benchmark


def _normalize_math_answer(text: str) -> str:
    """Normalize mathematical answer for comparison.

    Handles:
    - LaTeX box notation: \\boxed{answer}
    - LaTeX math mode: $...$
    - Fractions: \\frac{a}{b} -> a/b
    - Extra whitespace
    """
    text = text.strip()

    # Extract from \\boxed{...}
    boxed_match = re.search(r"\\boxed\{([^}]+)\}", text)
    if boxed_match:
        text = boxed_match.group(1)

    # Remove LaTeX math delimiters
    text = re.sub(r"\$+", "", text)

    # Simplify common LaTeX
    text = text.replace("\\frac", "frac")
    text = text.replace("\\left", "")
    text = text.replace("\\right", "")
    text = text.replace("\\cdot", "*")
    text = text.replace("\\times", "*")
    text = text.replace("\\div", "/")
    text = text.replace("\\pi", "pi")
    text = text.replace("\\sqrt", "sqrt")

    # Remove formatting
    text = re.sub(r"\\[a-zA-Z]+\{", "", text)
    text = text.replace("{", "").replace("}", "")

    # Normalize whitespace
    text = " ".join(text.split())

    return text.lower().strip()


def _extract_boxed_answer(text: str) -> str | None:
    """Extract answer from \\boxed{} notation.

    MATH problems typically have answers in \\boxed{} format.
    """
    # Standard boxed
    match = re.search(r"\\boxed\{([^{}]*(?:\{[^{}]*\}[^{}]*)*)\}", text)
    if match:
        return match.group(1).strip()

    # Alternative: answer after "the answer is"
    match = re.search(r"the\s+answer\s+is[:\s]*([^\n.]+)", text, re.IGNORECASE)
    if match:
        return match.group(1).strip()

    # Alternative: look for final expression
    match = re.search(r"=\s*([^\n=]+?)(?:\s*$|\s*\\)", text)
    if match:
        return match.group(1).strip()

    return None


def _parse_numeric(text: str) -> float | None:
    """Parse numeric value from mathematical expression."""
    normalized = _normalize_math_answer(text)

    # Try direct float parsing
    try:
        return float(normalized)
    except ValueError:
        pass

    # Handle fractions like "3/4"
    frac_match = re.match(r"(-?\d+)\s*/\s*(-?\d+)", normalized)
    if frac_match:
        numer = float(frac_match.group(1))
        denom = float(frac_match.group(2))
        if denom != 0:
            return numer / denom

    # Handle mixed numbers like "1 1/2"
    mixed_match = re.match(r"(-?\d+)\s+(\d+)\s*/\s*(\d+)", normalized)
    if mixed_match:
        whole = float(mixed_match.group(1))
        numer = float(mixed_match.group(2))
        denom = float(mixed_match.group(3))
        if denom != 0:
            sign = -1 if whole < 0 else 1
            return whole + sign * numer / denom

    # Extract first number
    num_match = re.search(r"-?\d+(?:\.\d+)?", normalized)
    if num_match:
        return float(num_match.group())

    return None


def _answers_equivalent(pred: str, ref: str) -> bool:
    """Check if two mathematical answers are equivalent.

    Handles symbolic and numeric comparisons.
    """
    pred_norm = _normalize_math_answer(pred)
    ref_norm = _normalize_math_answer(ref)

    # Exact string match
    if pred_norm == ref_norm:
        return True

    # Try numeric comparison
    pred_num = _parse_numeric(pred)
    ref_num = _parse_numeric(ref)

    if pred_num is not None and ref_num is not None:
        # Allow small tolerance for floating point
        return abs(pred_num - ref_num) < 1e-6

    return False


@register_benchmark("math500")
class MATH500Benchmark(Benchmark):
    """MATH-500 competition mathematics benchmark.

    500 diverse competition math problems covering:
    - Algebra
    - Geometry
    - Number Theory
    - Counting & Probability
    - Prealgebra
    - Precalculus
    - Intermediate Algebra

    Harder than GSM8K, uses \\boxed{} answer format.
    This benchmark is NOT in EB1 router training data.
    """

    name = "math500"
    version = "1.0"
    domain = BenchmarkDomain.MATH
    metric_type = MetricType.EXACT_MATCH
    description = "Competition mathematics requiring advanced reasoning"

    def __init__(self, max_examples: int | None = None):
        """Initialize MATH-500 benchmark.

        Args:
            max_examples: Maximum examples to load (None = all 500).
        """
        self._max_examples = max_examples

    def load(self, subset: str | None = None) -> Iterator[BenchmarkExample]:
        """Load MATH-500 examples via HuggingFace.

        Args:
            subset: Unused (single split).

        Yields:
            BenchmarkExample instances with problem and solution.
        """
        # MATH-500 on HuggingFace: lighteval/MATH-500
        source = HuggingFaceSource(
            "lighteval/MATH-500",
            split="test",
        )

        count = 0
        for batch in source.read_batches():
            for row in batch:
                if self._max_examples and count >= self._max_examples:
                    return

                problem = str(row.get("problem", ""))
                solution = str(row.get("solution", ""))
                answer = str(row.get("answer", ""))
                level = str(row.get("level", ""))
                category = str(row.get("type", row.get("subject", "")))

                # Extract boxed answer from solution if needed
                if not answer:
                    extracted = _extract_boxed_answer(solution)
                    answer = extracted if extracted else solution

                yield BenchmarkExample(
                    id=f"math500_{count}",
                    query=problem,
                    reference=answer,
                    metadata={
                        "domain": "math",
                        "source": "math500",
                        "level": level,
                        "category": category,
                        "full_solution": solution,
                        "difficulty": _map_level_to_difficulty(level),
                    },
                )
                count += 1

    def evaluate(
        self,
        prediction: str,
        example: BenchmarkExample,
    ) -> EvaluationScore:
        """Evaluate prediction by comparing extracted answers.

        Args:
            prediction: Model's output (should contain \\boxed{}).
            example: Ground truth with reference answer.

        Returns:
            EvaluationScore with equivalence result.
        """
        # Extract answer from prediction
        extracted = _extract_boxed_answer(prediction)
        if not extracted:
            # Fall back to last line
            lines = [ln.strip() for ln in prediction.split("\n") if ln.strip()]
            extracted = lines[-1] if lines else prediction

        # Get reference
        reference = (
            example.reference[0]
            if isinstance(example.reference, list)
            else example.reference
        )

        # Check equivalence
        correct = _answers_equivalent(extracted, reference)

        return EvaluationScore(
            correct=correct,
            score=1.0 if correct else 0.0,
            details={
                "metric": "math_equivalence",
                "extracted_answer": extracted,
                "expected_answer": reference,
                "normalized_pred": _normalize_math_answer(extracted),
                "normalized_ref": _normalize_math_answer(reference),
            },
        )


def _map_level_to_difficulty(level: str) -> str:
    """Map MATH difficulty level to standard difficulty."""
    level = level.lower().strip()

    if "1" in level:
        return "easy"
    elif "2" in level:
        return "easy"
    elif "3" in level:
        return "medium"
    elif "4" in level:
        return "hard"
    elif "5" in level:
        return "very_hard"

    return "unknown"
