"""Depth Anything 3 preprocessing: the package's ``InputProcessor`` (vendored in ``da3_input_processor.py``) behind the
transformers image-processor interface. Images (PIL / numpy / paths) -> ``pixel_values`` (1, N, 3, H, W), resized
("upper_bound_resize" to 504 by default, sides multiples of 14) and ImageNet-normalised; camera parameters given
with the images come back resized to match (``extrinsics`` / ``intrinsics``)."""

from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor

from .da3_input_processor import InputProcessor


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

    def __init__(self, process_res: int = 504, process_res_method: str = "upper_bound_resize", **kwargs):
        super().__init__(**kwargs)
        self.process_res = process_res
        self.process_res_method = process_res_method

    def preprocess(self, images, extrinsics=None, intrinsics=None, process_res=None, process_res_method=None,
                   return_tensors=None, **kwargs) -> BatchFeature:
        if not isinstance(images, (list, tuple)):
            images = [images]
        tensor, extrinsics, intrinsics = InputProcessor()(
            list(images),
            extrinsics.copy() if extrinsics is not None else None,
            intrinsics.copy() if intrinsics is not None else None,
            process_res or self.process_res,
            process_res_method or self.process_res_method,
            **kwargs,
        )
        # InputProcessor returns (N, 3, H, W); the package API adds the batch axis (_prepare_model_inputs)
        data = {"pixel_values": tensor[None].float()}
        if extrinsics is not None:
            data["extrinsics"] = extrinsics[None]
        if intrinsics is not None:
            data["intrinsics"] = intrinsics[None]
        return BatchFeature(data, tensor_type=return_tensors)


__all__ = ["DepthAnything3ImageProcessor"]