"""Depth Anything (v1) preprocessing: the transform of the official run.py (Resize lower_bound to 518 keeping the aspect
ratio, sides multiples of 14, cv2 bicubic; ImageNet normalisation; vendored in ``dav1_dpt.py``) behind the
transformers image-processor interface. A path is read with cv2 as run.py does; a PIL image / RGB array is used as RGB.
``post_process_depth_estimation`` resizes to the given sizes as run.py does (bilinear, align_corners=False -- not
transformers' bicubic)."""

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 .dav1_dpt import NormalizeImage, PrepareForNet, Resize


class DepthAnythingV1ImageProcessor(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, images, input_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature:
        if not isinstance(images, (list, tuple)):
            images = [images]
        size = input_size or self.input_size
        transform = Compose([
            Resize(width=size, height=size, 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(),
        ])
        tensors = []
        for image in images:
            if isinstance(image, str):
                rgb = cv2.cvtColor(cv2.imread(image), cv2.COLOR_BGR2RGB) / 255.0
            else:
                rgb = np.asarray(image.convert("RGB") if isinstance(image, Image.Image) else image) / 255.0
            tensors.append(torch.from_numpy(transform({"image": rgb})["image"]))
        if len({t.shape for t in tensors}) > 1:
            raise ValueError("images resize to different shapes; process them one at a time")
        return BatchFeature({"pixel_values": torch.stack(tensors)}, tensor_type=return_tensors)

    def post_process_depth_estimation(self, outputs, target_sizes=None) -> list:
        """``DepthEstimatorOutput`` -> ``[{"predicted_depth": (h, w)}]`` per image; target_sizes: (height, width) per
        image (e.g. the original sizes), 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:
                depth = F.interpolate(depth[None, None], (int(size[0]), int(size[1])), mode="bilinear", align_corners=False)[0, 0]
            results.append({"predicted_depth": depth})
        return results


__all__ = ["DepthAnythingV1ImageProcessor"]