"""
Video Preprocessing - Shared between training and inference.

This module provides the canonical preprocessing pipeline for video frames.
Both training and inference use this same code to ensure consistency.

Features:
    - Float-based frame sampling (accurate timing)
    - Direct resize to square (simple, no distortion modes)
    - ImageNet normalization on GPU
    - Incremental frame buffer for streaming inference
    - Debug utilities for tensor visualization
"""

import os
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from typing import List, Optional, Tuple


# ImageNet normalization constants
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

# Pre-allocated normalisation tensors (shape: 1,3,1,1 for batch broadcast)
# Avoids per-call tensor allocation inside hot DataLoader workers.
_NORM_MEAN = torch.tensor(IMAGENET_MEAN, dtype=torch.float32).view(1, 3, 1, 1)
_NORM_STD  = torch.tensor(IMAGENET_STD,  dtype=torch.float32).view(1, 3, 1, 1)


# =============================================================================
# FRAME SAMPLING
# =============================================================================

def calculate_frame_indices(
    total_frames: int,
    source_fps: float,
    target_fps: float,
) -> List[int]:
    """
    Calculate which frame indices to sample using float accumulation.

    This ensures accurate timing by avoiding integer truncation errors.

    Args:
        total_frames: Total frames in source video
        source_fps: Source video FPS (e.g., 30.0)
        target_fps: Target sampling FPS (e.g., 8.0)

    Returns:
        List of frame indices to sample

    Example:
        For 30fps video sampled at 8fps:
        - Interval = 30/8 = 3.75
        - Indices: [0, 3, 7, 11, 15, ...] (float accumulation)
    """
    if source_fps <= 0 or target_fps <= 0:
        raise ValueError(f"FPS must be positive: source={source_fps}, target={target_fps}")

    interval = source_fps / target_fps
    indices = []
    position = 0.0

    while position < total_frames:
        indices.append(int(position))
        position += interval

    return indices


def sample_frames_uniform(total_frames: int, num_frames: int) -> List[int]:
    """
    Sample frames uniformly using float accumulation.

    Args:
        total_frames: Total frames available
        num_frames: Number of frames to sample

    Returns:
        List of frame indices
    """
    if total_frames <= num_frames:
        return list(range(total_frames))

    interval = total_frames / num_frames
    indices = []
    position = 0.0

    for _ in range(num_frames):
        indices.append(int(position))
        position += interval

    return indices


# =============================================================================
# PREPROCESSING
# =============================================================================

def resize_frame(frame: np.ndarray, size: int) -> np.ndarray:
    """
    Resize frame to square dimensions.

    Args:
        frame: BGR or RGB image (H, W, 3)
        size: Target size (will be size x size)

    Returns:
        Resized image (size, size, 3)
    """
    return cv2.resize(frame, (size, size), interpolation=cv2.INTER_LINEAR)


def preprocess_frame_cpu(
    frame: np.ndarray,
    size: int,
    bgr_to_rgb: bool = True,
) -> np.ndarray:
    """
    Preprocess a single frame on CPU (resize and color conversion).

    Args:
        frame: Input image (H, W, 3)
        size: Target size
        bgr_to_rgb: Convert BGR to RGB (True if frame is from OpenCV)

    Returns:
        Preprocessed RGB image (size, size, 3) as uint8
    """
    resized = resize_frame(frame, size)
    if bgr_to_rgb:
        resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
    return resized


def preprocess_single_frame(
    frame,
    img_size: int = None,
) -> torch.Tensor:
    """
    Preprocess a single RGB frame to a normalized tensor.

    THIS IS THE CANONICAL PREPROCESSING FUNCTION.
    Both training and inference MUST use this function to ensure consistency.

    Accepts either a numpy uint8 HWC array or a float CHW tensor.
    When a tensor is passed, the resize and uint8-to-float conversion steps
    are skipped and only ImageNet normalization is applied.

    Pipeline (numpy input):
        1. Resize to square (img_size x img_size) — skipped if img_size is None
        2. Convert to float tensor
        3. HWC -> CHW
        4. Normalize to [0, 1]
        5. Apply ImageNet normalization

    Pipeline (tensor input):
        1. Apply ImageNet normalization

    Args:
        frame: RGB uint8 image (H, W, 3) as np.ndarray, or float (C, H, W) tensor
        img_size: Target size (will be img_size x img_size); ignored for tensor input

    Returns:
        Normalized tensor (3, H, W) on CPU, float32
    """
    if isinstance(frame, torch.Tensor):
        tensor = frame.float()
    else:
        # Resize
        if img_size is not None:
            frame = resize_frame(frame, img_size)

        # Convert to tensor and normalize to [0, 1]
        tensor = torch.from_numpy(frame).float()
        tensor = tensor.permute(2, 0, 1) / 255.0

    # ImageNet normalization — use pre-allocated module-level tensors (no per-call alloc)
    mean = _NORM_MEAN.view(3, 1, 1).to(tensor.device)
    std  = _NORM_STD.view(3, 1, 1).to(tensor.device)
    tensor = (tensor - mean) / std

    return tensor


def normalize_tensor(
    tensor: torch.Tensor,
    mean: Tuple[float, float, float] = IMAGENET_MEAN,
    std: Tuple[float, float, float] = IMAGENET_STD,
) -> torch.Tensor:
    """
    Apply ImageNet normalization to a tensor.

    Args:
        tensor: Input tensor (C, H, W) or (N, C, H, W), values in [0, 1]
        mean: Normalization mean
        std: Normalization std

    Returns:
        Normalized tensor
    """
    device = tensor.device
    dtype = tensor.dtype

    mean_t = torch.tensor(mean, device=device, dtype=dtype)
    std_t = torch.tensor(std, device=device, dtype=dtype)

    if tensor.dim() == 3:
        # (C, H, W)
        mean_t = mean_t.view(3, 1, 1)
        std_t = std_t.view(3, 1, 1)
    elif tensor.dim() == 4:
        # (N, C, H, W)
        mean_t = mean_t.view(1, 3, 1, 1)
        std_t = std_t.view(1, 3, 1, 1)

    return (tensor - mean_t) / std_t


# =============================================================================
# VIDEO PREPROCESSOR
# =============================================================================

class VideoPreprocessor:
    """
    GPU-accelerated video preprocessor with incremental frame buffer.

    Used by both training and inference for consistent preprocessing.

    Pipeline:
        1. Resize to square (CPU - OpenCV)
        2. BGR to RGB conversion
        3. Transfer to GPU
        4. Normalize to [0, 1]
        5. ImageNet normalization
        6. Store in frame buffer

    Args:
        img_size: Target image size (default 256)
        num_frames: Number of frames in buffer (default 16)
        device: Device for tensors (default 'cuda' if available)
        dtype: Output tensor dtype (default float32)

    Example:
        >>> preprocessor = VideoPreprocessor(img_size=256, num_frames=16)
        >>> for frame in video_frames:
        ...     if preprocessor.add_frame(frame):
        ...         tensor = preprocessor.get_tensor()
        ...         # tensor is (1, 16, 3, 256, 256)
    """

    def __init__(
        self,
        img_size: int = 256,
        num_frames: int = 16,
        device: str = None,
        dtype: torch.dtype = torch.float32,
    ):
        self.img_size = img_size
        self.num_frames = num_frames
        self.dtype = dtype

        # Set device
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = torch.device(device)

        # Double-length circular buffer (zero-copy contiguous slice in get_tensor)
        self._buffer: Optional[torch.Tensor] = None
        self._buffer_count = 0

        # Normalisation tensors pre-allocated on device — avoids per-call allocation
        self._norm_mean = torch.tensor(IMAGENET_MEAN, dtype=torch.float32, device=self.device).view(3, 1, 1)
        self._norm_std  = torch.tensor(IMAGENET_STD,  dtype=torch.float32, device=self.device).view(3, 1, 1)

        # CUDA stream pipelining: secondary stream for preprocessing concurrent with inference
        if self.device.type == 'cuda':
            self._preprocess_stream = torch.cuda.Stream(device=self.device)
            self._preprocess_event  = torch.cuda.Event()
        else:
            self._preprocess_stream = None
            self._preprocess_event  = None

    @property
    def is_ready(self) -> bool:
        """Whether buffer is full and ready for inference."""
        return self._buffer_count >= self.num_frames

    @property
    def buffer_count(self) -> int:
        """Number of frames in buffer."""
        return self._buffer_count

    def reset(self):
        """Clear the frame buffer."""
        self._buffer = None
        self._buffer_count = 0

    def _preprocess_single(self, frame: np.ndarray) -> torch.Tensor:
        """
        Preprocess a single BGR uint8 frame.

        GPU path: upload raw uint8, then do BGR→RGB / resize / normalize on GPU.
        CPU path: original cv2-based pipeline.

        Args:
            frame: BGR uint8 image (H, W, 3) from OpenCV

        Returns:
            Preprocessed tensor (3, img_size, img_size) on device
        """
        if self.device.type == 'cuda':
            # Upload raw uint8 to GPU (cheaper than uploading float32 after CPU work)
            frame_t = torch.from_numpy(np.ascontiguousarray(frame)).to(
                self.device, non_blocking=True
            )  # (H, W, 3) uint8 BGR on GPU

            # BGR→RGB, HWC→CHW, uint8→float32 [0,1] — all on GPU
            frame_t = frame_t.permute(2, 0, 1).flip(0).float().div_(255.0)  # (3, H, W)

            # GPU resize (only when source size differs from target)
            if frame_t.shape[1] != self.img_size or frame_t.shape[2] != self.img_size:
                frame_t = F.interpolate(
                    frame_t.unsqueeze(0),
                    size=(self.img_size, self.img_size),
                    mode='bilinear',
                    align_corners=False,
                    antialias=False,
                ).squeeze(0)  # (3, img_size, img_size)

            # ImageNet normalize — in-place, pre-cached GPU tensors, no allocation
            frame_t.sub_(self._norm_mean).div_(self._norm_std)

            return frame_t  # (3, img_size, img_size) float32 on GPU

        else:
            # CPU fallback: original pipeline
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            return preprocess_single_frame(rgb, self.img_size)  # (3, img_size, img_size) CPU

    def add_frame(self, frame: np.ndarray) -> bool:
        """
        Add a frame to the buffer.

        Args:
            frame: BGR uint8 image (H, W, 3) from OpenCV

        Returns:
            True if buffer is full (ready for inference)
        """
        processed = self._preprocess_single(frame)  # (3, img_size, img_size)

        if self._buffer is None:
            # Double-length buffer: slot pos and pos+num_frames are always kept in sync,
            # so get_tensor() can return buffer[read_ptr : read_ptr+num_frames] as a
            # zero-copy contiguous view — no torch.roll needed.
            self._buffer = torch.zeros(
                2 * self.num_frames, 3, self.img_size, self.img_size,
                device=self.device,
                dtype=self.dtype,
            )

        pos = self._buffer_count % self.num_frames
        self._buffer[pos] = processed
        self._buffer[pos + self.num_frames] = processed  # mirror write for zero-copy reads
        self._buffer_count += 1

        return self._buffer_count >= self.num_frames

    def get_tensor(self) -> Optional[torch.Tensor]:
        """
        Get the buffer as model input tensor.

        Returns:
            Tensor (1, num_frames, 3, H, W) or None if not ready
        """
        if not self.is_ready:
            return None

        # read_ptr is the oldest slot; [read_ptr, read_ptr+num_frames) is always
        # a valid contiguous window thanks to the double-write in add_frame.
        read_ptr = self._buffer_count % self.num_frames
        return self._buffer[read_ptr : read_ptr + self.num_frames].unsqueeze(0)
        # Returns a view — no allocation, no copy.

    def get_partial_tensor(self, n_frames: int) -> torch.Tensor:
        """Return a view of the first n_frames already in the buffer.

        Only valid during the startup phase (buffer_count < num_frames).
        During this phase frames are stored at positions 0..n_frames-1 in order.
        The caller must .clone() the result before async GPU submission.

        Args:
            n_frames: Number of frames to include (must be <= buffer_count)

        Returns:
            Tensor view (1, n_frames, 3, H, W).
        """
        return self._buffer[:n_frames].unsqueeze(0)

    def start_preprocess_async(self, frame: np.ndarray) -> bool:
        """
        Preprocess frame on a secondary CUDA stream, concurrent with inference.

        Runs the same GPU preprocessing as _preprocess_single inside a secondary
        CUDA stream so it overlaps with model inference on the default stream.
        Records a CUDA event when done; get_inference_tensor() waits on it.

        Falls back to synchronous add_frame() on CPU.

        Args:
            frame: BGR uint8 image (H, W, 3) from OpenCV

        Returns:
            True if buffer is full (ready for inference)
        """
        if self._preprocess_stream is None:
            return self.add_frame(frame)

        with torch.cuda.stream(self._preprocess_stream):
            # _preprocess_single uses non_blocking=True; GPU ops are queued on
            # whichever stream is current (preprocess_stream here).
            processed = self._preprocess_single(frame)

            if self._buffer is None:
                self._buffer = torch.zeros(
                    2 * self.num_frames, 3, self.img_size, self.img_size,
                    device=self.device, dtype=self.dtype,
                )

            pos = self._buffer_count % self.num_frames
            self._buffer[pos] = processed
            self._buffer[pos + self.num_frames] = processed  # mirror write
            self._buffer_count += 1

            # Signal that all writes to buffer are complete on preprocess_stream
            self._preprocess_event.record(self._preprocess_stream)

        return self._buffer_count >= self.num_frames

    def get_inference_tensor(self) -> Optional[torch.Tensor]:
        """
        Get a cloned buffer for pipelined inference.

        Makes the default (inference) CUDA stream wait for the preprocess stream
        event (GPU-side wait — CPU does not block). Then clones so preprocessing
        can safely overwrite the circular buffer while inference reads the clone.

        Use in place of get_tensor() when pipelining with start_preprocess_async().

        Returns:
            Cloned tensor (1, num_frames, 3, H, W) or None if not ready
        """
        if not self.is_ready:
            return None

        if self._preprocess_event is not None:
            # GPU-side sync: default stream waits for preprocess stream to finish
            torch.cuda.current_stream().wait_event(self._preprocess_event)
            return self.get_tensor().clone()  # clone so preprocess can reuse buffer

        return self.get_tensor()  # CPU path: no concurrency, view is safe

    def preprocess_frames(
        self,
        frames: List[np.ndarray],
        dtype: torch.dtype = None,
    ) -> torch.Tensor:
        """Preprocess a fixed list of BGR frames into a (1, T, C, H, W) tensor.

        Creates a temporary preprocessor so the rolling buffer of the main
        instance is never disturbed.  Safe to call at any point in a stream.

        Args:
            frames: List of BGR uint8 images
            dtype: Output dtype (default: self.dtype)

        Returns:
            Tensor (1, len(frames), 3, H, W)
        """
        if dtype is None:
            dtype = self.dtype

        temp = VideoPreprocessor(
            num_frames=len(frames),
            img_size=self.img_size,
            device=str(self.device),
            dtype=dtype,
        )
        for frame in frames:
            temp.add_frame(frame)
        return temp.get_tensor()

    def preprocess_window(self, frames: List[np.ndarray]) -> torch.Tensor:
        """
        Preprocess a window of frames for inference.

        Alias for preprocess_frames with default dtype for consistency
        with MultiBufferPreprocessor API.

        Args:
            frames: List of BGR uint8 images from OpenCV

        Returns:
            Tensor (1, len(frames), 3, H, W)
        """
        return self.preprocess_frames(frames, self.dtype)

    def get_input_tensor(self) -> Optional[torch.Tensor]:
        """
        Get the buffer as model input tensor.

        Alias for get_tensor() for consistency with MultiBufferPreprocessor API.

        Returns:
            Tensor (1, num_frames, 3, H, W) or None if not ready
        """
        return self.get_tensor()


# =============================================================================
# BATCH PREPROCESSING FOR TRAINING
# =============================================================================

def preprocess_video_frames(
    frames: List[np.ndarray],
    img_size: int = 256,
    device: str = None,
    dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """
    Preprocess a list of video frames for model input.

    Convenience function for training - creates a temporary preprocessor.

    Args:
        frames: List of BGR uint8 images from OpenCV
        img_size: Target size
        device: Device for output tensor
        dtype: Output dtype

    Returns:
        Tensor (num_frames, 3, H, W)
    """
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    preprocessor = VideoPreprocessor(
        img_size=img_size,
        num_frames=len(frames),
        device=device,
        dtype=dtype,
    )

    for frame in frames:
        preprocessor.add_frame(frame)

    tensor = preprocessor.get_tensor()
    return tensor.squeeze(0)  # Remove batch dim for training


def preprocess_frames_batch(
    frames: List[np.ndarray],
    img_size: int = 256,
) -> torch.Tensor:
    """
    Preprocess frames on CPU (for use with DataLoader).

    Returns tensor on CPU - DataLoader will handle device transfer.
    Uses bulk numpy→tensor ops: one allocation for mean/std, one stack op,
    one normalisation broadcast — instead of N per-frame allocations.

    Args:
        frames: List of RGB uint8 images (already converted from BGR)
        img_size: Target size

    Returns:
        Tensor (num_frames, 3, H, W) on CPU
    """
    # Stack frames: skip cv2.resize if they are already the target size
    h, w = frames[0].shape[:2]
    if h == img_size and w == img_size:
        stacked = np.stack(frames)  # (T, H, W, 3) — no copy needed for resize
    else:
        stacked = np.stack([
            cv2.resize(f, (img_size, img_size), interpolation=cv2.INTER_LINEAR)
            for f in frames
        ])
    # One tensor conversion + one broadcast normalisation using pre-allocated constants
    tensor = torch.from_numpy(np.ascontiguousarray(stacked)).float()
    tensor = tensor.permute(0, 3, 1, 2) / 255.0  # (T, 3, H, W)
    return (tensor - _NORM_MEAN) / _NORM_STD


# =============================================================================
# DEBUG UTILITIES
# =============================================================================

def denormalize_tensor(
    tensor: torch.Tensor,
    mean: Tuple[float, float, float] = IMAGENET_MEAN,
    std: Tuple[float, float, float] = IMAGENET_STD,
) -> List[np.ndarray]:
    """
    Reverse ImageNet normalization and convert tensor back to uint8 RGB frames.

    Args:
        tensor: (T, C, H, W) or (B, T, C, H, W) normalized tensor
        mean: Normalization mean used
        std: Normalization std used

    Returns:
        List of RGB uint8 numpy arrays (H, W, 3)
    """
    if tensor.dim() == 5:
        tensor = tensor.squeeze(0)  # Remove batch dim

    # Denormalize: x = (normalized * std) + mean
    mean_t = torch.tensor(mean, dtype=tensor.dtype).view(1, 3, 1, 1)
    std_t = torch.tensor(std, dtype=tensor.dtype).view(1, 3, 1, 1)

    tensor = tensor * std_t + mean_t  # Back to [0, 1]
    tensor = tensor * 255.0  # Back to [0, 255]
    tensor = tensor.clamp(0, 255).byte()

    # Convert to numpy frames
    frames = []
    for t in range(tensor.shape[0]):
        # CHW -> HWC
        frame = tensor[t].permute(1, 2, 0).cpu().numpy()
        frames.append(frame)

    return frames


def save_tensor_as_video(
    tensor: torch.Tensor,
    output_path: str,
    fps: int = 8,
    mean: Tuple[float, float, float] = IMAGENET_MEAN,
    std: Tuple[float, float, float] = IMAGENET_STD,
) -> str:
    """
    Denormalize a tensor and save as video file.

    Uses H.264 codec via ffmpeg for maximum compatibility (macOS, Windows, etc.)

    Args:
        tensor: (T, C, H, W) or (B, T, C, H, W) normalized tensor
        output_path: Path to save the video
        fps: Output video FPS
        mean: Normalization mean used
        std: Normalization std used

    Returns:
        Path to saved video
    """
    import subprocess

    frames = denormalize_tensor(tensor, mean, std)
    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)

    h, w = frames[0].shape[:2]

    # Use ffmpeg with H.264 codec for maximum compatibility
    cmd = [
        'ffmpeg', '-y',  # Overwrite output
        '-f', 'rawvideo',
        '-vcodec', 'rawvideo',
        '-s', f'{w}x{h}',
        '-pix_fmt', 'rgb24',
        '-r', str(fps),
        '-i', '-',  # Read from stdin
        '-c:v', 'libx264',  # H.264 codec
        '-pix_fmt', 'yuv420p',  # Required for QuickTime compatibility
        '-preset', 'fast',
        '-crf', '18',  # High quality
        output_path
    ]

    proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)

    for frame in frames:
        proc.stdin.write(frame.tobytes())

    proc.stdin.close()
    proc.wait()

    return output_path


def save_tensor_debug(
    tensor: torch.Tensor,
    output_dir: str,
    name: str,
    fps: int = 8,
) -> dict:
    """
    Save tensor for debugging: both as .pt file and as video.

    Args:
        tensor: (T, C, H, W) or (B, T, C, H, W) normalized tensor
        output_dir: Directory to save files
        name: Base name for files
        fps: Output video FPS

    Returns:
        Dict with paths to saved files
    """
    os.makedirs(output_dir, exist_ok=True)

    # Save raw tensor
    tensor_path = os.path.join(output_dir, f'{name}.pt')
    torch.save(tensor, tensor_path)

    # Save as video
    video_path = os.path.join(output_dir, f'{name}.mp4')
    save_tensor_as_video(tensor, video_path, fps)

    return {
        'tensor': tensor_path,
        'video': video_path,
    }
