"""Person tracker for the on-device detector.

Dependency-free on purpose: detect_onnx.py cannot be imported without cv2 +
onnxruntime, so keeping the tracker here is what makes it testable
(`python detector/test_tracker.py`).

Why this exists in its current shape
------------------------------------
Footfall ("visitors") is the count of DISTINCT people who arrived, and the
detector derives it from the tracker: one new person = one arrival. The original
tracker counted an arrival the first frame a box appeared, and matched boxes to
tracks by IOU alone. Both are wrong for a real camera:

* Inference runs at ~3 fps. A person walking across the frame moves further than
  `iou_thresh` between two inferences, so the box stops overlapping its own
  track, a brand-new id is minted, and the same person is counted again. On a
  driveway camera this inflated a handful of real people into 396 "visitors".
* A single-frame false positive (a shadow, a passing car mirror) also minted an
  id and was counted immediately.

So: match by IOU first, then fall back to centre distance for boxes that moved
too far to overlap, and only count an arrival once a track has been seen for
`min_hits` inference frames. `people_in_frame` still includes unconfirmed tracks,
so the live count stays responsive; only footfall waits for confirmation.
"""

def iou(a, b):
    ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
    ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
    inter = iw * ih
    ua = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / ua if ua > 0 else 0.0


def _center(b):
    return ((b[0] + b[2]) / 2.0, (b[1] + b[3]) / 2.0)


def _diag(b):
    w, h = b[2] - b[0], b[3] - b[1]
    return (w * w + h * h) ** 0.5


def _area(b):
    return max(1.0, float((b[2] - b[0]) * (b[3] - b[1])))


class Track:
    __slots__ = ("id", "box", "missed", "matched", "hits", "confirmed")

    def __init__(self, tid, box):
        self.id = tid
        self.box = box
        self.missed = 0
        self.matched = True
        self.hits = 1
        self.confirmed = False


class IouTracker:
    """Greedy IOU tracker with a centre-distance fallback and hit confirmation.

    `min_hits`     : inference frames a track must be seen before it counts as an
                     arrival. 1 restores the old count-on-first-sight behaviour.
    `center_gate`  : how far a box may travel between frames and still be matched
                     to the same track, as a fraction of the mean box diagonal.
    `size_ratio`   : reject a centre match whose box area changed more than this,
                     so a person is never matched onto a very different blob.
    """

    def __init__(self, iou_thresh=0.3, max_missed=30, min_hits=3,
                 center_gate=0.6, size_ratio=3.0):
        self.tracks = []
        self.next_id = 1
        self.iou_thresh = iou_thresh
        self.max_missed = max_missed
        self.min_hits = max(1, int(min_hits))
        self.center_gate = center_gate
        self.size_ratio = size_ratio
        # Track ids that reached `min_hits` on the most recent update(). The
        # caller turns these into footfall arrivals (after its ignore-mask check).
        self.newly_confirmed = set()

    def _center_score(self, b, t):
        """0..1 similarity for a box that failed the IOU gate, or None to reject."""
        bc, tc = _center(b), _center(t.box)
        dist = ((bc[0] - tc[0]) ** 2 + (bc[1] - tc[1]) ** 2) ** 0.5
        limit = self.center_gate * ((_diag(b) + _diag(t.box)) / 2.0)
        if limit <= 0 or dist > limit:
            return None
        ratio = _area(b) / _area(t.box)
        if ratio > self.size_ratio or ratio < 1.0 / self.size_ratio:
            return None
        return 1.0 - (dist / limit)

    def _greedy(self, candidates, used_boxes, used_tracks, box_to_track):
        candidates.sort(key=lambda p: p[0], reverse=True)
        for _, bi, ti in candidates:
            if bi in used_boxes or ti in used_tracks:
                continue
            used_boxes.add(bi)
            used_tracks.add(ti)
            box_to_track[bi] = ti

    def update(self, boxes):
        self.newly_confirmed = set()
        for t in self.tracks:
            t.matched = False

        used_boxes, used_tracks, box_to_track = set(), set(), {}

        # Pass 1: overlap. Strong evidence, so it always wins a contested box.
        iou_pairs = []
        for bi, b in enumerate(boxes):
            for ti, t in enumerate(self.tracks):
                v = iou(b, t.box)
                if v >= self.iou_thresh:
                    iou_pairs.append((v, bi, ti))
        self._greedy(iou_pairs, used_boxes, used_tracks, box_to_track)

        # Pass 2: proximity, for a person who moved too far between inferences to
        # still overlap. Without this the tracker mints a new id (a new "arrival")
        # every few frames for anyone actually walking.
        near_pairs = []
        for bi, b in enumerate(boxes):
            if bi in used_boxes:
                continue
            for ti, t in enumerate(self.tracks):
                if ti in used_tracks:
                    continue
                s = self._center_score(b, t)
                if s is not None:
                    near_pairs.append((s, bi, ti))
        self._greedy(near_pairs, used_boxes, used_tracks, box_to_track)

        out, new_tracks = [], []
        for bi, b in enumerate(boxes):
            if bi in box_to_track:
                t = self.tracks[box_to_track[bi]]
                t.box, t.missed, t.matched = b, 0, True
                t.hits += 1
                if not t.confirmed and t.hits >= self.min_hits:
                    t.confirmed = True
                    self.newly_confirmed.add(t.id)
                out.append((t.id, b))
            else:
                t = Track(self.next_id, b)
                self.next_id += 1
                if t.hits >= self.min_hits:  # min_hits == 1: confirm on sight
                    t.confirmed = True
                    self.newly_confirmed.add(t.id)
                new_tracks.append(t)
                out.append((t.id, b))

        survivors = []
        for t in self.tracks:
            if t.matched:
                survivors.append(t)
            else:
                t.missed += 1
                if t.missed <= self.max_missed:
                    survivors.append(t)
        self.tracks = survivors + new_tracks
        return out
