"""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',
        # es parity source: evalscope/competition_math (ModelScope) -- the
        # EleutherAI mirror shares NO questions with es's copy (0/199 text
        # overlap verified), same-question runs must use this source
        source='evalscope/competition_math',
        subset='Level 1',  # Level 1..5; override with --subset <subject>
        split='test',
        few_shot_split='train',
        few_shot_num=4,
        gen_config={'temperature': 0.0, 'max_tokens': 32768},
        prompt_style='imo_es',  # es: Problem: prefix + boxed suffix,
        task_type='math',
        params={'hub': 'modelscope', 'filter_column': 'level'},
        tags=['math'],
        description='MATH competition problems (Hendrycks). Target = \\boxed answer.',
    )
)
def competition_math():
    def to_sample(record: dict) -> Sample:
        solution = record.get('solution') or ''
        problem = record.get('problem') or record.get('input') or ''
        target = _extract_boxed(solution) or str(record.get('answer') or '').strip() or solution.strip()
        return Sample(
            input=problem,
            target=target,
            metadata={'es_math_fewshot': True, 'level': record.get('level') or '',
                      'type': record.get('type') or '', 'solution': solution},
        )

    return to_sample