"""MATH (official source: EleutherAI/hendrycks_math, the maintained Hendrycks MATH)."""

from ..sample import Sample
from ..registry import register_dataset
from ..spec import DatasetSpec


def _extract_boxed(text: str) -> str:
    """Extract the last \\boxed{...} content with brace balancing."""
    idx = text.rfind('\\boxed{')
    if idx < 0:
        return ''
    i = idx + len('\\boxed{')
    depth, start = 1, i
    while i < len(text) and depth:
        if text[i] == '{':
            depth += 1
        elif text[i] == '}':
            depth -= 1
        i += 1
    return text[start : i - 1] if depth == 0 else ''


@register_dataset(
    DatasetSpec(
        name='competition_math',
        source='EleutherAI/hendrycks_math',  # https://huggingface.co/datasets/EleutherAI/hendrycks_math
        subset='algebra',  # 7 subjects; override with --subset <subject>
        split='test',
        task_type='math',
        tags=['math'],
        description='MATH competition problems (Hendrycks). Target = \\boxed answer.',
    )
)
def competition_math():
    def to_sample(record: dict) -> Sample:
        solution = record['solution']
        return Sample(
            input=record['problem'],
            target=_extract_boxed(solution) or solution.strip(),
            metadata={'level': record.get('level'), 'type': record.get('type'), 'solution': solution},
        )

    return to_sample