"""
Lightweight ByteTrack Tracker for multi-person tracking.

Implements KalmanFilter-based tracking with IoU + confidence matching.
Track lifecycle: tentative (new) -> confirmed (3+ hits) -> removed (30 lost frames).
"""

from typing import List, Dict, Optional, Tuple
import numpy as np
from scipy.optimize import linear_sum_assignment
from collections import defaultdict


class KalmanFilter:
    """Simple Kalman filter for bounding box tracking."""

    def __init__(self):
        # State: [x, y, s, r, vx, vy, vs]
        # x, y = center, s = area, r = aspect ratio
        # vx, vy, vs = velocities
        self.dt = 1.0

        # State transition matrix
        self.F = np.array([
            [1, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 1, 0],
            [0, 0, 1, 0, 0, 0, 1],
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 0, 1],
        ], dtype=np.float32)

        # Measurement matrix
        self.H = np.array([
            [1, 0, 0, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0],
        ], dtype=np.float32)

        # Process noise
        self.Q = np.eye(7, dtype=np.float32) * 0.01
        self.Q[4:, 4:] *= 0.1

        # Measurement noise
        self.R = np.eye(4, dtype=np.float32) * 0.1

        # Error covariance
        self.P = np.eye(7, dtype=np.float32) * 10

        self.mean = np.zeros(7, dtype=np.float32)
        self.covariance = self.P.copy()

    def initiate(self, measurement: np.ndarray) -> None:
        """Initialize state from a detection measurement [x, y, s, r]."""
        self.mean[:4] = measurement
        self.mean[4:] = 0
        self.covariance = self.P.copy()

    def predict(self) -> np.ndarray:
        """Predict next state. Returns predicted [x, y, s, r]."""
        self.mean = self.F @ self.mean
        self.covariance = self.F @ self.covariance @ self.F.T + self.Q
        return self.mean[:4].copy()

    def update(self, measurement: np.ndarray) -> None:
        """Update state with measurement [x, y, s, r]."""
        innovation = measurement - self.H @ self.mean
        S = self.H @ self.covariance @ self.H.T + self.R
        K = self.covariance @ self.H.T @ np.linalg.inv(S)
        self.mean = self.mean + K @ innovation
        self.covariance = (np.eye(7) - K @ self.H) @ self.covariance

    def get_state(self) -> np.ndarray:
        """Get current [x, y, s, r] state estimate."""
        return self.mean[:4].copy()


def bbox_to_state(bbox: List[float]) -> np.ndarray:
    """Convert [x1, y1, x2, y2] to [cx, cy, area, aspect_ratio]."""
    x1, y1, x2, y2 = bbox
    w = max(x2 - x1, 1.0)
    h = max(y2 - y1, 1.0)
    cx = (x1 + x2) / 2.0
    cy = (y1 + y2) / 2.0
    area = w * h
    ar = w / h
    return np.array([cx, cy, area, ar], dtype=np.float32)


def state_to_bbox(state: np.ndarray) -> List[float]:
    """Convert [cx, cy, area, ar] to [x1, y1, x2, y2]."""
    cx, cy, area, ar = state
    w = np.sqrt(area * ar)
    h = area / w
    x1 = cx - w / 2
    y1 = cy - h / 2
    x2 = cx + w / 2
    y2 = cy + h / 2
    return [float(x1), float(y1), float(x2), float(y2)]


def compute_iou(bbox1: List[float], bbox2: List[float]) -> float:
    """Compute IoU between two bounding boxes."""
    x1 = max(bbox1[0], bbox2[0])
    y1 = max(bbox1[1], bbox2[1])
    x2 = min(bbox1[2], bbox2[2])
    y2 = min(bbox1[3], bbox2[3])

    inter = max(0, x2 - x1) * max(0, y2 - y1)
    if inter <= 0:
        return 0.0

    area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
    area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
    union = area1 + area2 - inter

    return inter / union if union > 0 else 0.0


class Track:
    """Single track with Kalman filtering and state management."""

    def __init__(self, track_id: int, bbox: List[float], confidence: float):
        self.track_id = track_id
        self.bbox = bbox
        self.confidence = confidence
        self.state = 'tentative'  # tentative, confirmed, lost, removed
        self.age = 0
        self.hits = 1  # Number of successful matches
        self.lost_frames = 0  # Consecutive frames without match
        self.kalman = KalmanFilter()
        self.kalman.initiate(bbox_to_state(bbox))
        self.prev_bbox = bbox.copy()
        self.speed = 0.0  # Movement speed in pixels/frame

    def predict(self) -> List[float]:
        """Predict next bbox position."""
        predicted_state = self.kalman.predict()
        self.bbox = state_to_bbox(predicted_state)
        return self.bbox

    def update(self, bbox: List[float], confidence: float) -> None:
        """Update track with new detection."""
        self.prev_bbox = self.bbox.copy()
        self.bbox = bbox
        self.confidence = confidence
        self.hits += 1
        self.lost_frames = 0
        self.age += 1

        # Compute speed
        dx = (self.bbox[0] + self.bbox[2]) / 2 - (self.prev_bbox[0] + self.prev_bbox[2]) / 2
        dy = (self.bbox[1] + self.bbox[3]) / 2 - (self.prev_bbox[1] + self.prev_bbox[3]) / 2
        self.speed = np.sqrt(dx**2 + dy**2)

        # Update Kalman
        measurement = bbox_to_state(bbox)
        self.kalman.update(measurement)

        # Promote to confirmed after 3 hits
        if self.state == 'tentative' and self.hits >= 3:
            self.state = 'confirmed'

    def mark_lost(self) -> None:
        """Mark track as lost (no matching detection this frame)."""
        self.lost_frames += 1
        self.age += 1
        if self.state == 'confirmed':
            self.state = 'lost'
        if self.lost_frames >= 30:
            self.state = 'removed'

    def is_removed(self) -> bool:
        """Check if track should be removed."""
        return self.state == 'removed' or self.lost_frames >= 30


class TrackManager:
    """Manages all tracks across frames."""

    def __init__(self, max_lost_frames: int = 30, min_hits: int = 3):
        self.tracks: Dict[int, Track] = {}
        self.next_id = 1
        self.max_lost_frames = max_lost_frames
        self.min_hits = min_hits

    def update(self, detections: List[Dict]) -> List[Dict]:
        """
        Update tracks with new detections.

        Args:
            detections: List of detection dicts with 'bbox', 'confidence'.

        Returns:
            List of track dicts with 'bbox', 'track_id', 'confidence', 'state'.
        """
        # Predict existing tracks
        for track in self.tracks.values():
            if track.state != 'removed':
                track.predict()

        # Remove dead tracks
        self.tracks = {
            tid: t for tid, t in self.tracks.items()
            if t.state != 'removed'
        }

        if not detections:
            # Mark all as lost
            for track in self.tracks.values():
                track.mark_lost()
            return self._get_active_tracks()

        # Separate high and low confidence detections
        high_dets = [d for d in detections if d['confidence'] >= 0.5]
        low_dets = [d for d in detections if d['confidence'] < 0.5]

        # Get confirmed and tentative tracks
        confirmed_tracks = {
            tid: t for tid, t in self.tracks.items()
            if t.state in ('confirmed', 'lost')
        }
        tentative_tracks = {
            tid: t for tid, t in self.tracks.items()
            if t.state == 'tentative'
        }
        lost_tracks = {
            tid: t for tid, t in self.tracks.items()
            if t.state == 'lost'
        }

        # First match: high score detections to confirmed tracks
        matched_dets = set()
        matched_tracks = set()

        if confirmed_tracks and high_dets:
            matched = self._match_tracks(confirmed_tracks, high_dets, min_iou=0.3)
            for track_id, det_idx in matched.items():
                t = self.tracks[track_id]
                t.update(high_dets[det_idx]['bbox'], high_dets[det_idx]['confidence'])
                matched_dets.add(det_idx)
                matched_tracks.add(track_id)

        # Second match: remaining high detections to tentative tracks
        remaining_high = [
            d for i, d in enumerate(high_dets) if i not in matched_dets
        ]
        remaining_tentative = {
            tid: t for tid, t in tentative_tracks.items()
            if tid not in matched_tracks
        }

        if remaining_tentative and remaining_high:
            matched = self._match_tracks(remaining_tentative, remaining_high, min_iou=0.5)
            for track_id, det_idx in matched.items():
                actual_idx = [i for i, d in enumerate(high_dets)
                              if d is remaining_high[det_idx]][0]
                t = self.tracks[track_id]
                t.update(high_dets[actual_idx]['bbox'], high_dets[actual_idx]['confidence'])
                matched_dets.add(actual_idx)
                matched_tracks.add(track_id)

        # Third match: low score detections to unmatched confirmed tracks
        remaining_low = [d for i, d in enumerate(low_dets)]
        unmatched_confirmed = {
            tid: t for tid, t in confirmed_tracks.items()
            if tid not in matched_tracks
        }

        if unmatched_confirmed and remaining_low:
            matched = self._match_tracks(unmatched_confirmed, remaining_low, min_iou=0.3)
            for track_id, det_idx in matched.items():
                t = self.tracks[track_id]
                t.update(low_dets[det_idx]['bbox'], low_dets[det_idx]['confidence'])
                matched_tracks.add(track_id)

        # Mark unmatched confirmed tracks as lost
        for tid, t in list(self.tracks.items()):
            if tid not in matched_tracks and t.state in ('confirmed', 'lost'):
                t.mark_lost()

        # Create new tracks for unmatched high-confidence detections
        for i, d in enumerate(high_dets):
            if i not in matched_dets:
                track = Track(self.next_id, d['bbox'], d['confidence'])
                self.tracks[self.next_id] = track
                self.next_id += 1

        # Clean removed tracks
        self.tracks = {
            tid: t for tid, t in self.tracks.items()
            if not t.is_removed()
        }

        return self._get_active_tracks()

    def _match_tracks(
        self,
        tracks: Dict[int, Track],
        detections: List[Dict],
        min_iou: float = 0.3
    ) -> Dict[int, int]:
        """
        Match tracks to detections using IoU-based Hungarian algorithm.

        Args:
            tracks: Dict of track_id -> Track.
            detections: List of detection dicts.
            min_iou: Minimum IoU threshold.

        Returns:
            Dict of track_id -> detection_index for matches.
        """
        if not tracks or not detections:
            return {}

        track_ids = list(tracks.keys())
        track_bboxes = [tracks[tid].bbox for tid in track_ids]
        det_bboxes = [d['bbox'] for d in detections]

        # Build IoU matrix
        iou_matrix = np.zeros((len(track_ids), len(det_bboxes)), dtype=np.float32)
        for i, tbox in enumerate(track_bboxes):
            for j, dbox in enumerate(det_bboxes):
                iou_matrix[i, j] = compute_iou(tbox, dbox)

        # Hungarian matching
        row_indices, col_indices = linear_sum_assignment(-iou_matrix)

        # Filter by IoU threshold
        matches = {}
        for i, j in zip(row_indices, col_indices):
            if iou_matrix[i, j] >= min_iou:
                matches[track_ids[i]] = j

        return matches

    def _get_active_tracks(self) -> List[Dict]:
        """Get list of active tracks (confirmed + tentative)."""
        results = []
        for track in self.tracks.values():
            if track.state in ('tentative', 'confirmed', 'lost'):
                results.append({
                    'track_id': track.track_id,
                    'bbox': track.bbox,
                    'confidence': track.confidence,
                    'state': track.state,
                    'speed': track.speed,
                    'lost_frames': track.lost_frames,
                })
        return results

    def get_track(self, track_id: int) -> Optional[Track]:
        """Get track by ID."""
        return self.tracks.get(track_id)

    def get_all_tracks(self) -> Dict[int, Track]:
        """Get all tracks."""
        return self.tracks

    def reset(self) -> None:
        """Reset all tracks."""
        self.tracks = {}
        self.next_id = 1


if __name__ == "__main__":
    # Quick test
    manager = TrackManager()

    # Simulate 5 frames of a person walking
    for frame_id in range(5):
        x = 100 + frame_id * 10
        detections = [{'bbox': [x, 100, x + 50, 200], 'confidence': 0.9}]
        tracks = manager.update(detections)
        print(f"Frame {frame_id}: {len(tracks)} tracks, IDs: {[t['track_id'] for t in tracks]}")
        for t in tracks:
            print(f"  Track {t['track_id']}: state={t['state']}, bbox={t['bbox']}")

    print("Tracker module OK!")