# SPDX-License-Identifier: Apache-2.0
"""
Minimal image processing utilities for FastVideo.
This module provides lightweight image preprocessing without external dependencies beyond PyTorch/NumPy/PIL.
"""

import numpy as np
import PIL.Image
import torch


class ImageProcessor:
    """
    Minimal image processor for video frame preprocessing.

    This is a lightweight alternative to diffusers.VideoProcessor that handles:
    - PIL image to tensor conversion
    - Resizing to specified dimensions
    - Normalization to [-1, 1] range

    Args:
        vae_scale_factor: The VAE scale factor used to ensure dimensions are multiples of this value.
    """

    def __init__(self, vae_scale_factor: int = 8) -> None:
        self.vae_scale_factor = vae_scale_factor

    def preprocess(
        self,
        image: PIL.Image.Image | np.ndarray | torch.Tensor,
        height: int | None = None,
        width: int | None = None,
    ) -> torch.Tensor:
        """
        Preprocess an image to a normalized torch tensor.

        Args:
            image: Input image (PIL Image, NumPy array, or torch tensor)
            height: Target height. If None, uses image's original height.
            width: Target width. If None, uses image's original width.

        Returns:
            torch.Tensor: Normalized tensor of shape (1, 3, height, width) or (1, 1, height, width) for grayscale,
                         with values in range [-1, 1].
        """
        # Handle different input types
        if isinstance(image, PIL.Image.Image):
            return self._preprocess_pil(image, height, width)
        elif isinstance(image, np.ndarray):
            return self._preprocess_numpy(image, height, width)
        elif isinstance(image, torch.Tensor):
            return self._preprocess_tensor(image, height, width)
        else:
            raise ValueError(f"Unsupported image type: {type(image)}. "
                             "Supported types: PIL.Image.Image, np.ndarray, torch.Tensor")

    def _preprocess_pil(
        self,
        image: PIL.Image.Image,
        height: int | None = None,
        width: int | None = None,
    ) -> torch.Tensor:
        """Preprocess a PIL image."""
        if height is None:
            height = image.height
        if width is None:
            width = image.width

        height = height - (height % self.vae_scale_factor)
        width = width - (width % self.vae_scale_factor)

        image = image.resize((width, height), resample=PIL.Image.Resampling.LANCZOS)

        image_np = np.array(image, dtype=np.float32) / 255.0

        if image_np.ndim == 2:  # Grayscale
            image_np = np.expand_dims(image_np, axis=-1)

        return self._normalize_to_tensor(image_np)

    def _preprocess_numpy(
        self,
        image: np.ndarray,
        height: int | None = None,
        width: int | None = None,
    ) -> torch.Tensor:
        """Preprocess a numpy array."""
        # Determine target dimensions if not provided
        if image.ndim == 3:
            img_height, img_width = image.shape[:2]
        elif image.ndim == 2:
            img_height, img_width = image.shape
        else:
            raise ValueError(f"Expected 2D or 3D array, got {image.ndim}D")

        if height is None:
            height = img_height
        if width is None:
            width = img_width

        height = height - (height % self.vae_scale_factor)
        width = width - (width % self.vae_scale_factor)

        if image.dtype == np.uint8:
            pil_image = PIL.Image.fromarray(image)
        else:
            # Assume normalized [0, 1] or similar
            image_uint8 = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
            pil_image = PIL.Image.fromarray(image_uint8)

        pil_image = pil_image.resize((width, height), resample=PIL.Image.Resampling.LANCZOS)
        image_np = np.array(pil_image, dtype=np.float32) / 255.0

        # Ensure 3D shape
        if image_np.ndim == 2:
            image_np = np.expand_dims(image_np, axis=-1)

        return self._normalize_to_tensor(image_np)

    def _preprocess_tensor(
        self,
        image: torch.Tensor,
        height: int | None = None,
        width: int | None = None,
    ) -> torch.Tensor:
        """Preprocess a torch tensor."""
        # Determine target dimensions
        if image.ndim == 3:  # (H, W, C) or (C, H, W)
            if image.shape[0] in (1, 3, 4):  # Likely (C, H, W)
                img_height, img_width = image.shape[1], image.shape[2]
            else:  # Likely (H, W, C)
                img_height, img_width = image.shape[0], image.shape[1]
        elif image.ndim == 2:  # (H, W)
            img_height, img_width = image.shape
        else:
            raise ValueError(f"Expected 2D or 3D tensor, got {image.ndim}D")

        if height is None:
            height = img_height
        if width is None:
            width = img_width

        height = height - (height % self.vae_scale_factor)
        width = width - (width % self.vae_scale_factor)

        if image.ndim == 2:
            image = image.unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
        elif image.ndim == 3:
            # (C, H, W) or (H, W, C) - rearrange to (1, C, H, W)
            image = image.unsqueeze(0) if image.shape[0] in (1, 3, 4) else image.permute(2, 0, 1).unsqueeze(0)

        image = torch.nn.functional.interpolate(image, size=(height, width), mode="bilinear", align_corners=False)

        if image.max() > 1.0:  # Assume [0, 255] range
            image = image / 255.0

        image = 2.0 * image - 1.0

        return image

    def _normalize_to_tensor(self, image_np: np.ndarray) -> torch.Tensor:
        """
        Convert normalized numpy array [0, 1] to torch tensor [-1, 1].

        Args:
            image_np: NumPy array with shape (H, W) or (H, W, C) with values in [0, 1]

        Returns:
            torch.Tensor: Shape (1, C, H, W) or (1, 1, H, W) with values in [-1, 1]
        """
        # Convert to tensor
        if image_np.ndim == 2:  # (H, W) - grayscale
            tensor = torch.from_numpy(image_np).unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
        elif image_np.ndim == 3:  # (H, W, C)
            tensor = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0)  # (1, C, H, W)
        else:
            raise ValueError(f"Expected 2D or 3D array, got {image_np.ndim}D")

        # Normalize to [-1, 1]
        tensor = 2.0 * tensor - 1.0

        return tensor

    def _preprocess_cosmos25(
        self,
        image: PIL.Image.Image | np.ndarray | torch.Tensor,
        height: int,
        width: int,
    ) -> torch.Tensor:
        """
        Cosmos-Predict2.5-style preprocessing for image2world:
        - aspect-preserving resize (scale so both dims >= target)
        - center crop to (height, width)
        - normalize to [-1, 1]

        Returns:
            torch.Tensor: (1, 3, height, width) in [-1, 1]
        """
        if not isinstance(image, PIL.Image.Image):
            # Reuse the generic preprocess for non-PIL inputs.
            return self.preprocess(image, height=height, width=width)

        # Ensure RGB (official uses torchvision.to_tensor on PIL, which yields C=3 for RGB inputs)
        image = image.convert("RGB")

        # Make sure target dims are multiples of VAE scale factor (FastVideo convention)
        height = height - (height % self.vae_scale_factor)
        width = width - (width % self.vae_scale_factor)

        orig_w, orig_h = image.size
        # Scale so that resized image fully covers the target box (like official resize_input()).
        scale = max(width / float(orig_w), height / float(orig_h))
        resized_w = int(np.ceil(scale * orig_w))
        resized_h = int(np.ceil(scale * orig_h))

        image = image.resize((resized_w, resized_h), resample=PIL.Image.Resampling.BILINEAR)

        # Center crop
        left = max(0, (resized_w - width) // 2)
        top = max(0, (resized_h - height) // 2)
        image = image.crop((left, top, left + width, top + height))

        image_np = np.array(image, dtype=np.float32) / 255.0  # [0,1]
        return self._normalize_to_tensor(image_np)  # -> [-1,1], (1,3,H,W)