TIR_PROJ/tests/test_perception.py

"""LiDAR simulation + perception pipeline + multi-target tracker."""

import math

import numpy as np
import pytest

from herding.perception.lidar_perception import (
    STATIC_REJECT, detections_from_scan,
)
from herding.perception.lidar_sim import (
    LIDAR_MAX_RANGE, LIDAR_N_RAYS, SHEEP_RADIUS, ray_angles, simulate_scan,
)
from herding.perception.sheep_tracker import (
    FORGET_STEPS, GATE_M, MAX_ACTIVE_TRACKS, REACQUIRE_GATE_M,
    REACQUIRE_MIN_AGE, SheepTracker,
)


# ---------------------------------------------------------------------------
# Sim
# ---------------------------------------------------------------------------

def test_simulate_scan_shape_and_dtype():
    ranges = simulate_scan(0.0, 0.0, 0.0, [(5.0, 0.0)], noise=0.0)
    assert ranges.shape == (LIDAR_N_RAYS,)
    assert ranges.dtype == np.float32


def test_simulate_scan_no_sheep_far_from_walls():
    # Dog at origin, no sheep, walls all ≥ 15 m away → all rays at max.
    ranges = simulate_scan(0.0, 0.0, 0.0, [], noise=0.0)
    # Walls (east/west at ±15) are beyond LIDAR_MAX_RANGE=12, so no hits.
    assert (ranges == LIDAR_MAX_RANGE).all()


def test_simulate_scan_sheep_in_front_returns_centre_hit():
    # Sheep dead ahead at 5 m. Centre ray should hit ~ 5 - SHEEP_RADIUS.
    ranges = simulate_scan(0.0, 0.0, 0.0, [(5.0, 0.0)], noise=0.0)
    centre = ranges[LIDAR_N_RAYS // 2]
    assert math.isclose(float(centre), 5.0 - SHEEP_RADIUS, abs_tol=0.01)


def test_simulate_scan_sheep_behind_dog_not_hit():
    ranges = simulate_scan(0.0, 0.0, 0.0, [(-5.0, 0.0)], noise=0.0)
    assert (ranges == LIDAR_MAX_RANGE).all()


def test_simulate_scan_wall_hit():
    # Dog 1 m south of the north wall, facing north → centre ray ≈ 1 m.
    ranges = simulate_scan(0.0, 14.0, math.pi / 2, [], noise=0.0)
    centre = ranges[LIDAR_N_RAYS // 2]
    assert math.isclose(float(centre), 1.0, abs_tol=0.01)


# ---------------------------------------------------------------------------
# Perception
# ---------------------------------------------------------------------------

def test_detections_recover_sheep_position():
    sheep = [(5.0, 0.0), (3.0, 1.0)]
    ranges = simulate_scan(0.0, 0.0, 0.0, sheep, noise=0.0)
    det = detections_from_scan(ranges, 0.0, 0.0, 0.0)
    assert len(det) == 2
    # Centroid bias is corrected to within ~5 cm.
    for truth in sheep:
        assert any(math.hypot(d[0] - truth[0], d[1] - truth[1]) < 0.1
                   for d in det)


def test_detections_filter_gate_post():
    # An empty scene at the dog right next to a gate post produces no
    # detections — the static-feature filter drops the post return.
    ranges = simulate_scan(11.5, -10.0, -math.pi / 2, [], noise=0.0)
    det = detections_from_scan(ranges, 11.5, -10.0, -math.pi / 2)
    for cx, cy in det:
        assert math.hypot(cx - 10.0, cy + 15.0) > STATIC_REJECT
        assert math.hypot(cx - 13.0, cy + 15.0) > STATIC_REJECT


def test_detections_empty_scan_returns_nothing():
    assert detections_from_scan(np.array([], dtype=np.float32),
                                0.0, 0.0, 0.0) == []


# ---------------------------------------------------------------------------
# Tracker
# ---------------------------------------------------------------------------

def test_tracker_creates_track_for_new_detection():
    t = SheepTracker()
    t.update([(5.0, 0.0)])
    assert t.n_active() == 1


def test_tracker_associates_close_detections():
    """A small movement within the gate keeps the same track."""
    t = SheepTracker()
    t.update([(5.0, 0.0)])
    t.update([(5.5, 0.0)])
    assert t.n_active() == 1


def test_tracker_spawns_new_track_far_detection():
    t = SheepTracker()
    t.update([(5.0, 0.0)])
    t.update([(-5.0, 0.0)])           # well outside the gate
    assert t.n_active() == 2


def test_tracker_reacquisition_for_stale_track():
    """A stale track within the wider re-acquisition gate rebinds rather
    than spawning a duplicate."""
    t = SheepTracker()
    t.update([(0.0, 0.0)])
    # Let it go stale.
    for _ in range(REACQUIRE_MIN_AGE):
        t.update([])
    # Re-emerges within REACQUIRE_GATE but outside the primary GATE.
    offset = (GATE_M + REACQUIRE_GATE_M) / 2.0
    t.update([(offset, 0.0)])
    assert t.n_active() == 1


def test_tracker_forgets_stale_tracks():
    t = SheepTracker()
    t.update([(0.0, 0.0)])
    for _ in range(FORGET_STEPS + 1):
        t.update([])
    assert t.n_active() == 0


def test_tracker_penned_position_promotes_track():
    t = SheepTracker()
    t.update([(11.5, -16.0)])         # spawn inside the pen column
    # is_penned_position is True for this point.
    assert t.n_penned() == 1
    assert t.n_active() == 0


def test_tracker_penned_tracks_persist():
    t = SheepTracker()
    t.update([(11.5, -16.0)])
    for _ in range(FORGET_STEPS * 2):
        t.update([])
    # Penned tracks are not forgotten.
    assert t.n_penned() == 1


def test_tracker_caps_active_set():
    t = SheepTracker()
    # Spawn more than the cap, each well outside the others' gates.
    for k in range(MAX_ACTIVE_TRACKS + 5):
        t.update([(k * (GATE_M + 1.0), 0.0)])
    assert t.n_active() <= MAX_ACTIVE_TRACKS


def test_tracker_reset_clears_state():
    t = SheepTracker()
    t.update([(0.0, 0.0)])
    t.reset()
    assert t.n_active() == 0
    assert t.step == 0