src/video.py

"""
VideoModule - temporal frame sampling + motion features for MORPH-AI v6.
Lazy-loads a video model when available; falls back to frame statistics.
"""

import json
import os
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class VideoFacts:
    duration: float = 0.0
    fps: float = 0.0
    frame_count: int = 0
    motion_score: float = 0.0
    scene_changes: List[float] = field(default_factory=list)
    embeddings: Optional[torch.Tensor] = None
    key_frames: List[str] = field(default_factory=list)

    def to_text(self) -> str:
        parts = [f"video {self.duration:.1f}s {self.fps:.1f}fps {self.frame_count}frames"]
        if self.motion_score > 0:
            parts.append(f"motion {self.motion_score:.2f}")
        if self.scene_changes:
            parts.append(f"scenes at {', '.join(f'{t:.1f}s' for t in self.scene_changes[:5])}")
        return " | ".join(parts)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "duration": self.duration,
            "fps": self.fps,
            "frame_count": self.frame_count,
            "motion_score": self.motion_score,
            "scene_changes": self.scene_changes,
        }


class VideoModule(nn.Module):
    """Temporal frame sampling + motion features for video understanding."""

    def __init__(self, config: MorphConfig, hidden_dim: int):
        super().__init__()
        self.max_frames = config.video_max_frames
        self.frame_proj = nn.Linear(hidden_dim, config.video_hidden)
        self.temporal_encoder = nn.GRU(
            config.video_hidden, config.video_hidden,
            batch_first=True, bidirectional=False
        )
        self.motion_proj = nn.Linear(config.video_hidden, hidden_dim)
        self.scene_detector = nn.Sequential(
            nn.Linear(hidden_dim, 128),
            nn.GELU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )
        nn.init.zeros_(self.motion_proj.weight)
        nn.init.zeros_(self.motion_proj.bias)

    def analyze(self, source) -> VideoFacts:
        """Analyze video: extract frames, compute motion, detect scenes."""
        facts = VideoFacts()
        try:
            import cv2
            import numpy as np

            cap = cv2.VideoCapture(source)
            if not cap.isOpened():
                return facts

            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = frame_count / fps if fps > 0 else 0

            facts.fps = fps
            facts.frame_count = frame_count
            facts.duration = duration

            frames = []
            prev_gray = None
            motion_scores = []
            scene_times = []

            sample_rate = max(1, frame_count // self.max_frames)
            for i in range(0, frame_count, sample_rate):
                cap.set(cv2.CAP_PROP_POS_FRAMES, i)
                ret, frame = cap.read()
                if not ret:
                    break

                small = cv2.resize(frame, (224, 224))
                gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)

                if prev_gray is not None:
                    diff = cv2.absdiff(prev_gray, gray)
                    motion = diff.mean() / 255.0
                    motion_scores.append(motion)
                    if motion > 0.3 and len(motion_scores) > 1:
                        scene_times.append(i / fps)
                prev_gray = gray
                frames.append(small)

                if len(frames) >= self.max_frames:
                    break

            cap.release()

            facts.motion_score = sum(motion_scores) / len(motion_scores) if motion_scores else 0
            facts.scene_changes = scene_times[:10]
            facts.key_frames = [f"frame_{i}" for i in range(len(frames))]

            if frames:
                frame_tensor = torch.tensor(frames, dtype=torch.float32).permute(0, 3, 1, 2) / 255.0
                facts.embeddings = frame_tensor

        except ImportError:
            facts.key_frames = ["[video analysis requires opencv-python: pip install opencv-python]"]
        except Exception as e:
            facts.key_frames = [f"[video analysis error: {e}]"]

        return facts

    def forward(self, hidden: torch.Tensor, frame_embeddings: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Project video frame embeddings into hidden space."""
        if frame_embeddings is None:
            return hidden

        B, T, H = hidden.shape
        frames = frame_embeddings.to(hidden.device)

        if frames.dim() == 4:
            frames = frames.mean(dim=[2, 3])

        frame_emb = self.frame_proj(frames)
        if frame_emb.dim() == 2:
            frame_emb = frame_emb.unsqueeze(0)

        _, last_hidden = self.temporal_encoder(frame_emb)
        motion = self.motion_proj(last_hidden.squeeze(0))
        return hidden + motion.unsqueeze(1)