TIR_PROJ/herding/perception/sheep_tracker.py

"""Multi-target tracker for LiDAR-detected sheep.

Three-stage greedy nearest-neighbour data association:

1. **Consensus promotion**. New detections start as *candidate* tracks
   invisible to ``get_positions``. They must accumulate ``consensus_k``
   matches within ``consensus_radius_m`` to promote; candidates that
   fail to re-confirm within ``consensus_max_age`` steps die. This
   filters one-shot LiDAR phantoms — wall returns, multi-cluster sheep
   splits, fast-moving sheep position jumps — at the cost of a small
   acquisition latency (~50 ms at the default ``consensus_k=3``).
   ``consensus_k=1`` disables the stage.
2. **Constant-velocity prediction**. Each track carries a smoothed
   ``(vx, vy)``. While a track is occluded its position is
   extrapolated for up to ``PREDICT_STEPS`` frames, then falls back to
   last-seen static memory until ``FORGET_STEPS`` deletes it.
3. **Pen latching**. A track whose estimated position crosses the gate
   plane south of ``is_penned_position`` is marked penned, excluded
   from ``get_positions``, and kept indefinitely.

Output of :meth:`SheepTracker.get_positions` is ``{name: (x, y)}`` —
Strömbom, Sequential and the BC observation builder consume it
directly.
"""

from __future__ import annotations

import math
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from herding.config import TrackerConfig

from herding.world.geometry import MAX_SHEEP, in_pen, is_penned_position


GATE_M = 2.5              # m — primary NN gate (recently observed tracks)
REACQUIRE_GATE_M = 4.5    # m — wider gate for re-binding stale tracks
REACQUIRE_MIN_AGE = 20    # steps — track must be this stale to use the wider gate
PENNED_GATE_M = 4.0       # m — gate for matching detections to existing penned tracks
FORGET_STEPS = 200        # ~3.2 s — delete stale active tracks (penned ones kept forever)
MAX_ACTIVE_TRACKS = MAX_SHEEP

# Predictive tracking constants.
PREDICT_STEPS = 120       # ~1.9 s — extrapolate velocity this many frames
VELOCITY_CLAMP = 1.0      # m/s — max predicted speed (sheep max is ~0.78 m/s)


class Track:
    """Single track with position, velocity, and age.

    Attributes
    ----------
    candidate
        ``True`` while the track has not yet accumulated enough
        consensus matches to be visible (``hit_count < consensus_k``).
        Candidates are excluded from :meth:`SheepTracker.get_positions`
        and from the active/penned counters.
    hit_count
        Number of detections this track has absorbed since spawn,
        used by the consensus filter.
    """

    __slots__ = ("x", "y", "vx", "vy", "last_seen", "penned",
                 "candidate", "hit_count")

    def __init__(
        self,
        x: float,
        y: float,
        step: int,
        penned: bool = False,
        candidate: bool = False,
    ):
        self.x = x
        self.y = y
        self.vx = 0.0
        self.vy = 0.0
        self.last_seen = step
        self.penned = penned
        self.candidate = candidate
        self.hit_count = 1

    @property
    def age(self) -> int:
        """Not-a-property in the hot loop — callers pass current step."""
        raise NotImplementedError

    def predicted_position(
        self,
        current_step: int,
        predict_steps: int = PREDICT_STEPS,
        velocity_clamp: float = VELOCITY_CLAMP,
    ) -> tuple[float, float]:
        """Extrapolated position using constant velocity, clamped."""
        dt = current_step - self.last_seen
        if dt <= 0 or dt > predict_steps:
            return self.x, self.y
        speed = math.hypot(self.vx, self.vy)
        if speed < 1e-4:
            return self.x, self.y
        # Clamp extrapolation distance.
        max_d = velocity_clamp * dt * 0.016  # steps → seconds
        d = min(speed * dt * 0.016, max_d)
        return (
            self.x + d * (self.vx / speed),
            self.y + d * (self.vy / speed),
        )

    def update(self, x: float, y: float, step: int) -> None:
        """Absorb a new detection and re-estimate velocity."""
        dt = step - self.last_seen
        if dt > 0:
            dt_s = dt * 0.016  # steps → seconds
            new_vx = (x - self.x) / dt_s
            new_vy = (y - self.y) / dt_s
            # Exponential smoothing on velocity.
            alpha = 0.6
            self.vx = alpha * new_vx + (1.0 - alpha) * self.vx
            self.vy = alpha * new_vy + (1.0 - alpha) * self.vy
        self.x = x
        self.y = y
        self.last_seen = step


class SheepTracker:
    """Online tracker with NN association, prediction, and forgetful memory.

    Each track is a :class:`Track` with position, velocity estimate,
    last-seen step, and penned flag.

    Pass a :class:`~herding.config.TrackerConfig` to override any
    module-level defaults without changing this file.
    """

    def __init__(
        self,
        gate: float = GATE_M,
        tracker_cfg: "TrackerConfig | None" = None,
    ):
        if tracker_cfg is not None:
            self.gate = tracker_cfg.gate_m
            self._reacquire_gate = tracker_cfg.reacquire_gate_m
            self._reacquire_min_age = tracker_cfg.reacquire_min_age
            self._penned_gate = tracker_cfg.penned_gate_m
            self._forget_steps = tracker_cfg.forget_steps
            self._predict_steps = tracker_cfg.predict_steps
            self._velocity_clamp = tracker_cfg.velocity_clamp
            self._max_new_per_step = tracker_cfg.max_new_tracks_per_step
            self._pen_latch_depth = tracker_cfg.pen_latch_depth
            self._consensus_k = tracker_cfg.consensus_k
            self._consensus_radius = tracker_cfg.consensus_radius_m
            self._consensus_max_age = tracker_cfg.consensus_max_age
        else:
            self.gate = gate
            self._reacquire_gate = REACQUIRE_GATE_M
            self._reacquire_min_age = REACQUIRE_MIN_AGE
            self._penned_gate = PENNED_GATE_M
            self._forget_steps = FORGET_STEPS
            self._predict_steps = PREDICT_STEPS
            self._velocity_clamp = VELOCITY_CLAMP
            self._max_new_per_step = MAX_ACTIVE_TRACKS
            self._pen_latch_depth = 0.0
            self._consensus_k = 1
            self._consensus_radius = 0.5
            self._consensus_max_age = 8
        self._tracks: dict[int, Track] = {}
        self._next_id = 0
        self.step = 0

    def reset(self) -> None:
        self._tracks.clear()
        self._next_id = 0
        self.step = 0

    def update(self, detections: list[tuple[float, float]]) -> dict[str, tuple[float, float]]:
        """Fold a new set of detections in and return active positions."""
        self.step += 1

        det_used: set[int] = set()
        updated_tids: set[int] = set()

        # Pass 1 — match promoted active tracks within the primary gate.
        # Use predicted positions for matching, oldest-first. Candidates
        # are excluded; they get their own (tighter) pass below so a
        # stray detection cannot rescue an already-stale candidate.
        active_tids = [tid for tid, t in self._tracks.items()
                       if not t.penned and not t.candidate]
        active_tids.sort(key=lambda tid: self._tracks[tid].last_seen)
        for tid in active_tids:
            track = self._tracks[tid]
            tx, ty = track.predicted_position(
                self.step, self._predict_steps, self._velocity_clamp)
            best_j, best_d = -1, self.gate
            for j, (dx, dy) in enumerate(detections):
                if j in det_used:
                    continue
                d = math.hypot(dx - tx, dy - ty)
                if d < best_d:
                    best_d = d
                    best_j = j
            if best_j >= 0:
                dx, dy = detections[best_j]
                track.update(dx, dy, self.step)
                track.hit_count += 1
                det_used.add(best_j)
                updated_tids.add(tid)

        # Pass 1b — re-acquisition with wider gate for stale tracks.
        for tid in active_tids:
            if tid in updated_tids:
                continue
            track = self._tracks[tid]
            if (self.step - track.last_seen) < self._reacquire_min_age:
                continue
            tx, ty = track.predicted_position(
                self.step, self._predict_steps, self._velocity_clamp)
            best_j, best_d = -1, self._reacquire_gate
            for j, (dx, dy) in enumerate(detections):
                if j in det_used:
                    continue
                d = math.hypot(dx - tx, dy - ty)
                if d < best_d:
                    best_d = d
                    best_j = j
            if best_j >= 0:
                dx, dy = detections[best_j]
                track.update(dx, dy, self.step)
                track.hit_count += 1
                det_used.add(best_j)
                updated_tids.add(tid)

        # Pass 1c — match remaining detections to candidate tracks within
        # the tight consensus radius. Each hit ages the candidate; once
        # hit_count reaches consensus_k it is promoted (handled below).
        candidate_tids = [tid for tid, t in self._tracks.items() if t.candidate]
        candidate_tids.sort(key=lambda tid: self._tracks[tid].last_seen)
        for tid in candidate_tids:
            track = self._tracks[tid]
            best_j, best_d = -1, self._consensus_radius
            for j, (dx, dy) in enumerate(detections):
                if j in det_used:
                    continue
                d = math.hypot(dx - track.x, dy - track.y)
                if d < best_d:
                    best_d = d
                    best_j = j
            if best_j >= 0:
                dx, dy = detections[best_j]
                track.update(dx, dy, self.step)
                track.hit_count += 1
                det_used.add(best_j)

        # Pass 2 — match remaining detections to penned tracks.
        penned_tids = [tid for tid, t in self._tracks.items() if t.penned]
        for tid in penned_tids:
            track = self._tracks[tid]
            best_j, best_d = -1, self._penned_gate
            for j, (dx, dy) in enumerate(detections):
                if j in det_used:
                    continue
                d = math.hypot(dx - track.x, dy - track.y)
                if d < best_d:
                    best_d = d
                    best_j = j
            if best_j >= 0:
                dx, dy = detections[best_j]
                track.update(dx, dy, self.step)
                track.hit_count += 1
                det_used.add(best_j)

        # Spawn tracks for still-unmatched detections.
        #
        # When ``consensus_k > 1`` every new track starts as a candidate
        # and remains invisible to ``get_positions`` until it accumulates
        # the required matches. Penned latching is deferred to after
        # promotion — otherwise gate-area phantoms could still skip the
        # consensus filter by landing inside the pen column and being
        # latched forever, which is exactly the failure mode the filter
        # is meant to eliminate. ``max_new_tracks_per_step`` continues
        # to rate-cap spawns.
        spawned = 0
        spawn_candidates = self._consensus_k > 1
        for j, (dx, dy) in enumerate(detections):
            if j in det_used:
                continue
            if spawned >= self._max_new_per_step:
                break
            if spawn_candidates:
                self._tracks[self._next_id] = Track(
                    dx, dy, self.step, penned=False, candidate=True)
            else:
                penned = self._is_penned(dx, dy)
                self._tracks[self._next_id] = Track(
                    dx, dy, self.step, penned=penned, candidate=False)
            self._next_id += 1
            spawned += 1

        # Promote candidates that have accumulated enough matches.
        for track in self._tracks.values():
            if track.candidate and track.hit_count >= self._consensus_k:
                track.candidate = False

        # Promote active tracks whose current estimate crosses the gate.
        # Candidates are deliberately excluded — a track that hasn't yet
        # earned visibility shouldn't be allowed to latch as penned
        # either (that path is exactly how south-wall FPs persisted
        # forever before the consensus filter existed).
        for track in self._tracks.values():
            if track.penned or track.candidate:
                continue
            px, py = track.predicted_position(
                self.step, self._predict_steps, self._velocity_clamp)
            if self._is_penned(px, py):
                track.penned = True

        # Forget stale tracks. Candidates have their own short timeout
        # (one window to confirm or die); promoted active tracks decay at
        # forget_steps; penned tracks decay 8× slower because real penned
        # sheep are still observed when the dog faces the pen.
        penned_forget = self._forget_steps * 8
        stale: list[int] = []
        for tid, t in self._tracks.items():
            age = self.step - t.last_seen
            if t.candidate:
                if age > self._consensus_max_age:
                    stale.append(tid)
            elif t.penned:
                if age > penned_forget:
                    stale.append(tid)
            else:
                if age > self._forget_steps:
                    stale.append(tid)
        for tid in stale:
            del self._tracks[tid]

        # Hard cap on the visible (promoted, not penned) active set —
        # drop the oldest-seen overflow. Candidates are not counted here:
        # they don't compete for slots until they earn promotion, and
        # rate-limiting their spawn is the job of ``max_new_per_step``.
        active = [(tid, t.last_seen) for tid, t in self._tracks.items()
                  if not t.penned and not t.candidate]
        if len(active) > MAX_ACTIVE_TRACKS:
            active.sort(key=lambda kv: kv[1])
            for tid, _ in active[: len(active) - MAX_ACTIVE_TRACKS]:
                del self._tracks[tid]

        return self.get_positions()

    def _is_penned(self, x: float, y: float) -> bool:
        """Check whether a position should be considered penned.

        Uses ``pen_latch_depth`` to require the position to be that many
        metres past the gate line before latching.  Increasing the depth
        prevents gate-area LiDAR false positives (gate hardware reflections
        at y ≈ -15) from being permanently latched as penned tracks.
        """
        from herding.world.geometry import GATE_Y
        # Apply depth threshold to both in_pen and is_penned_position so
        # that any position in the gate column must clear GATE_Y - depth.
        threshold = GATE_Y - self._pen_latch_depth
        return (in_pen(x, y) or is_penned_position(x, y)) and y <= threshold

    def get_positions(self, min_freshness: int | None = None) -> dict[str, tuple[float, float]]:
        """Promoted (non-candidate, non-penned) tracks as ``{name: (x, y)}``.

        For tracks currently being predicted (occluded but within
        predict_steps), returns the extrapolated position so the teacher
        sees a smooth estimate.

        Candidate tracks — those that have not yet accumulated
        ``consensus_k`` matches — are excluded so a one-shot phantom
        detection never reaches the policy/teacher.

        ``min_freshness`` (optional, deploy-only): drop tracks whose
        last_seen is older than ``step - min_freshness``. Real sheep in
        FOV are detected nearly every step; phantom tracks from sporadic
        Webots FPs stop being re-observed and decay. Default ``None``
        preserves training behaviour (extrapolated tracks visible).
        """
        result = {}
        for tid, track in self._tracks.items():
            if track.penned or track.candidate:
                continue
            if (min_freshness is not None
                    and self.step - track.last_seen > min_freshness):
                continue
            px, py = track.predicted_position(
                self.step, self._predict_steps, self._velocity_clamp)
            result[f"t{tid}"] = (px, py)
        return result

    def get_penned_set(self) -> set[str]:
        return {f"t{tid}" for tid, t in self._tracks.items() if t.penned}

    def n_active(self) -> int:
        """Number of promoted (non-candidate, non-penned) tracks."""
        return sum(1 for t in self._tracks.values()
                   if not t.penned and not t.candidate)

    def n_penned(self) -> int:
        return sum(1 for t in self._tracks.values() if t.penned)

    def n_candidate(self) -> int:
        """Number of unpromoted candidate tracks awaiting consensus."""
        return sum(1 for t in self._tracks.values() if t.candidate)

    def n_predicted(self) -> int:
        """Number of promoted active tracks currently being extrapolated (not directly observed)."""
        return sum(1 for t in self._tracks.values()
                   if not t.penned and not t.candidate
                   and (self.step - t.last_seen) > 0
                   and (self.step - t.last_seen) <= self._predict_steps)