"""
tracker.py — Single-target person tracker for saltybot person-following mode.

Strategy:
  - Select closest valid detection each frame (smallest depth within range)
  - Re-associate across frames using IoU matching with existing track
  - Hold last known track for `hold_duration` seconds when detections are lost
  - Assign monotonically increasing track IDs

Usage:
    tracker = SimplePersonTracker(hold_duration=2.0)
    track = tracker.update(detections)   # detections: list of (bbox, depth, conf)
    if track is not None:
        print(track.bbox, track.depth, track.track_id)
"""

import time
import numpy as np


class PersonTrack:
    """Single tracked person."""

    def __init__(self, bbox, depth, confidence, track_id):
        """
        Args:
            bbox: (x1, y1, x2, y2) in pixels
            depth: distance to person in metres
            confidence: detection confidence 0–1
            track_id: unique integer ID
        """
        self.bbox = bbox
        self.depth = depth
        self.confidence = confidence
        self.track_id = track_id
        self.position_3d = None       # (X, Y, Z) in camera frame — set by node
        self._first_seen = time.monotonic()
        self._last_seen = time.monotonic()

    def touch(self, bbox, depth, confidence):
        self.bbox = bbox
        self.depth = depth
        self.confidence = confidence
        self._last_seen = time.monotonic()

    @property
    def age(self):
        """Seconds since last detection."""
        return time.monotonic() - self._last_seen

    @property
    def is_stale(self):
        return self.age > 0     # Always check externally against hold_duration

    @property
    def center(self):
        x1, y1, x2, y2 = self.bbox
        return ((x1 + x2) / 2.0, (y1 + y2) / 2.0)

    @property
    def area(self):
        x1, y1, x2, y2 = self.bbox
        return max(0.0, (x2 - x1) * (y2 - y1))


class SimplePersonTracker:
    """
    Lightweight single-target person tracker.

    Maintains one active PersonTrack at a time. On each update:
      1. Filter detections to valid depth range.
      2. If a track is active, attempt IoU re-association.
      3. If re-association fails and track is within hold_duration, keep stale.
      4. If no active track, initialise from closest detection.
    """

    def __init__(
        self,
        hold_duration: float = 2.0,
        iou_threshold: float = 0.25,
        min_depth: float = 0.3,
        max_depth: float = 5.0,
    ):
        """
        Args:
            hold_duration: seconds to hold last known position after losing track
            iou_threshold: minimum IoU to accept a re-association
            min_depth: minimum valid depth in metres
            max_depth: maximum valid depth in metres
        """
        self._hold_duration = hold_duration
        self._iou_threshold = iou_threshold
        self._min_depth = min_depth
        self._max_depth = max_depth
        self._track: PersonTrack | None = None
        self._next_id = 1

    def update(self, detections):
        """
        Update tracker with new detections.

        Args:
            detections: list of (bbox, depth, confidence) where
                        bbox = (x1, y1, x2, y2) pixels,
                        depth = float metres (0 = invalid),
                        confidence = float 0–1

        Returns:
            PersonTrack or None.
            Returns a stale PersonTrack (track.age > 0) during hold period.
            Returns None after hold_duration or if never seen a person.
        """
        valid = [
            (b, d, c) for b, d, c in detections
            if self._min_depth < d < self._max_depth
        ]

        if not valid:
            # No valid detections this frame
            if self._track is not None:
                if self._track.age <= self._hold_duration:
                    return self._track      # hold last known
                else:
                    self._track = None
            return None

        if self._track is not None:
            # Attempt re-association by IoU
            best_iou = self._iou_threshold
            best = None
            for bbox, depth, conf in valid:
                iou = _iou(bbox, self._track.bbox)
                if iou > best_iou:
                    best_iou = iou
                    best = (bbox, depth, conf)

            if best is not None:
                self._track.touch(*best)
                return self._track

            # No IoU match — keep stale track within hold window
            if self._track.age <= self._hold_duration:
                return self._track
            else:
                self._track = None  # lost — start fresh

        # No active track: pick closest valid detection
        bbox, depth, conf = min(valid, key=lambda x: x[1])
        self._track = PersonTrack(bbox, depth, conf, self._next_id)
        self._next_id += 1
        return self._track

    def reset(self):
        """Drop current track."""
        self._track = None

    @property
    def active(self):
        return self._track is not None


def _iou(bbox_a, bbox_b):
    """Compute Intersection-over-Union of two bounding boxes (x1,y1,x2,y2)."""
    ax1, ay1, ax2, ay2 = bbox_a
    bx1, by1, bx2, by2 = bbox_b

    ix1 = max(ax1, bx1)
    iy1 = max(ay1, by1)
    ix2 = min(ax2, bx2)
    iy2 = min(ay2, by2)

    if ix2 <= ix1 or iy2 <= iy1:
        return 0.0

    intersection = (ix2 - ix1) * (iy2 - iy1)
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    union = area_a + area_b - intersection
    return intersection / union if union > 0 else 0.0
