"""
Canonical benchmark registry for the EEG Finetune Arena.

Maps internal benchmark keys to HuggingFace dataset IDs, task IDs,
and metadata used by both eval.yaml generation and eval results pushing.
"""

from dataclasses import dataclass, field
from typing import Dict


@dataclass(frozen=True)
class BenchmarkInfo:
    """Metadata for a single EEG benchmark dataset."""

    dataset_id: str
    task_id: str
    display_name: str
    description: str
    category: str
    num_classes: int
    accuracy_field: str
    split: str = "test"
    config: str = "default"


EEG_BENCHMARKS: Dict[str, BenchmarkInfo] = {
    "bcic2a": BenchmarkInfo(
        dataset_id="braindecode/bcic2a",
        task_id="motor_imagery_4class",
        display_name="BCIC-IV-2a",
        description="BCI Competition IV Dataset 2a - 4-class motor imagery classification.",
        category="Motor Imagery",
        num_classes=4,
        accuracy_field="bcic2a_accuracy",
    ),
    "physionet": BenchmarkInfo(
        dataset_id="braindecode/physionet",
        task_id="motor_imagery_4class",
        display_name="PhysioNet-MI",
        description="PhysioNet Motor Imagery dataset - 4-class motor imagery classification.",
        category="Motor Imagery",
        num_classes=4,
        accuracy_field="physionet_accuracy",
    ),
    "bcic2020_3": BenchmarkInfo(
        dataset_id="braindecode/bcic2020-3",
        task_id="imagined_speech_5class",
        display_name="BCIC-2020-3",
        description="BCI Competition 2020 Track #3 - 5-class imagined speech classification.",
        category="Speech Imagery",
        num_classes=5,
        accuracy_field="bcic2020_3_accuracy",
    ),
    "isruc_sleep": BenchmarkInfo(
        dataset_id="braindecode/isruc-sleep",
        task_id="sleep_staging_5class",
        display_name="ISRUC",
        description="ISRUC-SLEEP dataset - 5-class sleep staging classification.",
        category="Sleep Staging",
        num_classes=5,
        accuracy_field="isruc_sleep_accuracy",
    ),
    "tuab": BenchmarkInfo(
        dataset_id="braindecode/tuab",
        task_id="pathology_binary",
        display_name="TUAB",
        description="Temple University Abnormal EEG corpus - binary pathology detection.",
        category="Pathology Detection",
        num_classes=2,
        accuracy_field="tuab_accuracy",
    ),
    "tuev": BenchmarkInfo(
        dataset_id="braindecode/tuev",
        task_id="event_classification_6class",
        display_name="TUEV",
        description="Temple University EEG Events corpus - 6-class event classification.",
        category="Event Classification",
        num_classes=6,
        accuracy_field="tuev_accuracy",
    ),
    "mdd_mumtaz2016": BenchmarkInfo(
        dataset_id="braindecode/mdd_mumtaz2016",
        task_id="pathology_binary",
        display_name="Mumtaz",
        description="Mumtaz 2016 MDD dataset - binary depression vs. healthy classification.",
        category="Pathology Detection",
        num_classes=2,
        accuracy_field="mdd_mumtaz2016_accuracy",
    ),
    "arithmetic_zyma2019": BenchmarkInfo(
        dataset_id="braindecode/arithmetic_zyma2019",
        task_id="cognitive_load_binary",
        display_name="MAT",
        description="EEG During Mental Arithmetic Tasks (Zyma 2019) - binary rest vs. task classification.",
        category="Cognitive Load",
        num_classes=2,
        accuracy_field="arithmetic_zyma2019_accuracy",
    ),
    "faced": BenchmarkInfo(
        dataset_id="braindecode/faced",
        task_id="emotion_recognition_9class",
        display_name="FACED",
        description="FACED dataset - 9-class emotion recognition from EEG.",
        category="Emotion Recognition",
        num_classes=9,
        accuracy_field="faced_accuracy",
    ),
    "seed_v": BenchmarkInfo(
        dataset_id="braindecode/seed-v",
        task_id="emotion_recognition_5class",
        display_name="SEED-V",
        description="SEED-V dataset - 5-class emotion recognition from EEG.",
        category="Emotion Recognition",
        num_classes=5,
        accuracy_field="seed_v_accuracy",
    ),
}

EVALUATION_FRAMEWORK = "eeg-finetune-arena"