"""Advanced jitter reduction and smoothing for tracking."""

import numpy as np
from typing import Dict, Optional

try:
    from filterpy.kalman import KalmanFilter as _KalmanFilter
    _HAS_FILTERPY = True
except ImportError:
    _HAS_FILTERPY = False


class KalmanSmoothing:
    """Kalman filter for smooth tracking (jitter reduction).
    
    Use for: Smooth camera movement, reduce box jitter with predictive modeling.
    """
    
    def __init__(
        self,
        process_noise: float = 0.01,
        measurement_noise: float = 0.1
    ):
        """
        Args:
            process_noise: How much the object can move between frames (lower = smoother)
            measurement_noise: Trust in measurements (lower = trust more)
        """
        if not _HAS_FILTERPY:
            raise ImportError(
                "KalmanSmoothing requires filterpy. "
                "Install with: pip install filterpy"
            )
        self.process_noise = process_noise
        self.measurement_noise = measurement_noise
        self.filters: Dict[int, '_KalmanFilter'] = {}
    
    def update(self, tracker_id: int, box: np.ndarray) -> np.ndarray:
        """Smooth box using Kalman filter.
        
        Args:
            tracker_id: Object ID
            box: [x1, y1, x2, y2]
        
        Returns:
            Smoothed box
        """
        if tracker_id not in self.filters:
            self._init_filter(tracker_id, box)
        
        kf = self.filters[tracker_id]
        
        # Measurement: [x_center, y_center, width, height]
        cx = (box[0] + box[2]) / 2
        cy = (box[1] + box[3]) / 2
        w = box[2] - box[0]
        h = box[3] - box[1]
        measurement = np.array([cx, cy, w, h])
        
        # Predict and update
        kf.predict()
        kf.update(measurement)
        
        # Get smoothed state
        state = kf.x.flatten()
        smoothed_cx, smoothed_cy, smoothed_w, smoothed_h = state[:4]
        
        # Convert back to xyxy
        return np.array([
            smoothed_cx - smoothed_w / 2,
            smoothed_cy - smoothed_h / 2,
            smoothed_cx + smoothed_w / 2,
            smoothed_cy + smoothed_h / 2
        ])
    
    def _init_filter(self, tracker_id: int, box: np.ndarray):
        """Initialize Kalman filter for new track."""
        kf = _KalmanFilter(dim_x=8, dim_z=4)
        
        # State: [cx, cy, w, h, vcx, vcy, vw, vh]
        cx = (box[0] + box[2]) / 2
        cy = (box[1] + box[3]) / 2
        w = box[2] - box[0]
        h = box[3] - box[1]
        kf.x = np.array([cx, cy, w, h, 0, 0, 0, 0]).reshape(8, 1)
        
        # State transition matrix
        kf.F = np.eye(8)
        kf.F[:4, 4:] = np.eye(4)
        
        # Measurement matrix
        kf.H = np.zeros((4, 8))
        kf.H[:4, :4] = np.eye(4)
        
        # Covariance matrices
        kf.P *= 10
        kf.R *= self.measurement_noise
        kf.Q[4:, 4:] *= self.process_noise
        
        self.filters[tracker_id] = kf
    
    def reset(self, tracker_id: int = None):
        """Reset filter(s)."""
        if tracker_id is None:
            self.filters.clear()
        elif tracker_id in self.filters:
            del self.filters[tracker_id]


class ExponentialSmoothing:
    """Exponential moving average for simple smoothing.
    
    Use for: Quick smoothing, less computational overhead than Kalman.
    """
    
    def __init__(self, alpha: float = 0.3):
        """
        Args:
            alpha: Smoothing factor (0-1)
                   Lower = smoother but slower response
                   Higher = faster response but less smooth
        """
        self.alpha = alpha
        self.previous: Dict[int, np.ndarray] = {}
    
    def update(self, tracker_id: int, box: np.ndarray) -> np.ndarray:
        """Smooth box using exponential moving average.
        
        Args:
            tracker_id: Object ID
            box: [x1, y1, x2, y2]
        
        Returns:
            Smoothed box
        """
        if tracker_id not in self.previous:
            self.previous[tracker_id] = box
            return box
        
        # EMA: smoothed = alpha * new + (1 - alpha) * previous
        smoothed = self.alpha * box + (1 - self.alpha) * self.previous[tracker_id]
        self.previous[tracker_id] = smoothed
        
        return smoothed
    
    def reset(self, tracker_id: int = None):
        """Reset smoothing history."""
        if tracker_id is None:
            self.previous.clear()
        elif tracker_id in self.previous:
            del self.previous[tracker_id]


def smooth_detections(
    detections,
    smoother,
    inplace: bool = False
):
    """Apply smoothing to all detections.
    
    Args:
        detections: Detections object
        smoother: KalmanSmoothing or ExponentialSmoothing instance
        inplace: Modify in place (faster) or return new
    
    Returns:
        Smoothed detections
    """
    if detections.tracker_id is None or len(detections) == 0:
        return detections
    
    smoothed_boxes = np.empty_like(detections.xyxy)
    
    for i in range(len(detections)):
        tid = detections.tracker_id[i]
        smoothed_boxes[i] = smoother.update(tid, detections.xyxy[i])
    
    if inplace:
        object.__setattr__(detections, 'xyxy', smoothed_boxes)
        return detections
    else:
        from eye.detection.core import Detections
        return Detections(
            xyxy=smoothed_boxes,
            mask=detections.mask,
            confidence=detections.confidence,
            class_id=detections.class_id,
            tracker_id=detections.tracker_id,
            data=detections.data
        )
