"""GPQA Diamond.

Official source Idavidrein/gpqa is gated on HF (needs auth); we load the
nmayorga7 CSV export 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='nmayorga7/gpqa_diamond',  # HF CSV export of the gated official
        split='train',
        gen_config={'temperature': 1.0, 'max_tokens': 8192},
        prompt_style='cot_letter',  # CoT then ANSWER:X (es GPQA template)
        task_type='mcq',
        tags=['knowledge', 'science'],
        description='GPQA diamond split, graduate-level science MCQ (official content).',
        params={'hub': 'hf_raw'},
        # order_policy: 'sha256' (deterministic per-question shuffle -- position
        # bias protection) | 'es-dump:<path>' (pin the exact option order es
        # used in a specific run, for same-order alignment) | 'official'
        # (keep the CSV's raw order: incorrect 1-3 then correct)
        # set via spec params at get_dataset time or the default below.
    )
)
def gpqa_diamond(order_policy: str = 'sha256', order_dump: str = ''):
    import hashlib
    import random as _rnd
    import json as _json
    import os as _os

    dump = {}
    if order_policy.startswith('es-dump'):
        path = order_dump or order_policy.split(':', 1)[1] if ':' in order_policy else order_dump
        path = path or _os.environ.get('EVALHARNESS_CACHE', '') + '/../gpqa_es_order.json'
        if _os.path.exists(path):
            dump = _json.load(open(path))

    def to_sample(record: dict) -> Sample:
        choices = [
            str(record['Incorrect Answer 1'] or '').strip(),
            str(record['Incorrect Answer 2'] or '').strip(),
            str(record['Incorrect Answer 3'] or '').strip(),
            str(record['Correct Answer'] or '').strip(),
        ]
        q = str(record['Question']).strip()
        if q in dump:
            # pinned order from an es run dump
            choices = list(dump[q]['order'])
            target = dump[q]['target']
        elif order_policy == 'sha256':
            seed = int.from_bytes(
                hashlib.sha256(q.encode('utf-8')).digest()[:8], 'big')
            _rnd.Random(seed).shuffle(choices)
            target = 'ABCD'[choices.index(str(record['Correct Answer'] or '').strip())]
        else:  # 'official': raw order, correct is D
            target = 'D'
        return Sample(
            input=record['Question'],
            choices=choices,
            target=target,
            metadata={'subdomain': record.get('Subdomain')},
        )

    return to_sample