"""
Ball tracking module with buffer-based filtering and trajectory smoothing
"""
import numpy as np
from typing import List, Dict, Optional, Tuple
from collections import deque
import cv2


class BallTracker:
    """
    Tracks ball across frames using buffer-based filtering
    Handles occlusions and false detections
    """

    def __init__(
        self,
        buffer_size: int = 10,
        max_distance_threshold: int = 100,
        min_confidence: float = 0.3
    ):
        """
        Initialize ball tracker

        Args:
            buffer_size: Number of recent positions to store for filtering
            max_distance_threshold: Maximum pixel distance between frames to consider same ball
            min_confidence: Minimum confidence threshold for detections
        """
        self.buffer_size = buffer_size
        self.max_distance_threshold = max_distance_threshold
        self.min_confidence = min_confidence

        # Buffer for storing recent ball positions
        self.position_buffer = deque(maxlen=buffer_size)

        # Track ball state
        self.current_position = None
        self.lost_frames = 0
        self.max_lost_frames = 10  # Maximum frames to interpolate when ball is lost

        # Trajectory history
        self.trajectory = []

    def update(self, detections: List[Dict], frame_number: int) -> Optional[Dict]:
        """
        Update tracker with new detections from current frame

        Args:
            detections: List of ball detections from detector
            frame_number: Current frame number

        Returns:
            Best ball detection (filtered), or None if no valid detection
        """
        # Filter low confidence detections
        valid_detections = [
            det for det in detections
            if det['confidence'] >= self.min_confidence
        ]

        if len(valid_detections) == 0:
            # No detections - try to interpolate if recently had detection
            return self._handle_missing_detection(frame_number)

        # Select best detection
        best_detection = self._select_best_detection(valid_detections)

        if best_detection is None:
            return self._handle_missing_detection(frame_number)

        # Update buffer and state
        self.position_buffer.append(best_detection['center'])
        self.current_position = best_detection['center']
        self.lost_frames = 0

        # Add to trajectory
        self.trajectory.append({
            'frame': frame_number,
            'position': best_detection['center'],
            'pixel_coords': best_detection['center'],
            'real_coords': None,  # Will be filled by processor
            'confidence': best_detection['confidence']
        })

        return best_detection

    def _select_best_detection(self, detections: List[Dict]) -> Optional[Dict]:
        """
        Select the most likely ball detection from multiple candidates

        Args:
            detections: List of valid detections

        Returns:
            Best detection, or None
        """
        if len(detections) == 0:
            return None

        if len(detections) == 1:
            return detections[0]

        # If we have position history, use it to filter
        if len(self.position_buffer) > 0:
            return self._select_by_proximity(detections)
        else:
            # No history - return highest confidence
            return max(detections, key=lambda d: d['confidence'])

    def _select_by_proximity(self, detections: List[Dict]) -> Optional[Dict]:
        """
        Select detection closest to predicted position based on buffer

        Args:
            detections: List of detections

        Returns:
            Detection closest to predicted position
        """
        # Calculate average position from buffer
        avg_position = np.mean(self.position_buffer, axis=0)

        # Find detection closest to average
        min_distance = float('inf')
        best_detection = None

        for det in detections:
            distance = np.linalg.norm(
                np.array(det['center']) - avg_position
            )

            # Check if within threshold
            if distance < self.max_distance_threshold and distance < min_distance:
                min_distance = distance
                best_detection = det

        # If no detection within threshold, return highest confidence
        if best_detection is None:
            best_detection = max(detections, key=lambda d: d['confidence'])

        return best_detection

    def _handle_missing_detection(self, frame_number: int) -> Optional[Dict]:
        """
        Handle case when no valid detection in current frame

        Args:
            frame_number: Current frame number

        Returns:
            Interpolated detection if possible, None otherwise
        """
        self.lost_frames += 1

        # If lost for too long, stop interpolating
        if self.lost_frames > self.max_lost_frames:
            self.current_position = None
            return None

        # Try to interpolate based on recent positions
        if len(self.position_buffer) >= 2:
            interpolated = self._interpolate_position()

            # Add interpolated position to trajectory (marked as interpolated)
            self.trajectory.append({
                'frame': frame_number,
                'position': interpolated,
                'pixel_coords': interpolated,
                'real_coords': None,
                'confidence': 0.0,  # Mark as interpolated
                'interpolated': True
            })

            return {
                'center': interpolated,
                'bbox': self._estimate_bbox(interpolated),
                'confidence': 0.0,
                'interpolated': True
            }

        return None

    def _interpolate_position(self) -> Tuple[float, float]:
        """
        Interpolate ball position based on recent trajectory

        Returns:
            Estimated (x, y) position
        """
        if len(self.position_buffer) < 2:
            return self.current_position

        # Simple linear interpolation based on last two positions
        positions = np.array(self.position_buffer)

        # Calculate velocity
        velocity = positions[-1] - positions[-2]

        # Predict next position
        predicted = positions[-1] + velocity

        return tuple(predicted)

    def _estimate_bbox(self, center: Tuple[float, float], size: int = 20) -> List[float]:
        """
        Estimate bounding box for interpolated position

        Args:
            center: Center position
            size: Estimated ball size in pixels

        Returns:
            [x1, y1, x2, y2] bbox
        """
        cx, cy = center
        half_size = size / 2
        return [cx - half_size, cy - half_size, cx + half_size, cy + half_size]

    def get_trajectory(self) -> List[Dict]:
        """
        Get full trajectory history

        Returns:
            List of trajectory points
        """
        return self.trajectory

    def get_smoothed_trajectory(self, window_size: int = 5) -> List[Dict]:
        """
        Get smoothed trajectory using moving average

        Args:
            window_size: Size of smoothing window

        Returns:
            Smoothed trajectory
        """
        if len(self.trajectory) < window_size:
            return self.trajectory

        smoothed = []

        for i, point in enumerate(self.trajectory):
            # Get window
            start = max(0, i - window_size // 2)
            end = min(len(self.trajectory), i + window_size // 2 + 1)
            window = self.trajectory[start:end]

            # Calculate average position (only non-interpolated points)
            valid_positions = [
                p['position'] for p in window
                if not p.get('interpolated', False)
            ]

            if len(valid_positions) > 0:
                avg_position = np.mean(valid_positions, axis=0)
                smoothed_point = point.copy()
                smoothed_point['smoothed_position'] = tuple(avg_position)
                smoothed.append(smoothed_point)
            else:
                smoothed.append(point)

        return smoothed

    def reset(self):
        """Reset tracker state"""
        self.position_buffer.clear()
        self.current_position = None
        self.lost_frames = 0
        self.trajectory = []

    def draw_trajectory(self, frame: np.ndarray, max_points: int = 30) -> np.ndarray:
        """
        Draw ball trajectory on frame

        Args:
            frame: Input frame
            max_points: Maximum number of trajectory points to draw

        Returns:
            Frame with trajectory overlay
        """
        if len(self.trajectory) == 0:
            return frame

        overlay = frame.copy()

        # Get recent trajectory points
        recent_trajectory = self.trajectory[-max_points:]

        # Draw trajectory line
        points = [point['position'] for point in recent_trajectory]
        if len(points) > 1:
            pts = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
            cv2.polylines(overlay, [pts], False, (0, 255, 255), 2)

        # Draw current position
        if self.current_position is not None:
            cv2.circle(
                overlay,
                (int(self.current_position[0]), int(self.current_position[1])),
                8,
                (0, 0, 255),
                -1
            )

        return overlay


class MultiObjectTracker:
    """
    Extended tracker for tracking multiple objects (ball, players, paddles)
    Uses simple centroid-based tracking
    """

    def __init__(self, max_disappeared: int = 10, max_distance: int = 100):
        """
        Initialize multi-object tracker

        Args:
            max_disappeared: Max frames object can disappear before removing
            max_distance: Max distance to associate detection with existing object
        """
        self.next_object_id = 0
        self.objects = {}  # object_id -> centroid
        self.disappeared = {}  # object_id -> num_frames_disappeared
        self.max_disappeared = max_disappeared
        self.max_distance = max_distance

    def register(self, centroid: Tuple[float, float]) -> int:
        """Register new object"""
        object_id = self.next_object_id
        self.objects[object_id] = centroid
        self.disappeared[object_id] = 0
        self.next_object_id += 1
        return object_id

    def deregister(self, object_id: int):
        """Remove object from tracking"""
        del self.objects[object_id]
        del self.disappeared[object_id]

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

        Args:
            detections: List of detections with 'center' key

        Returns:
            Dict mapping object_id to detection
        """
        # If no detections, mark all as disappeared
        if len(detections) == 0:
            for object_id in list(self.disappeared.keys()):
                self.disappeared[object_id] += 1
                if self.disappeared[object_id] > self.max_disappeared:
                    self.deregister(object_id)
            return {}

        centroids = np.array([det['center'] for det in detections])

        # If no objects being tracked, register all
        if len(self.objects) == 0:
            result = {}
            for i, det in enumerate(detections):
                object_id = self.register(centroids[i])
                result[object_id] = det
            return result

        # Associate detections with existing objects
        object_ids = list(self.objects.keys())
        object_centroids = np.array([self.objects[oid] for oid in object_ids])

        # Calculate distance matrix
        D = np.linalg.norm(
            object_centroids[:, np.newaxis] - centroids,
            axis=2
        )

        # Find best matches
        rows = D.min(axis=1).argsort()
        cols = D.argmin(axis=1)[rows]

        used_rows = set()
        used_cols = set()
        result = {}

        for row, col in zip(rows, cols):
            if row in used_rows or col in used_cols:
                continue

            if D[row, col] > self.max_distance:
                continue

            object_id = object_ids[row]
            self.objects[object_id] = centroids[col]
            self.disappeared[object_id] = 0
            result[object_id] = detections[col]

            used_rows.add(row)
            used_cols.add(col)

        # Handle disappeared objects
        unused_rows = set(range(D.shape[0])) - used_rows
        for row in unused_rows:
            object_id = object_ids[row]
            self.disappeared[object_id] += 1
            if self.disappeared[object_id] > self.max_disappeared:
                self.deregister(object_id)

        # Register new objects
        unused_cols = set(range(D.shape[1])) - used_cols
        for col in unused_cols:
            object_id = self.register(centroids[col])
            result[object_id] = detections[col]

        return result
