tests/test_segmentation.py

"""Unit tests for shot segmentation (slice A1, BRIEF 13/08/2026) -- on a
SYNTHETIC track.

They check the boundary-cutting CORE (``segment_shots`` -- pure numpy, runs
on the app venv with no cv2 and no torch): a single shot, several shots,
"every ball is still" rather than the cue alone, coasting and lost detections
mid-shot, a lost scene, a static video that must not invent phantom shots,
and the boundary padding. They do NOT check YOLO or decoding (``segment_video``
was accepted on a real VOD at gate A1.2) -- the same split of duties as
test_broadcast_analytics against gate G-24.3.

The physical thresholds in these scenarios come from the set agreed on 11/08
(imported from broadcast): a shot runs faster than V_INIT 0.25 m/s, and a
still ball jitters ~2mm/frame (the p99 apparent speed of 55 mm/s at 30fps --
exactly the measurement noise seen in P0).
"""

from __future__ import annotations

import numpy as np
import pytest

from poolcoach_cv.broadcast import V_INIT_MPS, V_STILL_MPS
from poolcoach_cv.segmentation import (BLIND_MAX_S, PAD_POST_S, PAD_PRE_S,
                                       STILL_WIN_S, segment_shots)

FPS = 30.0


def _pos_at(t: float, p0, segs):
    """Position from a motion schedule: ``segs`` = [(t_a, t_b, p_a, p_b)],
    each a constant-speed straight line; outside a segment the ball rests at
    the nearest end point."""
    x, y = p0
    for t_a, t_b, p_a, p_b in segs:
        if t < t_a:
            break
        if t >= t_b:
            x, y = p_b
        else:
            f = (t - t_a) / (t_b - t_a)
            x = p_a[0] + (p_b[0] - p_a[0]) * f
            y = p_a[1] + (p_b[1] - p_a[1]) * f
    return x, y


def build(dur_s, cue_p0, cue_segs=(), balls=(), cue_missing=(), blind=(),
          jitter=0.0, seed=7):
    """A synthetic 30fps track -> (rows, others) in exactly the shape
    segment_shots expects.

    ``balls``: a tuple (x, y) for a still ball, or a dict
    {"p0": (x,y), "segs": [...]} for a ball with its own schedule.
    ``cue_missing``/``blind``: lists of (t_a, t_b) -- the stretch where the
    cue has no detection / the stretch where NO ball is seen (a lost scene).
    """
    rng = np.random.default_rng(seed)
    rows, others = [], []
    n = int(round(dur_s * FPS)) + 1
    for i in range(n):
        t = i / FPS
        name = f"{i:05d}"
        is_blind = any(a <= t < b for a, b in blind)
        cue_gone = is_blind or any(a <= t < b for a, b in cue_missing)
        cx, cy = _pos_at(t, cue_p0, cue_segs)
        if jitter:
            cx += rng.normal(0, jitter)
            cy += rng.normal(0, jitter)
        rows.append({
            "frame_file": name, "t_s": round(t, 6),
            "covered": 0 if cue_gone else 1,
            "table_x_m": "" if cue_gone else cx,
            "table_y_m": "" if cue_gone else cy,
            # img_diff is a constant 1.0 (a "live" table) so the duplicate
            # frame rule of find_dup_frames never fires -- that one has its
            # own test over in broadcast
            "img_diff": 1.0 if i else -1.0,
        })
        if is_blind:
            continue
        for spec in balls:
            if isinstance(spec, dict):
                bx, by = _pos_at(t, spec["p0"], spec["segs"])
            else:
                bx, by = spec
            if jitter:
                bx += rng.normal(0, jitter)
                by += rng.normal(0, jitter)
            others.append({"frame_file": name, "t_s": round(t, 6),
                           "x_m": bx, "y_m": by})
    return rows, others


STILL_BALLS = [(0.4, 2.0), (0.9, 1.8), (0.6, 0.4)]


# ------------------------------------------------------------- a single shot

def test_a_single_shot_is_ok_with_the_right_boundaries():
    rows, others = build(
        6.0, cue_p0=(0.3, 0.5),
        cue_segs=[(1.0, 2.0, (0.3, 0.5), (0.9, 1.7))],   # ~1.34 m/s
        balls=STILL_BALLS)
    out = segment_shots(rows, others)
    assert [s["status"] for s in out["shots"]] == ["ok"]
    s = out["shots"][0]
    assert s["idx"] == 1
    # the onset follows the V_INIT_RUN rule (a step above V_INIT) -- caught
    # at the moment motion starts
    assert s["t_onset_s"] == pytest.approx(1.0, abs=0.15)
    # settle lags the true stop (2.0s) by at most one chord window
    assert 2.0 <= s["t_settle_s"] <= 2.0 + 2 * STILL_WIN_S
    # padding at both ends, still inside the video
    assert s["t_start_s"] == pytest.approx(s["t_onset_s"] - PAD_PRE_S,
                                           abs=0.05)
    assert s["t_settle_s"] < s["t_end_s"] <= s["t_settle_s"] + PAD_POST_S
    assert 0.0 <= s["t_start_s"] < s["t_onset_s"]


def test_a_static_video_invents_no_phantom_shot():
    """A still table for 10s with 2mm/frame detection jitter (about the p99
    apparent speed of 55mm/s in P0) -- not one shot may come out."""
    rows, others = build(10.0, cue_p0=(0.5, 1.0), balls=STILL_BALLS,
                         jitter=0.002)
    out = segment_shots(rows, others)
    assert out["shots"] == []


def test_an_empty_track_raises_valueerror():
    with pytest.raises(ValueError):
        segment_shots([])


# ------------------------------------------------------------ several shots

def test_three_shots_come_out_in_order():
    segs = [(1.0, 2.0, (0.3, 0.5), (0.9, 1.7)),
            (5.0, 6.2, (0.9, 1.7), (0.4, 0.8)),
            (10.0, 10.8, (0.4, 0.8), (1.0, 2.2))]
    rows, others = build(14.0, cue_p0=(0.3, 0.5), cue_segs=segs,
                         balls=STILL_BALLS)
    out = segment_shots(rows, others)
    shots = out["shots"]
    assert [s["status"] for s in shots] == ["ok", "ok", "ok"]
    assert [s["idx"] for s in shots] == [1, 2, 3]
    for s, (t_a, _b, _pa, _pb) in zip(shots, segs):
        assert s["t_onset_s"] == pytest.approx(t_a, abs=0.15)
    # the coarse boundaries increase and never overlap
    for a, b in zip(shots, shots[1:]):
        assert a["t_settle_s"] <= b["t_onset_s"]
        # the padding must not swallow the neighbouring shot's motion
        assert a["t_end_s"] <= b["t_onset_s"] + 1e-6
        assert b["t_start_s"] >= a["t_settle_s"] - 1e-6


def test_two_shots_back_to_back_still_separate():
    """Shot 2 breaks the moment the table has settled -- after closing a shot
    the detector must arm again immediately (not wait for another still
    window) or it misses the next onset."""
    segs = [(1.0, 2.0, (0.3, 0.5), (0.9, 1.7)),
            (3.2, 4.0, (0.9, 1.7), (0.5, 0.9))]
    rows, others = build(7.0, cue_p0=(0.3, 0.5), cue_segs=segs,
                         balls=STILL_BALLS)
    out = segment_shots(rows, others)
    assert [s["status"] for s in out["shots"]] == ["ok", "ok"]
    assert out["shots"][1]["t_onset_s"] == pytest.approx(3.2, abs=0.15)


# --------------------------------------- "EVERY ball is still", not just cue

def test_a_still_rolling_target_ball_keeps_the_shot_open():
    """The cue stops at 2.0 but the target ball rolls on until 4.0 -- settle
    has to wait for EVERY ball (the boundary definition of design section 4)
    and must not close early on the cue alone."""
    rolling = {"p0": (0.9, 1.8),
               "segs": [(2.0, 4.0, (0.9, 1.8), (0.5, 1.2))]}   # ~0.36 m/s
    rows, others = build(
        7.0, cue_p0=(0.3, 0.5),
        cue_segs=[(1.0, 2.0, (0.3, 0.5), (0.9, 1.7))],
        balls=[rolling, (0.4, 2.0)])
    out = segment_shots(rows, others)
    assert [s["status"] for s in out["shots"]] == ["ok"]
    assert out["shots"][0]["t_settle_s"] >= 4.0 - 0.1


# ---------------------------------------- coasting / lost detection mid-shot

def test_a_short_cue_detection_gap_does_not_split_the_shot():
    rows, others = build(
        6.0, cue_p0=(0.3, 0.5),
        cue_segs=[(1.0, 3.0, (0.3, 0.5), (1.1, 2.3))],
        balls=STILL_BALLS, cue_missing=[(1.8, 2.1)])
    out = segment_shots(rows, others)
    assert [s["status"] for s in out["shots"]] == ["ok"]
    assert out["shots"][0]["t_settle_s"] >= 3.0 - 0.1


def test_a_cue_gap_longer_than_the_window_still_does_not_close_early():
    """The cue disappears for a WHOLE chord window while in flight (someone
    standing in front of it -- the shot_07 lesson): every other ball is still,
    but that is no licence to declare "settled" while the main witness is
    absent."""
    rows, others = build(
        6.0, cue_p0=(0.3, 0.5),
        cue_segs=[(1.0, 3.0, (0.3, 0.5), (1.1, 2.3))],
        balls=STILL_BALLS, cue_missing=[(1.5, 2.5)])
    out = segment_shots(rows, others)
    assert [s["status"] for s in out["shots"]] == ["ok"]
    s = out["shots"][0]
    # it must not close while the cue is invisible (1.5-2.5)
    assert s["t_settle_s"] >= 3.0 - 0.1


# ------------------------------------------------- a lost scene (section 5)

def test_a_lost_scene_mid_shot_is_reported_and_not_swallowed():
    rows, others = build(
        9.0, cue_p0=(0.3, 0.5),
        cue_segs=[(1.0, 2.0, (0.3, 0.5), (0.9, 1.7)),
                  (6.0, 6.8, (0.9, 1.7), (0.5, 1.0))],
        balls=STILL_BALLS, blind=[(1.5, 4.0)])
    out = segment_shots(rows, others)
    shots = out["shots"]
    assert len(shots) == 2
    assert shots[0]["status"] == "error"
    assert "scene was lost" in shots[0]["reason"]
    # once the scene is back and the table settles, the next shot is caught
    # as usual
    assert shots[1]["status"] == "ok"
    assert shots[1]["t_onset_s"] == pytest.approx(6.0, abs=0.15)


def test_a_blind_gap_under_the_threshold_does_not_break_the_shot():
    """A short detection gap (< BLIND_MAX_S) mid-shot -- a stutter in the
    video or a brief occlusion -- must not split the shot in two or report a
    lost scene."""
    assert BLIND_MAX_S >= 0.5
    rows, others = build(
        6.0, cue_p0=(0.3, 0.5),
        cue_segs=[(1.0, 3.0, (0.3, 0.5), (1.1, 2.3))],
        balls=STILL_BALLS, blind=[(1.8, 2.1)])
    out = segment_shots(rows, others)
    assert [s["status"] for s in out["shots"]] == ["ok"]


def test_a_video_that_ends_mid_shot_is_reported():
    rows, others = build(
        4.5, cue_p0=(0.3, 0.5),
        cue_segs=[(4.0, 6.0, (0.3, 0.5), (0.9, 1.7))],
        balls=STILL_BALLS)
    out = segment_shots(rows, others)
    assert len(out["shots"]) == 1
    assert out["shots"][0]["status"] == "error"
    assert "video ends mid-shot" in out["shots"][0]["reason"]


# ---------------------------------------------------------- boundary padding

def test_the_padding_never_runs_past_the_video():
    rows, others = build(
        3.0, cue_p0=(0.3, 0.5),
        cue_segs=[(0.8, 1.6, (0.3, 0.5), (0.9, 1.7))],
        balls=STILL_BALLS)
    out = segment_shots(rows, others)
    assert len(out["shots"]) == 1
    s = out["shots"][0]
    assert s["t_start_s"] >= 0.0
    assert s["t_end_s"] <= 3.0 + 1e-6
    assert s["t_start_s"] < s["t_onset_s"] < s["t_settle_s"] <= s["t_end_s"]


def test_the_agreed_thresholds_are_not_redefined():
    """Pins the "use the thresholds agreed on 11/08 as they are" contract:
    segmentation has to read V_STILL/V_INIT from broadcast rather than keep
    its own copy, which could drift."""
    import poolcoach_cv.broadcast as bc
    import poolcoach_cv.segmentation as seg

    assert seg.V_STILL_MPS is bc.V_STILL_MPS
    assert seg.V_INIT_MPS is bc.V_INIT_MPS
    assert seg.VMAX_MPS is bc.VMAX_MPS
    assert V_STILL_MPS == 0.055 and V_INIT_MPS == 0.25