from typing import Sequence

from kinematic_tracker.core.loose_items import get_loose_indices
from kinematic_tracker.types import FMT, IVT

from .track import NdKkfTrack


def upd_for_loose_tracks(
    tracks: list[NdKkfTrack], targets: IVT, num_tracks_before: int, num_misses_max: int
) -> None:
    """Update the list of tracks for loose tracks (i.e., tracks without association).

    Args:
        tracks: list of tracks to be updated.
        targets: Indices of associated tracks.
        num_tracks_before: Number of tracks before association.
        num_misses_max: Number of misses to tolerate before removal of a track.
    """
    tracks_wo_correspondence = get_loose_indices(targets, num_tracks_before)
    for t in tracks_wo_correspondence:
        tracks[t].bump_num_misses()

    to_remove = [t for t in tracks if t.num_miss > num_misses_max]
    for t in to_remove:
        tracks.remove(t)


# fmt: off
def correct_associated(tracks: Sequence[NdKkfTrack], reports: IVT, targets: IVT, metric_rt: FMT, det_rz: FMT, ids_r: IVT) -> None:
# fmt: on
    """Correct tracks based on the given association matches.

    Args:
        tracks: Tracks to be eventually corrected.
        reports: Indices of associated detection reports.
        targets: Indices of associated tracked targets.
        metric_rt: Metric values for report-target pairs.
        det_rz: Detection vectors to be used eventually for correction step.
        ids_r: Detection IDs to be used eventually for assigning the update IDs.
    """
    for r, t in zip(reports, targets):
        track = tracks[t]
        track.kkf.correct(det_rz[r])
        track.bump_num_detections(float(metric_rt[r, t]), int(ids_r[r]))
