formscout/analysis/relevant_joints.py

"""
Per-test relevant-joint and relevant-angle definitions.

Drives the "always describe just the joint relevant to the screening action"
requirement: each FMS test names the joints and the joint-angles that matter for
its rubric, plus the single primary angle used for the headline angle-over-time
graph. Angles are COCO (a, b, c) triplets — the angle is measured at joint b.
"""
from __future__ import annotations

# COCO-17 joint indices
NOSE = 0
L_SHOULDER, R_SHOULDER = 5, 6
L_ELBOW, R_ELBOW = 7, 8
L_WRIST, R_WRIST = 9, 10
L_HIP, R_HIP = 11, 12
L_KNEE, R_KNEE = 13, 14
L_ANKLE, R_ANKLE = 15, 16

COCO_NAMES = {
    0: "nose", 1: "left_eye", 2: "right_eye", 3: "left_ear", 4: "right_ear",
    5: "left_shoulder", 6: "right_shoulder", 7: "left_elbow", 8: "right_elbow",
    9: "left_wrist", 10: "right_wrist", 11: "left_hip", 12: "right_hip",
    13: "left_knee", 14: "right_knee", 15: "left_ankle", 16: "right_ankle",
}

# Each test: the joints that matter, the named angles (a, b, c) measured at b,
# and the primary angle for the headline graph.
RELEVANT: dict[str, dict] = {
    "deep_squat": {
        "joints": [L_SHOULDER, R_SHOULDER, L_HIP, R_HIP, L_KNEE, R_KNEE, L_ANKLE, R_ANKLE],
        "angles": {
            "left_knee_flexion": (L_HIP, L_KNEE, L_ANKLE),
            "right_knee_flexion": (R_HIP, R_KNEE, R_ANKLE),
            "left_hip_flexion": (L_SHOULDER, L_HIP, L_KNEE),
            "right_hip_flexion": (R_SHOULDER, R_HIP, R_KNEE),
        },
        "primary_angle": "left_knee_flexion",
    },
    "hurdle_step": {
        "joints": [L_HIP, R_HIP, L_KNEE, R_KNEE, L_ANKLE, R_ANKLE],
        "angles": {
            "left_knee_flexion": (L_HIP, L_KNEE, L_ANKLE),
            "right_knee_flexion": (R_HIP, R_KNEE, R_ANKLE),
            "left_hip_flexion": (L_SHOULDER, L_HIP, L_KNEE),
            "right_hip_flexion": (R_SHOULDER, R_HIP, R_KNEE),
        },
        "primary_angle": "left_hip_flexion",
    },
    "inline_lunge": {
        "joints": [L_HIP, R_HIP, L_KNEE, R_KNEE, L_ANKLE, R_ANKLE],
        "angles": {
            "left_knee_flexion": (L_HIP, L_KNEE, L_ANKLE),
            "right_knee_flexion": (R_HIP, R_KNEE, R_ANKLE),
        },
        "primary_angle": "left_knee_flexion",
    },
    "shoulder_mobility": {
        "joints": [L_SHOULDER, R_SHOULDER, L_ELBOW, R_ELBOW, L_WRIST, R_WRIST],
        "angles": {
            "left_shoulder_angle": (L_ELBOW, L_SHOULDER, L_HIP),
            "right_shoulder_angle": (R_ELBOW, R_SHOULDER, R_HIP),
            "left_elbow_flexion": (L_SHOULDER, L_ELBOW, L_WRIST),
            "right_elbow_flexion": (R_SHOULDER, R_ELBOW, R_WRIST),
        },
        "primary_angle": "left_shoulder_angle",
    },
    "active_slr": {
        "joints": [L_HIP, R_HIP, L_KNEE, R_KNEE, L_ANKLE, R_ANKLE],
        "angles": {
            "left_hip_flexion": (L_SHOULDER, L_HIP, L_KNEE),
            "right_hip_flexion": (R_SHOULDER, R_HIP, R_KNEE),
            "left_knee_flexion": (L_HIP, L_KNEE, L_ANKLE),
            "right_knee_flexion": (R_HIP, R_KNEE, R_ANKLE),
        },
        "primary_angle": "left_hip_flexion",
    },
    "trunk_stability_pushup": {
        "joints": [L_SHOULDER, R_SHOULDER, L_ELBOW, R_ELBOW, L_HIP, R_HIP, L_ANKLE, R_ANKLE],
        "angles": {
            "left_elbow_flexion": (L_SHOULDER, L_ELBOW, L_WRIST),
            "right_elbow_flexion": (R_SHOULDER, R_ELBOW, R_WRIST),
            "left_hip_line": (L_SHOULDER, L_HIP, L_ANKLE),
            "right_hip_line": (R_SHOULDER, R_HIP, R_ANKLE),
        },
        "primary_angle": "left_hip_line",
    },
    "rotary_stability": {
        "joints": [L_SHOULDER, R_SHOULDER, L_ELBOW, R_ELBOW, L_HIP, R_HIP, L_KNEE, R_KNEE],
        "angles": {
            "left_hip_line": (L_SHOULDER, L_HIP, L_KNEE),
            "right_hip_line": (R_SHOULDER, R_HIP, R_KNEE),
        },
        "primary_angle": "left_hip_line",
    },
}


def relevant_joints(test_name: str) -> list[int]:
    """COCO joint indices relevant to this test (empty for unknown tests)."""
    return list(RELEVANT.get(test_name, {}).get("joints", []))


def relevant_angles(test_name: str) -> dict[str, tuple]:
    """Named (a, b, c) angle triplets relevant to this test."""
    return dict(RELEVANT.get(test_name, {}).get("angles", {}))


def primary_angle(test_name: str) -> str | None:
    """Name of the headline angle for the angle-over-time graph, or None."""
    return RELEVANT.get(test_name, {}).get("primary_angle")


def openness_label(angle_deg: float) -> str:
    """Describe how open/closed a joint is from its interior angle in degrees."""
    if angle_deg >= 160:
        return "open / extended"
    if angle_deg >= 110:
        return "mid-range"
    if angle_deg >= 60:
        return "flexed"
    return "deeply flexed / closed"