tests/test_datagen_discovery.py

"""Synthetic, GPU-free tests for fpgm.datagen.discovery.

No TAPNext++/SAM3/PointWorld h5 dependency anywhere here -- every test builds
its own small :class:`~fpgm.types.Track2D` and/or robot-seg array by hand, per
the repo's "synthetic, no GPU, no network" test convention.
"""

from __future__ import annotations

import numpy as np
import pytest

from fpgm.config import DatagenConfig
from fpgm.datagen.discovery import (
    DiscoveredObject,
    DiscoveryResult,
    DiscoveryStats,
    build_cluster_features,
    build_densify_queries,
    cluster_tracks,
    compute_displacement,
    compute_object_bbox_px,
    compute_robot_overlap,
    densify_undersampled_clusters,
    discover_objects,
    filter_robot_tracks,
    filter_static_tracks,
    merge_consistent_clusters,
    read_track2d,
    sample_dense_grid_in_bbox,
    sample_grid_points,
    write_track2d,
)
from fpgm.types import DataError, Track2D


def _make_track(
    uv: np.ndarray, visible: np.ndarray | None = None, resolution=(200, 200)
) -> Track2D:
    t, q, _ = uv.shape
    if visible is None:
        visible = np.ones((t, q), dtype=bool)
    return Track2D(
        point_id=np.arange(q, dtype=np.int32),
        frames=np.arange(t, dtype=np.int32),
        uv=uv.astype(np.float32),
        visible=visible.astype(bool),
        resolution=resolution,
    )


def _linear_track(starts: np.ndarray, deltas: np.ndarray, n_frames: int = 5) -> Track2D:
    """Q points moving linearly from `starts` to `starts + deltas` over n_frames."""
    uv = np.stack(
        [starts + deltas * (t / (n_frames - 1)) for t in range(n_frames)], axis=0
    )
    return _make_track(uv)


# --------------------------------------------------------------------------- #
# Grid sampler
# --------------------------------------------------------------------------- #


class TestSampleGridPoints:
    def test_count_and_bounds(self):
        pts = sample_grid_points(width=1280, height=720, n_cols=40, n_rows=24)
        assert pts.shape == (40 * 24, 2)
        assert pts.dtype == np.float32
        assert np.all(pts[:, 0] >= 0) and np.all(pts[:, 0] < 1280)
        assert np.all(pts[:, 1] >= 0) and np.all(pts[:, 1] < 720)

    def test_no_jitter_is_a_regular_lattice(self):
        pts = sample_grid_points(width=100, height=100, n_cols=10, n_rows=10)
        xs = sorted(set(np.round(pts[:, 0], 4)))
        assert len(xs) == 10

    def test_deterministic_under_a_seed(self):
        rng1 = np.random.default_rng(42)
        rng2 = np.random.default_rng(42)
        a = sample_grid_points(200, 200, 10, 10, jitter_frac=0.5, rng=rng1)
        b = sample_grid_points(200, 200, 10, 10, jitter_frac=0.5, rng=rng2)
        np.testing.assert_array_equal(a, b)

    def test_jitter_changes_output(self):
        a = sample_grid_points(200, 200, 10, 10, jitter_frac=0.5, rng=np.random.default_rng(1))
        b = sample_grid_points(200, 200, 10, 10, jitter_frac=0.0)
        assert not np.array_equal(a, b)

    def test_rejects_nonpositive_args(self):
        with pytest.raises(ValueError):
            sample_grid_points(0, 100, 5, 5)
        with pytest.raises(ValueError):
            sample_grid_points(100, 100, 0, 5)


# --------------------------------------------------------------------------- #
# Robot subtraction
# --------------------------------------------------------------------------- #


class TestRobotSubtraction:
    def test_drops_track_at_query_frame_inside_robot(self):
        # point 0 sits inside the robot box every frame; point 1 never does.
        uv = np.zeros((3, 2, 2), dtype=np.float32)
        uv[:, 0] = [10, 10]
        uv[:, 1] = [150, 150]
        track = _make_track(uv)

        robot_seg = np.zeros((3, 200, 200), dtype=np.uint8)
        robot_seg[:, 5:15, 5:15] = 1

        keep = filter_robot_tracks(track, robot_seg, query_frame_idx=0)
        assert keep.tolist() == [False, True]

    def test_drops_track_by_majority_life_fraction_even_off_query_frame(self):
        # point 0 is OUTSIDE the robot at the query frame (t=0) but spends 3/4
        # of its visible life inside it afterwards -> still dropped.
        uv = np.zeros((4, 1, 2), dtype=np.float32)
        uv[0, 0] = [150, 150]  # query frame: outside
        uv[1, 0] = [10, 10]
        uv[2, 0] = [10, 10]
        uv[3, 0] = [10, 10]
        track = _make_track(uv)

        robot_seg = np.zeros((4, 200, 200), dtype=np.uint8)
        robot_seg[:, 5:15, 5:15] = 1

        at_query, life_fraction = compute_robot_overlap(track, robot_seg, query_frame_idx=0)
        assert at_query.tolist() == [False]
        assert life_fraction[0] == pytest.approx(0.75)

        keep = filter_robot_tracks(track, robot_seg, query_frame_idx=0, life_fraction_gate=0.5)
        assert keep.tolist() == [False]

    def test_keeps_track_never_touching_robot(self):
        uv = np.full((3, 1, 2), 150.0, dtype=np.float32)
        track = _make_track(uv)
        robot_seg = np.zeros((3, 200, 200), dtype=np.uint8)
        robot_seg[:, 5:15, 5:15] = 1
        keep = filter_robot_tracks(track, robot_seg, query_frame_idx=0)
        assert keep.tolist() == [True]


# --------------------------------------------------------------------------- #
# Static rejection
# --------------------------------------------------------------------------- #


class TestStaticRejection:
    def test_zero_displacement_is_dropped(self):
        uv = np.full((5, 1, 2), 50.0, dtype=np.float32)
        track = _make_track(uv)
        _, magnitude = compute_displacement(track)
        assert magnitude[0] == pytest.approx(0.0)
        keep = filter_static_tracks(track, displacement_gate_px=6.0)
        assert keep.tolist() == [False]

    def test_large_displacement_is_kept(self):
        track = _linear_track(np.array([[50.0, 50.0]]), np.array([[20.0, 0.0]]))
        _, magnitude = compute_displacement(track)
        assert magnitude[0] == pytest.approx(20.0, abs=1e-3)
        keep = filter_static_tracks(track, displacement_gate_px=6.0)
        assert keep.tolist() == [True]

    def test_below_threshold_is_dropped(self):
        track = _linear_track(np.array([[50.0, 50.0]]), np.array([[3.0, 0.0]]))
        keep = filter_static_tracks(track, displacement_gate_px=6.0)
        assert keep.tolist() == [False]

    def test_never_visible_track_has_zero_displacement(self):
        uv = np.zeros((5, 1, 2), dtype=np.float32)
        visible = np.zeros((5, 1), dtype=bool)
        track = _make_track(uv, visible=visible)
        _, magnitude = compute_displacement(track)
        assert magnitude[0] == 0.0


# --------------------------------------------------------------------------- #
# Clustering
# --------------------------------------------------------------------------- #


class TestClusterTracks:
    def test_two_well_separated_groups_give_two_clusters(self):
        rng = np.random.default_rng(0)
        group_a = rng.normal(loc=[0.0, 0.0, 0.5, 0.0, 0.0], scale=0.003, size=(10, 5))
        group_b = rng.normal(loc=[1.0, 1.0, 1.0, 0.05, 0.0], scale=0.003, size=(10, 5))
        features = np.concatenate([group_a, group_b], axis=0)

        labels = cluster_tracks(features, eps_m=0.05, min_samples=4)
        non_noise = set(labels.tolist()) - {-1}
        assert len(non_noise) == 2
        assert (labels[:10] == labels[0]).all()
        assert (labels[10:] == labels[10]).all()
        assert labels[0] != labels[10]

    def test_one_group_gives_one_cluster(self):
        rng = np.random.default_rng(1)
        features = rng.normal(loc=[0.2, 0.2, 0.6, 0.01, 0.0], scale=0.003, size=(12, 5))
        labels = cluster_tracks(features, eps_m=0.05, min_samples=4)
        non_noise = set(labels.tolist()) - {-1}
        assert len(non_noise) == 1

    def test_pure_noise_yields_no_clusters(self):
        rng = np.random.default_rng(2)
        features = rng.uniform(low=0.0, high=10.0, size=(30, 5))
        labels = cluster_tracks(features, eps_m=0.05, min_samples=6)
        non_noise = set(labels.tolist()) - {-1}
        assert len(non_noise) == 0, "clustering hallucinated a cluster out of uniform noise"

    def test_empty_features(self):
        labels = cluster_tracks(np.zeros((0, 5)), eps_m=0.05, min_samples=4)
        assert labels.shape == (0,)


class TestBuildClusterFeatures:
    def test_shape_and_metric_conversion(self):
        world_xyz = np.array([[0.0, 0.0, 1.0], [1.0, 1.0, 2.0]])
        net_vector_px = np.array([[100.0, 0.0], [0.0, 50.0]])
        z_cam = np.array([1.0, 2.0])
        features = build_cluster_features(world_xyz, net_vector_px, z_cam, fx=500.0, fy=500.0)
        assert features.shape == (2, 5)
        # d_m = d_px * z / f
        assert features[0, 3] == pytest.approx(100.0 * 1.0 / 500.0)
        assert features[1, 4] == pytest.approx(50.0 * 2.0 / 500.0)


# --------------------------------------------------------------------------- #
# End-to-end discover_objects (still synthetic: no GPU, no h5)
# --------------------------------------------------------------------------- #


class TestDiscoverObjects:
    def _two_object_scene(self):
        n_frames = 5
        # group A: 6 points near pixel (100, 100), moving +30px in x
        a_start = np.array([[100.0 + 2 * i, 100.0 + 2 * i] for i in range(6)])
        a_delta = np.tile([30.0, 0.0], (6, 1))
        # group B: 6 points near pixel (500, 400), moving +30px in y
        b_start = np.array([[500.0 + 2 * i, 400.0 + 2 * i] for i in range(6)])
        b_delta = np.tile([0.0, 30.0], (6, 1))

        starts = np.concatenate([a_start, b_start], axis=0)
        deltas = np.concatenate([a_delta, b_delta], axis=0)
        track = _linear_track(starts, deltas, n_frames=n_frames)

        robot_seg = np.zeros((n_frames, 720, 1280), dtype=np.uint8)

        rng_a, rng_b = np.random.default_rng(0), np.random.default_rng(1)
        world_xyz = np.concatenate(
            [
                np.tile([0.0, 0.0, 0.5], (6, 1)) + rng_a.normal(scale=0.002, size=(6, 3)),
                np.tile([1.0, 1.0, 1.0], (6, 1)) + rng_b.normal(scale=0.002, size=(6, 3)),
            ],
            axis=0,
        )
        z_cam = world_xyz[:, 2].copy()
        return track, robot_seg, world_xyz, z_cam

    def test_two_clusters_found(self):
        track, robot_seg, world_xyz, z_cam = self._two_object_scene()
        cfg = DatagenConfig(dbscan_eps_m=0.05, dbscan_min_samples=4)
        result = discover_objects(
            track, robot_seg, query_frame_idx=0, world_xyz=world_xyz, z_cam=z_cam,
            fx=500.0, fy=500.0, cfg=cfg,
        )
        assert result.stats.n_clusters == 2
        assert len(result.objects) == 2
        total_tracks = sum(obj.n_tracks for obj in result.objects)
        assert total_tracks == 12
        for obj in result.objects:
            assert obj.n_tracks == 6
            assert np.all(np.isfinite(obj.mean_3d_position))
            assert obj.total_displacement_px == pytest.approx(30.0, abs=1.0)

    def test_robot_filter_removes_a_whole_object(self):
        track, robot_seg, world_xyz, z_cam = self._two_object_scene()
        # Robot covers group A's whole neighbourhood for every frame.
        robot_seg[:, 90:120, 90:120] = 1
        cfg = DatagenConfig(dbscan_eps_m=0.05, dbscan_min_samples=4)
        result = discover_objects(
            track, robot_seg, query_frame_idx=0, world_xyz=world_xyz, z_cam=z_cam,
            fx=500.0, fy=500.0, cfg=cfg,
        )
        assert result.stats.n_after_robot_filter == 6
        assert result.stats.n_clusters == 1
        assert result.objects[0].n_tracks == 6

    def test_everything_filtered_gives_no_objects_not_a_crash(self):
        n_frames = 3
        uv = np.full((n_frames, 2, 2), 50.0, dtype=np.float32)  # static
        track = _make_track(uv)
        robot_seg = np.zeros((n_frames, 200, 200), dtype=np.uint8)
        world_xyz = np.zeros((2, 3))
        z_cam = np.ones(2)
        cfg = DatagenConfig()
        result = discover_objects(
            track, robot_seg, query_frame_idx=0, world_xyz=world_xyz, z_cam=z_cam,
            fx=500.0, fy=500.0, cfg=cfg,
        )
        assert result.objects == ()
        assert result.stats.n_clusters == 0
        assert result.track.uv.shape[1] == 0


# --------------------------------------------------------------------------- #
# Dataclass validation
# --------------------------------------------------------------------------- #


class TestValidation:
    def test_discovered_object_rejects_bad_query_uv_shape(self):
        with pytest.raises(DataError):
            DiscoveredObject(
                cluster_id=0,
                query_uv=np.zeros((3,), dtype=np.float32),  # wrong: not (n, 2)
                centroid_uv_per_frame=np.zeros((5, 2), dtype=np.float32),
                n_tracks=3,
                total_displacement_px=10.0,
                mean_3d_position=np.zeros(3),
            )

    def test_discovery_result_rejects_misaligned_labels(self):
        track = _make_track(np.zeros((2, 4, 2), dtype=np.float32))
        with pytest.raises(DataError):
            DiscoveryResult(
                objects=(),
                track=track,
                cluster_labels=np.zeros((3,), dtype=np.int32),  # should be (4,)
                stats=None,  # not touched before the shape check raises
                seed_frame_idx=0,
            )


# --------------------------------------------------------------------------- #
# Track2D (de)serialization
# --------------------------------------------------------------------------- #


class TestTrack2DRoundTrip:
    def test_round_trip(self, tmp_path):
        starts = np.array([[10.0, 20.0], [30.0, 40.0]])
        deltas = np.array([[5.0, 0.0], [0.0, 5.0]])
        track = _linear_track(starts, deltas)
        path = tmp_path / "track.npz"
        write_track2d(path, track)
        loaded = read_track2d(path)
        np.testing.assert_array_equal(loaded.point_id, track.point_id)
        np.testing.assert_array_equal(loaded.frames, track.frames)
        np.testing.assert_array_almost_equal(loaded.uv, track.uv)
        np.testing.assert_array_equal(loaded.visible, track.visible)
        assert loaded.resolution == track.resolution


# --------------------------------------------------------------------------- #
# Cluster merging (over-segmentation cleanup)
# --------------------------------------------------------------------------- #


class TestMergeConsistentClusters:
    def _three_cluster_scene(self):
        """Two spatially-close, same-direction "drawer patch" groups (A, B) that
        DBSCAN keeps apart at a tight eps, plus one far, differently-moving
        "brick" group (C) that must never merge with anything.
        """
        n_frames = 5
        # Group A: 4 points near pixel (100, 100), net +10px in x.
        a_start = np.array([[100.0 + 2 * i, 100.0 + 2 * i] for i in range(4)])
        a_delta = np.tile([10.0, 0.0], (4, 1))
        # Group B: 4 points near pixel (300, 100) -- close to A in *world*
        # position (0.05 m, see world_xyz below) and moving almost exactly
        # the same way (+10.5px x, +0.3px y).
        b_start = np.array([[300.0 + 2 * i, 100.0 + 2 * i] for i in range(4)])
        b_delta = np.tile([10.5, 0.3], (4, 1))
        # Group C: 4 points far away, moving in a different direction.
        c_start = np.array([[500.0 + 2 * i, 400.0 + 2 * i] for i in range(4)])
        c_delta = np.tile([0.0, 10.0], (4, 1))

        starts = np.concatenate([a_start, b_start, c_start], axis=0)
        deltas = np.concatenate([a_delta, b_delta, c_delta], axis=0)
        track = _linear_track(starts, deltas, n_frames=n_frames)
        robot_seg = np.zeros((n_frames, 720, 1280), dtype=np.uint8)

        rng = np.random.default_rng(0)
        world_xyz = np.concatenate(
            [
                np.tile([0.0, 0.0, 0.5], (4, 1)) + rng.normal(scale=0.001, size=(4, 3)),
                np.tile([0.05, 0.0, 0.5], (4, 1)) + rng.normal(scale=0.001, size=(4, 3)),
                np.tile([1.0, 1.0, 1.0], (4, 1)) + rng.normal(scale=0.001, size=(4, 3)),
            ],
            axis=0,
        )
        z_cam = np.ones(12)  # z=1 everywhere -> dx_m = dx_px / fx exactly
        return track, robot_seg, world_xyz, z_cam

    def _discover(self, eps_m=0.03, min_samples=3):
        track, robot_seg, world_xyz, z_cam = self._three_cluster_scene()
        cfg = DatagenConfig(dbscan_eps_m=eps_m, dbscan_min_samples=min_samples)
        return discover_objects(
            track, robot_seg, query_frame_idx=0, world_xyz=world_xyz, z_cam=z_cam,
            fx=500.0, fy=500.0, cfg=cfg,
        )

    def test_raw_dbscan_keeps_three_clusters(self):
        # Sanity check on the fixture itself: at a tight eps, A and B (0.05 m
        # apart) are NOT joined by DBSCAN alone -- merging is this function's job.
        result = self._discover()
        assert result.stats.n_clusters == 3

    def test_close_same_direction_clusters_merge(self):
        result = self._discover()
        merged, groups = merge_consistent_clusters(result)
        assert merged.stats.n_clusters == 2
        sizes = sorted(o.n_tracks for o in merged.objects)
        assert sizes == [4, 8]
        # the far/differently-moving group must be a singleton in the merge map
        singleton_groups = [g for g in groups if len(g) == 1]
        assert len(singleton_groups) == 1

    def test_merged_object_stats_are_track_weighted_means(self):
        result = self._discover()
        merged, groups = merge_consistent_clusters(result)
        big = next(o for o in merged.objects if o.n_tracks == 8)
        # exact weighted mean of A (n=4, x=0.0) and B (n=4, x=0.05): 0.025
        assert big.mean_3d_position[0] == pytest.approx(0.025, abs=5e-3)
        assert big.mean_displacement_m[0] > 0  # still pointing the same (+x) way

    def test_no_spurious_merge_when_nothing_agrees(self):
        # Re-run with min_samples=1 so every point is its own tiny cluster's
        # worth of "signal", but drop A/B close enough only by chance -- use
        # the original three well-separated-by-direction groups from
        # TestDiscoverObjects to confirm merge is a no-op when it should be.
        track, robot_seg, world_xyz, z_cam = TestDiscoverObjects()._two_object_scene()
        cfg = DatagenConfig(dbscan_eps_m=0.05, dbscan_min_samples=4)
        result = discover_objects(
            track, robot_seg, query_frame_idx=0, world_xyz=world_xyz, z_cam=z_cam,
            fx=500.0, fy=500.0, cfg=cfg,
        )
        assert result.stats.n_clusters == 2
        merged, groups = merge_consistent_clusters(result)
        assert merged.stats.n_clusters == 2
        assert all(len(g) == 1 for g in groups)

    def test_merge_is_a_noop_on_zero_or_one_cluster(self):
        empty = DiscoveryResult(
            objects=(),
            track=_make_track(np.zeros((2, 0, 2), dtype=np.float32)),
            cluster_labels=np.zeros((0,), dtype=np.int32),
            stats=DiscoveryStats(0, 0, 0, 0, 0, 0),
            seed_frame_idx=0,
        )
        merged, groups = merge_consistent_clusters(empty)
        assert merged.stats.n_clusters == 0
        assert groups == ()


# --------------------------------------------------------------------------- #
# Second-pass densification
# --------------------------------------------------------------------------- #


class TestComputeObjectBboxPx:
    def test_pads_around_a_tight_cluster(self):
        query_uv = np.array([[10.0, 10.0], [14.0, 14.0]])
        x0, y0, x1, y1 = compute_object_bbox_px(query_uv, width=200, height=200)
        # raw extent (4, 4) is floored to the 8px minimum, then padded 1x each side
        assert x1 - x0 == pytest.approx(16.0, abs=1e-6)
        assert y1 - y0 == pytest.approx(16.0, abs=1e-6)
        assert x0 <= 10.0 and x1 >= 14.0

    def test_clamped_to_frame_bounds(self):
        query_uv = np.array([[1.0, 1.0], [2.0, 2.0]])
        x0, y0, x1, y1 = compute_object_bbox_px(query_uv, width=50, height=50)
        assert x0 >= 0.0 and y0 >= 0.0
        assert x1 <= 49.0 and y1 <= 49.0

    def test_rejects_empty_query_uv(self):
        with pytest.raises(ValueError):
            compute_object_bbox_px(np.zeros((0, 2)), width=100, height=100)


class TestSampleDenseGridInBbox:
    def test_count_and_containment(self):
        pts = sample_dense_grid_in_bbox((10.0, 10.0, 30.0, 30.0), n_cols=4, n_rows=4)
        assert pts.shape == (16, 2)
        assert np.all(pts[:, 0] >= 10.0) and np.all(pts[:, 0] <= 30.0)
        assert np.all(pts[:, 1] >= 10.0) and np.all(pts[:, 1] <= 30.0)


class TestBuildDensifyQueries:
    def _two_object_result(self, small_n=3, big_n=50):
        small = DiscoveredObject(
            cluster_id=0,
            query_uv=np.array([[100.0, 100.0], [104.0, 100.0], [100.0, 104.0]][:small_n]),
            centroid_uv_per_frame=np.zeros((5, 2), dtype=np.float32),
            n_tracks=small_n,
            total_displacement_px=15.0,
            mean_3d_position=np.array([0.0, 0.0, 0.5]),
            mean_displacement_m=np.array([0.02, 0.0]),
        )
        big = DiscoveredObject(
            cluster_id=1,
            query_uv=np.tile([500.0, 400.0], (big_n, 1)).astype(np.float32),
            centroid_uv_per_frame=np.zeros((5, 2), dtype=np.float32),
            n_tracks=big_n,
            total_displacement_px=20.0,
            mean_3d_position=np.array([1.0, 1.0, 1.0]),
            mean_displacement_m=np.array([0.0, 0.02]),
        )
        return DiscoveryResult(
            objects=(small, big),
            track=_make_track(np.zeros((5, 0, 2), dtype=np.float32)),
            cluster_labels=np.zeros((0,), dtype=np.int32),
            stats=DiscoveryStats(0, 0, 0, 0, 2, 0),
            seed_frame_idx=0,
        )

    def test_only_undersampled_cluster_gets_reseeded(self):
        result = self._two_object_result()
        uv, source = build_densify_queries(
            result, width=1280, height=720, min_tracks=20, grid=(3, 3)
        )
        assert uv.shape == (9, 2)
        assert set(source.tolist()) == {0}

    def test_no_reseed_when_everything_is_dense_enough(self):
        result = self._two_object_result(small_n=25)
        uv, source = build_densify_queries(
            result, width=1280, height=720, min_tracks=20, grid=(3, 3)
        )
        assert uv.shape == (0, 2)
        assert source.shape == (0,)


class TestDensifyUndersampledClusters:
    def _sparse_result(self):
        query_uv = np.array([[100.0, 100.0], [104.0, 100.0]], dtype=np.float32)
        track = _make_track(
            np.stack([query_uv, query_uv + [5, 0], query_uv + [10, 0]], axis=0)
        )
        obj = DiscoveredObject(
            cluster_id=0,
            query_uv=query_uv,
            centroid_uv_per_frame=track.uv.mean(axis=1),
            n_tracks=2,
            total_displacement_px=10.0,
            mean_3d_position=np.array([0.0, 0.0, 0.5]),
            mean_displacement_m=np.array([0.02, 0.0]),
        )
        result = DiscoveryResult(
            objects=(obj,),
            track=track,
            cluster_labels=np.zeros(2, dtype=np.int32),
            stats=DiscoveryStats(
                n_query_points=100, n_after_robot_filter=50, n_after_static_filter=2,
                n_with_3d_lift=2, n_clusters=1, n_noise=0,
            ),
            seed_frame_idx=0,
        )
        return result

    def test_surviving_reseed_points_grow_the_cluster(self):
        result = self._sparse_result()
        n_frames = result.track.uv.shape[0]
        dense_uv = np.array([[102.0, 102.0], [98.0, 98.0], [101.0, 99.0]], dtype=np.float32)
        dense_track = _make_track(
            np.stack([dense_uv, dense_uv + [5, 0], dense_uv + [10, 0]], axis=0)
        )
        assert dense_track.uv.shape[0] == n_frames
        dense_world_xyz = np.tile([0.0, 0.0, 0.5], (3, 1)) + np.array(
            [[0.001, 0.0, 0.0], [-0.001, 0.0, 0.0], [0.0005, 0.0, 0.0]]
        )
        dense_z_cam = np.ones(3)
        source_cluster_id = np.zeros(3, dtype=np.int32)
        robot_seg = np.zeros((n_frames, 720, 1280), dtype=np.uint8)

        merged = densify_undersampled_clusters(
            result, dense_track, dense_world_xyz, dense_z_cam, source_cluster_id,
            robot_seg, query_frame_idx=0,
        )
        assert merged.objects[0].n_tracks == 5  # 2 original + 3 new
        assert merged.track.uv.shape[1] == 5
        assert merged.cluster_labels.shape == (5,)
        # exact weighted-mean identity: (2*0.0 + 3*mean(new x)) / 5
        expected_x = (2 * 0.0 + 3 * np.mean([0.001, -0.001, 0.0005])) / 5
        assert merged.objects[0].mean_3d_position[0] == pytest.approx(expected_x, abs=1e-6)
        # funnel stats accumulate the second pass on top of the first
        assert merged.stats.n_query_points == 100 + 3
        assert merged.stats.n_after_robot_filter == 50 + 3

    def test_points_failing_the_robot_gate_are_dropped_not_reassigned(self):
        result = self._sparse_result()
        n_frames = result.track.uv.shape[0]
        dense_uv = np.array([[10.0, 10.0]], dtype=np.float32)  # inside the robot box below
        dense_track = _make_track(np.tile(dense_uv, (n_frames, 1, 1)))
        dense_world_xyz = np.array([[0.0, 0.0, 0.5]])
        dense_z_cam = np.array([1.0])
        source_cluster_id = np.zeros(1, dtype=np.int32)
        robot_seg = np.zeros((n_frames, 720, 1280), dtype=np.uint8)
        robot_seg[:, 0:20, 0:20] = 1  # covers the reseed point's query frame

        merged = densify_undersampled_clusters(
            result, dense_track, dense_world_xyz, dense_z_cam, source_cluster_id,
            robot_seg, query_frame_idx=0,
        )
        assert merged.objects[0].n_tracks == 2  # unchanged -- the reseed point was dropped
        assert merged.stats.n_query_points == 100 + 1
        assert merged.stats.n_after_robot_filter == 50 + 0

    def test_empty_dense_track_is_a_noop(self):
        result = self._sparse_result()
        empty_dense = _make_track(np.zeros((result.track.uv.shape[0], 0, 2), dtype=np.float32))
        merged = densify_undersampled_clusters(
            result, empty_dense, np.zeros((0, 3)), np.zeros(0), np.zeros(0, dtype=np.int32),
            np.zeros((3, 720, 1280), dtype=np.uint8), query_frame_idx=0,
        )
        assert merged is result