"""GPQA Diamond.

Official source Idavidrein/gpqa is gated on HF (needs auth); we load the
ModelScope mirror (AI-ModelScope/gpqa_diamond) which keeps the official
column layout (Question / Correct Answer / Incorrect Answer 1-3).

Note: choices are stored with the correct answer first (target 'A').
Option shuffling is an eval-time concern (the evaluator should shuffle
choices and remap the target, like Dataset.shuffle_choices would).
"""

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


@register_dataset(
    DatasetSpec(
        name='gpqa_diamond',
        source='AI-ModelScope/gpqa_diamond',  # ModelScope mirror of the gated official
        split='train',
        task_type='mcq',
        tags=['knowledge', 'science'],
        description='GPQA diamond split, graduate-level science MCQ (official content).',
        params={'hub': 'modelscope'},
    )
)
def gpqa_diamond():
    def to_sample(record: dict) -> Sample:
        choices = [
            record['Correct Answer'],
            record['Incorrect Answer 1'],
            record['Incorrect Answer 2'],
            record['Incorrect Answer 3'],
        ]
        return Sample(
            input=record['Question'],
            choices=choices,
            target='A',  # correct answer is first; shuffle at eval time
            metadata={'subdomain': record.get('Subdomain'), 'unshuffled': True},
        )

    return to_sample