models/motion_detector.py

"""
ShortSmith v2 - Motion Detector Module

Motion analysis using optical flow for:
- Detecting action-heavy segments
- Identifying camera movement vs subject movement
- Dynamic FPS scaling based on motion level

Uses RAFT (Recurrent All-Pairs Field Transforms) for high-quality
optical flow, with fallback to Farneback for speed.
"""

from pathlib import Path
from typing import List, Optional, Tuple, Union
from dataclasses import dataclass
import numpy as np

from utils.logger import get_logger, LogTimer
from utils.helpers import ModelLoadError, InferenceError
from config import get_config, ModelConfig

logger = get_logger("models.motion_detector")


@dataclass
class MotionScore:
    """Motion analysis result for a frame pair."""
    timestamp: float           # Timestamp of second frame
    magnitude: float           # Average motion magnitude (0-1 normalized)
    direction: float           # Dominant motion direction (radians)
    uniformity: float          # How uniform the motion is (1 = all same direction)
    is_camera_motion: bool     # Likely camera motion vs subject motion

    @property
    def is_high_motion(self) -> bool:
        """Check if this is a high-motion segment."""
        return self.magnitude > 0.3

    @property
    def is_action(self) -> bool:
        """Check if this likely contains action (non-uniform motion)."""
        return self.magnitude > 0.2 and self.uniformity < 0.7


class MotionDetector:
    """
    Motion detection using optical flow.

    Supports:
    - RAFT optical flow (high quality, GPU)
    - Farneback optical flow (faster, CPU)
    - Motion magnitude scoring
    - Camera vs subject motion detection
    """

    def __init__(
        self,
        config: Optional[ModelConfig] = None,
        use_raft: bool = True,
    ):
        """
        Initialize motion detector.

        Args:
            config: Model configuration
            use_raft: Whether to use RAFT (True) or Farneback (False)
        """
        self.config = config or get_config().model
        self.use_raft = use_raft
        self.raft_model = None

        if use_raft:
            self._load_raft()

        logger.info(f"MotionDetector initialized (RAFT={use_raft})")

    def _load_raft(self) -> None:
        """Load RAFT optical flow model."""
        try:
            import torch
            from torchvision.models.optical_flow import raft_small, Raft_Small_Weights

            logger.info("Loading RAFT optical flow model...")

            weights = Raft_Small_Weights.DEFAULT
            self.raft_model = raft_small(weights=weights)

            if self.config.device == "cuda" and torch.cuda.is_available():
                self.raft_model = self.raft_model.cuda()

            self.raft_model.eval()

            # Store preprocessing transforms
            self._raft_transforms = weights.transforms()

            logger.info("RAFT model loaded successfully")

        except Exception as e:
            logger.warning(f"Failed to load RAFT model, using Farneback: {e}")
            self.use_raft = False
            self.raft_model = None

    def compute_flow(
        self,
        frame1: np.ndarray,
        frame2: np.ndarray,
    ) -> np.ndarray:
        """
        Compute optical flow between two frames.

        Args:
            frame1: First frame (BGR or RGB, HxWxC)
            frame2: Second frame (BGR or RGB, HxWxC)

        Returns:
            Optical flow array (HxWx2), flow[y,x] = (dx, dy)
        """
        if self.use_raft and self.raft_model is not None:
            return self._compute_raft_flow(frame1, frame2)
        else:
            return self._compute_farneback_flow(frame1, frame2)

    def _compute_raft_flow(
        self,
        frame1: np.ndarray,
        frame2: np.ndarray,
    ) -> np.ndarray:
        """Compute flow using RAFT."""
        import torch

        try:
            # Convert to RGB if BGR
            if frame1.shape[2] == 3:
                frame1_rgb = frame1[:, :, ::-1].copy()
                frame2_rgb = frame2[:, :, ::-1].copy()
            else:
                frame1_rgb = frame1
                frame2_rgb = frame2

            # Convert to tensors
            img1 = torch.from_numpy(frame1_rgb).permute(2, 0, 1).float().unsqueeze(0)
            img2 = torch.from_numpy(frame2_rgb).permute(2, 0, 1).float().unsqueeze(0)

            if self.config.device == "cuda" and torch.cuda.is_available():
                img1 = img1.cuda()
                img2 = img2.cuda()

            # Compute flow
            with torch.no_grad():
                flow_predictions = self.raft_model(img1, img2)
                flow = flow_predictions[-1]  # Use final prediction

            # Convert back to numpy
            flow = flow[0].permute(1, 2, 0).cpu().numpy()

            return flow

        except Exception as e:
            logger.warning(f"RAFT flow failed, using Farneback: {e}")
            return self._compute_farneback_flow(frame1, frame2)

    def _compute_farneback_flow(
        self,
        frame1: np.ndarray,
        frame2: np.ndarray,
    ) -> np.ndarray:
        """Compute flow using Farneback algorithm."""
        import cv2

        # Convert to grayscale
        if len(frame1.shape) == 3:
            gray1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
            gray2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        else:
            gray1 = frame1
            gray2 = frame2

        # Compute Farneback optical flow
        flow = cv2.calcOpticalFlowFarneback(
            gray1, gray2,
            None,
            pyr_scale=0.5,
            levels=3,
            winsize=15,
            iterations=3,
            poly_n=5,
            poly_sigma=1.2,
            flags=0,
        )

        return flow

    def analyze_motion(
        self,
        frame1: np.ndarray,
        frame2: np.ndarray,
        timestamp: float = 0.0,
    ) -> MotionScore:
        """
        Analyze motion between two frames.

        Args:
            frame1: First frame
            frame2: Second frame
            timestamp: Timestamp of second frame

        Returns:
            MotionScore with analysis results
        """
        flow = self.compute_flow(frame1, frame2)

        # Compute magnitude and direction
        magnitude = np.sqrt(flow[:, :, 0]**2 + flow[:, :, 1]**2)
        direction = np.arctan2(flow[:, :, 1], flow[:, :, 0])

        # Average magnitude (normalized by image diagonal)
        h, w = frame1.shape[:2]
        diagonal = np.sqrt(h**2 + w**2)
        avg_magnitude = float(np.mean(magnitude) / diagonal)

        # Dominant direction
        # Weight by magnitude to get dominant direction
        weighted_direction = np.average(direction, weights=magnitude + 1e-8)

        # Uniformity: how consistent is the motion direction?
        # High uniformity = likely camera motion
        dir_std = float(np.std(direction))
        uniformity = 1.0 / (1.0 + dir_std)

        # Detect camera motion (uniform direction across frame)
        is_camera = uniformity > 0.7 and avg_magnitude > 0.05

        return MotionScore(
            timestamp=timestamp,
            magnitude=min(1.0, avg_magnitude * 10),  # Scale up
            direction=float(weighted_direction),
            uniformity=uniformity,
            is_camera_motion=is_camera,
        )

    def analyze_video_segment(
        self,
        frames: List[np.ndarray],
        timestamps: List[float],
    ) -> List[MotionScore]:
        """
        Analyze motion across a video segment.

        Args:
            frames: List of frames
            timestamps: Timestamps for each frame

        Returns:
            List of MotionScore objects (one per frame pair)
        """
        if len(frames) < 2:
            return []

        scores = []

        with LogTimer(logger, f"Analyzing motion in {len(frames)} frames"):
            for i in range(1, len(frames)):
                try:
                    score = self.analyze_motion(
                        frames[i-1],
                        frames[i],
                        timestamps[i],
                    )
                    scores.append(score)
                except Exception as e:
                    logger.warning(f"Motion analysis failed for frame {i}: {e}")

        return scores

    def get_motion_heatmap(
        self,
        frame1: np.ndarray,
        frame2: np.ndarray,
    ) -> np.ndarray:
        """
        Get motion magnitude heatmap.

        Args:
            frame1: First frame
            frame2: Second frame

        Returns:
            Heatmap of motion magnitude (HxW, values 0-255)
        """
        flow = self.compute_flow(frame1, frame2)
        magnitude = np.sqrt(flow[:, :, 0]**2 + flow[:, :, 1]**2)

        # Normalize to 0-255
        max_mag = np.percentile(magnitude, 99)  # Robust max
        if max_mag > 0:
            normalized = np.clip(magnitude / max_mag * 255, 0, 255)
        else:
            normalized = np.zeros_like(magnitude)

        return normalized.astype(np.uint8)

    def compute_aggregate_motion(
        self,
        scores: List[MotionScore],
    ) -> float:
        """
        Compute aggregate motion score for a segment.

        Args:
            scores: List of MotionScore objects

        Returns:
            Aggregate motion score (0-1)
        """
        if not scores:
            return 0.0

        # Weight by non-camera motion
        weighted_sum = sum(
            s.magnitude * (0.3 if s.is_camera_motion else 1.0)
            for s in scores
        )

        return weighted_sum / len(scores)

    def identify_high_motion_segments(
        self,
        scores: List[MotionScore],
        threshold: float = 0.3,
        min_duration: int = 3,
    ) -> List[Tuple[float, float, float]]:
        """
        Identify segments with high motion.

        Args:
            scores: List of MotionScore objects
            threshold: Minimum motion magnitude
            min_duration: Minimum number of consecutive frames

        Returns:
            List of (start_time, end_time, avg_motion) tuples
        """
        if not scores:
            return []

        segments = []
        in_segment = False
        segment_start = 0.0
        segment_scores = []

        for score in scores:
            if score.magnitude >= threshold:
                if not in_segment:
                    in_segment = True
                    segment_start = score.timestamp
                    segment_scores = [score.magnitude]
                else:
                    segment_scores.append(score.magnitude)
            else:
                if in_segment:
                    if len(segment_scores) >= min_duration:
                        segments.append((
                            segment_start,
                            score.timestamp,
                            sum(segment_scores) / len(segment_scores),
                        ))
                    in_segment = False

        # Handle segment at end
        if in_segment and len(segment_scores) >= min_duration:
            segments.append((
                segment_start,
                scores[-1].timestamp,
                sum(segment_scores) / len(segment_scores),
            ))

        logger.info(f"Found {len(segments)} high-motion segments")
        return segments


# Export public interface
__all__ = ["MotionDetector", "MotionScore"]