"""Camera Depth Model preprocessing: cdm/infer.py ``load_images`` (raw depth / depth_scale, values above max_depth set
to 0, inverse depth where valid) and ``RGBDDepth.image2tensor`` (Resize lower_bound to 518 keeping the aspect ratio,
sides multiples of 14, cv2 bicubic for the image and nearest for the depth; ImageNet normalisation; vendored in
``cdm_dpt.py``) behind the transformers image-processor interface.

As in infer.py and the README API, the image is given as RGB and ``image2tensor`` applies ``cv2.COLOR_BGR2RGB`` to it,
so the network sees the channels swapped; this is kept unchanged. ``post_process_depth_estimation`` resizes to the
given sizes as ``infer_image`` does (nearest). The model predicts inverse depth: meters = 1 / prediction (infer.py)."""

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


class CameraDepthModelImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values", "prompt_depth"]

    def __init__(self, input_size: int = 518, depth_scale: float = 1000.0, max_depth: float = 25.0, **kwargs):
        super().__init__(**kwargs)
        self.input_size = input_size
        self.depth_scale = depth_scale
        self.max_depth = max_depth  # infer.py calls load_images(..., 25.0)

    def load_inverse_depth(self, depth, depth_scale: float, max_depth: float):
        # infer.py load_images; a path is the raw camera image (cv2 IMREAD_UNCHANGED, / depth_scale), an array is meters
        if isinstance(depth, str):
            depth_low_res = cv2.imread(depth, cv2.IMREAD_UNCHANGED)
            if depth_low_res is None:
                raise ValueError(f"Could not load depth image from {depth}")
            depth_low_res = np.asarray(depth_low_res).astype(np.float32) / depth_scale
        else:
            depth_low_res = np.array(depth, dtype=np.float32)
        depth_low_res[depth_low_res > max_depth] = 0.0  # Remove values beyond max range
        simi_depth_low_res = np.zeros_like(depth_low_res)
        simi_depth_low_res[depth_low_res > 0] = 1 / depth_low_res[depth_low_res > 0]
        return simi_depth_low_res

    def preprocess(self, images, depth=None, input_size: int | None = None, depth_scale: float | None = None,
                   max_depth: float | None = None, return_tensors=None, **kwargs) -> BatchFeature:
        """images: one image or a list (paths, PIL images or RGB uint8 arrays); depth: one raw depth per image (paths
        of the camera's depth png, or arrays in meters). pixel_values (B, 3, H, W); prompt_depth (B, 1, H, W): the
        inverse depth resized with the image."""
        if not isinstance(images, (list, tuple)):
            images = [images]
        if not isinstance(depth, (list, tuple)):
            depth = [depth]
        if len(depth) != len(images):
            raise ValueError("give one depth per image")
        size = input_size or self.input_size
        transform = Compose([
            Resize(width=size, height=size, resize_target=True, 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(),
        ])
        pixel_values, prompt_depth = [], []
        for image, raw_depth in zip(images, depth):
            if isinstance(image, str):
                raw_image = np.asarray(cv2.imread(image)[:, :, ::-1])  # infer.py load_images
            else:
                raw_image = np.asarray(image.convert("RGB") if isinstance(image, Image.Image) else image)
            simi_depth = self.load_inverse_depth(raw_depth, depth_scale or self.depth_scale, max_depth or self.max_depth)
            # RGBDDepth.image2tensor
            prepared = transform({"image": cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0, "depth": simi_depth})
            pixel_values.append(torch.from_numpy(prepared["image"]))
            prompt_depth.append(torch.from_numpy(prepared["depth"]).unsqueeze(0))
        if len({t.shape for t in pixel_values}) > 1:
            raise ValueError("images resize to different shapes; process them one at a time")
        return BatchFeature({"pixel_values": torch.stack(pixel_values), "prompt_depth": torch.stack(prompt_depth)},
                            tensor_type=return_tensors)

    def post_process_depth_estimation(self, outputs, target_sizes=None) -> list:
        """``DepthEstimatorOutput`` -> ``[{"predicted_depth": (h, w)}]`` (inverse depth) 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:  # infer_image
                depth = F.interpolate(depth[None, None], (int(size[0]), int(size[1])), mode="nearest")[0, 0]
            results.append({"predicted_depth": depth})
        return results


__all__ = ["CameraDepthModelImageProcessor"]