15_07_expt/scripts/labels.py

"""Canonical label taxonomy and source-dataset label mappings."""

from __future__ import annotations

import re
from typing import Optional


TAXONOMY = {
    0: "3Toeloop",
    1: "3Loop",
    2: "2Axel",
    3: "CamelSpin",
    4: "SitSpin",
    5: "UprightSpin",
    6: "2Salchow",
    7: "2Toeloop",
    8: "3Salchow",
    9: "3Axel",
    10: "3Flip",
    11: "3Lutz",
    12: "NoBasic",
    13: "2Lutz",
    14: "4Salchow",
    15: "4Flip",
    16: "4Toeloop",
    17: "4Lutz",
    18: "4Loop",
    19: "2Flip",
    20: "2Loop",
    21: "1Axel",
    22: "1Loop",
    23: "1Salchow",
    24: "1Toeloop",
    25: "1Flip",
    26: "1Lutz",
    27: "Sequence",
}

LABEL_TO_IDX = {label: idx for idx, label in TAXONOMY.items()}


# --- Coarse ("action-level") taxonomy: drop the rotation count from jumps ---
# 2Axel/3Axel/1Axel -> Axel, 3Toeloop/2Toeloop/4Toeloop -> Toeloop, etc. Spins, Sequence and
# NoBasic map to themselves. Used for high-level classification; derived from TAXONOMY so the
# two stay in sync. FINE_TO_COARSE_IDX is the single authoritative fine-index -> coarse-index map.
_JUMP_TYPES = ("Axel", "Toeloop", "Salchow", "Loop", "Flip", "Lutz")
_COARSE_ORDER = _JUMP_TYPES + ("CamelSpin", "SitSpin", "UprightSpin", "Sequence", "NoBasic")
COARSE_TAXONOMY = {idx: name for idx, name in enumerate(_COARSE_ORDER)}
COARSE_LABEL_TO_IDX = {name: idx for idx, name in COARSE_TAXONOMY.items()}


def coarse_name_for_fine(fine_name: str) -> str:
    """Map a fine class name to its action-level group ('3Axel' -> 'Axel'; spins/Sequence unchanged)."""
    m = re.fullmatch(r"[1-4](" + "|".join(_JUMP_TYPES) + r")", fine_name)
    return m.group(1) if m else fine_name


FINE_TO_COARSE_IDX = {
    fine_idx: COARSE_LABEL_TO_IDX[coarse_name_for_fine(fine_name)]
    for fine_idx, fine_name in TAXONOMY.items()
}


_NORMALIZED_LABELS = {
    re.sub(r"[^a-z0-9]", "", label.lower()): label for label in LABEL_TO_IDX
}
_NORMALIZED_LABELS.update(
    {
        "3toe": "3Toeloop",
        "3t": "3Toeloop",
        "2toe": "2Toeloop",
        "2t": "2Toeloop",
        "4toe": "4Toeloop",
        "4t": "4Toeloop",
        "1toe": "1Toeloop",
        "1t": "1Toeloop",
        "nobasic": "NoBasic",
        "none": "NoBasic",
        "background": "NoBasic",
        "stepseq": "Sequence",
        "stepsequence": "Sequence",
        "choreseq": "Sequence",
        "choresequence": "Sequence",
    }
)

CURRENT_DATASET_DEFAULTS = {
    "Axel": "2Axel",
    "Flip": "3Flip",
    "Loop": "3Loop",
    "Lutz": "3Lutz",
    "Salchow": "3Salchow",
    "Toeloop": "3Toeloop",
    "ToeLoop": "3Toeloop",
    "Toe Loop": "3Toeloop",
    "CamelSpin": "CamelSpin",
    "ChCamelSp3": "CamelSpin",
    "ChCamelSp4": "CamelSpin",
    "SitSpin": "SitSpin",
    "ChSitSp4": "SitSpin",
    "FlySitSp3": "SitSpin",
    "UprightSpin": "UprightSpin",
    "LaybackSp3": "UprightSpin",
    "LaybackSp4": "UprightSpin",
    "ChComboSp2": "UprightSpin",
    "ChComboSp3": "UprightSpin",
    "ChComboSp4": "UprightSpin",
    "StepSeq1": "Sequence",
    "StepSeq2": "Sequence",
    "StepSeq3": "Sequence",
    "StepSeq4": "Sequence",
    "ChoreSeq": "Sequence",
    "ChoreSeq1": "Sequence",
    "NoBasic": "NoBasic",
}


def _canonicalize(label: str) -> Optional[str]:
    key = re.sub(r"[^a-z0-9]", "", str(label).lower())
    return _NORMALIZED_LABELS.get(key)


def _with_index(label: str) -> tuple[str, int]:
    return label, LABEL_TO_IDX[label]


def map_skatingverse_label(label_idx: int | str) -> tuple[str, int]:
    """Map a SkatingVerse label index or label string to the canonical taxonomy."""
    if isinstance(label_idx, str) and not label_idx.strip().isdigit():
        canonical = _canonicalize(label_idx)
        if canonical is None:
            raise ValueError(f"Unknown SkatingVerse label: {label_idx!r}")
        return _with_index(canonical)

    idx = int(label_idx)
    if idx not in TAXONOMY:
        raise ValueError(f"SkatingVerse label index {idx} is outside 0..{len(TAXONOMY) - 1}")
    return TAXONOMY[idx], idx


def map_mmfs_label(spatial_label: str, temporal_label: str = "") -> tuple[str, int] | None:
    """Map an MMFS spatial/temporal category to the canonical taxonomy.

    MMFS labels are commonly expressed as compound spatial and temporal labels. This mapper
    searches both strings for a recognizable element. Unrecognized quality-only or broad action
    labels return None so callers can discard them.
    """
    candidates = [
        spatial_label,
        temporal_label,
        f"{temporal_label}{spatial_label}",
        f"{spatial_label}{temporal_label}",
    ]
    for candidate in candidates:
        canonical = _canonicalize(candidate)
        if canonical is not None:
            return _with_index(canonical)

    combined = f"{spatial_label} {temporal_label}".lower()
    if "camel" in combined:
        return _with_index("CamelSpin")
    if "sit" in combined:
        return _with_index("SitSpin")
    if "upright" in combined or "layback" in combined or "combo" in combined:
        return _with_index("UprightSpin")
    if "step" in combined or "sequence" in combined or "choreo" in combined:
        return _with_index("Sequence")
    return None


def map_current_dataset_label(label: str, ambiguous_defaults: bool = True) -> tuple[str, int] | None:
    """Map the existing 19-class schema to the 28-class canonical taxonomy.

    Labels that already include a rotation count map exactly. Broad jump labels such as
    "Axel" are ambiguous; with ambiguous_defaults=True they are mapped to the common defaults
    in CURRENT_DATASET_DEFAULTS so existing data can still be used as noisy labels.
    """
    canonical = _canonicalize(label)
    if canonical is not None:
        return _with_index(canonical)

    if ambiguous_defaults and label in CURRENT_DATASET_DEFAULTS:
        return _with_index(CURRENT_DATASET_DEFAULTS[label])
    return None