"""AI2 ARC (official source: allenai/ai2_arc, ARC-Easy / ARC-Challenge)."""

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


@register_dataset(
    DatasetSpec(
        name='arc',
        source='allenai/ai2_arc',  # official: https://huggingface.co/datasets/allenai/ai2_arc
        subset='ARC-Easy',  # or ARC-Challenge
        split='test',
        task_type='mcq',
        tags=['knowledge', 'science'],
        description='AI2 Reasoning Challenge (official).',
    )
)
def arc():
    def to_sample(record: dict) -> Sample:
        choices = record['choices']  # {'text': [...], 'label': [...]}
        return Sample(
            input=record['question'],
            choices=list(choices['text']),
            target=str(record['answerKey']).strip(),  # letter label
            metadata={'id': record.get('id'), 'labels': list(choices['label'])},
        )

    return to_sample