"""
Object tracking capabilities for FMUS-VID.

This module provides AI-powered object tracking.
"""

import numpy as np
from typing import List, Dict, Tuple, Optional, Set
import logging
from dataclasses import dataclass
from collections import deque
import cv2

from .detection import Detection

logger = logging.getLogger(__name__)

@dataclass
class Track:
    """Represents a tracked object."""
    id: int
    label: str
    bbox: Tuple[float, float, float, float]  # x1, y1, x2, y2 (normalized)
    confidence: float
    history: deque  # Recent positions

class ObjectTracker:
    """Multi-object tracking using detection and motion prediction."""

    def __init__(self, max_history: int = 30, iou_threshold: float = 0.3):
        """
        Initialize object tracker.

        Args:
            max_history: Maximum track history length
            iou_threshold: IoU threshold for track matching
        """
        self.max_history = max_history
        self.iou_threshold = iou_threshold
        self.tracks: List[Track] = []
        self.next_id = 0
        self.frame_count = 0

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

        Args:
            detections: List of new detections

        Returns:
            List of active tracks
        """
        self.frame_count += 1

        # Convert detections to array format for matching
        if detections:
            det_boxes = np.array([d.bbox for d in detections])
            det_scores = np.array([d.confidence for d in detections])
            det_labels = [d.label for d in detections]
        else:
            det_boxes = np.empty((0, 4))
            det_scores = np.empty(0)
            det_labels = []

        # Get existing track boxes
        if self.tracks:
            track_boxes = np.array([t.bbox for t in self.tracks])
        else:
            track_boxes = np.empty((0, 4))

        # Compute IoU between all tracks and detections
        if len(track_boxes) > 0 and len(det_boxes) > 0:
            iou_matrix = self._compute_iou(track_boxes, det_boxes)

            # Match tracks to detections
            matched_indices = self._match_detections(iou_matrix)
        else:
            matched_indices = []

        # Update matched tracks
        matched_track_indices = set()
        matched_det_indices = set()

        for track_idx, det_idx in matched_indices:
            self._update_track(
                self.tracks[track_idx],
                det_boxes[det_idx],
                det_scores[det_idx],
                det_labels[det_idx]
            )
            matched_track_indices.add(track_idx)
            matched_det_indices.add(det_idx)

        # Create new tracks for unmatched detections
        for i in range(len(detections)):
            if i not in matched_det_indices:
                self._create_track(
                    det_boxes[i],
                    det_scores[i],
                    det_labels[i]
                )

        # Remove unmatched tracks
        self.tracks = [t for i, t in enumerate(self.tracks)
                      if i in matched_track_indices]

        return self.tracks

    def _compute_iou(self, boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
        """Compute IoU between two sets of boxes."""
        x11, y11, x12, y12 = np.split(boxes1, 4, axis=1)
        x21, y21, x22, y22 = np.split(boxes2, 4, axis=1)

        xA = np.maximum(x11, np.transpose(x21))
        yA = np.maximum(y11, np.transpose(y21))
        xB = np.minimum(x12, np.transpose(x22))
        yB = np.minimum(y12, np.transpose(y22))

        interArea = np.maximum((xB - xA), 0) * np.maximum((yB - yA), 0)

        box1Area = (x12 - x11) * (y12 - y11)
        box2Area = (x22 - x21) * (y22 - y21)

        iou = interArea / (box1Area + np.transpose(box2Area) - interArea)

        return iou

    def _match_detections(self, iou_matrix: np.ndarray) -> List[Tuple[int, int]]:
        """Match tracks to detections using IoU."""
        matches = []

        # Find matches greedily
        while True:
            # Get highest IoU
            if iou_matrix.size == 0:
                break

            max_iou = np.max(iou_matrix)
            if max_iou < self.iou_threshold:
                break

            # Get indices of best match
            track_idx, det_idx = np.unravel_index(
                np.argmax(iou_matrix),
                iou_matrix.shape
            )

            # Add match
            matches.append((track_idx, det_idx))

            # Remove matched entries
            iou_matrix[track_idx, :] = 0
            iou_matrix[:, det_idx] = 0

        return matches

    def _create_track(self, bbox: np.ndarray, score: float, label: str) -> None:
        """Create a new track."""
        track = Track(
            id=self.next_id,
            label=label,
            bbox=tuple(bbox),
            confidence=score,
            history=deque(maxlen=self.max_history)
        )
        track.history.append(bbox)
        self.tracks.append(track)
        self.next_id += 1

    def _update_track(self, track: Track, bbox: np.ndarray,
                     score: float, label: str) -> None:
        """Update an existing track."""
        track.bbox = tuple(bbox)
        track.confidence = score
        track.label = label
        track.history.append(bbox)

    def annotate_frame(self, frame: np.ndarray) -> np.ndarray:
        """
        Draw tracking information on frame.

        Args:
            frame: Input frame

        Returns:
            Annotated frame
        """
        result = frame.copy()
        height, width = frame.shape[:2]

        for track in self.tracks:
            # Convert normalized coordinates to pixel coordinates
            x1 = int(track.bbox[0] * width)
            y1 = int(track.bbox[1] * height)
            x2 = int(track.bbox[2] * width)
            y2 = int(track.bbox[3] * height)

            # Draw box
            cv2.rectangle(result, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # Draw label with ID
            label = f"{track.label} #{track.id} {track.confidence:.2f}"
            cv2.putText(result, label, (x1, y1-10),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            # Draw motion trail
            if len(track.history) > 1:
                points = [(int(x[0] * width), int(x[1] * height))
                         for x in track.history]
                for i in range(len(points)-1):
                    cv2.line(result, points[i], points[i+1], (0, 255, 0), 2)

        return result
