tests/test_sampling.py

"""Unit tests for fpgm.tracking.sampling -- pure numpy/scipy, no GPU or weights."""

from __future__ import annotations

import logging

import numpy as np
import pytest
from scipy import ndimage

from fpgm.tracking.sampling import mask_centroid, points_in_mask, sample_points_in_mask


def _disk_mask(height: int, width: int, cx: float, cy: float, radius: float) -> np.ndarray:
    yy, xx = np.mgrid[0:height, 0:width]
    return ((xx - cx) ** 2 + (yy - cy) ** 2) <= radius**2


def _pairwise_min_distance(points: np.ndarray) -> float:
    """Smallest distance between any two distinct points -- a spread-out set has a large one."""
    diffs = points[:, None, :] - points[None, :, :]
    dists = np.sqrt(np.sum(diffs**2, axis=-1))
    np.fill_diagonal(dists, np.inf)
    return float(dists.min())


class TestMaskCentroidAndPointsInMask:
    def test_mask_centroid_of_square(self) -> None:
        mask = np.zeros((20, 20), dtype=bool)
        mask[5:15, 8:12] = True  # x in [8, 11], y in [5, 14]
        centroid = mask_centroid(mask)
        assert centroid == pytest.approx([9.5, 9.5], abs=1e-6)

    def test_mask_centroid_empty_raises(self) -> None:
        with pytest.raises(ValueError):
            mask_centroid(np.zeros((10, 10), dtype=bool))

    def test_points_in_mask(self) -> None:
        mask = np.zeros((10, 10), dtype=bool)
        mask[2:5, 2:5] = True
        uv = np.array([[3, 3], [0, 0], [-1, -1], [100, 100], [3.4, 2.6]])
        result = points_in_mask(mask, uv)
        assert result.tolist() == [True, False, False, False, True]


@pytest.mark.parametrize("method", ["farthest", "grid", "random"])
class TestSamplePointsInMask:
    def test_points_lie_inside_mask(self, method: str) -> None:
        mask = _disk_mask(120, 120, cx=60, cy=60, radius=45)
        rng = np.random.default_rng(0)
        points = sample_points_in_mask(
            mask, n_points=25, method=method, boundary_erosion_px=3, rng=rng
        )
        assert points_in_mask(mask, points).all()

    def test_count_is_respected(self, method: str) -> None:
        mask = _disk_mask(120, 120, cx=60, cy=60, radius=45)
        rng = np.random.default_rng(1)
        n_points = 30
        points = sample_points_in_mask(
            mask, n_points=n_points, method=method, boundary_erosion_px=3, rng=rng
        )
        assert points.shape == (n_points, 2)
        assert points.dtype == np.float32

    def test_empty_mask_raises(self, method: str) -> None:
        mask = np.zeros((30, 30), dtype=bool)
        with pytest.raises(ValueError):
            sample_points_in_mask(mask, n_points=5, method=method)


class TestBoundaryErosion:
    def test_erosion_pulls_points_away_from_boundary(self) -> None:
        mask = _disk_mask(150, 150, cx=75, cy=75, radius=60)
        # Euclidean distance-to-background: eroded samples should sit farther
        # from the boundary, on average, than a generic mask pixel does.
        dist_to_boundary = ndimage.distance_transform_edt(mask)
        all_mask_pixels_mean = dist_to_boundary[mask].mean()

        rng = np.random.default_rng(2)
        points = sample_points_in_mask(
            mask, n_points=40, method="farthest", boundary_erosion_px=8, rng=rng
        )
        sampled_u = np.round(points[:, 0]).astype(int)
        sampled_v = np.round(points[:, 1]).astype(int)
        sampled_dist = dist_to_boundary[sampled_v, sampled_u]

        assert sampled_dist.mean() > all_mask_pixels_mean
        assert sampled_dist.min() >= 1.0  # none sit exactly on the boundary pixel

    def test_erosion_emptying_mask_falls_back_with_warning(
        self, caplog: pytest.LogCaptureFixture
    ) -> None:
        # A mask thinner than 2 * boundary_erosion_px erodes to nothing.
        mask = np.zeros((30, 30), dtype=bool)
        mask[10:13, 5:25] = True  # 3px-tall strip
        caplog.set_level(logging.WARNING, logger="fpgm.tracking.sampling")

        rng = np.random.default_rng(3)
        points = sample_points_in_mask(
            mask, n_points=5, method="random", boundary_erosion_px=5, rng=rng
        )

        assert points_in_mask(mask, points).all()
        assert any("emptied the mask" in record.message for record in caplog.records)


class TestFarthestPointSampling:
    def test_farthest_spreads_points_more_than_random(self) -> None:
        mask = _disk_mask(200, 200, cx=100, cy=100, radius=80)
        n_points = 30

        farthest = sample_points_in_mask(
            mask,
            n_points=n_points,
            method="farthest",
            boundary_erosion_px=0,
            rng=np.random.default_rng(0),
        )
        random_pts = sample_points_in_mask(
            mask,
            n_points=n_points,
            method="random",
            boundary_erosion_px=0,
            rng=np.random.default_rng(0),
        )

        assert _pairwise_min_distance(farthest) > _pairwise_min_distance(random_pts)


class TestSmallMask:
    def test_mask_smaller_than_n_points_returns_fewer(
        self, caplog: pytest.LogCaptureFixture
    ) -> None:
        mask = np.zeros((20, 20), dtype=bool)
        mask[10, 10] = True
        mask[10, 11] = True
        mask[11, 10] = True  # exactly 3 pixels
        caplog.set_level(logging.WARNING, logger="fpgm.tracking.sampling")

        points = sample_points_in_mask(
            mask,
            n_points=10,
            method="farthest",
            boundary_erosion_px=0,
            rng=np.random.default_rng(4),
        )

        assert points.shape[0] == 3
        assert points_in_mask(mask, points).all()
        assert any(
            "only" in record.message or "requested" in record.message for record in caplog.records
        )

    def test_grid_method_small_mask_does_not_crash(self) -> None:
        mask = np.zeros((20, 20), dtype=bool)
        mask[10:12, 10:12] = True  # 4 pixels
        points = sample_points_in_mask(
            mask, n_points=50, method="grid", boundary_erosion_px=0, rng=np.random.default_rng(5)
        )
        assert 0 < points.shape[0] <= 50
        assert points_in_mask(mask, points).all()


def test_invalid_method_raises() -> None:
    mask = np.ones((10, 10), dtype=bool)
    with pytest.raises(ValueError):
        sample_points_in_mask(mask, n_points=5, method="bogus")


def test_non_positive_n_points_raises() -> None:
    mask = np.ones((10, 10), dtype=bool)
    with pytest.raises(ValueError):
        sample_points_in_mask(mask, n_points=0)