"""Math evaluation sets and calibration data.

The leaderboard set is hidden and probably postdates the model, so we keep two
tiers deliberately separate:

* **gate**  — small, fast, run on every recipe. Cheap signal for iteration.
* **holdout** — recent competitions we never tune against. The honest estimate.

AIME 2024 is deliberately excluded from the holdout: it is measurably
contaminated (inflating scores 10-20 points over clean contests), so it flatters
every recipe equally and discriminates between none of them.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterable, Sequence


@dataclass
class MathExample:
    example_id: str
    problem: str
    answer: str
    source: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class DatasetSpec:
    """How to pull one benchmark off the Hub.

    Field names differ between mirrors of the same benchmark, so each role lists
    candidate column names tried in order.
    """

    name: str
    hf_id: str
    split: str = "test"
    config: str | None = None
    problem_fields: Sequence[str] = ("problem", "Problem", "question", "Question")
    answer_fields: Sequence[str] = ("answer", "Answer", "solution", "expected_answer")
    tier: str = "gate"
    filters: tuple[tuple[str, str, Any], ...] = ()
    max_examples: int | None = None
    note: str = ""


# Three tiers, by how often we run them and how much tuning pressure they can
# absorb before their numbers stop meaning anything.
#
#   gate       every experiment. Cheap, tuned against freely.
#   checkpoint before a weekly leaderboard submission. Moderate tuning risk.
#   holdout    the two graded checkpoints only. NEVER tuned against — these are
#              post-release contests and the closest proxy we have for a hidden
#              eval that "is not available in public domain today".
REGISTRY: dict[str, DatasetSpec] = {
    "math500_hard": DatasetSpec(
        name="math500_hard",
        hf_id="HuggingFaceH4/MATH-500",
        split="test",
        tier="gate",
        filters=(("level", "gte", 4),),
        max_examples=100,
        note="MATH-500 levels 4-5, deterministic 100-problem subsample. The "
        "fast regression signal: sensitive enough to catch damage, cheap "
        "enough to run on every recipe.",
    ),
    "math500": DatasetSpec(
        name="math500",
        hf_id="HuggingFaceH4/MATH-500",
        split="test",
        tier="checkpoint",
        note="Full 500. Largely saturated for this model (~84.5 bf16), so it "
        "detects collapse but not subtle reasoning damage.",
    ),
    "aime25": DatasetSpec(
        name="aime25",
        hf_id="MathArena/aime_2025",
        split="train",
        tier="checkpoint",
        note="30 problems. Hard tail — where quantization damage actually shows.",
    ),
    "hmmt_feb25": DatasetSpec(
        name="hmmt_feb25",
        hf_id="MathArena/hmmt_feb_2025",
        split="train",
        tier="checkpoint",
        note="30 problems. Reported on the model card (74.0), so we have a "
        "published bf16 reference to validate our harness against.",
    ),
    "aime26": DatasetSpec(
        name="aime26",
        hf_id="MathArena/aime_2026",
        split="train",
        tier="holdout",
        note="30 problems, Feb 2026 contest. Post-dates most training data.",
    ),
    "hmmt_feb26": DatasetSpec(
        name="hmmt_feb26",
        hf_id="MathArena/hmmt_feb_2026",
        split="train",
        tier="holdout",
        note="33 problems, Feb 2026 contest. Cleanest proxy for the hidden eval.",
    ),
    "aime24": DatasetSpec(
        name="aime24",
        hf_id="Maxwell-Jia/AIME_2024",
        split="train",
        tier="diagnostic",
        note="Measurably contaminated — inflates scores 10-20 points over clean "
        "contests. Diagnostic only, never for recipe selection.",
    ),
}

# Training pools — problems with known answers, used to generate our own
# reasoning traces. Disjoint from every eval set above: MATH-500 is drawn from
# the MATH *test* split, so the MATH train split cannot leak into it.
TRAIN_REGISTRY: dict[str, DatasetSpec] = {
    "math_train": DatasetSpec(
        name="math_train",
        hf_id="EleutherAI/hendrycks_math",
        split="train",
        config="algebra",
        answer_fields=("solution",),  # gold answer is the \boxed{} in the solution
        tier="train",
        note="MATH train split. Pass --config to pick a subject.",
    ),
    "openr1": DatasetSpec(
        name="openr1",
        hf_id="open-r1/OpenR1-Math-220k",
        split="train",
        answer_fields=("answer", "solution"),
        tier="train",
        note="220k competition problems with verified answers.",
    ),
}

MATH_SUBJECTS = (
    "algebra", "counting_and_probability", "geometry", "intermediate_algebra",
    "number_theory", "prealgebra", "precalculus",
)

SUITES: dict[str, list[str]] = {
    "gate": ["math500_hard"],
    "checkpoint": ["math500", "aime25", "hmmt_feb25"],
    "holdout": ["aime26", "hmmt_feb26"],
}


def _passes_filters(row: dict[str, Any], filters: Sequence[tuple[str, str, Any]]) -> bool:
    for field_name, op, value in filters:
        actual = row.get(field_name)
        if actual is None:
            return False
        if op == "gte" and not actual >= value:
            return False
        if op == "lte" and not actual <= value:
            return False
        if op == "eq" and actual != value:
            return False
        if op == "in" and actual not in value:
            return False
    return True


def _subsample(examples: list[MathExample], n: int, seed: int = 0) -> list[MathExample]:
    """Deterministic subsample, stable across runs and machines.

    Shuffles with a fixed seed rather than taking a prefix, because these sets
    are ordered by subject/difficulty and a prefix would be badly skewed. The
    gate set must be identical across every recipe or the comparison is
    meaningless.
    """
    import random

    if len(examples) <= n:
        return examples
    indices = sorted(range(len(examples)))
    random.Random(seed).shuffle(indices)
    return [examples[i] for i in sorted(indices[:n])]


def _resolve_field(row: dict[str, Any], candidates: Iterable[str]) -> str | None:
    lowered = {k.lower(): k for k in row}
    for candidate in candidates:
        key = lowered.get(candidate.lower())
        if key is not None and row[key] is not None:
            return str(row[key])
    return None


def load_dataset_examples(
    spec: DatasetSpec | str,
    limit: int | None = None,
    cache_dir: str | None = None,
) -> list[MathExample]:
    """Load one benchmark into ``MathExample`` records.

    Raises with the observed column names when a field cannot be resolved, so a
    schema change on the Hub produces an actionable error instead of silently
    empty problems.
    """
    from datasets import load_dataset

    if isinstance(spec, str):
        table = {**REGISTRY, **TRAIN_REGISTRY}
        if spec not in table:
            raise KeyError(f"Unknown dataset {spec!r}. Known: {sorted(table)}")
        spec = table[spec]

    kwargs: dict[str, Any] = {"split": spec.split}
    if spec.config:
        kwargs["name"] = spec.config
    if cache_dir:
        kwargs["cache_dir"] = cache_dir

    dataset = load_dataset(spec.hf_id, **kwargs)

    examples: list[MathExample] = []
    for i, row in enumerate(dataset):
        if not _passes_filters(row, spec.filters):
            continue
        problem = _resolve_field(row, spec.problem_fields)
        answer = _resolve_field(row, spec.answer_fields)
        if problem is None or answer is None:
            raise ValueError(
                f"{spec.name}: could not resolve problem/answer fields. "
                f"Available columns: {sorted(row)}. "
                f"Tried problem={list(spec.problem_fields)}, answer={list(spec.answer_fields)}."
            )
        if "\\boxed" in answer:
            from .answers import extract_boxed

            boxed = extract_boxed(answer)
            if boxed is None:
                continue  # unparseable gold: drop rather than train on it
            answer = boxed

        examples.append(
            MathExample(
                example_id=f"{spec.name}:{i}",
                problem=problem,
                answer=answer,
                source=spec.name,
                metadata={
                    k: row[k]
                    for k in ("level", "subject", "type", "url", "id", "problem_idx")
                    if k in row
                },
            )
        )

    # Spec cap first (defines the canonical set), then the ad-hoc --limit.
    if spec.max_examples is not None:
        examples = _subsample(examples, spec.max_examples)
    if limit is not None:
        examples = examples[:limit]
    return examples


def load_suite(
    names: Sequence[str],
    limit: int | None = None,
    cache_dir: str | None = None,
) -> list[MathExample]:
    """Load and concatenate several benchmarks. ``limit`` applies per dataset.

    Accepts tier names (``gate``/``checkpoint``/``holdout``) as shorthand for
    the datasets in that tier.
    """
    resolved: list[str] = []
    for name in names:
        resolved.extend(SUITES[name] if name in SUITES else [name])

    out: list[MathExample] = []
    for name in resolved:
        out.extend(load_dataset_examples(name, limit=limit, cache_dir=cache_dir))
    return out


def describe_registry() -> str:
    lines = []
    for tier in ("gate", "checkpoint", "holdout", "diagnostic"):
        members = [s for s in REGISTRY.values() if s.tier == tier]
        if not members:
            continue
        lines.append(f"[{tier}]")
        for spec in members:
            cap = f" (capped at {spec.max_examples})" if spec.max_examples else ""
            lines.append(f"  {spec.name:<14} {spec.hf_id}{cap}")
            lines.append(f"  {'':<14} {spec.note}")
    return "\n".join(lines)


MATH_PROMPT = (
    "Solve the following math problem. Put your final answer inside "
    "\\boxed{{}} on the last line.\n\n"
    "Problem:\n{problem}"
)


def build_prompt(example: MathExample) -> str:
    return MATH_PROMPT.format(problem=example.problem.strip())