experiment-2107/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


# --- FS-Jump3D: its own dedicated taxonomy, kept separate from TAXONOMY/COARSE_TAXONOMY ---
# FS-Jump3D's folder-name labels are jump-type only (no rotation count, e.g. "Axel" not "2Axel"),
# plus a "Comb" (combination-jump) category the other sources don't have. Rather than force
# these onto the canonical class scheme -- which would mean guessing a rotation count that isn't
# in the data -- they get their own small taxonomy for now, used as-is. Not merged with anything.
FS_JUMP3D_TAXONOMY = {
    0: "Axel",
    1: "Comb",
    2: "Flip",
    3: "Loop",
    4: "Lutz",
    5: "Salchow",
    6: "Toeloop",
}
FS_JUMP3D_LABEL_TO_IDX = {name: idx for idx, name in FS_JUMP3D_TAXONOMY.items()}


def map_fs_jump3d_label(label: str) -> tuple[str, int]:
    """Map an FS-Jump3D folder-name label to its own dedicated taxonomy, as-is (no rotation-count
    guessing, no merging into TAXONOMY/COARSE_TAXONOMY)."""
    name = str(label)
    if name not in FS_JUMP3D_LABEL_TO_IDX:
        raise ValueError(f"Unknown FS-Jump3D label: {label!r}")
    return name, FS_JUMP3D_LABEL_TO_IDX[name]


# Singles-only variant of FS_JUMP3D_TAXONOMY with "Comb" dropped, for the combo-decomposition
# experiment: train a classifier on the 6 known single-jump types only, then run it with a
# sliding window over the (held-out, never-trained-on) Comb clips to see what sequence of known
# jump types it thinks it's looking at. Comb samples never appear in this taxonomy's label space
# at all -- they're inference-only inputs, not a class to predict.
FS_JUMP3D_SINGLES_TAXONOMY = {
    idx: name for idx, name in enumerate(
        name for name in FS_JUMP3D_TAXONOMY.values() if name != "Comb"
    )
}
FS_JUMP3D_SINGLES_LABEL_TO_IDX = {name: idx for idx, name in FS_JUMP3D_SINGLES_TAXONOMY.items()}


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