"""CMMLU dataset plugin (mirror: evalscope/cmmlu on ModelScope).

The mirror packs all 67 subjects into one parquet with a ``category``
column, so the loader filters by ``subset`` (filter_column convention).
The official HF repo (haonan-li/cmmlu) is script-based; use --subset to
pick a subject, 'all' for everything.
"""

import re

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


@register_dataset(
    DatasetSpec(
        name='cmmlu',
        source='evalscope/cmmlu',  # ModelScope parquet mirror; the HF original is script-based
        subset='anatomy',  # 67 subjects; override with --subset <subject> or 'all'
        split='test',
        few_shot_split='dev',
        few_shot_num=0,   # es default is 0-shot (docstring says 5 but code says 0)
        task_type='mcq',
        tags=['zh', 'knowledge'],
        description='Chinese multiple-choice QA (official content, ModelScope mirror).',
        params={'hub': 'modelscope', 'filter_column': 'category'},
    )
)
def cmmlu():
    def to_sample(record: dict) -> Sample:
        # mirror layout: question/choices(['(A) ...', ...])/answer('(B) ...'); official: Question/A-D/Answer
        if 'Question' in record:
            choices = [record[k] for k in ('A', 'B', 'C', 'D') if record.get(k) is not None]
            return Sample(
                input=record['Question'],
                choices=choices,
                target=str(record.get('Answer', '')).strip(),
                metadata={'category': record.get('Subject')},
            )
        choices = [re.sub(r'^\([A-J]\)\s*', '', c) for c in record['choices']]
        answer = str(record.get('answer', ''))
        m = re.match(r'^\(?([A-J])\)?', answer)
        target = m.group(1) if m and len(answer) > 1 else answer
        return Sample(
            input=record['question'],
            choices=choices,
            target=target,
            metadata={'category': record.get('category'), 'id': record.get('id')},
        )

    return to_sample