utils/tracker.py

import logging
import numpy as np
from scipy.optimize import linear_sum_assignment
import scipy.linalg

from utils.schemas import AssessmentStatus


class KalmanFilter:
    """
    A simple Kalman Filter for tracking bounding boxes in image space.
    The 8-dimensional state space is (x, y, a, h, vx, vy, va, vh), where
    x, y is the center position, a is the aspect ratio, and h is the height.
    """

    def __init__(self):
        ndim, dt = 4, 1.0

        # Create Kalman filter model matrices.
        self._motion_mat = np.eye(2 * ndim, 2 * ndim)
        for i in range(ndim):
            self._motion_mat[i, ndim + i] = dt
        self._update_mat = np.eye(ndim, 2 * ndim)

        # Motion and observation uncertainty are chosen relative to the current
        # state estimate. These weights control the amount of uncertainty in
        # the model. This is a bit heuristic.
        self._std_weight_position = 1.0 / 20
        self._std_weight_velocity = 1.0 / 160

    def initiate(self, measurement):
        """Create track from unassociated measurement.
        
        Parameters
        ----------
        measurement : dbo
            Bounding box coordinates (x1, y1, x2, y2) with confidence score.
            
        Returns
        -------
        (mean, covariance)
            Returns the mean vector (8 dimensional) and covariance matrix (8x8)
            of the new track.
        """
        mean_pos = self._xyah_from_xyxy(measurement)
        mean = np.r_[mean_pos, np.zeros_like(mean_pos)]

        std = [
            2 * self._std_weight_position * mean_pos[3],
            2 * self._std_weight_position * mean_pos[3],
            1e-2,
            2 * self._std_weight_position * mean_pos[3],
            10 * self._std_weight_velocity * mean_pos[3],
            10 * self._std_weight_velocity * mean_pos[3],
            1e-5,
            10 * self._std_weight_velocity * mean_pos[3],
        ]
        covariance = np.diag(np.square(std))
        return mean, covariance

    def predict(self, mean, covariance):
        """Run Kalman filter prediction step.

        Parameters
        ----------
        mean : ndarray
            The 8 dimensional mean vector of the object state at the previous
            time step.
        covariance : ndarray
            The 8x8 dimensional covariance matrix of the object state at the
            previous time step.

        Returns
        -------
        (mean, covariance)
            Returns the mean vector and covariance matrix of the predicted
            state.
        """
        std_pos = [
            self._std_weight_position * mean[3],
            self._std_weight_position * mean[3],
            1e-2,
            self._std_weight_position * mean[3],
        ]
        
        std_vel = [
            self._std_weight_velocity * mean[3],
            self._std_weight_velocity * mean[3],
            1e-5,
            self._std_weight_velocity * mean[3],
        ]
        
        motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
        mean = np.dot(self._motion_mat, mean)
        covariance = (
            np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T))
            + motion_cov
        )
        return mean, covariance

    def project(self, mean, covariance):
        """Project state distribution to measurement space.

        Parameters
        ----------
        mean : ndarray
            The state's mean vector (8 dimensional).
        covariance : ndarray
            The state's covariance matrix (8x8 dimensional).

        Returns
        -------
        (mean, covariance)
            Returns the projected mean and covariance matrix of the given state
            estimate.
        """
        std = [
            self._std_weight_position * mean[3],
            self._std_weight_position * mean[3],
            1e-1,
            self._std_weight_position * mean[3],
        ]
        
        innovation_cov = np.diag(np.square(std))
        mean = np.dot(self._update_mat, mean)
        covariance = np.linalg.multi_dot((self._update_mat, covariance, self._update_mat.T))
        return mean, covariance + innovation_cov

    def update(self, mean, covariance, measurement):
        """Run Kalman filter correction step.

        Parameters
        ----------
        mean : ndarray
            The predicted state's mean vector (8 dimensional).
        covariance : ndarray
            The state's covariance matrix (8x8 dimensional).
        measurement : ndarray
            The 4 dimensional measurement vector (x, y, a, h), where (x, y)
            is the center position, a the aspect ratio, and h the height.

        Returns
        -------
        (mean, covariance)
            Returns the measurement-corrected state distribution.
        """
        projected_mean, projected_cov = self.project(mean, covariance)
        chol_factor, lower = scipy.linalg.cho_factor(
            projected_cov, lower=True, check_finite=False
        )
        kalman_gain = scipy.linalg.cho_solve(
            (chol_factor, lower),
            np.dot(covariance, self._update_mat.T).T,
            check_finite=False,
        ).T
        innovation = measurement - projected_mean
        new_mean = mean + np.dot(innovation, kalman_gain.T)
        new_covariance = covariance - np.linalg.multi_dot(
            (kalman_gain, projected_cov, kalman_gain.T)
        )
        return new_mean, new_covariance

    def gating_distance(self, mean, covariance, measurements, only_position=False, metric="mahalanobis"):
        """Compute gating distance between state distribution and measurements."""
        mean, covariance = self.project(mean, covariance)
        if only_position:
            mean, covariance = mean[:2], covariance[:2, :2]
            measurements = measurements[:, :2]

        d = measurements - mean
        if metric == "gaussian":
            return np.sum(d * d, axis=1)
        elif metric == "mahalanobis":
            cholesky_factor = np.linalg.cholesky(covariance)
            z = scipy.linalg.solve_triangular(
                cholesky_factor, d.T, lower=True, check_finite=False, overwrite_b=True
            )
            squared_maha = np.sum(z * z, axis=0)
            return squared_maha
        else:
            raise ValueError("invalid distance metric")

    def _xyah_from_xyxy(self, xyxy):
        """Convert bounding box to format `(center x, center y, aspect ratio,
        height)`, where the aspect ratio is `width / height`.
        """
        bbox = np.asarray(xyxy).copy()
        cx = (bbox[0] + bbox[2]) / 2.0
        cy = (bbox[1] + bbox[3]) / 2.0
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]
        
        ret = np.zeros(4, dtype=bbox.dtype)
        ret[0] = cx
        ret[1] = cy
        ret[2] = w / h
        ret[3] = h
        return ret


# Default staleness threshold: GPT metadata older than this many frames is flagged STALE
MAX_STALE_FRAMES = 300

GPT_SYNC_KEYS = frozenset({
    # Legacy / polyfilled fields (consumed by frontend cards)
    "gpt_distance_m", "gpt_direction", "gpt_description", "gpt_raw",
    "threat_level_score", "distance_m", "direction", "description",
    # Universal schema fields
    "object_type", "size", "visible_weapons", "weapon_readiness",
    "motion_status", "range_estimate", "bearing",
    "threat_level", "threat_classification", "tactical_intent",
    "dynamic_features",
    # Provenance and temporal validity
    "assessment_frame_index", "assessment_status",
    # Mission relevance
    "mission_relevant", "relevance_reason",
})


class STrack:
    """
    Single object track. Wrapper around KalmanFilter state.
    """
    
    def __init__(self, tlwh, score, label):
        # wait, input is xyxy usually in our pipeline
        # ByteTrack usually uses tlwh internally.
        # Let's standardize to input xyxy.

        self._tlwh = np.asarray(self._tlwh_from_xyxy(tlwh), dtype=np.float32)
        self.is_activated = False
        self.track_id = 0
        self.state = 1 # 1: New, 2: Tracked, 3: Lost, 4: Removed

        self.score = score
        self.label = label
        self.start_frame = 0
        self.frame_id = 0
        self.time_since_update = 0
        self.hits = 0
        
        # Multi-frame history
        self.history = []
        
        # Kalman Filter
        self.kalman_filter = None
        self.mean = None
        self.covariance = None
        
        # GPT attributes (persistent)
        self.gpt_data = {}

    def _tlwh_from_xyxy(self, xyxy):
        """Convert xyxy to tlwh."""
        w = xyxy[2] - xyxy[0]
        h = xyxy[3] - xyxy[1]
        return [xyxy[0], xyxy[1], w, h]
        
    def _xyxy_from_tlwh(self, tlwh):
        """Convert tlwh to xyxy."""
        x1 = tlwh[0]
        y1 = tlwh[1]
        x2 = x1 + tlwh[2]
        y2 = y1 + tlwh[3]
        return [x1, y1, x2, y2]
    
    @property
    def tlwh(self):
        """Get current position in bounding box format `(top left x, top left y,
        width, height)`.
        """
        if self.mean is None:
            return self._tlwh.copy()
        ret = self.mean[:4].copy()
        ret[2] *= ret[3]
        ret[:2] -= ret[2:] / 2
        return ret

    @property
    def tlbr(self):
        """Get current position in bounding box format `(min x, min y, max x,
        max y)`.
        """
        ret = self.tlwh.copy()
        ret[2:] += ret[:2]
        return ret

    def activate(self, kalman_filter, frame_id):
        """Start a new track (tentative until min_hits reached)."""
        self.kalman_filter = kalman_filter
        self.track_id = self.next_id()
        self.mean, self.covariance = self.kalman_filter.initiate(self.tlbr) # Initiate needs xyxy

        self.state = 2 # Tracked
        self.frame_id = frame_id
        self.start_frame = frame_id
        self.is_activated = False  # Tentative until min_hits reached
        self.hits = 1

    def re_activate(self, new_track, frame_id, new_id=False):
        """Reactivate a lost track with a new detection."""
        self.mean, self.covariance = self.kalman_filter.update(
            self.mean, self.covariance, self._xyah_from_xyxy(new_track.tlbr)
        )
        self.time_since_update = 0
        self.state = 2 # Tracked
        self.frame_id = frame_id
        self.score = new_track.score
        
        if new_id:
            self.track_id = self.next_id()

    def update(self, new_track, frame_id, min_hits=3):
        """Update a tracked object with a new detection."""
        self.frame_id = frame_id
        self.time_since_update = 0
        self.score = new_track.score
        self.hits += 1

        self.mean, self.covariance = self.kalman_filter.update(
            self.mean, self.covariance, self._xyah_from_xyxy(new_track.tlbr)
        )
        self.state = 2 # Tracked
        if self.hits >= min_hits:
            self.is_activated = True  # Confirmed after N consecutive hits
        
    def predict(self):
        """Propagate tracking state distribution one time step forward."""
        if self.mean is None: return
        if self.state != 2: # Only predict if tracked? ByteTrack predicts always?
             # Standard implementation predicts for all active/lost tracks
             pass
        self.mean, self.covariance = self.kalman_filter.predict(self.mean, self.covariance)

    def _xyah_from_xyxy(self, xyxy):
        """Internal helper for measurement conversion."""
        bbox = np.asarray(xyxy).copy()
        cx = (bbox[0] + bbox[2]) / 2.0
        cy = (bbox[1] + bbox[3]) / 2.0
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]
        
        ret = np.zeros(4, dtype=bbox.dtype)
        ret[0] = cx
        ret[1] = cy
        ret[2] = w / h
        ret[3] = h
        return ret

    @staticmethod
    def next_id():
        # Global counter
        if not hasattr(STrack, "_count"):
            STrack._count = 0
        STrack._count += 1
        return STrack._count

    @staticmethod
    def reset_count():
        """Reset track ID counter so the next video starts at T01."""
        STrack._count = 0


class ByteTracker:
    def __init__(self, track_high_thresh=0.3, track_low_thresh=0.1,
                 new_track_thresh=0.3, track_buffer=60, match_thresh=0.8,
                 frame_rate=30, min_hits=3):
        STrack.reset_count()
        self.track_high_thresh = track_high_thresh
        self.track_low_thresh = track_low_thresh
        self.new_track_thresh = new_track_thresh
        self.track_buffer = track_buffer
        self.match_thresh = match_thresh
        self.min_hits = min_hits
        self.frame_id = 0

        self.tracked_stracks = []  # Type: List[STrack]
        self.lost_stracks = []     # Type: List[STrack]
        self.removed_stracks = []  # Type: List[STrack]

        self.kalman_filter = KalmanFilter()

    def update(self, detections_list):
        """
        Update the tracker with a list of detections.
        
        Args:
            detections_list: List of dicts, each having:
                - bbox: [x1, y1, x2, y2]
                - score: float
                - label: str
                - (optional) other keys preserved
        
        Returns:
            List of dicts with 'track_id' added/updated.
        """
        self.frame_id += 1
        
        # 0. STrack Conversion using generic interface
        activated_stracks = []
        refind_stracks = []
        lost_stracks = []
        removed_stracks = []

        scores = [d['score'] for d in detections_list]
        bboxes = [d['bbox'] for d in detections_list]
        
        # Split into high and low confidence
        detections = []
        detections_second = []
        
        # Need to keep mapping to original dict to populate results later
        # We wrap original dict in STrack
        
        for d in detections_list:
            score = d['score']
            if score < self.track_low_thresh:
                continue  # Background noise — discard entirely

            t = STrack(d['bbox'], score, d['label'])
            t.original_data = d  # Link back

            if score >= self.track_high_thresh:
                detections.append(t)
            else:
                detections_second.append(t)

        # 1. Prediction
        unconfirmed = []
        tracked_stracks = []  # Type: List[STrack]
        for track in self.tracked_stracks:
            if not track.is_activated:
                unconfirmed.append(track)
            else:
                tracked_stracks.append(track)

        strack_pool = join_stracks(tracked_stracks, self.lost_stracks)
        # Predict the current location with KF
        STrack.multi_predict(strack_pool, self.kalman_filter)

        # 2. First association (High score)
        dists = iou_distance(strack_pool, detections)
        dists = fuse_score(dists, detections) # Optional? ByteTrack uses it
        matches, u_track, u_detection = linear_assignment(dists, thresh=self.match_thresh)

        for itracked, idet in matches:
            track = strack_pool[itracked]
            det = detections[idet]
            if track.state == 2:
                track.update(det, self.frame_id, min_hits=self.min_hits)
                activated_stracks.append(track)
            else:
                track.re_activate(det, self.frame_id, new_id=False)
                refind_stracks.append(track)

            # Persist data
            self._sync_data(track, det)

        # 3. Second association (Low score)
        # Match unmatched tracks to low score detections
        r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state == 2]
        dists = iou_distance(r_tracked_stracks, detections_second)
        matches, u_track, u_detection_second = linear_assignment(dists, thresh=0.5)

        for itracked, idet in matches:
            track = r_tracked_stracks[itracked]
            det = detections_second[idet]
            if track.state == 2:
                track.update(det, self.frame_id, min_hits=self.min_hits)
                activated_stracks.append(track)
            else:
                track.re_activate(det, self.frame_id, new_id=False)
                refind_stracks.append(track)

            self._sync_data(track, det)

        for it in u_track:
            track = r_tracked_stracks[it]
            if not track.state == 3: # If not already lost
                track.state = 3 # Lost
                lost_stracks.append(track)

        # 3.5 Associate unconfirmed tracks with remaining unmatched detections
        if unconfirmed and u_detection:
            remaining_dets = [detections[i] for i in u_detection]
            dists = iou_distance(unconfirmed, remaining_dets)
            matches_unc, u_unconfirmed, u_det_remaining = linear_assignment(dists, thresh=0.7)

            for itracked, idet in matches_unc:
                track = unconfirmed[itracked]
                det = remaining_dets[idet]
                track.update(det, self.frame_id, min_hits=self.min_hits)
                activated_stracks.append(track)
                self._sync_data(track, det)

            # Update u_detection to only contain indices not matched to unconfirmed
            matched_det_indices = set(u_detection[idet] for _, idet in matches_unc) if len(matches_unc) > 0 else set()
            u_detection = [i for i in u_detection if i not in matched_det_indices]

            # Unconfirmed tracks that didn't match → remove (too noisy to keep)
            for it in u_unconfirmed:
                track = unconfirmed[it]
                track.state = 4  # Removed
                removed_stracks.append(track)

        elif unconfirmed:
            # No detections left to match — remove all unconfirmed
            for track in unconfirmed:
                track.state = 4  # Removed
                removed_stracks.append(track)

        # 4. Init new tracks from unmatched high score detections
        # Note: Unmatched low score detections are ignored (noise)
        unmatched_dets = [detections[i] for i in u_detection]
        rejected_by_thresh = 0
        for track in unmatched_dets:
            if track.score < self.new_track_thresh:
                rejected_by_thresh += 1
                continue  # Not confident enough to start a new track

            track.activate(self.kalman_filter, self.frame_id)
            activated_stracks.append(track)
            self._sync_data(track, track)  # Sync self

        if rejected_by_thresh > 0 and self.frame_id <= 5:
            logging.warning(
                "Tracker frame %d: %d detections rejected by new_track_thresh=%.2f (scores: %s)",
                self.frame_id, rejected_by_thresh, self.new_track_thresh,
                [f"{t.score:.3f}" for t in unmatched_dets if t.score < self.new_track_thresh]
            )

        # 5. Update state
        self.tracked_stracks = [t for t in self.tracked_stracks if t.state == 2]
        self.tracked_stracks = join_stracks(self.tracked_stracks, activated_stracks)
        self.tracked_stracks = join_stracks(self.tracked_stracks, refind_stracks)
        self.lost_stracks = sub_stracks(self.lost_stracks, self.tracked_stracks)
        self.lost_stracks.extend(lost_stracks)
        self.lost_stracks = sub_stracks(self.lost_stracks, self.removed_stracks)
        self.removed_stracks.extend(removed_stracks)
        self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(self.tracked_stracks, self.lost_stracks)

        # 6. Age out lost tracks
        for track in self.lost_stracks:
            if self.frame_id - track.frame_id > self.track_buffer:
                self.removed_stracks.append(track)
        self.lost_stracks = [t for t in self.lost_stracks if self.frame_id - t.frame_id <= self.track_buffer]

        # 7. Final Output Construction
        # We need to update the original dictionaries in detections_list IN PLACE, 
        # or return a new list. The logic in inference.py expects us to modify detections dicts
        # or we might want to return the tracked ones.
        # But wait, we iterate `detections_list` at start.
        # We want to return ONLY the currently tracked/active objects? 
        # Usually inference pipeline draws ALL detections, but standard tracking ONLY output active tracks.
        # If we only output active tracks, we might suppress valid high-confidence detections that just started?
        # No, activated_stracks includes new ones.
        
        # Let's collect all active tracks
        output_stracks = [t for t in self.tracked_stracks if t.is_activated]
        
        results = []
        for track in output_stracks:
            d_out = track.original_data.copy() if hasattr(track, 'original_data') else {}

            tracked_bbox = track.tlbr
            d_out['bbox'] = [float(x) for x in tracked_bbox]
            d_out['track_id'] = f"T{str(track.track_id).zfill(2)}"

            # Restore GPT data if track has it and current detection didn't
            for k, v in track.gpt_data.items():
                if k not in d_out:
                    d_out[k] = v

            # --- Temporal validity check (INV-5, INV-11) ---
            assessment_frame = d_out.get('assessment_frame_index')
            if assessment_frame is not None:
                frames_since = self.frame_id - assessment_frame
                if frames_since > MAX_STALE_FRAMES:
                    d_out['assessment_status'] = AssessmentStatus.STALE
                    d_out['assessment_age_frames'] = frames_since
            elif d_out.get('assessment_status') != AssessmentStatus.ASSESSED:
                # INV-6: Unassessed objects get explicit UNASSESSED status
                d_out['assessment_status'] = AssessmentStatus.UNASSESSED

            # Update history
            if 'history' not in track.gpt_data:
                track.gpt_data['history'] = []
            track.gpt_data['history'].append(d_out['bbox'])
            if len(track.gpt_data['history']) > 30:
                track.gpt_data['history'].pop(0)
            d_out['history'] = track.gpt_data['history']

            results.append(d_out)

        if self.frame_id % 30 == 0:
            logging.info(
                "Tracker frame %d: %d tracked (%d activated), %d lost, %d new this frame, output=%d",
                self.frame_id,
                len(self.tracked_stracks),
                sum(1 for t in self.tracked_stracks if t.is_activated),
                len(self.lost_stracks),
                len(activated_stracks),
                len(results),
            )

        return results

    def _sync_data(self, track, det_source):
        """Propagate attributes like GPT data between track and detection."""
        # 1. From Source to Track (Update)
        source_data = det_source.original_data if hasattr(det_source, 'original_data') else {}
        for k in GPT_SYNC_KEYS:
            if k in source_data:
                track.gpt_data[k] = source_data[k]

        # 2. From Track to Source (Forward fill logic handled in output construction)

    def inject_metadata(self, tracked_dets):
        """Push metadata from post-processed detection dicts back into internal STrack objects.

        Needed because GPT results are added to detection dicts *after* tracker.update()
        returns, so the tracker's internal state doesn't have GPT data unless we
        explicitly push it back in.

        Records assessment_frame_index for temporal validity tracking (INV-5).
        """
        meta_by_tid = {}
        for d in tracked_dets:
            tid = d.get('track_id')
            if not tid:
                continue
            meta = {k: d[k] for k in GPT_SYNC_KEYS if k in d}
            if meta:
                # Ensure assessment_frame_index is recorded
                if "assessment_frame_index" not in meta and any(
                    k in meta for k in ("threat_level_score", "gpt_raw", "object_type")
                ):
                    meta["assessment_frame_index"] = self.frame_id
                    meta["assessment_status"] = AssessmentStatus.ASSESSED
                meta_by_tid[tid] = meta
        for track in self.tracked_stracks:
            tid_str = f"T{str(track.track_id).zfill(2)}"
            if tid_str in meta_by_tid:
                track.gpt_data.update(meta_by_tid[tid_str])


# --- Helper Functions ---

def linear_assignment(cost_matrix, thresh):
    """Linear assignment with threshold using scipy."""
    if cost_matrix.size == 0:
        return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
        
    matches, unmatched_a, unmatched_b = [], [], []
    
    # Scipy linear_sum_assignment finds min cost
    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    
    for r, c in zip(row_ind, col_ind):
        if cost_matrix[r, c] <= thresh:
            matches.append((r, c))
        else:
            unmatched_a.append(r)
            unmatched_b.append(c)
            
    # Add accumulation of indices that weren't selected
    # (scipy returns perfect matching for square, but partial for rectangular)
    # Actually scipy matches rows to cols. Any row not in row_ind is unmatched?
    # No, row_ind covers all rows if N < M.
    
    if cost_matrix.shape[0] > cost_matrix.shape[1]: # More rows than cols
         unmatched_a += list(set(range(cost_matrix.shape[0])) - set(row_ind))
    elif cost_matrix.shape[0] < cost_matrix.shape[1]: # More cols than rows
         unmatched_b += list(set(range(cost_matrix.shape[1])) - set(col_ind))
         
    # Also filter out threshold failures
    for r, c in zip(row_ind, col_ind):
        if cost_matrix[r, c] > thresh:
             if r not in unmatched_a: unmatched_a.append(r)
             if c not in unmatched_b: unmatched_b.append(c)

    # Clean up
    matches = np.array(matches) if len(matches) > 0 else np.empty((0, 2), dtype=int)
    return matches, unmatched_a, unmatched_b


def iou_distance(atracks, btracks):
    """Compute IOU cost matrix between tracks and detections."""
    if (len(atracks) == 0 and len(btracks) == 0) or len(atracks) == 0 or len(btracks) == 0:
        return np.zeros((len(atracks), len(btracks)), dtype=float)
        
    atlbrs = [track.tlbr for track in atracks]
    btlbrs = [track.tlbr for track in btracks]
    
    _ious = bbox_ious(np.array(atlbrs), np.array(btlbrs))
    cost_matrix = 1 - _ious
    return cost_matrix

def bbox_ious(boxes1, boxes2):
    """IOU matrix."""
    b1_x1, b1_y1, b1_x2, b1_y2 = boxes1[:, 0], boxes1[:, 1], boxes1[:, 2], boxes1[:, 3]
    b2_x1, b2_y1, b2_x2, b2_y2 = boxes2[:, 0], boxes2[:, 1], boxes2[:, 2], boxes2[:, 3]

    inter_rect_x1 = np.maximum(b1_x1[:, None], b2_x1)
    inter_rect_y1 = np.maximum(b1_y1[:, None], b2_y1)
    inter_rect_x2 = np.minimum(b1_x2[:, None], b2_x2)
    inter_rect_y2 = np.minimum(b1_y2[:, None], b2_y2)
    
    inter_area = np.maximum(inter_rect_x2 - inter_rect_x1, 0) * np.maximum(inter_rect_y2 - inter_rect_y1, 0)
    
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
    
    iou = inter_area / (b1_area[:, None] + b2_area - inter_area + 1e-6)
    return iou


def fuse_score(cost_matrix, detections):
    """Refine cost matrix with detection scores."""
    if cost_matrix.size == 0: return cost_matrix
    iou_sim = 1 - cost_matrix
    det_scores = np.array([d.score for d in detections])
    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


# STrack collection helpers

def join_stracks(tlist_a, tlist_b):
    exists = {}
    res = []
    for t in tlist_a:
        exists[t.track_id] = 1
        res.append(t)
    for t in tlist_b:
        tid = t.track_id
        if not exists.get(tid, 0):
            exists[tid] = 1
            res.append(t)
    return res

def sub_stracks(tlist_a, tlist_b):
    stracks = {}
    for t in tlist_a:
        stracks[t.track_id] = t
    for t in tlist_b:
        tid = t.track_id
        if stracks.get(tid, 0):
            del stracks[tid]
    return list(stracks.values())

def remove_duplicate_stracks(stracksa, stracksb):
    pdist = iou_distance(stracksa, stracksb)
    pairs = np.where(pdist < 0.15)
    dupa, dupb = list(pairs[0]), list(pairs[1])
    for a, b in zip(dupa, dupb):
        time_a = stracksa[a].frame_id - stracksa[a].start_frame
        time_b = stracksb[b].frame_id - stracksb[b].start_frame
        if time_a > time_b:
            dupb.append(b) # Bug in orig ByteTrack? It assumes removing from list. 
                           # We mark for removal.
        else:
            dupa.append(a)
            
    res_a = [t for i, t in enumerate(stracksa) if not i in dupa]
    res_b = [t for i, t in enumerate(stracksb) if not i in dupb]
    return res_a, res_b


# Monkey patch for multi_predict since STrack is not in a module
def multi_predict(stracks, kalman_filter):
    for t in stracks:
        if t.state != 2:
            t.mean[7] = 0 # reset velocity h if lost
        t.mean, t.covariance = kalman_filter.predict(t.mean, t.covariance)

STrack.multi_predict = static_method_multi_predict = multi_predict