test_analysis.py

"""Headless tests for the movement math — no camera, no MediaPipe required.

``analysis.py`` deliberately imports only ``numpy``/``math`` (MediaPipe stays
isolated in ``pose_estimator.py``), so these run anywhere the deps are installed:

    python test_analysis.py        # plain asserts, exits non-zero on failure
    pytest test_analysis.py        # also works
"""
from __future__ import annotations

import analysis
from analysis import (
    RepCounter,
    SQUAT_KNEE_DOWN,
    SQUAT_KNEE_UP,
    SYMMETRY_TOL,
    VIS_THRESHOLD,
    knee_angle,
)


def _world(overrides):
    """A list of 33 fully-visible 3D world landmarks at the origin, overrides applied.

    World landmarks are ``(x, y, z, visibility)`` (what the angle math consumes).
    ``overrides`` maps a MediaPipe index to ``(x, y, z)`` or ``(x, y, z, vis)``.
    """
    pts = [(0.0, 0.0, 0.0, 1.0) for _ in range(33)]
    for idx, val in overrides.items():
        x, y, z = val[0], val[1], val[2]
        vis = val[3] if len(val) > 3 else 1.0
        pts[idx] = (float(x), float(y), float(z), float(vis))
    return pts


# --- rep counter -----------------------------------------------------------

# Each side's angle is EMA-smoothed (the squat uses an ABSOLUTE knee threshold),
# so each "pose" must be held a handful of frames for the smoother to settle —
# mirroring real video.
HOLD = 30


def _hold(*values):
    """Expand (value, value, ...) into a per-frame stream, each held HOLD frames."""
    seq = []
    for v in values:
        seq.extend([v] * HOLD)
    return seq


def _count(values):
    rc = RepCounter()
    for a in _hold(*values):
        rc.update(a, a)  # symmetric: both knees at the same angle
    return rc.reps


def test_one_full_squat_counts_one():
    # stand -> deep squat (both knees below SQUAT_KNEE_DOWN) -> stand: one rep.
    assert 80 < SQUAT_KNEE_DOWN
    assert _count([170, 80, 170]) == 1


def test_shallow_dip_counts_zero():
    # A dip that stays above SQUAT_KNEE_DOWN is not a squat -> no rep.
    assert 160 > SQUAT_KNEE_DOWN
    assert _count([170, 160, 170]) == 0


def test_front_facing_compressed_squat_not_counted():
    # Accepted limitation of the absolute threshold: a front-facing squat whose
    # knee only compresses to ~125 deg never crosses SQUAT_KNEE_DOWN, so it does
    # not count (side-on views, where the 3D angle is true, do).
    assert 125 > SQUAT_KNEE_DOWN
    assert _count([150, 125, 150]) == 0


def test_two_squats_count_two():
    assert _count([170, 80, 170, 75, 165]) == 2


def test_partial_recovery_does_not_double_count():
    # Bend below SQUAT_KNEE_DOWN, rise only into the hysteresis band (110: above
    # SQUAT_KNEE_DOWN but below SQUAT_KNEE_UP, so not yet "standing"), bend again,
    # then fully stand: a single completed rep, not two.
    assert SQUAT_KNEE_DOWN < 110 < SQUAT_KNEE_UP
    assert _count([170, 80, 110, 80, 170]) == 1


def test_posture_labels():
    rc = RepCounter()
    posture = "--"
    for a in _hold(170):  # standing: both knees extended
        posture = rc.update(a, a)
    assert posture == "Neutral"
    for a in [80] * HOLD:  # drop below SQUAT_KNEE_DOWN
        posture = rc.update(a, a)
    assert posture == "Squatting"
    assert rc.update(None, None) == "--"


def test_reset_clears_state_and_count():
    rc = RepCounter()
    for a in _hold(170, 80, 170):
        rc.update(a, a)
    assert rc.reps == 1
    rc.reset()
    assert rc.reps == 0
    assert rc._machine._ema_l is None
    # After reset a fresh cycle counts cleanly.
    for a in _hold(170, 80, 170):
        rc.update(a, a)
    assert rc.reps == 1


def test_symmetric_squat_within_tolerance_counts():
    # Both knees below SQUAT_KNEE_DOWN and within SYMMETRY_TOL of each other -> counts.
    assert 88 < SQUAT_KNEE_DOWN and 88 - 75 < SYMMETRY_TOL
    rc = RepCounter()
    for l, r in _hold((170, 170), (75, 88), (170, 170)):
        rc.update(l, r)
    assert rc.reps == 1


def test_asymmetric_squat_does_not_count():
    # Both knees dip below the threshold but by very different amounts (a lopsided
    # squat, > SYMMETRY_TOL apart): the symmetry gate rejects it.
    assert 55 < SQUAT_KNEE_DOWN and 88 < SQUAT_KNEE_DOWN and 88 - 55 > SYMMETRY_TOL
    rc = RepCounter()
    for l, r in _hold((170, 170), (55, 88), (170, 170)):
        rc.update(l, r)
    assert rc.reps == 0


def test_one_knee_hidden_pauses_squat():
    # A missing side -> posture "--" and no count (the rule needs both knees).
    rc = RepCounter()
    posture = "--"
    for l, r in _hold((170, 170), (80, None)):
        posture = rc.update(l, r)
    assert posture == "--"
    assert rc.reps == 0


# --- angle geometry --------------------------------------------------------

def test_knee_angle_right_angle():
    # hip above knee, ankle to the side => 90 deg at the knee.
    lm = _world(
        {analysis.L_HIP: (0, 0, 0), analysis.L_KNEE: (0, 1, 0), analysis.L_ANKLE: (1, 1, 0)}
    )
    assert abs(knee_angle(lm, "left") - 90.0) < 1e-6


def test_knee_angle_uses_depth():
    # Foreshortened in 2D (hip/knee/ankle share x,y) but bent in depth (z):
    # the 3D angle still reads 90 deg where a 2D projection would collapse.
    lm = _world(
        {analysis.L_HIP: (0, 0, 0), analysis.L_KNEE: (0, 1, 0), analysis.L_ANKLE: (0, 1, 1)}
    )
    assert abs(knee_angle(lm, "left") - 90.0) < 1e-6


def test_knee_angle_straight_leg():
    # collinear hip-knee-ankle => 180 deg (fully extended).
    lm = _world(
        {analysis.L_HIP: (0, 0, 0), analysis.L_KNEE: (0, 1, 0), analysis.L_ANKLE: (0, 2, 0)}
    )
    assert abs(knee_angle(lm, "left") - 180.0) < 1e-6


def test_low_visibility_knee_returns_none():
    lm = _world({analysis.L_KNEE: (0, 1, 0, VIS_THRESHOLD - 0.1)})
    assert knee_angle(lm, "left") is None


# --- activity recognition helpers ------------------------------------------

import math


def _pixel(overrides):
    """A list of 33 fully-visible pixel landmarks at the origin, overrides applied.

    Pixel landmarks are ``(x, y, visibility)`` (centroid / torso / JJ signals).
    ``overrides`` maps a MediaPipe index to ``(x, y)`` or ``(x, y, vis)``.
    """
    pts = [(0.0, 0.0, 1.0) for _ in range(33)]
    for idx, val in overrides.items():
        x, y = val[0], val[1]
        vis = val[2] if len(val) > 2 else 1.0
        pts[idx] = (float(x), float(y), float(vis))
    return pts


def _jj_pose(arms_up, legs_apart):
    """A jumping-jack frame. Shoulders at y=100, hips at y=200 (torso = 100 px).
    Arms up => wrists above the shoulders; legs apart => ankles spread wide."""
    wrist_y = 20 if arms_up else 200
    la_x, ra_x = (0, 120) if legs_apart else (45, 55)
    return _pixel({
        analysis.L_SHOULDER: (40, 100), analysis.R_SHOULDER: (60, 100),
        analysis.L_HIP: (40, 200), analysis.R_HIP: (60, 200),
        analysis.L_WRIST: (50, wrist_y), analysis.R_WRIST: (50, wrist_y),
        analysis.L_ANKLE: (la_x, 200), analysis.R_ANKLE: (ra_x, 200),
    })


def _pushup_world(bent, horizontal=True):
    """A push-up frame in world coords. ``horizontal`` lays the torso flat (gate
    open); ``bent`` flexes BOTH elbows to ~90 deg, else extends them to ~180.
    Both elbows are set because the counter now requires both sides."""
    s = (1.0, 0.0, 0.0) if horizontal else (0.0, 1.0, 0.0)
    h = (0.0, 0.0, 0.0)
    e = (s[0], s[1] - 1, s[2])
    w = (s[0] + 1, s[1] - 1, s[2]) if bent else (s[0], s[1] - 2, s[2])
    return _world({
        analysis.L_SHOULDER: s, analysis.R_SHOULDER: s,
        analysis.L_HIP: h, analysis.R_HIP: h,
        analysis.L_ELBOW: e, analysis.L_WRIST: w,
        analysis.R_ELBOW: e, analysis.R_WRIST: w,
    })


def _knee_world(angle_deg):
    """World landmarks whose BOTH knees subtend ``angle_deg`` (hip->knee->ankle),
    symmetrically — the squat counter now requires both sides."""
    rad = math.radians(angle_deg)
    ankle = (math.sin(rad), math.cos(rad), 0)
    return _world({
        analysis.L_HIP: (0, 1, 0), analysis.L_KNEE: (0, 0, 0), analysis.L_ANKLE: ankle,
        analysis.R_HIP: (0, 1, 0), analysis.R_KNEE: (0, 0, 0), analysis.R_ANKLE: ankle,
    })


class _FullRepChecker:
    """Stand-in with the ``.active`` / ``.reps`` surface ActivityClassifier reads.

    ``reps`` is the count of *completed full reps* (a whole down-and-up cycle) — the
    arbiter only labels an exercise once that count reaches ``MIN_REPS_TO_LABEL``, so
    these tests drive it with the full-rep total a real counter would expose.
    """

    def __init__(self, active=False, reps=0):
        self.active = active
        self.reps = reps


# --- jumping jacks ---------------------------------------------------------

def test_one_jumping_jack_counts_one():
    jj = analysis.JumpingJackCounter()
    closed, opened = _jj_pose(False, False), _jj_pose(True, True)
    for lm in [closed] * HOLD + [opened] * HOLD + [closed] * HOLD:
        jj.update(lm)
    assert jj.reps == 1


def test_jumping_jack_active_holds_then_decays():
    # After a jack, `active` stays True for a short hold (so the label doesn't
    # flicker to Walking on a brief range dip mid-set), then decays once the
    # joints go missing for longer than _ACTIVE_HOLD.
    jj = analysis.JumpingJackCounter()
    closed, opened = _jj_pose(False, False), _jj_pose(True, True)
    for lm in [closed] * HOLD + [opened] * HOLD + [closed] * HOLD:
        jj.update(lm)
    assert jj.active
    hidden = _pixel({  # a wrist drops below the visibility threshold
        analysis.L_SHOULDER: (40, 100), analysis.R_SHOULDER: (60, 100),
        analysis.L_HIP: (40, 200), analysis.R_HIP: (60, 200),
        analysis.L_WRIST: (50, 200, 0.0), analysis.R_WRIST: (50, 200),
        analysis.L_ANKLE: (45, 200), analysis.R_ANKLE: (55, 200),
    })
    jj.update(hidden)
    assert jj.active  # one missing frame doesn't drop the label
    for _ in range(analysis._ACTIVE_HOLD):
        jj.update(hidden)
    assert not jj.active  # sustained absence finally decays it


def test_arm_only_motion_not_counted():
    # Arms cycle but the legs never spread -> not a jumping jack.
    jj = analysis.JumpingJackCounter()
    closed, arms = _jj_pose(False, False), _jj_pose(True, False)
    for lm in [closed] * HOLD + [arms] * HOLD + [closed] * HOLD:
        jj.update(lm)
    assert jj.reps == 0


def test_small_fidget_not_counted():
    # Sub-threshold swings in both signals stay below JJ_MIN_*_RANGE.
    jj = analysis.JumpingJackCounter()
    a = _pixel({
        analysis.L_SHOULDER: (40, 100), analysis.R_SHOULDER: (60, 100),
        analysis.L_HIP: (40, 200), analysis.R_HIP: (60, 200),
        analysis.L_WRIST: (50, 190), analysis.R_WRIST: (50, 190),  # arm_raise -0.9
        analysis.L_ANKLE: (45, 200), analysis.R_ANKLE: (55, 200),  # spread 0.1
    })
    b = _pixel({
        analysis.L_SHOULDER: (40, 100), analysis.R_SHOULDER: (60, 100),
        analysis.L_HIP: (40, 200), analysis.R_HIP: (60, 200),
        analysis.L_WRIST: (50, 170), analysis.R_WRIST: (50, 170),  # arm_raise -0.7 (span 0.2)
        analysis.L_ANKLE: (43, 200), analysis.R_ANKLE: (57, 200),  # spread 0.14 (span 0.04)
    })
    for lm in [a] * HOLD + [b] * HOLD + [a] * HOLD:
        jj.update(lm)
    assert jj.reps == 0


def test_one_arm_raised_not_counted():
    # Only the left arm goes up (legs spread normally). min() over the two arms
    # is held down by the right arm, so the pose never reads as "open".
    jj = analysis.JumpingJackCounter()
    closed = _jj_pose(False, False)
    one_arm = _pixel({
        analysis.L_SHOULDER: (40, 100), analysis.R_SHOULDER: (60, 100),
        analysis.L_HIP: (40, 200), analysis.R_HIP: (60, 200),
        analysis.L_WRIST: (50, 20), analysis.R_WRIST: (50, 200),  # only left up
        analysis.L_ANKLE: (0, 200), analysis.R_ANKLE: (120, 200),
    })
    for lm in [closed] * HOLD + [one_arm] * HOLD + [closed] * HOLD:
        jj.update(lm)
    assert jj.reps == 0


# --- push-ups --------------------------------------------------------------

def test_pushup_horizontal_gate_and_rep():
    pu = analysis.PushupCounter()
    up, down = _pushup_world(bent=False), _pushup_world(bent=True)
    for w in [up] * HOLD + [down] * HOLD + [up] * HOLD:
        pu.update(w)
    assert pu.reps == 1


def _hide(world, idxs):
    """Copy of ``world`` with the given landmark indices marked low-visibility."""
    out = list(world)
    for i in idxs:
        x, y, z, _ = out[i]
        out[i] = (x, y, z, VIS_THRESHOLD - 0.1)
    return out


def test_pushup_counts_with_far_arm_occluded():
    # Side-on push-up (the usual filming angle): the far (left) arm is occluded by
    # the body at the bottom of each rep, so only the right elbow is reliable there.
    # The counter must still count via the visible side (require_both=False), unlike
    # the squat which holds when a knee is hidden.
    pu = analysis.PushupCounter()
    up = _pushup_world(bent=False)
    down = _hide(_pushup_world(bent=True), [analysis.L_ELBOW, analysis.L_WRIST])
    for w in [up] * HOLD + [down] * HOLD + [up] * HOLD + [down] * HOLD + [up] * HOLD:
        pu.update(w)
    assert pu.reps == 2


def test_vertical_torso_blocks_pushup():
    # Same elbow flexion, but an upright torso fails the horizontal gate.
    pu = analysis.PushupCounter()
    up = _pushup_world(bent=False, horizontal=False)
    down = _pushup_world(bent=True, horizontal=False)
    for w in [up] * HOLD + [down] * HOLD + [up] * HOLD:
        pu.update(w)
    assert pu.reps == 0


def test_torso_incline_horizontal_vs_vertical():
    assert analysis.torso_incline(_pushup_world(bent=False, horizontal=True)) > 80
    assert analysis.torso_incline(_pushup_world(bent=False, horizontal=False)) < 10


# --- elbow geometry --------------------------------------------------------

def test_elbow_angle_right_angle():
    lm = _world({
        analysis.L_SHOULDER: (0, 1, 0),
        analysis.L_ELBOW: (0, 0, 0),
        analysis.L_WRIST: (1, 0, 0),
    })
    assert abs(analysis.elbow_angle(lm, "left") - 90.0) < 1e-6


# --- locomotion ------------------------------------------------------------

def test_neutral_low_speed():
    loco = analysis.LocomotionClassifier()
    label = "Neutral"
    for _ in range(HOLD):
        label = loco.update(0.1)
    assert label == "Neutral"


def test_walking_band():
    loco = analysis.LocomotionClassifier()
    label = "Neutral"
    for _ in range(HOLD):
        label = loco.update(1.0)  # between WALK_ENTER and RUN_ENTER
    assert label == "Walking"


def test_running_high_speed():
    loco = analysis.LocomotionClassifier()
    label = "Neutral"
    for _ in range(HOLD):
        label = loco.update(3.0)
    assert label == "Running"


def test_loco_hysteresis_no_flicker():
    loco = analysis.LocomotionClassifier()
    label = "Neutral"
    for _ in range(HOLD):
        label = loco.update(1.0)
    assert label == "Walking"
    # Dip between WALK_EXIT and WALK_ENTER: stays Walking (wouldn't enter from Neutral).
    assert analysis.WALK_EXIT < 0.4 < analysis.WALK_ENTER
    for _ in range(HOLD):
        label = loco.update(0.4)
    assert label == "Walking"


# --- activity arbitration --------------------------------------------------

def test_pushup_beats_locomotion():
    ac = analysis.ActivityClassifier()
    label = reps = None
    for _ in range(analysis.ACTIVITY_DEBOUNCE + 2):
        label, reps = ac.choose(
            _FullRepChecker(), _FullRepChecker(), _FullRepChecker(active=True, reps=3), "Walking"
        )
    assert label == "Push-ups" and reps == 3


def test_debounce_prevents_flicker():
    ac = analysis.ActivityClassifier()
    idle = (_FullRepChecker(), _FullRepChecker(), _FullRepChecker())
    label = None
    for _ in range(analysis.ACTIVITY_DEBOUNCE + 1):
        label, _ = ac.choose(*idle, "Neutral")
    assert label == "Neutral"
    for _ in range(analysis.ACTIVITY_DEBOUNCE - 1):  # too brief to commit
        label, _ = ac.choose(*idle, "Walking")
    assert label == "Neutral"


def test_activity_reps_tracks_selected_exercise():
    ac = analysis.ActivityClassifier()
    squat = _FullRepChecker(active=True, reps=5)
    label = reps = None
    for _ in range(analysis.ACTIVITY_DEBOUNCE + 1):
        label, reps = ac.choose(squat, _FullRepChecker(), _FullRepChecker(), "Neutral")
    assert label == "Squat" and reps == 5


def test_active_below_min_reps_shows_neutral_not_locomotion():
    # An exercise active but below MIN_REPS_TO_LABEL must show neither the exercise
    # nor a locomotion label from in-place bounce -> it reads neutral "Neutral".
    ac = analysis.ActivityClassifier()
    squat = _FullRepChecker(active=True, reps=analysis.MIN_REPS_TO_LABEL - 1)
    label = None
    for _ in range(analysis.ACTIVITY_DEBOUNCE + 2):
        label, _ = ac.choose(squat, _FullRepChecker(), _FullRepChecker(), "Running")
    assert label == "Neutral"


def test_idle_passes_locomotion_through():
    # With no exercise active, the locomotion label flows through unchanged.
    ac = analysis.ActivityClassifier()
    label = None
    for _ in range(analysis.ACTIVITY_DEBOUNCE + 2):
        label, _ = ac.choose(_FullRepChecker(), _FullRepChecker(), _FullRepChecker(), "Running")
    assert label == "Running"


# --- end-to-end regression -------------------------------------------------

def test_squat_still_counts_via_analyzer():
    # Two full squats through MovementAnalyzer count and label as "Squat" (the
    # label turns on once MIN_REPS_TO_LABEL reps are reached).
    an = analysis.MovementAnalyzer()
    px = _pixel({
        analysis.L_HIP: (40, 200), analysis.R_HIP: (60, 200),
        analysis.L_SHOULDER: (40, 100), analysis.R_SHOULDER: (60, 100),
    })
    out = {}
    for i, angle in enumerate(_hold(170, 80, 170, 80, 170)):
        out = an.update(px, _knee_world(angle), i / 30.0)
    assert out["reps"] == 2
    assert out["activity"] == "Squat"
    assert out["activity_reps"] == out["reps"]


def _run_all():
    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    failures = 0
    for t in tests:
        try:
            t()
            print(f"PASS {t.__name__}")
        except AssertionError as e:
            failures += 1
            print(f"FAIL {t.__name__}: {e}")
    print(f"\n{len(tests) - failures}/{len(tests)} passed")
    return failures


if __name__ == "__main__":
    raise SystemExit(1 if _run_all() else 0)