"""Prompt Depth Anything preprocessing: the official ``load_image`` / ``load_depth`` (promptda/utils/io_wrapper.py @
9161547c) behind the transformers image-processor interface. That module also imports open3d / matplotlib for
visualisation, so its two loaders are re-typed here line by line (checked bit-exact against the originals).

Image: RGB in [0, 1] (the model normalises it), longer side at most ``max_size`` (cv2 INTER_AREA, sides floored to
multiples of 14). A path is read with imageio, as the original. Prompt depth: meters, (h, w); a .png path is
millimeters (/ 1000), a .npz path holds "depth", as the original."""

import cv2
import imageio
import numpy as np
import torch
from PIL import Image
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor


def _to_tensor(arr):  # io_wrapper.to_tensor_func
    if arr.ndim == 2:
        arr = arr[:, :, np.newaxis]
    return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)


def _ensure_multiple_of(x, multiple_of=14):  # io_wrapper.ensure_multiple_of
    return int(x // multiple_of * multiple_of)


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

    def __init__(self, max_size: int = 1008, **kwargs):
        super().__init__(**kwargs)
        self.max_size = max_size

    def load_image(self, image, max_size: int, multiple_of: int = 14):
        # io_wrapper.load_image, also taking a PIL image / RGB uint8 array
        if isinstance(image, str):
            image = np.asarray(imageio.imread(image)).astype(np.float32)
        else:
            image = np.asarray(image.convert("RGB") if isinstance(image, Image.Image) else image).astype(np.float32)
        image = image / 255.

        max_size = max_size // multiple_of * multiple_of
        if max(image.shape) > max_size:
            h, w = image.shape[:2]
            scale = max_size / max(h, w)
            tar_h = _ensure_multiple_of(h * scale)
            tar_w = _ensure_multiple_of(w * scale)
            image = cv2.resize(image, (tar_w, tar_h), interpolation=cv2.INTER_AREA)
        return _to_tensor(image)

    @staticmethod
    def load_depth(depth):
        # io_wrapper.load_depth, also taking an array / tensor in meters
        if isinstance(depth, str):
            if depth.endswith('.png'):
                depth = np.asarray(imageio.imread(depth)).astype(np.float32)
                depth = depth / 1000.
            elif depth.endswith('.npz'):
                depth = np.load(depth)['depth']
            else:
                raise ValueError(f"Unsupported depth format: {depth}")
        elif isinstance(depth, torch.Tensor):
            depth = depth.cpu().numpy()
        return _to_tensor(np.asarray(depth))

    def preprocess(self, images, prompt_depth=None, max_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature:
        """images: one image or a list (paths, PIL images or RGB uint8 arrays); prompt_depth: one depth map per image
        (paths, or arrays in meters). pixel_values (B, 3, H, W) in [0, 1]; prompt_depth (B, 1, h, w)."""
        if not isinstance(images, (list, tuple)):
            images = [images]
        tensors = [self.load_image(image, max_size or self.max_size)[0] for image in images]
        if len({t.shape for t in tensors}) > 1:
            raise ValueError("images resize to different shapes; process them one at a time")
        data = {"pixel_values": torch.stack(tensors)}
        if prompt_depth is not None:
            if not isinstance(prompt_depth, (list, tuple)):
                prompt_depth = [prompt_depth]
            if len(prompt_depth) != len(images):
                raise ValueError("give one prompt depth per image")
            data["prompt_depth"] = torch.stack([self.load_depth(depth)[0] for depth in prompt_depth])
        return BatchFeature(data, tensor_type=return_tensors)


__all__ = ["PromptDepthAnythingImageProcessor"]