"""ByteTrack: Multi-Object Tracking by Associating Every Detection Box
Based on: https://arxiv.org/abs/2110.06864
"""

import numpy as np
from typing import List, Tuple
try:
    import lap
    LAP_AVAILABLE = True
except ImportError:
    LAP_AVAILABLE = False


class KalmanBoxTracker:
    """Kalman Filter for tracking bounding boxes."""
    count = 0
    
    def __init__(self, bbox):
        """Initialize tracker with bounding box [x1, y1, x2, y2]."""
        # State: [x_center, y_center, area, aspect_ratio, vx, vy, va, vr]
        self.kf_x = np.zeros(8)
        self.kf_x[:4] = self._convert_bbox_to_z(bbox)
        
        # Covariance matrix
        self.kf_P = np.eye(8) * 10.
        self.kf_P[4:, 4:] *= 1000.  # High uncertainty for velocities
        
        # State transition matrix
        self.kf_F = np.eye(8)
        for i in range(4):
            self.kf_F[i, i+4] = 1.0  # Position += velocity
        
        # Measurement matrix
        self.kf_H = np.eye(4, 8)
        
        # Process noise
        self.kf_Q = np.eye(8)
        self.kf_Q[4:, 4:] *= 0.01
        
        # Measurement noise
        self.kf_R = np.eye(4)
        
        self.time_since_update = 0
        self.id = KalmanBoxTracker.count
        KalmanBoxTracker.count += 1
        self.history = []
        self.hits = 0
        self.hit_streak = 0
        self.age = 0
        
    def _convert_bbox_to_z(self, bbox):
        """Convert [x1,y1,x2,y2] to [x,y,s,r] where s=area, r=aspect_ratio."""
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]
        x = bbox[0] + w/2.
        y = bbox[1] + h/2.
        s = w * h
        r = w / float(h) if h != 0 else 1.0
        return np.array([x, y, s, r])
    
    def _convert_x_to_bbox(self, x):
        """Convert [x,y,s,r] to [x1,y1,x2,y2]."""
        area_ratio = max(x[2] * x[3], 1e-6)  # clamp to avoid sqrt(negative)
        w = np.sqrt(area_ratio)
        h = x[2] / w if w > 1e-6 else 1.0
        return np.array([
            x[0] - w/2.,
            x[1] - h/2.,
            x[0] + w/2.,
            x[1] + h/2.
        ])
    
    def update(self, bbox):
        """Update state with observed bbox."""
        self.time_since_update = 0
        self.history = []
        self.hits += 1
        self.hit_streak += 1
        
        z = self._convert_bbox_to_z(bbox)
        
        # Kalman update (H is 4x8, x is 8, P is 8x8)
        y = z - self.kf_H @ self.kf_x                    # (4,)
        S = self.kf_H @ self.kf_P @ self.kf_H.T + self.kf_R  # (4,4)
        K = self.kf_P @ self.kf_H.T @ np.linalg.inv(S)       # (8,4)
        self.kf_x = self.kf_x + K @ y                         # (8,)
        I_KH = np.eye(len(self.kf_x)) - K @ self.kf_H         # (8,8)
        self.kf_P = I_KH @ self.kf_P                           # (8,8)
        
    def predict(self):
        """Predict next state."""
        # Kalman predict
        self.kf_x = np.dot(self.kf_F, self.kf_x)
        self.kf_P = np.dot(np.dot(self.kf_F, self.kf_P), self.kf_F.T) + self.kf_Q
        
        self.age += 1
        if self.time_since_update > 0:
            self.hit_streak = 0
        self.time_since_update += 1
        self.history.append(self._convert_x_to_bbox(self.kf_x))
        return self.history[-1]
    
    def get_state(self):
        """Return current bounding box estimate."""
        return self._convert_x_to_bbox(self.kf_x)


def iou_batch(bb_test, bb_gt):
    """Compute IoU between two sets of boxes."""
    bb_gt = np.expand_dims(bb_gt, 0)
    bb_test = np.expand_dims(bb_test, 1)
    
    xx1 = np.maximum(bb_test[..., 0], bb_gt[..., 0])
    yy1 = np.maximum(bb_test[..., 1], bb_gt[..., 1])
    xx2 = np.minimum(bb_test[..., 2], bb_gt[..., 2])
    yy2 = np.minimum(bb_test[..., 3], bb_gt[..., 3])
    w = np.maximum(0., xx2 - xx1)
    h = np.maximum(0., yy2 - yy1)
    wh = w * h
    
    area_test = (bb_test[..., 2] - bb_test[..., 0]) * (bb_test[..., 3] - bb_test[..., 1])
    area_gt = (bb_gt[..., 2] - bb_gt[..., 0]) * (bb_gt[..., 3] - bb_gt[..., 1])
    
    iou = wh / (area_test + area_gt - wh)
    return iou


def linear_assignment(cost_matrix):
    """Solve linear assignment problem."""
    if not LAP_AVAILABLE:
        # Fallback: greedy assignment
        matches = []
        cost = cost_matrix.copy()
        for _ in range(min(cost.shape)):
            i, j = np.unravel_index(cost.argmin(), cost.shape)
            if cost[i, j] < 1e9:
                matches.append([i, j])
                cost[i, :] = 1e9
                cost[:, j] = 1e9
        
        matches = np.array(matches) if matches else np.empty((0, 2), dtype=int)
        unmatched_a = list(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
        unmatched_b = list(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
        return matches, unmatched_a, unmatched_b
    
    # Use lap for optimal assignment
    _, x, y = lap.lapjv(cost_matrix, extend_cost=True)
    matches = [[i, x[i]] for i in range(len(x)) if x[i] >= 0]
    unmatched_a = [i for i in range(len(x)) if x[i] < 0]
    unmatched_b = [j for j in range(len(y)) if y[j] < 0]
    return np.array(matches), unmatched_a, unmatched_b


class ByteTrackTracker:
    """ByteTrack tracker with high/low confidence detection association."""
    
    def __init__(
        self,
        max_age: int = 30,
        min_hits: int = 3,
        iou_threshold: float = 0.3,
        high_threshold: float = 0.5,
        low_threshold: float = 0.1
    ):
        """
        Args:
            max_age: Maximum frames to keep lost tracks
            min_hits: Minimum hits to confirm track
            iou_threshold: IoU threshold for matching
            high_threshold: High confidence threshold
            low_threshold: Low confidence threshold
        """
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold
        self.high_threshold = high_threshold
        self.low_threshold = low_threshold
        self.trackers = []
        self.frame_count = 0
    
    def update(self, detections):
        """
        Args:
            detections: nx5 array [x1, y1, x2, y2, score]
        
        Returns:
            nx5 array [x1, y1, x2, y2, track_id]
        """
        self.frame_count += 1
        
        # Get predictions from existing trackers
        trks = np.zeros((len(self.trackers), 5))
        to_del = []
        for t, trk in enumerate(trks):
            pos = self.trackers[t].predict()
            trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
            if np.any(np.isnan(pos)):
                to_del.append(t)
        
        trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
        for t in reversed(to_del):
            self.trackers.pop(t)
        
        # Split detections by confidence
        if len(detections) > 0:
            high_dets = detections[detections[:, 4] >= self.high_threshold]
            low_dets = detections[(detections[:, 4] >= self.low_threshold) & 
                                 (detections[:, 4] < self.high_threshold)]
        else:
            high_dets = np.empty((0, 5))
            low_dets = np.empty((0, 5))
        
        # First association: high confidence detections
        matched, unmatched_dets, unmatched_trks = self._associate_detections_to_trackers(
            high_dets, trks, self.iou_threshold
        )
        
        # Update matched trackers
        for m in matched:
            self.trackers[m[1]].update(high_dets[m[0], :4])
        
        # Second association: low confidence detections with unmatched tracks
        if len(low_dets) > 0 and len(unmatched_trks) > 0:
            unmatched_trks_boxes = trks[unmatched_trks]
            matched_low, unmatched_dets_low, unmatched_trks_low = self._associate_detections_to_trackers(
                low_dets, unmatched_trks_boxes, self.iou_threshold
            )
            
            for m in matched_low:
                self.trackers[unmatched_trks[m[1]]].update(low_dets[m[0], :4])
            
            unmatched_trks = [unmatched_trks[i] for i in unmatched_trks_low]
        
        # Create new trackers for unmatched high confidence detections
        for i in unmatched_dets:
            trk = KalmanBoxTracker(high_dets[i, :4])
            self.trackers.append(trk)
        
        # Remove dead trackers
        i = len(self.trackers)
        ret = []
        for trk in reversed(self.trackers):
            d = trk.get_state()
            if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits):
                ret.append(np.concatenate((d, [trk.id])).reshape(1, -1))
            i -= 1
            if trk.time_since_update > self.max_age:
                self.trackers.pop(i)
        
        if len(ret) > 0:
            return np.concatenate(ret)
        return np.empty((0, 5))
    
    def _associate_detections_to_trackers(self, detections, trackers, iou_threshold):
        """Associate detections to tracked objects."""
        n_dets = len(detections)
        n_trks = len(trackers)

        if n_trks == 0:
            return np.empty((0, 2), dtype=int), list(range(n_dets)), []

        iou_matrix = iou_batch(detections[:, :4], trackers[:, :4])

        # Default: everything unmatched
        unmatched_a = list(range(n_dets))
        unmatched_b = list(range(n_trks))

        if min(iou_matrix.shape) > 0:
            a = (iou_matrix > iou_threshold).astype(np.int32)
            if a.sum(1).max() == 1 and a.sum(0).max() == 1:
                matched_indices = np.stack(np.where(a), axis=1)
            else:
                matched_indices, unmatched_a, unmatched_b = linear_assignment(
                    -iou_matrix
                )
                # Ensure lists (linear_assignment may return np arrays)
                unmatched_a = list(unmatched_a)
                unmatched_b = list(unmatched_b)
        else:
            matched_indices = np.empty((0, 2), dtype=int)

        # Filter out matched with low IoU
        matches = []
        for m in matched_indices:
            d_idx, t_idx = int(m[0]), int(m[1])
            if iou_matrix[d_idx, t_idx] < iou_threshold:
                if d_idx not in unmatched_a:
                    unmatched_a.append(d_idx)
                if t_idx not in unmatched_b:
                    unmatched_b.append(t_idx)
            else:
                matches.append(m.reshape(1, 2))
                # Remove from unmatched if present
                if d_idx in unmatched_a:
                    unmatched_a.remove(d_idx)
                if t_idx in unmatched_b:
                    unmatched_b.remove(t_idx)

        if len(matches) == 0:
            matches = np.empty((0, 2), dtype=int)
        else:
            matches = np.concatenate(matches, axis=0)

        return matches, unmatched_a, unmatched_b
