# ------------------------------------------------------------------------
# Trackers
# Copyright (c) 2026 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------

"""Unit tests for deterministic unmatched-index ordering in _get_associated_indices.

All trackers must return sorted lists for unmatched tracks and unmatched detections so that tracker-ID assignment is
stable across CPython versions and implementations.
"""

from __future__ import annotations

from collections.abc import Callable
from unittest.mock import MagicMock

import numpy as np
import pytest
import supervision as sv

from trackers.core.botsort.tracker import BoTSORTTracker
from trackers.core.bytetrack.tracker import ByteTrackTracker
from trackers.core.ocsort.tracker import OCSORTTracker
from trackers.core.sort.tracker import SORTTracker


def _call_sort(
    n_tracks: int, n_detections: int, iou_matrix: np.ndarray
) -> tuple[list[tuple[int, int]], list[int], list[int]]:
    """Call SORTTracker._get_associated_indices with n_tracks dummy tracks."""
    tracker = SORTTracker()
    tracker.tracks = [None] * n_tracks  # type: ignore[list-item]
    detection_boxes = np.zeros((n_detections, 4))
    return tracker._get_associated_indices(iou_matrix, detection_boxes)


def _call_bytetrack(
    n_tracks: int, n_detections: int, iou_matrix: np.ndarray
) -> tuple[list[tuple[int, int]], list[int], list[int]]:
    """Call ByteTrackTracker._get_associated_indices with the given matrix."""
    tracker = ByteTrackTracker()
    return tracker._get_associated_indices(iou_matrix, tracker.minimum_iou_threshold)


def _call_ocsort(
    n_tracks: int, n_detections: int, iou_matrix: np.ndarray
) -> tuple[list[tuple[int, int]], list[int], list[int]]:
    """Call OCSORTTracker._get_associated_indices with zeros for direction matrix."""
    tracker = OCSORTTracker()
    direction_matrix = np.zeros_like(iou_matrix)
    return tracker._get_associated_indices(iou_matrix, direction_matrix)


def _call_botsort(
    n_tracks: int, n_detections: int, iou_matrix: np.ndarray
) -> tuple[list[tuple[int, int]], list[int], list[int]]:
    """Call BoTSORTTracker._get_associated_indices with default similarity threshold."""
    tracker = BoTSORTTracker(enable_cmc=False)
    matched, unmatched_tracks, unmatched_detections = tracker._get_associated_indices(iou_matrix, 0.1)
    return matched, unmatched_tracks, unmatched_detections


_TRACKER_CALL_FNS: list[
    pytest.param  # type: ignore[type-arg]
] = [
    pytest.param(_call_sort, id="sort"),
    pytest.param(_call_bytetrack, id="bytetrack"),
    pytest.param(_call_ocsort, id="ocsort"),
    pytest.param(_call_botsort, id="botsort"),
]

CallFn = Callable[[int, int, np.ndarray], tuple[list[tuple[int, int]], list[int], list[int]]]


class TestGetAssociatedIndicesSortedOutput:
    """_get_associated_indices returns sorted unmatched lists across all trackers."""

    @pytest.mark.parametrize("call_fn", _TRACKER_CALL_FNS)
    def test_all_unmatched_detections_are_sorted(self, call_fn: CallFn) -> None:
        """All detections are unmatched when no tracks exist; result is sorted."""
        n_detections = 4
        iou_matrix = np.zeros((0, n_detections))
        _, unmatched_tracks, unmatched_detections = call_fn(0, n_detections, iou_matrix)
        assert unmatched_tracks == []
        assert unmatched_detections == list(range(n_detections))

    @pytest.mark.parametrize("call_fn", _TRACKER_CALL_FNS)
    def test_non_contiguous_unmatched_indices_are_sorted(self, call_fn: CallFn) -> None:
        """Non-contiguous unmatched indices are returned in ascending order.

        Scenario: 4 tracks, 5 detections; track 1 matches detection 2 (IoU=0.9).
        Expected unmatched_tracks  = [0, 2, 3]  (not {0, 2, 3} in arbitrary set order).
        Expected unmatched_detections = [0, 1, 3, 4].
        """
        n_tracks, n_detections = 4, 5
        iou_matrix = np.zeros((n_tracks, n_detections))
        iou_matrix[1, 2] = 0.9  # only match: track 1 ↔ detection 2

        _, unmatched_tracks, unmatched_detections = call_fn(n_tracks, n_detections, iou_matrix)

        assert unmatched_tracks == [0, 2, 3]
        assert unmatched_detections == [0, 1, 3, 4]
        assert unmatched_tracks == sorted(unmatched_tracks)
        assert unmatched_detections == sorted(unmatched_detections)


def test_ocsort_none_direction_matrix_matches_explicit_zero_matrix() -> None:
    """OCSORT weight-0 path (``None`` direction matrix) associates identically to a zero matrix.

    The direction-consistency matrix is finite and bounded, so the historical ``iou_matrix + weight * matrix`` reduced
    to ``iou_matrix`` whenever the matrix was all zeros. Passing ``None`` (the weight-0 early-out) must be bit-identical
    to passing an explicit zero matrix. ``None`` is only valid at weight 0, so the tracker is constructed with
    ``direction_consistency_weight=0``.
    """
    tracker = OCSORTTracker(direction_consistency_weight=0.0)
    iou_matrix = np.zeros((4, 5))
    iou_matrix[1, 2] = 0.9
    iou_matrix[3, 0] = 0.5

    with_none = tracker._get_associated_indices(iou_matrix, None)
    with_zeros = tracker._get_associated_indices(iou_matrix, np.zeros_like(iou_matrix))

    assert with_none == with_zeros


def test_ocsort_none_direction_matrix_with_nonzero_weight_raises() -> None:
    """Passing ``None`` while direction_consistency_weight != 0 raises — None is the weight-0 sentinel only."""
    tracker = OCSORTTracker(direction_consistency_weight=0.2)
    iou_matrix = np.zeros((2, 3))
    iou_matrix[0, 1] = 0.9

    with pytest.raises(ValueError, match="only valid when the weight is 0"):
        tracker._get_associated_indices(iou_matrix, None)


@pytest.mark.parametrize("call_fn", _TRACKER_CALL_FNS)
def test_all_trackers_associated_indices_are_deterministically_sorted(
    call_fn: CallFn,
) -> None:
    """All trackers return stable ascending unmatched indices on same pattern."""
    n_tracks, n_detections = 4, 5
    similarity_matrix = np.zeros((n_tracks, n_detections), dtype=np.float32)
    similarity_matrix[1, 2] = 0.9

    _, unmatched_tracks, unmatched_detections = call_fn(n_tracks, n_detections, similarity_matrix)

    assert unmatched_tracks == [0, 2, 3]
    assert unmatched_detections == [0, 1, 3, 4]


@pytest.mark.parametrize(
    ("weight", "expect_matrix_built"),
    [
        pytest.param(0.0, False, id="weight-zero-skips-matrix"),
        pytest.param(0.2, True, id="weight-nonzero-builds-matrix"),
    ],
)
def test_ocsort_direction_matrix_built_only_when_weight_nonzero(
    weight: float, expect_matrix_built: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
    """OCSORT builds the direction-consistency matrix iff direction_consistency_weight != 0.

    Exercises the weight-0 early-out itself: with weight 0 the matrix build is
    skipped entirely (not merely cancelled by an explicit zero matrix), while a
    nonzero weight still triggers it. A track is established on the first frame so
    the matrix would be built on the second frame whenever the weight allows it.
    """
    tracker = OCSORTTracker(direction_consistency_weight=weight)
    detections = sv.Detections(
        xyxy=np.array([[10.0, 10.0, 50.0, 50.0]], dtype=np.float64),
        confidence=np.array([0.9]),
    )
    tracker.update(detections)  # frame 1: establish a track
    spy = MagicMock(wraps=tracker._compute_direction_consistency_matrix)
    monkeypatch.setattr(tracker, "_compute_direction_consistency_matrix", spy)

    tracker.update(detections)  # frame 2: track present; matrix built iff weight != 0

    assert spy.called is expect_matrix_built
