import cv2
import numpy as np
import torch
import os
from typing import List, Tuple, Optional

def get_video_properties(video_path: str) -> dict:
    """
    Get video properties.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return {}
    
    props = {
        'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
        'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
        'fps': cap.get(cv2.CAP_PROP_FPS),
        'frame_count': int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    }
    cap.release()
    return props

def load_video_frames(
    video_path: str, 
    num_frames: int = 30, 
    frame_size: Tuple[int, int] = (224, 224),
    frame_skip: int = 1
) -> Optional[np.ndarray]:
    """
    Load frames from a video with robust preprocessing:
    1. Center crop to square (min dimension).
    2. Resize to frame_size.
    3. Sample frames uniformly.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error opening video: {video_path}")
        return None
        
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total_frames <= 0:
        cap.release()
        return None
        
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # Calculate crop coordinates for center square crop
    min_dim = min(width, height)
    start_x = (width - min_dim) // 2
    start_y = (height - min_dim) // 2
    
    # Calculate frame indices to sample
    # We want 'num_frames' frames. 
    # Strategy: evenly space them across the video duration we look at.
    # But for simplicity and consistency, let's just grab them with a stride, 
    # or if video is short, grab all and pad.
    
    # Let's try to span as much of the video as possible? 
    # Or just stick to the requested architecture of sampling segments.
    # The prompt asked for "preprocess ... not based on aspect ratio".
    
    # Simple strategy: Read frames with skip, up to num_frames.
    # If video is too short, loop/pad? 
    # Better: Reservoir sampling or Linspace if we want fixed count?
    # Let's stick to the user's likely need: Fixed number of frames.
    
    sampled_frames = []
    
    # We'll seek effectively.
    # But looping is safer for some codecs.
    
    # Improved sampling: pick indices using linspace if we want to span whole video?
    # Or just sequential for temporal consistency (RNN prefer sequences).
    # Let's do sequential with skip.
    
    frame_idx = 0
    frames_collected = 0
    
    while frames_collected < num_frames:
        ret, frame = cap.read()
        if not ret:
            break
            
        if frame_idx % frame_skip == 0:
            # Preprocess Frame
            
            # 1. Center Crop
            crop = frame[start_y:start_y+min_dim, start_x:start_x+min_dim]
            
            # 2. Resize
            resized = cv2.resize(crop, frame_size)
            
            # 3. Convert BGR to RGB
            rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
            
            sampled_frames.append(rgb)
            frames_collected += 1
            
        frame_idx += 1
        
    cap.release()
    
    # Handle insufficient frames
    if len(sampled_frames) == 0:
        return None
        
    if len(sampled_frames) < num_frames:
        # Pad with last frame or zeros?
        # Let's pad with zeros (black frames) or loop?
        # Zero padding is safer to avoid motion artifacts.
        padding = [np.zeros((frame_size[1], frame_size[0], 3), dtype=np.uint8)] * (num_frames - len(sampled_frames))
        sampled_frames.extend(padding)
        
    return np.array(sampled_frames)

def normalize_frames(frames: np.ndarray) -> np.ndarray:
    """
    Normalize frames to [0, 1] and then standard ImageNet mean/std.
    Frames input: (N, H, W, C) in RGB, uint8 [0,255]
    """
    # Convert to float32 [0, 1]
    frames_norm = frames.astype(np.float32) / 255.0
    
    # Standard ImageNet mean and std
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    
    # Apply normalization
    # frames is (N, H, W, C), mean/std are (3,)
    # We allow broadcasting on the last dimension
    frames_norm = (frames_norm - mean) / std
    
    return frames_norm

def validate_video(video_path: str) -> bool:
    """
    Check if video is valid and openable.
    """
    if not os.path.exists(video_path):
        return False
    try:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            return False
        # Read one frame to be sure
        ret, _ = cap.read()
        cap.release()
        return ret
    except Exception:
        return False