"""Trace annotator for drawing object paths."""

import cv2
import numpy as np
from typing import Optional, Dict
from collections import deque
from ..core.detections import Detections
from ..core.colors import ColorPalette


class TraceAnnotator:
    """Draw tracking trails/paths for objects.
    
    Innovation: Fading trails, variable thickness, trail smoothing.
    """
    
    def __init__(
        self,
        color_palette: Optional[ColorPalette] = None,
        thickness: int = 2,
        trace_length: int = 50,
        fade_trail: bool = True,
        smooth_trail: bool = False
    ):
        """
        Args:
            color_palette: Color palette
            thickness: Line thickness
            trace_length: Maximum number of points in trail
            fade_trail: Gradually fade older trail points
            smooth_trail: Apply smoothing to trail
        """
        from ..core.colors import PredefinedPalettes
        self.color_palette = color_palette or PredefinedPalettes.bright()
        self.thickness = thickness
        self.trace_length = trace_length
        self.fade_trail = fade_trail
        self.smooth_trail = smooth_trail
        
        # Track history for each object
        self.history: Dict[int, deque] = {}
        # Position smoothing: base alpha (0-1), lower = more smoothing
        self.position_smoothing_alpha = 0.6
        # If movement speed (pixels/frame) exceeds this, use stronger smoothing
        self.position_speed_threshold = 8.0
        # Alpha to use when speed > threshold
        self.fast_position_smoothing_alpha = 0.3
    
    def annotate(
        self,
        image: np.ndarray,
        detections: Detections
    ) -> np.ndarray:
        """Draw trails on image."""
        annotated = image.copy()
        
        if detections.tracker_id is None:
            return annotated
        
        # Update history
        active_ids = set()
        for i in range(len(detections)):
            tracker_id = detections.tracker_id[i]
            active_ids.add(tracker_id)
            center = detections.center[i]
            
            if tracker_id not in self.history:
                self.history[tracker_id] = deque(maxlen=self.trace_length)
            # Apply optional temporal smoothing when appending to history
            if len(self.history[tracker_id]) > 0 and self.position_smoothing_alpha is not None:
                last = np.array(self.history[tracker_id][-1], dtype=float)
                newc = np.array(center, dtype=float)
                dist = float(np.linalg.norm(newc - last))
                if self.position_speed_threshold is not None and dist > float(self.position_speed_threshold):
                    alpha = float(self.fast_position_smoothing_alpha)
                else:
                    alpha = float(self.position_smoothing_alpha)
                center = alpha * newc + (1.0 - alpha) * last

            self.history[tracker_id].append(center)
        
        # Remove old tracks
        inactive_ids = set(self.history.keys()) - active_ids
        for tid in inactive_ids:
            if len(self.history[tid]) > 0:
                # Keep trail visible for a bit after object leaves
                pass
            else:
                del self.history[tid]
        
        # Draw trails
        for i in range(len(detections)):
            tracker_id = detections.tracker_id[i]
            if tracker_id not in self.history or len(self.history[tracker_id]) < 2:
                continue
            
            points = list(self.history[tracker_id])
            
            # Smooth if requested
            if self.smooth_trail and len(points) >= 3:
                points = self._smooth_points(points)
            
            # Get color
            if detections.class_id is not None:
                color = self.color_palette.by_id(detections.class_id[i])
            else:
                color = self.color_palette.by_id(tracker_id)
            
            # Draw trail
            for j in range(len(points) - 1):
                if self.fade_trail:
                    # Calculate alpha based on position in trail
                    alpha = (j + 1) / len(points)
                    faded_color = tuple(int(c * alpha) for c in color.as_bgr())
                else:
                    faded_color = color.as_bgr()
                
                pt1 = tuple(points[j].astype(int))
                pt2 = tuple(points[j + 1].astype(int))
                
                cv2.line(annotated, pt1, pt2, faded_color, self.thickness)
        
        return annotated
    
    @staticmethod
    def _smooth_points(points, window=3):
        """Apply moving average smoothing to points."""
        points = np.array(points)
        smoothed = []
        for i in range(len(points)):
            start = max(0, i - window // 2)
            end = min(len(points), i + window // 2 + 1)
            smoothed.append(np.mean(points[start:end], axis=0))
        return smoothed
    
    def reset(self):
        """Clear all tracking history."""
        self.history.clear()


# Alias for easier understanding
PathAnnotator = TraceAnnotator
