from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import numpy.typing as npt
from scipy.optimize import linear_sum_assignment

from supervision.detection.utils.iou_and_nms import box_iou_batch

if TYPE_CHECKING:
    from supervision.tracker.byte_tracker.single_object_track import STrack


def indices_to_matches(
    cost_matrix: npt.NDArray[np.float32], indices: npt.NDArray[np.int_], thresh: float
) -> tuple[npt.NDArray[np.int_], tuple[int, ...], tuple[int, ...]]:
    matched_cost = cost_matrix[tuple(zip(*indices))]
    matched_mask = matched_cost <= thresh

    matches = indices[matched_mask]
    unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
    unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
    return matches, unmatched_a, unmatched_b


def linear_assignment(
    cost_matrix: npt.NDArray[np.float32], thresh: float
) -> tuple[npt.NDArray[np.int_], tuple[int, ...], tuple[int, ...]]:
    if cost_matrix.size == 0:
        return (
            np.empty((0, 2), dtype=int),
            tuple(range(cost_matrix.shape[0])),
            tuple(range(cost_matrix.shape[1])),
        )

    cost_matrix[cost_matrix > thresh] = thresh + 1e-4
    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    indices = np.column_stack((row_ind, col_ind))

    return indices_to_matches(cost_matrix, indices, thresh)


def iou_distance(
    atracks: list[STrack] | list[npt.NDArray[np.float32]],
    btracks: list[STrack] | list[npt.NDArray[np.float32]],
) -> npt.NDArray[np.float32]:
    if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
        len(btracks) > 0 and isinstance(btracks[0], np.ndarray)
    ):
        atlbrs = atracks
        btlbrs = btracks
    else:
        atlbrs = [track.tlbr for track in atracks]
        btlbrs = [track.tlbr for track in btracks]

    _ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
    if _ious.size != 0:
        _ious = box_iou_batch(np.asarray(atlbrs), np.asarray(btlbrs))
    cost_matrix = 1 - _ious

    return cost_matrix


def fuse_score(
    cost_matrix: npt.NDArray[np.float32], stracks: list[STrack]
) -> npt.NDArray[np.float32]:
    if cost_matrix.size == 0:
        return cost_matrix
    iou_sim = 1 - cost_matrix
    det_scores = np.array([strack.score for strack in stracks])
    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
    fuse_sim = iou_sim * det_scores
    fuse_cost = 1 - fuse_sim
    return fuse_cost