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

import logging
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import cast

import numpy as np
import supervision as sv
from scipy.optimize import linear_sum_assignment

from trackers.core.base import BaseTracker
from trackers.core.botsort.utils import _fuse_score
from trackers.core.masks.base import MaskOutput, TrackletSnapshot
from trackers.core.masks.manager import (
    MASK_CREATION_BBOX_OVERLAP_THRESHOLD,
    MaskManager,
)
from trackers.core.mcbyte.mask_association import (
    MINIMUM_MASK_AVERAGE_CONFIDENCE,
    MINIMUM_MASK_COVERAGE,
    MINIMUM_MASK_FILL_RATIO,
    condition_similarity_with_masks,
)
from trackers.core.mcbyte.tracklet import McByteTracklet
from trackers.core.mcbyte.utils import _get_alive_tracklets
from trackers.utils.cmc import CMC, CMCConfig, CMCMethod
from trackers.utils.detections import default_confidences
from trackers.utils.iou import BaseIoU, IoU
from trackers.utils.state_representations import (
    BaseStateEstimator,
    XCYCWHStateEstimator,
)

logger = logging.getLogger(__name__)

# Number of consecutive out-of-memory failures (CUDA or MPS) in the mask
# pipeline after which mask-conditioned association is disabled for the
# remainder of the run.
_MAX_CONSECUTIVE_MASK_FAILURES = 3

# Detections with confidence at or below this floor are discarded entirely.
# Those above it but below ``high_conf_det_threshold`` are treated as
# low-confidence and used only in the second association stage.
_MINIMUM_DETECTION_CONFIDENCE = 0.1


def _is_out_of_memory(exc: BaseException) -> bool:
    """Return whether an exception represents a torch out-of-memory error.

    Detected without importing ``torch`` (an optional McByte dependency): a
    ``torch.cuda.OutOfMemoryError`` is a ``RuntimeError`` subclass, so it is
    recognised by class name, and plain out-of-memory ``RuntimeError``
    instances are recognised by their canonical message. This also covers the
    MPS backend, whose OOM is a plain ``RuntimeError`` with an
    ``"MPS backend out of memory"`` message.

    Args:
        exc: Exception raised by the mask pipeline.

    Returns:
        ``True`` when ``exc`` is a torch out-of-memory error (CUDA or MPS),
        otherwise ``False``.
    """
    if any(klass.__name__ == "OutOfMemoryError" for klass in type(exc).__mro__):
        return True
    return "out of memory" in str(exc).lower()


@dataclass(frozen=True)
class McByteMaskConfig:
    """Configuration for McByte's SAM and Cutie mask pipeline.

    The configuration is used only when ``McByteTracker`` automatically creates
    its default real ``MaskManager``. It is ignored when a custom manager is
    supplied directly.

    Args:
        device: Device shared by SAM and Cutie, for example ``"cuda"``,
            ``"cuda:0"``, ``"mps"``, or ``"cpu"``. The default ``"auto"``
            resolves to CUDA when available, otherwise CPU. Apple MPS is
            never auto-selected (measured roughly an order of magnitude
            slower than CPU for this pipeline); pass ``device="mps"``
            explicitly to use it.
        sam_checkpoint_path: Optional SAM checkpoint path. When omitted, the
            default checkpoint for ``sam_model_type`` is used and downloaded
            automatically when necessary.
        sam_model_type: SAM model variant used for box-prompted mask generation.
        cutie_weights_path: Optional Cutie checkpoint path. When omitted, the
            default checkpoint for ``cutie_model_type`` is used and downloaded
            automatically when necessary.
        cutie_model_type: Cutie model variant used for temporal propagation.
        cutie_config_path: Optional Cutie Hydra configuration directory. When
            omitted, it is inferred from the installed Cutie package.
        cutie_config_name: Hydra configuration name loaded by Cutie.
        cutie_use_amp: Whether Cutie may use automatic mixed precision. AMP is
            activated only when Cutie runs on a CUDA device. Disabled by
            default so that default runs use full fp32 precision on every
            backend; opt in explicitly after validating tracking-quality
            parity on your hardware.
        cutie_max_internal_size: Maximum shortest side of frames processed
            internally by Cutie. Larger frames are downscaled before the
            encoder and the propagated masks are restored to the original
            resolution, matching Cutie's own streaming preset. Use ``-1`` to
            propagate at full input resolution (Cutie's offline-benchmark
            behavior; substantially slower on high-resolution streams).
        cutie_mem_every: How often, in frames, Cutie updates its working
            memory. Higher values speed up processing. ``None`` keeps the
            value from the loaded Cutie configuration.
        cutie_use_long_term: Whether Cutie uses bounded long-term memory,
            recommended for videos longer than roughly one minute. ``None``
            keeps the value from the loaded Cutie configuration.
        cutie_channels_last: Opt-in ``channels_last`` memory format for the
            Cutie model. Off by default so default runs are unchanged; it may
            alter kernel selection, so validate fp32 tracking-quality parity on
            your backend before enabling. Primarily helps CUDA.
        cutie_compile: Opt-in ``torch.compile`` of Cutie's shape-stable
            per-frame encoder path. Off by default; incurs first-call warmup
            and may alter numerics, so validate fp32 parity before enabling.
            ``torch.compile`` support on MPS is experimental.
        mask_creation_bbox_overlap_threshold: Bounding-box overlap fraction at
            or above which mask creation is delayed by ``MaskManager``.
    """

    device: str = "auto"

    sam_checkpoint_path: str | Path | None = None
    sam_model_type: str = "vit_b"

    cutie_weights_path: str | Path | None = None
    cutie_model_type: str = "base-mega"
    cutie_config_path: str | Path | None = None
    cutie_config_name: str = "eval_config"
    cutie_use_amp: bool = False
    cutie_max_internal_size: int = 480
    cutie_mem_every: int | None = 10
    cutie_use_long_term: bool | None = True
    cutie_channels_last: bool = False
    cutie_compile: bool = False

    mask_creation_bbox_overlap_threshold: float = MASK_CREATION_BBOX_OVERLAP_THRESHOLD


def _build_default_mask_manager(
    config: McByteMaskConfig,
) -> MaskManager:
    """Create McByte's standard SAM + Cutie mask-management pipeline."""
    from trackers.core.masks.cutie import CutieMaskPropagator
    from trackers.core.masks.sam import SAMBoxMaskGenerator

    mask_generator = SAMBoxMaskGenerator(
        checkpoint_path=config.sam_checkpoint_path,
        model_type=config.sam_model_type,
        device=config.device,
    )

    mask_propagator = CutieMaskPropagator(
        weights_path=config.cutie_weights_path,
        model_type=config.cutie_model_type,
        config_path=config.cutie_config_path,
        config_name=config.cutie_config_name,
        device=config.device,
        use_amp=config.cutie_use_amp,
        max_internal_size=config.cutie_max_internal_size,
        mem_every=config.cutie_mem_every,
        use_long_term=config.cutie_use_long_term,
        channels_last=config.cutie_channels_last,
        compile_model=config.cutie_compile,
    )

    return MaskManager(
        mask_generator=mask_generator,
        mask_propagator=mask_propagator,
        mask_creation_bbox_overlap_threshold=(config.mask_creation_bbox_overlap_threshold),
    )


class McByteTracker(BaseTracker):
    """McByte multi-object tracker with optional mask-conditioned association.

    McByte extends a ByteTrack-style multi-stage tracking pipeline with
    clear-match locking, reduced assignment, optional camera motion
    compensation, and optional propagated-mask evidence.

    The tracker can operate in two configurations:

    - without a ``MaskManager``, association uses the McByte clear-match locking
      and reduced-assignment procedure with IoU-based similarities;
    - with a ``MaskManager`` (full McByte), masks are additionally used to condition
      ambiguous associations and, when enabled, isolated positive-IoU associations
      below the normal stage threshold.

    When ``enable_mask_manager=True``, the default mask pipeline initializes
    masks from detection boxes using SAM and propagates them temporally using
    Cutie. A custom ``MaskManager`` may instead be supplied directly, for
    example to inject alternative mask components or lightweight test doubles.

    Mask processing follows the original McByte timing. At frame ``t``, masks
    are updated before association using the frame, visible tracklets, newly
    created tracklets, and removed-tracklet events stored after processing frame
    ``t - 1``. Temporarily lost but still active tracklets retain their masks.
    Masks are removed only after the corresponding tracklets are terminated
    during tracker pruning.

    Input frames are expected in RGB channel order. A frame is required when
    mask management is enabled and is also needed for camera motion
    compensation. When no frame is supplied, those frame-dependent operations
    are skipped.

    Args:
        lost_track_buffer: Time buffer, expressed as a number of frames at
            30 FPS, for retaining unmatched tracks before deletion. The value
            is scaled according to ``frame_rate``.
        frame_rate: Sequence frame rate used to scale ``lost_track_buffer``.
        track_activation_threshold: Minimum detection confidence required to
            create a new tracklet.
        minimum_consecutive_frames: Number of successful tracklet updates
            required before assigning a confirmed non-negative tracker ID.
        minimum_iou_threshold_first_assoc: Minimum association similarity for
            matching high-confidence detections to confirmed and lost tracks.
            The default of ``0.1`` follows ByteTrack's deliberately low
            first-association threshold: a broad candidate set is admitted and
            resolved by fused IoU and detection score. In the default mode
            (``enable_mask_manager=False``) this ByteTrack parity is the sole
            safety net. When mask management is enabled, the same broad
            candidate set additionally lets mask-conditioned association resolve
            ambiguities and optional isolations.
        minimum_iou_threshold_second_assoc: Minimum association similarity for
            matching low-confidence detections to remaining tracked tracks.
        minimum_iou_threshold_unconfirmed_assoc: Minimum association similarity
            for matching unconfirmed tracks to remaining high-confidence
            detections.
        high_conf_det_threshold: Confidence threshold separating high- and
            low-confidence detections. Detections with confidence at or below
            0.1 are discarded.
        enable_cmc: Whether to apply camera motion compensation before
            association.
        cmc_method: Camera motion compensation method.
        cmc_downscale: Image downscale factor used during camera motion
            estimation. Defaults to 6 based on the 1280x720 SportsMOT
            validation aggregate-performance criterion; it is not a strict
            per-clip guarantee. Tune it for other workloads or pass 2 to
            preserve the previous conservative behavior.
        instant_first_frame_activation: Whether tracklets created on the first
            frame receive confirmed tracker IDs immediately.
        state_estimator_class: State estimator class used by newly created
            ``McByteTracklet`` instances.
        iou: IoU implementation used to compute association similarities. When
            omitted, the default ``IoU`` implementation is used.
        enable_mask_manager: Whether to construct McByte's default SAM and
            Cutie mask pipeline. It is disabled by default to avoid loading
            optional heavyweight models when mask-conditioned tracking is not
            requested.
        mask_manager: Optional custom ``MaskManager``. When supplied, it is used
            directly regardless of ``enable_mask_manager``, and automatic
            SAM/Cutie construction is skipped.
        mask_config: Configuration for automatic construction of the default
            SAM/Cutie pipeline. It requires ``enable_mask_manager=True`` and
            cannot be combined with a custom ``mask_manager``.
        minimum_mask_average_confidence: Minimum average confidence of a
            propagated mask before it may influence association.
        minimum_mask_coverage: Minimum fraction of the visible tracklet mask
            that must lie inside a candidate detection box.
        minimum_mask_fill_ratio: Minimum fraction of a candidate detection-box
            area that must be occupied by the tracklet mask.
        enable_isolated_mask_matching: Whether mask evidence may rescue an
            isolated candidate with positive IoU whose association similarity
            is below the normal stage threshold.
        minimum_mask_creation_frames: Number of consecutive frames a confirmed
            tracklet must remain visible in tracker output before its mask is
            created (the SAM prompt plus Cutie ``add_masks``). This defers the
            per-appearance mask encode for very-short-lived tracklets: those
            that terminate before reaching the threshold never pay the encode,
            at the cost of running IoU-only association for those tracklets
            until their mask exists. Because mask conditioning is deferred, this
            can alter tracking output and must be validated for CLEAR/HOTA/
            Identity parity on the target workload. Use ``1`` to create masks on
            a tracklet's first visible frame (the original immediate-creation
            timing). A deferred tracklet is withheld from the mask pipeline
            entirely, including Cutie's initial mask set: when every confirmed
            tracklet is still inside its defer window the first masks are
            produced only once at least one tracklet reaches the threshold.
    """

    tracker_id = "mcbyte"

    def __init__(
        self,
        lost_track_buffer: int = 30,
        frame_rate: float = 30.0,
        track_activation_threshold: float = 0.7,
        minimum_consecutive_frames: int = 2,
        minimum_iou_threshold_first_assoc: float = 0.1,
        minimum_iou_threshold_second_assoc: float = 0.5,
        minimum_iou_threshold_unconfirmed_assoc: float = 0.3,
        high_conf_det_threshold: float = 0.6,
        enable_cmc: bool = True,
        cmc_method: CMCMethod = "sparseOptFlow",
        cmc_downscale: int = 6,
        instant_first_frame_activation: bool = True,
        state_estimator_class: type[BaseStateEstimator] = XCYCWHStateEstimator,
        iou: BaseIoU | None = None,
        enable_mask_manager: bool = False,
        mask_manager: MaskManager | None = None,
        mask_config: McByteMaskConfig | None = None,
        minimum_mask_average_confidence: float = MINIMUM_MASK_AVERAGE_CONFIDENCE,
        minimum_mask_coverage: float = MINIMUM_MASK_COVERAGE,
        minimum_mask_fill_ratio: float = MINIMUM_MASK_FILL_RATIO,
        enable_isolated_mask_matching: bool = False,
        minimum_mask_creation_frames: int = 3,
    ) -> None:
        # Calculate maximum frames without update based on lost_track_buffer and
        # frame_rate. This scales the buffer based on the frame rate to ensure
        # consistent time-based tracking across different frame rates.
        self.maximum_frames_without_update = self._compute_maximum_frames_without_update(
            lost_track_buffer=lost_track_buffer,
            frame_rate=frame_rate,
        )
        self.maximum_time_without_update: float = lost_track_buffer / 30.0
        self.minimum_consecutive_frames = minimum_consecutive_frames
        if minimum_mask_creation_frames < 1:
            raise ValueError(
                "minimum_mask_creation_frames must be at least 1 "
                f"(1 creates masks immediately). Got {minimum_mask_creation_frames}."
            )
        self.minimum_mask_creation_frames = minimum_mask_creation_frames
        self.minimum_iou_threshold_first_assoc = minimum_iou_threshold_first_assoc
        self.minimum_iou_threshold_second_assoc = minimum_iou_threshold_second_assoc
        self.minimum_iou_threshold_unconfirmed_assoc = minimum_iou_threshold_unconfirmed_assoc
        self.track_activation_threshold = track_activation_threshold
        self.high_conf_det_threshold = high_conf_det_threshold
        self.instant_first_frame_activation = instant_first_frame_activation
        self.minimum_mask_average_confidence = minimum_mask_average_confidence
        self.minimum_mask_coverage = minimum_mask_coverage
        self.minimum_mask_fill_ratio = minimum_mask_fill_ratio
        self.enable_isolated_mask_matching = enable_isolated_mask_matching
        if high_conf_det_threshold <= _MINIMUM_DETECTION_CONFIDENCE:
            raise ValueError(
                "high_conf_det_threshold must be greater than the detection discard "
                f"floor ({_MINIMUM_DETECTION_CONFIDENCE}); otherwise detections at or "
                f"below the floor would be treated as high-confidence. "
                f"Got {high_conf_det_threshold}."
            )
        self.tracks: list[McByteTracklet] = []
        self.state_estimator_class = state_estimator_class
        self.iou = iou if iou is not None else IoU()
        self.frame_id: int = 0
        self._reset_id_allocator()

        self.enable_cmc = enable_cmc
        self.cmc = CMC(CMCConfig(method=cmc_method, downscale=cmc_downscale)) if enable_cmc else None

        self._init_timestamp_state(frame_rate)

        self.mask_manager: MaskManager | None

        if mask_manager is not None and mask_config is not None:
            raise ValueError("mask_config cannot be used together with a custom mask_manager.")
        if mask_config is not None and not enable_mask_manager:
            raise ValueError("mask_config requires enable_mask_manager=True when no custom mask_manager is supplied.")
        if mask_manager is not None:
            self.mask_manager = mask_manager
        elif enable_mask_manager:
            self.mask_manager = _build_default_mask_manager(
                mask_config if mask_config is not None else McByteMaskConfig()
            )
        else:
            self.mask_manager = None

        # Retain the original manager so reset() can re-attach it. The active
        # ``mask_manager`` is nulled after repeated CUDA out-of-memory failures
        # (see ``_run_mask_manager``); without this reference a subsequent
        # reset() — the documented new-video boundary where GPU memory is
        # typically freed — could never restore mask-conditioned association,
        # and a user-supplied custom manager would be lost permanently.
        self._mask_manager_original = self.mask_manager

        self._previous_frame: np.ndarray | None = None
        self._previous_tracklets: list[TrackletSnapshot] = []
        self._last_mask_output: MaskOutput | None = None
        self._previous_new_tracklets: list[TrackletSnapshot] = []
        self._previous_removed_tracklet_ids: list[int] = []
        self._mask_tracklet_ids: set[int] = set()
        # Consecutive-visible-frame counts for confirmed tracklets that are
        # awaiting mask creation under ``minimum_mask_creation_frames``. Only
        # populated when the threshold exceeds 1.
        self._mask_pending_ages: dict[int, int] = {}
        self._consecutive_mask_failures: int = 0
        self._warned_mask_manager_dynamic_rate = False

    def update(
        self,
        detections: sv.Detections,
        frame: np.ndarray | None = None,
        timestamp: float | None = None,
    ) -> sv.Detections:
        """Update the tracker with detections from the current frame.

        This is the main per-frame entry point. If a mask manager is configured and a
        frame is provided, masks are updated before association using tracker lifecycle
        events stored from the previous call. After association, the method stores the
        current frame's visible tracklets, newly created tracklets, and explicitly
        terminated tracklet IDs for the next frame's mask update.

        Args:
            detections: Supervision detections for the current frame. Must include
                ``.xyxy``. Confidence (`detections.confidence`) is optional but
                recommended. This method does not mutate the input detections; it
                returns a new ``sv.Detections`` with ``tracker_id`` assigned.
            frame: Current frame in **RGB** channel order, shape ``(H, W, 3)``.
                Note this deviates from the ``BaseTracker.update`` default of BGR:
                McByte's SAM and Cutie mask backends consume the frame as RGB
                without any internal channel conversion, so a BGR frame silently
                degrades mask quality (and therefore mask-conditioned
                association). Required for camera motion compensation and for
                mask-manager propagation.
            timestamp: Absolute time of the current frame in seconds, or ``None``
                for fixed-rate mode (``frame_step = 1.0`` per call). When provided,
                capture times must be non-decreasing; elapsed seconds are converted
                to Kalman frame units for prediction and used directly for
                lost-track pruning.

        Returns:
            New sv.Detections with tracker_id assigned for each output detection.
            Confirmed tracks have tracker_id >= 0; unmatched/unconfirmed detections have
            tracker_id of -1. When the update is skipped (backwards or non-finite
            timestamp), all ``tracker_id`` values are ``-1``.

        Warns:
            UserWarning: If ``timestamp`` is earlier than the previous call
                (backwards order); the whole update is skipped and all output
                IDs are ``-1``. If ``timestamp`` equals the previous call
                (duplicate); predict is skipped but association still runs on
                the last state.
        """
        timing = self._predict_timing(timestamp)
        if timing.skip_update:
            return self._detections_for_skipped_update(detections)

        self.frame_id += 1

        # For the convenience and better understanding. McByte processes uses previous
        # frame and current frame. It is better to keep the method argument as "frame",
        # as in case of the other trackers.
        current_frame = frame
        terminated_tracklet_ids: list[int] = []

        if timing.skip_predict:
            # Duplicate timestamp: the Kalman predict step is skipped and
            # association runs on the last state (see update() "Warns"). Masks
            # must not advance either — stepping the backend one phantom frame
            # would match freshly propagated masks against un-advanced track
            # state (a one-frame mask/state desync). Retain the previous mask
            # output so masks stay in lockstep with the state used here.
            pass
        elif self.mask_manager is not None and current_frame is not None:
            if timing.uses_elapsed_time and not self._warned_mask_manager_dynamic_rate:
                warnings.warn(
                    "enable_mask_manager=True with timestamp-based (dynamic-rate) "
                    "updates: the mask pipeline advances one step per update() call "
                    "regardless of elapsed time, while Kalman prediction and "
                    "lost-track pruning scale by timestamp. Mask propagation can "
                    "drift out of sync with track state across timestamp gaps.",
                    UserWarning,
                    stacklevel=2,
                )
                self._warned_mask_manager_dynamic_rate = True
            self._last_mask_output = self._run_mask_manager(self.mask_manager, current_frame)
        else:
            self._last_mask_output = None

        if len(self.tracks) == 0 and len(detections) == 0:
            result = sv.Detections.empty()
            result.tracker_id = np.array([], dtype=int)
            self._store_previous_mask_inputs(
                frame=current_frame,
                detections=result,
                removed_tracklet_ids=terminated_tracklet_ids,
            )
            return result

        out_det_indices: list[int] = []
        out_tracker_ids: list[int] = []

        # Predict new locations for existing tracks
        self._predict_tracklets(self.tracks, timing)

        # Ghost-ID prevention: budget-only filter before association.
        # Keeps immature tracks alive for matching; full lifecycle prune runs after.
        _budget = self._lost_track_time_budget(timing, self.maximum_time_without_update)
        self._prune_lost_tracks(timing)

        detection_boxes = detections.xyxy
        confidences = default_confidences(detections)

        # Split indices into high / low / discarded by confidence
        high_mask = confidences >= self.high_conf_det_threshold
        low_mask = (confidences > _MINIMUM_DETECTION_CONFIDENCE) & (~high_mask)

        high_indices = np.where(high_mask)[0]
        low_indices = np.where(low_mask)[0]

        high_boxes = detection_boxes[high_indices]
        low_boxes = detection_boxes[low_indices]
        high_scores = confidences[high_indices]

        # Split tracks into confirmed, unconfirmed, and lost.
        # After predict(), time_since_update == 1 means the track was matched in
        # the previous frame ("tracked"), while time_since_update > 1 means the
        # track has been unmatched for multiple frames ("lost").
        confirmed_tracks: list[McByteTracklet] = []
        unconfirmed_tracks: list[McByteTracklet] = []
        lost_tracks: list[McByteTracklet] = []
        for track in self.tracks:
            if track.time_since_update > 1:
                lost_tracks.append(track)
            elif track.tracker_id != -1 or track.number_of_successful_updates >= self.minimum_consecutive_frames:
                # Maturity is sticky: a track that already holds a real
                # tracker_id (e.g. an instant-activated first-frame track) stays
                # confirmed even before it reaches minimum_consecutive_frames.
                # On a miss it is kept as a confirmed (then eventually lost)
                # track rather than discarded as an unconfirmed one.
                confirmed_tracks.append(track)
            else:
                unconfirmed_tracks.append(track)

        # CMC: apply to all predicted tracks before association
        if self.enable_cmc and self.cmc is not None and current_frame is not None:
            mask_boxes = high_boxes if len(high_boxes) > 0 else None
            H = self.cmc.estimate(current_frame, mask_boxes)
            CMC.apply_batch(H, self.tracks)

        # Cache each tracklet's predicted state bbox once per update. Track state
        # is unchanged across the three association stages: a track matched in an
        # earlier stage never re-enters a later one, and the CMC adjustment above
        # is already applied. Recomputing get_state_bbox() per stage would be
        # redundant, so all stages read boxes from this map keyed by ``id()``.
        predicted_state_boxes = {id(track): track.get_state_bbox() for track in self.tracks}

        # Step 1: associate high-confidence detections to confirmed + lost tracks.
        # Lost tracks are included here (following the original ByteTrack), and
        # IoU is fused with detection scores.
        strack_pool = confirmed_tracks + lost_tracks
        raw_iou_similarity = self._get_iou_matrix(
            strack_pool,
            high_boxes,
            predicted_state_boxes,
        )
        association_similarity = _fuse_score(
            self.iou.normalize_for_fusion(raw_iou_similarity.copy()),
            high_scores,
        )

        matched, unmatched_pool, unmatched_high = self._get_mask_conditioned_associated_indices(
            similarity_matrix=association_similarity,
            raw_iou_similarity=raw_iou_similarity,
            tracklets=strack_pool,
            detection_boxes=high_boxes,
            min_similarity_thresh=self.minimum_iou_threshold_first_assoc,
        )

        for row, col in matched:
            track = strack_pool[row]
            track.update(high_boxes[col])
            if track.number_of_successful_updates >= self.minimum_consecutive_frames and track.tracker_id == -1:
                track.tracker_id = self._allocate_tracker_id()
            out_det_indices.append(int(high_indices[col]))
            out_tracker_ids.append(track.tracker_id)

        # Step 2: associate low-confidence detections to remaining *tracked* tracks
        # only (excluding lost tracks, following the original ByteTrack).
        # No score fusing in second association.
        remaining_tracked = [strack_pool[i] for i in unmatched_pool if strack_pool[i].time_since_update == 1]
        raw_iou_similarity = self._get_iou_matrix(
            remaining_tracked,
            low_boxes,
            predicted_state_boxes,
        )

        # There is no score fusion in stage 2, so the assignment matrix
        # and raw-IoU matrix are the same.
        matched, _, unmatched_low = self._get_mask_conditioned_associated_indices(
            similarity_matrix=raw_iou_similarity,
            raw_iou_similarity=raw_iou_similarity,
            tracklets=remaining_tracked,
            detection_boxes=low_boxes,
            min_similarity_thresh=self.minimum_iou_threshold_second_assoc,
        )

        for row, col in matched:
            track = remaining_tracked[row]
            track.update(low_boxes[col])
            if track.number_of_successful_updates >= self.minimum_consecutive_frames and track.tracker_id == -1:
                track.tracker_id = self._allocate_tracker_id()
            out_det_indices.append(int(low_indices[col]))
            out_tracker_ids.append(track.tracker_id)

        # Unmatched low-confidence detections
        for det_local_idx in sorted(unmatched_low):
            out_det_indices.append(int(low_indices[det_local_idx]))
            out_tracker_ids.append(-1)

        # Step 3: match unconfirmed tracks with remaining unmatched high-confidence
        # detections (with score fusing, following the original ByteTrack).
        # Unmatched unconfirmed tracks are removed (not kept as lost).
        unmatched_high_list = sorted(unmatched_high)
        unmatched_uc_indices: list[int] = list(range(len(unconfirmed_tracks)))

        if len(unconfirmed_tracks) > 0 and len(unmatched_high_list) > 0:
            uh_boxes = high_boxes[unmatched_high_list]
            uh_scores = high_scores[unmatched_high_list]

            raw_iou_similarity = self._get_iou_matrix(
                unconfirmed_tracks,
                uh_boxes,
                predicted_state_boxes,
            )
            association_similarity = _fuse_score(
                self.iou.normalize_for_fusion(raw_iou_similarity.copy()),
                uh_scores,
            )

            matched_uc, unmatched_uc_indices, remaining_uh = self._get_mask_conditioned_associated_indices(
                similarity_matrix=association_similarity,
                raw_iou_similarity=raw_iou_similarity,
                tracklets=unconfirmed_tracks,
                detection_boxes=uh_boxes,
                min_similarity_thresh=self.minimum_iou_threshold_unconfirmed_assoc,
            )

            for row, col in matched_uc:
                track = unconfirmed_tracks[row]
                orig_high_idx = unmatched_high_list[col]
                track.update(high_boxes[orig_high_idx])
                if track.number_of_successful_updates >= self.minimum_consecutive_frames and track.tracker_id == -1:
                    track.tracker_id = self._allocate_tracker_id()
                out_det_indices.append(int(high_indices[orig_high_idx]))
                out_tracker_ids.append(track.tracker_id)

            # Only remaining unmatched high-conf dets proceed to spawning
            unmatched_high = [unmatched_high_list[i] for i in remaining_uh]

        # Remove unmatched unconfirmed tracks (following original ByteTrack,
        # which marks them as removed rather than keeping them as lost).
        if len(unmatched_uc_indices) > 0:
            remove_ids = {id(unconfirmed_tracks[i]) for i in unmatched_uc_indices}
            self.tracks = [t for t in self.tracks if id(t) not in remove_ids]

        # Spawn new tracks from unmatched high-confidence detections
        self._spawn_new_tracks(
            detection_boxes,
            confidences,
            unmatched_high,
            high_indices,
            out_det_indices,
            out_tracker_ids,
            is_first_frame=(self.frame_id == 1),
        )

        # Kill terminated tracks. Temporarily lost tracks remain alive and keep masks.
        tracklet_ids_before_pruning = {int(track.tracker_id) for track in self.tracks if track.tracker_id >= 0}
        self.tracks = _get_alive_tracklets(
            tracklets=self.tracks,
            maximum_frames_without_update=self.maximum_frames_without_update,
            minimum_consecutive_frames=self.minimum_consecutive_frames,
            maximum_time_without_update=_budget,
        )
        tracklet_ids_after_pruning = {int(track.tracker_id) for track in self.tracks if track.tracker_id >= 0}
        terminated_tracklet_ids = sorted(tracklet_ids_before_pruning - tracklet_ids_after_pruning)

        # Build final detections
        if not out_det_indices:
            result = sv.Detections.empty()
            result.tracker_id = np.array([], dtype=int)
            self._store_previous_mask_inputs(
                frame=current_frame,
                detections=result,
                removed_tracklet_ids=terminated_tracklet_ids,
            )
            return result

        idx = np.array(out_det_indices)
        result = cast(sv.Detections, detections[idx])
        result.tracker_id = np.array(out_tracker_ids, dtype=int)
        self._store_previous_mask_inputs(
            frame=current_frame,
            detections=result,
            removed_tracklet_ids=terminated_tracklet_ids,
        )
        return result

    def _run_mask_manager(
        self,
        mask_manager: MaskManager,
        current_frame: np.ndarray,
    ) -> MaskOutput | None:
        """Run the mask pipeline for this frame, degrading gracefully on CUDA OOM.

        The SAM/Cutie step can exhaust GPU memory mid-sequence. A CUDA
        out-of-memory failure must not crash the whole ``update()`` call, so it
        is caught, logged, and treated as "no mask evidence for this frame":
        association then falls back to IoU only. After
        ``_MAX_CONSECUTIVE_MASK_FAILURES`` consecutive out-of-memory failures the
        mask manager is disabled for the remainder of the run. Any non-OOM error
        (for example ``ValueError`` or a non-OOM ``RuntimeError``) propagates
        unchanged so genuine bugs are not silently swallowed.

        Args:
            mask_manager: The active mask manager (already narrowed to non-None
                by the caller).
            current_frame: Current RGB frame passed to the mask pipeline.

        Returns:
            The propagated ``MaskOutput`` on success, or ``None`` when the
            pipeline raised a CUDA out-of-memory error for this frame.
        """
        try:
            mask_output = mask_manager.get_updated_masks(
                frame=current_frame,
                previous_frame=self._previous_frame,
                previous_tracklets=self._previous_tracklets,
                new_tracklets=self._previous_new_tracklets,
                removed_tracklet_ids=self._previous_removed_tracklet_ids,
            )
        except RuntimeError as exc:
            if not _is_out_of_memory(exc):
                raise
            self._consecutive_mask_failures += 1
            logger.warning(
                "McByte mask pipeline raised a CUDA out-of-memory error on frame %d; "
                "falling back to IoU-only association for this frame (failure %d/%d): %s",
                self.frame_id,
                self._consecutive_mask_failures,
                _MAX_CONSECUTIVE_MASK_FAILURES,
                exc,
            )
            if self._consecutive_mask_failures >= _MAX_CONSECUTIVE_MASK_FAILURES:
                logger.error(
                    "McByte mask pipeline hit %d consecutive out-of-memory failures; "
                    "disabling mask-conditioned association for the remainder of this "
                    "run. Association continues using IoU only.",
                    self._consecutive_mask_failures,
                )
                self.mask_manager = None
            return None

        self._consecutive_mask_failures = 0
        return mask_output

    def _detections_to_tracklet_snapshots(
        self,
        detections: sv.Detections,
    ) -> list[TrackletSnapshot]:
        """Convert tracker output detections into mask-manager tracklet snapshots.

        Only detections with valid non-negative tracker IDs are converted. The returned snapshots contain the tracker ID
        and ``xyxy`` box needed by mask generators.
        """
        if detections.tracker_id is None:
            return []

        return [
            TrackletSnapshot(
                tracker_id=int(tracker_id),
                xyxy=xyxy.astype(np.float32),
            )
            for xyxy, tracker_id in zip(detections.xyxy, detections.tracker_id)
            if tracker_id >= 0
        ]

    def _store_previous_mask_inputs(
        self,
        frame: np.ndarray | None,
        detections: sv.Detections,
        removed_tracklet_ids: list[int],
    ) -> None:
        """Store tracker outputs and mask lifecycle events for the next frame.

        The mask manager consumes these values at the beginning of the next ``update()`` call. New tracklets are
        detected among current visible outputs that do not yet have masks. Removed tracklets are provided explicitly
        from tracker pruning, so temporarily lost but still alive tracklets keep their masks.
        """
        if self.mask_manager is None or frame is None:
            self._previous_frame = None
            self._previous_tracklets = []
            self._previous_new_tracklets = []
            self._previous_removed_tracklet_ids = []
            # Preserve the already-masked tracklet-ID set across an occasional
            # frame=None call (e.g. a dropped/corrupt frame while the manager is
            # still active). Clearing it would make the next real frame treat
            # every visible tracklet as new: SAM would re-prompt all masks and
            # Cutie's add_masks would raise on tracklets that already own an
            # object. Only a fully absent mask manager warrants a full clear.
            if self.mask_manager is None:
                self._mask_tracklet_ids = set()
            return

        # Convert current output detections into TrackletSnapshots.
        # Only valid tracker IDs are kept.
        current_tracklets = self._detections_to_tracklet_snapshots(detections)

        # Remove from the “tracks that already have masks” set
        # only the IDs that were truly terminated/pruned.
        removed_tracklet_id_set = set(removed_tracklet_ids)
        self._mask_tracklet_ids -= removed_tracklet_id_set

        # Find current visible tracklets that do not yet have masks and have
        # been visible long enough to warrant a mask. These will be passed to
        # SAM/Cutie on the next frame.
        new_tracklets = self._collect_new_masked_tracklets(current_tracklets)

        # Mark those new tracklets as now mask-managed, so if they disappear temporarily
        # and later reappear, they are not treated as new again.
        self._mask_tracklet_ids.update(tracklet.tracker_id for tracklet in new_tracklets)

        # Store lifecycle events from this frame. At the next update(),
        # MaskManager receives these and calls add_masks() / remove_masks().
        self._previous_new_tracklets = new_tracklets
        self._previous_removed_tracklet_ids = removed_tracklet_ids

        # Stores the current frame and current visible tracklets as “previous”
        # inputs for the next frame. Only mask-eligible tracklets (already masked
        # or promoted this frame) are exposed to the mask manager: tracklets
        # still inside their ``minimum_mask_creation_frames`` defer window must
        # be invisible to the pipeline so they are neither masked at Cutie
        # initialization nor double-added once they cross the threshold. With
        # ``minimum_mask_creation_frames == 1`` every visible tracklet is
        # eligible immediately, so this equals ``current_tracklets``.
        self._previous_frame = frame
        if self.minimum_mask_creation_frames <= 1:
            self._previous_tracklets = current_tracklets
        else:
            self._previous_tracklets = [
                tracklet for tracklet in current_tracklets if tracklet.tracker_id in self._mask_tracklet_ids
            ]

    def _collect_new_masked_tracklets(
        self,
        current_tracklets: list[TrackletSnapshot],
    ) -> list[TrackletSnapshot]:
        """Select not-yet-masked tracklets ready for mask creation next frame.

        A confirmed tracklet becomes eligible for mask creation only after it has
        appeared in tracker output for ``minimum_mask_creation_frames`` consecutive
        frames. Very-short-lived tracklets that terminate before reaching the
        threshold are never handed to SAM/Cutie, saving the per-appearance mask
        encode. A tracklet that disappears before reaching the threshold restarts
        its count if it later reappears, so only stably visible tracklets are
        masked. With ``minimum_mask_creation_frames == 1`` every not-yet-masked
        tracklet is returned on its first visible frame, matching the original
        immediate-creation timing.

        Args:
            current_tracklets: Tracklet snapshots visible in the current tracker
                output (already filtered to valid, non-negative tracker IDs).

        Returns:
            Snapshots whose masks should be created on the next frame.
        """
        already_masked = self._mask_tracklet_ids
        if self.minimum_mask_creation_frames <= 1:
            return [tracklet for tracklet in current_tracklets if tracklet.tracker_id not in already_masked]

        current_ids = {tracklet.tracker_id for tracklet in current_tracklets}
        # Drop ages for tracklets no longer visible so a reappearing tracklet
        # must again accumulate consecutive visible frames before it is masked.
        self._mask_pending_ages = {
            tracker_id: age for tracker_id, age in self._mask_pending_ages.items() if tracker_id in current_ids
        }

        new_tracklets: list[TrackletSnapshot] = []
        for tracklet in current_tracklets:
            tracker_id = tracklet.tracker_id
            if tracker_id in already_masked:
                continue
            age = self._mask_pending_ages.get(tracker_id, 0) + 1
            if age >= self.minimum_mask_creation_frames:
                new_tracklets.append(tracklet)
                self._mask_pending_ages.pop(tracker_id, None)
            else:
                self._mask_pending_ages[tracker_id] = age
        return new_tracklets

    def _get_iou_matrix(
        self,
        tracklets: list[McByteTracklet],
        detections: np.ndarray,
        tracklet_boxes_by_id: dict[int, np.ndarray],
    ) -> np.ndarray:
        """Compute IoU similarity between tracklet states and detection boxes.

        Returns an ``(N, M)`` matrix where ``N`` is the number of tracklets and ``M`` is
        the number of detections. Empty inputs are handled by returning an empty matrix
        with the expected shape.

        Args:
            tracklets: Tracklets forming the rows of the returned matrix.
            detections: Detection boxes in ``xyxy`` format forming the columns.
            tracklet_boxes_by_id: Mapping from ``id(track)`` to the track's
                predicted state bbox, computed once per ``update()`` and reused
                across association stages to avoid recomputing ``get_state_bbox``.

        Raises:
            KeyError: If a tracklet passed in is absent from ``tracklet_boxes_by_id``
                — an internal-invariant violation, since the map is built from
                ``self.tracks`` and every tracklet here is drawn from it.
        """
        if len(tracklets) == 0:
            tracklet_boxes = np.empty((0, 4))
        else:
            try:
                tracklet_boxes = np.array([tracklet_boxes_by_id[id(tracklet)] for tracklet in tracklets])
            except KeyError as exc:
                raise KeyError(
                    f"tracklet id {exc.args[0]} missing from the per-frame decode-once box cache; "
                    "tracklet_boxes_by_id must contain every tracklet passed to this helper "
                    "(it is built from self.tracks once per update())"
                ) from exc
        return self.iou.compute(tracklet_boxes, detections)

    def _get_mask_conditioned_associated_indices(
        self,
        similarity_matrix: np.ndarray,
        raw_iou_similarity: np.ndarray,
        tracklets: list[McByteTracklet],
        detection_boxes: np.ndarray,
        min_similarity_thresh: float,
    ) -> tuple[list[tuple[int, int]], list[int], list[int]]:
        """Associate tracklets and detections using McByte mask conditioning.

        Clear threshold-valid pairs are locked before assignment when they are the
        only eligible candidate in both their row and column. The remaining
        association problem is conditioned with propagated-mask evidence for
        ambiguous pairs and, optionally, isolated positive-IoU pairs below the
        normal threshold.

        Hungarian assignment is applied only to the remaining reduced matrix.
        Reduced row and column indices are then mapped back to the original
        ``tracklets`` and ``detection_boxes`` index spaces and combined with the
        locked matches.

        When no propagated mask output is available, mask-based score updates are
        skipped. Clear-match locking and assignment of the remaining problem still
        follow the McByte association pipeline.

        Args:
            similarity_matrix: Stage-specific association similarity matrix with
                shape ``(num_tracklets, num_detections)``. This is score-fused IoU
                for the first and unconfirmed association stages, and raw IoU for
                the second association stage.
            raw_iou_similarity: Unfused IoU similarity matrix with the same shape.
                It is used to determine optional isolated geometric candidates.
            tracklets: Tracklets corresponding, in order, to the rows of both
                similarity matrices.
            detection_boxes: Detection boxes in ``xyxy`` format corresponding, in
                order, to the columns of both similarity matrices.
            min_similarity_thresh: Minimum stage-specific similarity required for
                a valid association.

        Returns:
            A tuple containing:

            - matched original ``(tracklet_index, detection_index)`` pairs;
            - sorted original indices of unmatched tracklets;
            - sorted original indices of unmatched detections.
        """
        conditioned_association = condition_similarity_with_masks(
            similarity=similarity_matrix,
            raw_iou_similarity=raw_iou_similarity,
            tracklet_ids=[int(tracklet.tracker_id) for tracklet in tracklets],
            detection_boxes=detection_boxes,
            mask_output=self._last_mask_output,
            minimum_similarity=min_similarity_thresh,
            minimum_mask_average_confidence=self.minimum_mask_average_confidence,
            minimum_mask_coverage=self.minimum_mask_coverage,
            minimum_mask_fill_ratio=self.minimum_mask_fill_ratio,
            enable_isolated_mask_matching=self.enable_isolated_mask_matching,
        )

        (
            reduced_matches,
            reduced_unmatched_track_indices,
            reduced_unmatched_detection_indices,
        ) = self._get_associated_indices(
            similarity_matrix=conditioned_association.conditioned_similarity,
            min_similarity_thresh=min_similarity_thresh,
        )

        remapped_matches = [
            (
                conditioned_association.remaining_track_indices[reduced_track_index],
                conditioned_association.remaining_detection_indices[reduced_detection_index],
            )
            for reduced_track_index, reduced_detection_index in reduced_matches
        ]

        matched = sorted(conditioned_association.locked_matches + remapped_matches)

        unmatched_tracks = sorted(
            conditioned_association.remaining_track_indices[reduced_track_index]
            for reduced_track_index in reduced_unmatched_track_indices
        )

        unmatched_detections = sorted(
            conditioned_association.remaining_detection_indices[reduced_detection_index]
            for reduced_detection_index in reduced_unmatched_detection_indices
        )

        return matched, unmatched_tracks, unmatched_detections

    def _get_associated_indices(
        self,
        similarity_matrix: np.ndarray,
        min_similarity_thresh: float,
    ) -> tuple[list[tuple[int, int]], list[int], list[int]]:
        """Associate detections to tracks based on Similarity (IoU) using the Jonker-Volgenant algorithm approach with
        no initialization instead of the Hungarian algorithm as mentioned in the SORT paper, but it solves the
        assignment problem in an optimal way.

        Args:
            similarity_matrix: Similarity matrix between tracks (rows) and detections
            (columns). min_similarity_thresh: Minimum similarity threshold for a valid
            match.

        Returns:
            matched: List of ``(tracker_idx, detection_idx)`` tuples for
                associations that meet the similarity threshold.
            unmatched_tracks: Sorted list of track indices not matched to any
                detection.
            unmatched_detections: Sorted list of detection indices not matched
                to any track.
        """
        matched_indices = []
        n_tracks, n_detections = similarity_matrix.shape
        unmatched_tracks = set(range(n_tracks))
        unmatched_detections = set(range(n_detections))

        if n_tracks > 0 and n_detections > 0:
            row_indices, col_indices = linear_sum_assignment(similarity_matrix, maximize=True)
            for row, col in zip(row_indices, col_indices):
                if similarity_matrix[row, col] >= min_similarity_thresh:
                    matched_indices.append((row, col))
                    unmatched_tracks.remove(row)
                    unmatched_detections.remove(col)

        # Return sorted lists for deterministic order across Python runtimes.
        return matched_indices, sorted(unmatched_tracks), sorted(unmatched_detections)

    def _spawn_new_tracks(
        self,
        detection_boxes: np.ndarray,
        confidences: np.ndarray,
        unmatched_high_local: list[int],
        high_indices: np.ndarray,
        out_det_indices: list[int],
        out_tracker_ids: list[int],
        is_first_frame: bool = False,
    ) -> None:
        """Create new tracklets from unmatched high-confidence detections.

        On the very first frame, new tracklets are immediately activated with a
        real tracker ID, following the original ByteTrack convention where
        ``activate()`` sets ``is_activated = True`` only when
        ``frame_id == 1``.
        """
        for det_local_idx in unmatched_high_local:
            global_idx = int(high_indices[det_local_idx])
            conf = float(confidences[global_idx])

            # Every unmatched high-confidence detection is returned, with
            # tracker_id -1 when it does not spawn a track. This matches the
            # documented update() contract ("unmatched/unconfirmed detections
            # have tracker_id of -1"), the symmetric handling of unmatched
            # low-confidence detections, and sibling ByteTrack. Only detections
            # clearing the activation threshold spawn a new tracklet; a
            # first-frame instant-activated spawn is emitted with its real ID.
            tracker_id_out = -1
            if conf >= self.track_activation_threshold:
                tracklet = McByteTracklet(
                    initial_bbox=detection_boxes[global_idx],
                    state_estimator_class=self.state_estimator_class,
                )
                if is_first_frame and self.instant_first_frame_activation:
                    tracklet.tracker_id = self._allocate_tracker_id()
                self.tracks.append(tracklet)
                tracker_id_out = tracklet.tracker_id

            out_det_indices.append(global_idx)
            out_tracker_ids.append(tracker_id_out)

    def reset(self) -> None:
        """Reset tracker, camera-motion, and mask-manager state.

        This clears active tracklets, resets the global McByte track ID counter, clears stored mask lifecycle inputs,
        and resets optional camera motion compensation and mask-manager components. Call this when switching to a new
        video or scene.
        """
        self.tracks = []
        self.frame_id = 0
        self._last_timestamp = None
        self._reset_id_allocator()
        self._previous_frame = None
        self._previous_tracklets = []
        self._last_mask_output = None
        # Re-attach a mask manager that was auto-disabled after repeated CUDA
        # out-of-memory failures. reset() marks the documented new-video
        # boundary where GPU memory has typically been freed.
        if self.mask_manager is None and self._consecutive_mask_failures >= _MAX_CONSECUTIVE_MASK_FAILURES:
            self.mask_manager = self._mask_manager_original
        if self.mask_manager is not None:
            self.mask_manager.reset()
        if self.cmc is not None:
            self.cmc.reset()
        self._previous_new_tracklets = []
        self._previous_removed_tracklet_ids = []
        self._mask_tracklet_ids = set()
        self._mask_pending_ages = {}
        self._consecutive_mask_failures = 0

    def apply_cmc_batch(self, H: np.ndarray | None) -> None:
        """Apply camera motion compensation to all active tracks.

        Convenience wrapper around :meth:`CMC.apply_batch` for callers that hold
        a tracker instance and an affine transform. ``update()`` applies CMC
        directly and does not rely on this method.

        Args:
            H: 2x3 affine transform matrix returned by CMC.estimate().
                If None, this method is a no-op.

        Examples:
            >>> tracker = McByteTracker()
            >>> tracker.apply_cmc_batch(None)  # no-op
        """
        CMC.apply_batch(H, self.tracks)
