from dataclasses import dataclass
from enum import Enum


@dataclass
class Metric:
    key: str
    col_name: str
    description: str = ""


@dataclass
class Task:
    benchmark: str
    metrics: list[Metric]
    col_name: str
    description: str
    task_type: str
    test_size: int
    dataset_url: str = ""
    primary_metric: str = ""

    def __post_init__(self):
        if not self.primary_metric:
            self.primary_metric = self.metrics[0].key

    @property
    def metric(self) -> str:
        """Backward-compatible primary metric key."""
        return self.primary_metric


MMRED_GROUP_DEFS = [
    ("mmred_dc_sa_c", "DC-SA-C", "Кто провёл больше/меньше всего времени в одиночестве?"),
    ("mmred_dc_sr_i", "DC-SR-I", "Сколько шагов X провёл в комнате Y?"),
    ("mmred_dc_cc_i", "DC-CC-I", "Сколько шагов в комнатах существовала толпа (3+ человек)?"),
    ("mmred_dc_ws_r", "DC-WS-R", "В какой комнате X провёл больше/меньше всего времени?"),
    ("mmred_dc_whs_c", "DC-WHS-C", "Кто провёл больше/меньше всего времени в комнате X?"),
]

MMRED_LENGTHS = [32, 64, 128]

MMRED_GROUP_KEYS = {group_key for group_key, _, _ in MMRED_GROUP_DEFS}

MMRED_SUBTASK_KEYS = {
    f"{group_key}_{length}"
    for group_key, _, _ in MMRED_GROUP_DEFS
    for length in MMRED_LENGTHS
}


def _mmred_metrics() -> list[Metric]:
    metrics = [
        Metric(
            "em.dc_aggregate",
            "Overall",
            "Итоговая оценка MMReD: гармоническое среднее пяти взвешенных EM по типам вопросов.",
        ),
    ]
    for group_key, group_label, group_desc in MMRED_GROUP_DEFS:
        metrics.append(
            Metric(
                f"{group_key}::em.dc_aggregate",
                group_label,
                f"{group_desc} Взвешенный EM по длинам 32/64/128.",
            )
        )
    for group_key, group_label, _ in MMRED_GROUP_DEFS:
        for length in MMRED_LENGTHS:
            metrics.append(
                Metric(
                    f"{group_key}_{length}::exact_match",
                    f"{group_label} × {length}",
                    f"Exact Match на подзадаче длины {length}.",
                )
            )
    return metrics


# Select your tasks here
# ---------------------------------------------------
class Tasks(Enum):
    ruaime = Task(
        benchmark="ruaime",
        metrics=[
            Metric(
                "exact_match",
                "Exact match",
                "Максимум из exact_match,flexible-extract и exact_match,strict-match.",
            ),
        ],
        col_name="ruAIME",
        dataset_url="https://huggingface.co/datasets/MERA-evaluation/ruAIME",
        description=(
            "Русскоязычный набор задач AIME (1983–2025): олимпиадная математика с числовыми "
            "ответами, требующая многошагового рассуждения."
        ),
        task_type="Математика, рассуждение",
        test_size=724,
        primary_metric="exact_match",
    )
    t_math = Task(
        benchmark="t_math",
        metrics=[
            Metric(
                "exact_match",
                "Exact match",
                "Математическая эквивалентность предсказания и эталона (math_verify).",
            ),
        ],
        col_name="T-math",
        dataset_url="https://huggingface.co/datasets/MERA-evaluation/T-math",
        description=(
            "310 олимпиадных математических задач на русском из Всероссийской олимпиады "
            "школьников и Московской олимпиады (2005–2025)."
        ),
        task_type="Математика, олимпиадные задачи",
        test_size=310,
        primary_metric="exact_match",
    )
    luzitania = Task(
        benchmark="luzitania",
        metrics=[
            Metric(
                "exact_match",
                "Exact match",
                "Максимум из exact_match,flexible-extract и exact_match,strict-match.",
            ),
        ],
        col_name="Luzitania",
        dataset_url="https://huggingface.co/datasets/MERA-evaluation/Luzitania",
        description=(
            "Набор для оценки многошагового математического рассуждения на русском языке: "
            "олимпиадные и продвинутые задачи с целочисленными ответами."
        ),
        task_type="Логика, рассуждение",
        test_size=251,
        primary_metric="exact_match",
    )
    mmred = Task(
        benchmark="mmred",
        metrics=_mmred_metrics(),
        col_name="MMReD",
        dataset_url="https://huggingface.co/datasets/MERA-evaluation/MMReD",
        description=(
            "Multi-Modal Reasoning in Dense Context — синтетический бенчмарк рассуждения "
            "на длинном плотном контексте. 5 типов вопросов × 3 длины (32, 64, 128) = 750 вопросов."
        ),
        task_type="Длинный контекст, рассуждение",
        test_size=750,
        primary_metric="em.dc_aggregate",
    )


NUM_FEWSHOT = 0
# ---------------------------------------------------