metrics_config.py

"""
Scenario metrics: CSV column keys, labels, and grouping for leaderboard + analytics.

Canonical scenario keys are far-field oriented. Legacy keys (wer_clean, …) are migrated
on CSV load in ``init.normalize_legacy_csv_row``.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from typing import Literal

Status = Literal["live", "planned"]


@dataclass(frozen=True)
class ScenarioMetric:
    """One eval scenario (column `key` in results CSV)."""

    key: str
    label: str
    short: str
    description: str
    group: str
    status: Status


# Order defines CSV column order for scenario WER columns.
SCENARIO_METRICS: tuple[ScenarioMetric, ...] = (
    ScenarioMetric(
        key="wer_anechoic_speech",
        label="Near Field Speech (close-mic)",
        short="Near Field Speech",
        description="Close-mic / anechoic near-field speech; easy baseline.",
        group="Core benchmarks",
        status="live",
    ),
    ScenarioMetric(
        key="wer_lab_measured",
        label="Lab Measured",
        short="Lab Measured",
        description="Controlled lab recordings with measured acoustics / additive noise.",
        group="Core benchmarks",
        status="live",
    ),
    ScenarioMetric(
        key="wer_lab_simulated",
        label="Lab Simulated",
        short="Lab Simulated",
        description="Simulated acoustics (e.g. simulated room / Treble-style reverb).",
        group="Core benchmarks",
        status="live",
    ),
    ScenarioMetric(
        key="wer_realistic_high_snr",
        label="Realistic conditions: High SNR",
        short="High SNR",
        description="Realistic far-field / room conditions at relatively high SNR.",
        group="Realistic conditions",
        status="live",
    ),
    ScenarioMetric(
        key="wer_realistic_mid_snr",
        label="Realistic conditions: Mid SNR",
        short="Mid SNR",
        description="Realistic far-field mixtures at mid SNR.",
        group="Realistic conditions",
        status="live",
    ),
    ScenarioMetric(
        key="wer_realistic_low_snr",
        label="Realistic conditions: Low SNR",
        short="Low SNR",
        description="Hard realistic mixtures (low SNR, strong reverb, distortion).",
        group="Realistic conditions",
        status="live",
    ),
    ScenarioMetric(
        key="wer_moving_low",
        label="Moving Low SNR",
        short="Moving Low SNR",
        description="Moving talker / varying geometry (packed moving_low.pt split).",
        group="Moving conditions",
        status="live",
    ),
    ScenarioMetric(
        key="wer_moving_mid",
        label="Moving Mid SNR",
        short="Moving Mid SNR",
        description="Moving talker / varying geometry (packed moving_mid.pt split).",
        group="Moving conditions",
        status="live",
    ),
    ScenarioMetric(
        key="wer_moving_high",
        label="Moving High SNR",
        short="Moving High SNR",
        description="Moving talker / varying geometry (packed moving_high.pt split).",
        group="Moving conditions",
        status="live",
    ),
)

SCENARIO_KEYS: tuple[str, ...] = tuple(m.key for m in SCENARIO_METRICS)

# Heatmap column order: each moving SNR split immediately follows its static peer.
HEATMAP_SCENARIO_KEYS: tuple[str, ...] = (
    "wer_anechoic_speech",
    "wer_lab_measured",
    "wer_lab_simulated",
    "wer_realistic_high_snr",
    "wer_moving_high",
    "wer_realistic_mid_snr",
    "wer_moving_mid",
    "wer_realistic_low_snr",
    "wer_moving_low",
)

LIVE_SCENARIO_KEYS: tuple[str, ...] = tuple(m.key for m in SCENARIO_METRICS if m.status == "live")

PLANNED_SCENARIO_KEYS: tuple[str, ...] = tuple(m.key for m in SCENARIO_METRICS if m.status == "planned")


def metric_by_key(key: str) -> ScenarioMetric | None:
    for m in SCENARIO_METRICS:
        if m.key == key:
            return m
    return None


def labels_for_keys(keys: list[str]) -> list[str]:
    out = []
    for k in keys:
        m = metric_by_key(k)
        out.append(m.short if m else k)
    return out


# Compact axis labels for the Analysis heatmap (avoids overlapping headers).
HEATMAP_COLUMN_LABELS: dict[str, str] = {
    "wer_anechoic_speech": "Near Field",
    "wer_lab_measured": "Lab Meas.",
    "wer_lab_simulated": "Lab Sim.",
    "wer_realistic_high_snr": "High SNR",
    "wer_realistic_mid_snr": "Mid SNR",
    "wer_realistic_low_snr": "Low SNR",
    "wer_moving_low": "Mov. Low",
    "wer_moving_mid": "Mov. Mid",
    "wer_moving_high": "Mov. High",
}


def heatmap_label_for_key(key: str) -> str:
    m = metric_by_key(key)
    return HEATMAP_COLUMN_LABELS.get(key) or (m.short if m else key)


def resolve_scenario_metric_keys(
    selected: Sequence[str] | None,
    choices: Sequence[tuple[str, str]] | None = None,
) -> list[str]:
    """Map CheckboxGroup values (canonical keys or display labels) to scenario keys."""
    if not selected:
        return []
    label_to_key: dict[str, str] = {}
    valid_keys: set[str] = set(SCENARIO_KEYS)
    for item in choices or ():
        if isinstance(item, (tuple, list)) and len(item) >= 2:
            label_to_key[str(item[0])] = str(item[1])
            valid_keys.add(str(item[1]))
        elif isinstance(item, str):
            valid_keys.add(item)
    out: list[str] = []
    seen: set[str] = set()
    for item in selected:
        s = str(item).strip()
        if not s:
            continue
        key = s if s in valid_keys else label_to_key.get(s)
        if key and key not in seen:
            out.append(key)
            seen.add(key)
    return out