"""Video Depth Anything preprocessing: the per-frame transform of the official ``infer_video_depth`` (Resize lower_bound
to 518 keeping the aspect ratio, sides multiples of 14, cv2 bicubic; ImageNet normalisation; vendored in ``vda_dpt.py``),
including its rule for wide videos (aspect ratio > 1.78: smaller input size), behind the transformers image-processor
interface. Frames are RGB uint8 arrays (H, W, 3), as the official ``read_video_frames`` returns them.
``post_process_depth_estimation`` resizes to the given sizes as ``infer_video_depth`` does (float32, bilinear,
align_corners=True)."""

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision.transforms import Compose
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor

from .vda_dpt import NormalizeImage, PrepareForNet, Resize


class VideoDepthAnythingImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def __init__(self, input_size: int = 518, **kwargs):
        super().__init__(**kwargs)
        self.input_size = input_size

    def preprocess(self, videos, input_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature:
        """videos: one video or a list of videos of the same size; a video is an array (T, H, W, 3) or a list of
        frames (RGB arrays or PIL images). pixel_values: (B, T, 3, H', W')."""
        is_frame = lambda x: isinstance(x, Image.Image) or isinstance(x, np.ndarray) and x.ndim == 3  # noqa: E731
        if isinstance(videos, np.ndarray) and videos.ndim == 4 or is_frame(videos[0]):
            videos = [videos]
        size = input_size or self.input_size
        batch = []
        for video in videos:
            frames = [np.asarray(f.convert("RGB")) if isinstance(f, Image.Image) else f for f in video]
            frame_height, frame_width = frames[0].shape[:2]
            # infer_video_depth: wide videos get a smaller input size
            input_size_video = size
            ratio = max(frame_height, frame_width) / min(frame_height, frame_width)
            if ratio > 1.78:
                input_size_video = int(input_size_video * 1.777 / ratio)
                input_size_video = round(input_size_video / 14) * 14
            transform = Compose([
                Resize(width=input_size_video, height=input_size_video, resize_target=False, keep_aspect_ratio=True,
                       ensure_multiple_of=14, resize_method="lower_bound", image_interpolation_method=cv2.INTER_CUBIC),
                NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
                PrepareForNet(),
            ])
            batch.append(torch.stack([torch.from_numpy(transform({"image": f.astype(np.float32) / 255.0})["image"]) for f in frames]))
        if len({t.shape for t in batch}) > 1:
            raise ValueError("videos resize to different shapes; process them one at a time")
        return BatchFeature({"pixel_values": torch.stack(batch)}, tensor_type=return_tensors)

    def post_process_depth_estimation(self, outputs, target_sizes=None) -> list:
        """``DepthEstimatorOutput`` (predicted_depth (B, T, H, W)) -> ``[{"predicted_depth": (T, h, w)}]`` per video;
        target_sizes: (height, width) per video (e.g. the frame size), None keeps the model resolution."""
        predicted_depth = outputs.predicted_depth
        if target_sizes is not None and len(predicted_depth) != len(target_sizes):
            raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the predicted depth")
        results = []
        for depth, size in zip(predicted_depth, target_sizes if target_sizes is not None else [None] * len(predicted_depth)):
            if size is not None:  # infer_video_depth: depth.to(float32 input dtype), then upsample
                depth = F.interpolate(depth.float().unsqueeze(1), size=(int(size[0]), int(size[1])), mode="bilinear", align_corners=True)[:, 0]
            results.append({"predicted_depth": depth})
        return results


__all__ = ["VideoDepthAnythingImageProcessor"]