"""Depth Anything V2 (DepthAnything/Depth-Anything-V2 @ a561b849) as a transformers depth-estimation model.

``DepthAnythingV2ForDepthEstimation`` holds the submodules of the official ``DepthAnythingV2`` under the same names
(``pretrained``: DINOv2, ``depth_head``: DPT head), so the released ``depth_anything_v2_*.pth`` tensors load by name,
and computes depth with that class's own ``forward`` -- the relative-depth class (``dav2_dpt.py``), or with
``config.max_depth`` set the metric-depth class (``dav2_metric_dpt.py``, metric_depth/ of the repo). It returns the
standard ``DepthEstimatorOutput`` (``predicted_depth``). The official ``infer_image`` / ``image2tensor`` helpers are
kept. Unlike transformers' ``DepthAnythingForDepthEstimation``, the DINOv2 position-embedding interpolation is the
original one, so non-square inputs give the original outputs.
"""

import torch
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import DepthEstimatorOutput

from .configuration_depth_anything_v2 import DepthAnythingV2Config
from .dav2_dpt import DepthAnythingV2 as RelativeDepthAnythingV2
from .dav2_metric_dpt import DepthAnythingV2 as MetricDepthAnythingV2


class DepthAnythingV2PreTrainedModel(PreTrainedModel):
    config_class = DepthAnythingV2Config
    config: DepthAnythingV2Config
    base_model_prefix = "depth_anything_v2"
    main_input_name = "pixel_values"


class DepthAnythingV2ForDepthEstimation(DepthAnythingV2PreTrainedModel):
    def __init__(self, config: DepthAnythingV2Config):
        super().__init__(config)
        kwargs = dict(encoder=config.encoder, features=config.features, out_channels=list(config.out_channels),
                      use_bn=config.use_bn, use_clstoken=config.use_clstoken)
        if config.max_depth is None:
            self._official = RelativeDepthAnythingV2
            net = RelativeDepthAnythingV2(**kwargs)
        else:
            self._official = MetricDepthAnythingV2
            net = MetricDepthAnythingV2(**kwargs, max_depth=config.max_depth)
            self.max_depth = net.max_depth
        # the official module's attributes, so its methods run on this object unchanged
        self.intermediate_layer_idx = net.intermediate_layer_idx
        self.encoder = net.encoder
        self.pretrained = net.pretrained
        self.depth_head = net.depth_head
        self.post_init()

    def forward(
        self,
        pixel_values: torch.Tensor,
        labels: torch.Tensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ) -> tuple | DepthEstimatorOutput:
        """pixel_values (B, 3, H, W), H and W multiples of 14, ImageNet-normalised (``DepthAnythingV2Processor``).
        ``predicted_depth`` (B, H, W): relative inverse depth, or metric depth in meters when ``config.max_depth``
        is set. ``output_hidden_states``: the DINOv2 features the DPT head reads (per layer, patch tokens)."""
        if labels is not None:
            raise NotImplementedError("Training is not implemented")
        if output_attentions:
            raise NotImplementedError("Attention weights are not returned (the original model does not expose them)")
        return_dict = return_dict if return_dict is not None else self.config.return_dict
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        predicted_depth = self._official.forward(self, pixel_values)
        hidden_states = None
        if output_hidden_states:
            features = self.pretrained.get_intermediate_layers(pixel_values, self.intermediate_layer_idx[self.encoder])
            hidden_states = tuple(features)
        if not return_dict:
            return (predicted_depth,) + ((hidden_states,) if hidden_states is not None else ())
        return DepthEstimatorOutput(predicted_depth=predicted_depth, hidden_states=hidden_states)

    # official helpers: raw BGR uint8 image (cv2.imread) -> depth at the original size (numpy)
    @torch.no_grad()
    def infer_image(self, raw_image, input_size: int = 518):
        # body of DepthAnythingV2.infer_image, calling the official forward (self.forward returns a DepthEstimatorOutput)
        image, (h, w) = self.image2tensor(raw_image, input_size)
        depth = self._official.forward(self, image)
        depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0]
        return depth.cpu().numpy()

    def image2tensor(self, raw_image, input_size: int = 518):
        return self._official.image2tensor(self, raw_image, input_size)


__all__ = ["DepthAnythingV2ForDepthEstimation", "DepthAnythingV2PreTrainedModel"]