tests/test_static_span.py

"""Synthetic tests for fpgm.datagen.static_span -- no GPU, no data files.

Covers both signals :class:`~fpgm.datagen.static_span.StaticSpanEstimator` exposes:
the primary ``from_tracks`` (tracked-point median displacement) and the weaker
``from_masks`` fallback (mask centroid) -- see that module's own docstring for why
tracks are preferred (the mask centroid moves under partial occlusion even when the
object does not).
"""

from __future__ import annotations

import numpy as np

from fpgm.config_datagen import StaticSpanConfig
from fpgm.datagen.static_span import StaticSpanEstimator

# max_step_px=1.5, max_total_drift_px=4.0, min_frames=4, max_frames=30
_DEFAULT_CFG = StaticSpanConfig()


def _tracks(
    n_frames: int, n_points: int, xy_per_frame: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """``xy_per_frame`` is ``(T, 2)`` -- every point sits at the same per-frame offset,
    all points visible every frame. Callers that need per-point jitter perturb the
    returned ``uv`` afterward.
    """
    uv = np.tile(xy_per_frame[:, None, :], (1, n_points, 1)).astype(np.float64)
    visible = np.ones((n_frames, n_points), dtype=bool)
    return uv, visible


class TestFromTracksStaticThenMotion:
    def test_static_run_then_real_motion_is_bounded_at_the_motion(self) -> None:
        rng = np.random.default_rng(0)
        n_frames, n_points = 20, 12
        base = np.zeros((n_frames, 2))
        # Frames 0-9: sub-pixel jitter only (well under max_step_px/max_total_drift_px).
        base[:10] = np.array([100.0, 100.0])
        # Frames 10+: a real, fast slide -- well past both thresholds per step.
        for t in range(10, n_frames):
            base[t] = base[9] + np.array([20.0, 0.0]) * (t - 9)
        uv, visible = _tracks(n_frames, n_points, base)
        uv += rng.uniform(-0.05, 0.05, size=uv.shape)  # sub-pixel per-point jitter

        result = StaticSpanEstimator(_DEFAULT_CFG).from_tracks(uv, visible)

        assert result.ok
        assert result.source == "tracks"
        assert result.span[0] == 0
        # Must stop at or before the real motion begins (frame 10) -- never later.
        assert result.span[1] <= 10
        assert not result.capped
        assert result.max_step_px < _DEFAULT_CFG.max_step_px

    def test_to_json_round_trips_the_measured_numbers(self) -> None:
        n_frames, n_points = 10, 10
        base = np.tile(np.array([50.0, 60.0]), (n_frames, 1))
        uv, visible = _tracks(n_frames, n_points, base)

        result = StaticSpanEstimator(_DEFAULT_CFG).from_tracks(uv, visible)
        payload = result.to_json()

        assert payload["span"] == list(result.span)
        assert payload["n_frames"] == result.n_frames
        assert payload["source"] == "tracks"
        assert payload["capped"] == result.capped


class TestFromTracksSteadyCreepRejectedByTotalDrift:
    def test_1px_per_frame_creep_passes_every_step_but_is_capped_by_drift(self) -> None:
        """Every single-frame step is 1.0 px (< max_step_px=1.5) -- a per-step-only
        check would happily accept the entire 30-frame run. Total drift crosses
        max_total_drift_px=4.0 well before that, so the returned span must be much
        shorter than the full run, capped by drift, not by max_frames.
        """
        n_frames, n_points = 30, 10
        base = np.stack([np.arange(n_frames, dtype=np.float64), np.zeros(n_frames)], axis=1)
        base += np.array([200.0, 150.0])  # arbitrary origin
        uv, visible = _tracks(n_frames, n_points, base)

        result = StaticSpanEstimator(_DEFAULT_CFG).from_tracks(uv, visible)

        assert result.ok
        # every individual step legitimately passed the per-step gate:
        assert result.max_step_px <= _DEFAULT_CFG.max_step_px + 1e-9
        assert result.n_frames < n_frames  # did NOT run the full synthetic creep
        assert not result.capped  # stopped by drift, not by hitting max_frames
        assert result.total_drift_px <= _DEFAULT_CFG.max_total_drift_px + 1e-9


class TestFromMasksEmptyLeadingFramesSkipped:
    def test_leading_empty_masks_are_skipped_and_reported(self) -> None:
        h, w = 64, 64
        n_frames = 12
        n_leading_empty = 3
        masks = np.zeros((n_frames, h, w), dtype=bool)
        # A small, barely-moving blob starting only at frame `n_leading_empty`.
        for t in range(n_leading_empty, n_frames):
            cx = 20 + t - n_leading_empty  # 1 px/frame drift -- stays under both gates briefly
            masks[t, 20:24, cx:cx + 4] = True

        result = StaticSpanEstimator(_DEFAULT_CFG).from_masks(masks)

        assert result.n_leading_empty == n_leading_empty
        if result.ok:
            assert result.span[0] == n_leading_empty
        assert result.source == "mask_centroid"


class TestNeverVisibleReturnsNone:
    def test_from_tracks_never_enough_covisible_points(self) -> None:
        n_frames, n_points = 10, 10
        uv = np.zeros((n_frames, n_points, 2))
        visible = np.zeros((n_frames, n_points), dtype=bool)  # nothing ever visible

        result = StaticSpanEstimator(_DEFAULT_CFG).from_tracks(uv, visible)

        assert result.span is None
        assert not result.ok
        assert result.n_frames == 0
        assert "visible" in result.reason

    def test_from_masks_object_never_visible(self) -> None:
        masks = np.zeros((10, 32, 32), dtype=bool)  # never any mask pixels at all

        result = StaticSpanEstimator(_DEFAULT_CFG).from_masks(masks)

        assert result.span is None
        assert not result.ok
        assert result.reason == "object never visible"


class TestMeasurePrefersTracksOverMasks:
    def test_label_with_a_track_uses_tracks_even_if_masks_are_also_given(self) -> None:
        n_frames, n_points = 10, 10
        base = np.tile(np.array([10.0, 10.0]), (n_frames, 1))
        uv, visible = _tracks(n_frames, n_points, base)
        masks = np.zeros((n_frames, 16, 16), dtype=bool)
        masks[:, :4, :4] = True  # would also measure "static" via the mask fallback

        results = StaticSpanEstimator(_DEFAULT_CFG).measure(
            tracks_by_label={"fixture": (uv, visible)}, masks_by_label={"fixture": masks}
        )

        assert results["fixture"].source == "tracks"

    def test_label_with_only_masks_falls_back(self) -> None:
        n_frames = 10
        masks = np.zeros((n_frames, 16, 16), dtype=bool)
        masks[:, :4, :4] = True

        results = StaticSpanEstimator(_DEFAULT_CFG).measure(masks_by_label={"fixture": masks})

        assert results["fixture"].source == "mask_centroid"

    def test_empty_inputs_yield_empty_results(self) -> None:
        results = StaticSpanEstimator(_DEFAULT_CFG).measure()
        assert results == {}