"""Camera Depth Model (ByteDance-Seed/manip-as-in-sim-suite @ 757b2ea2) as a transformers depth-estimation model.

``CameraDepthModelForDepthEstimation`` holds the submodules of the official ``RGBDDepth`` under the same names
(``pretrained``: the RGB DINOv2, ``depth_pretrained``: the depth DINOv2, ``crossAtts``, ``depth_head_rgbd``), so the
released checkpoints load by name, and computes depth with that class's own ``forward`` on the RGB + inverse-depth
input, returned as the standard ``DepthEstimatorOutput`` (``predicted_depth``: inverse depth). The official
``infer_image`` / ``image2tensor`` helpers are kept.
"""

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

from .cdm_dpt import RGBDDepth
from .configuration_camera_depth_model import CameraDepthModelConfig


class CameraDepthModelPreTrainedModel(PreTrainedModel):
    config_class = CameraDepthModelConfig
    config: CameraDepthModelConfig
    base_model_prefix = "camera_depth_model"
    main_input_name = "pixel_values"


class CameraDepthModelForDepthEstimation(CameraDepthModelPreTrainedModel):
    def __init__(self, config: CameraDepthModelConfig):
        super().__init__(config)
        net = RGBDDepth(encoder=config.encoder, features=config.features, out_channels=list(config.out_channels),
                        use_bn=config.use_bn, use_clstoken=config.use_clstoken, max_depth=config.max_depth)
        # the official module's attributes, so its methods run on this object unchanged
        self.intermediate_layer_idx = net.intermediate_layer_idx
        self.max_depth = net.max_depth
        self.encoder = net.encoder
        self.pretrained = net.pretrained
        self.depth_pretrained = net.depth_pretrained
        self.depth_head_rgbd = net.depth_head_rgbd
        self.crossAtts = net.crossAtts
        self.post_init()

    def forward(
        self,
        pixel_values: torch.Tensor,
        prompt_depth: torch.Tensor | None = None,
        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; prompt_depth (B, 1, H, W): the
        camera's inverse depth (1 / meters, 0 where invalid) at the same size; both from ``CameraDepthModelProcessor``.
        ``predicted_depth`` (B, H, W): inverse depth (meters = 1 / predicted_depth, as infer.py).
        ``output_hidden_states``: the DINOv2 features of the RGB and the depth branch (per layer, (patch tokens, class
        token)), RGB layers first."""
        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)")
        if prompt_depth is None:
            raise ValueError("prompt_depth (the camera's inverse depth) is required")
        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
        x = torch.cat((pixel_values, prompt_depth), dim=1)  # image2tensor
        predicted_depth = RGBDDepth.forward(self, x)
        hidden_states = None
        if output_hidden_states:
            layers = self.intermediate_layer_idx[self.encoder]
            hidden_states = tuple(self.pretrained.get_intermediate_layers(pixel_values, layers, return_class_token=True)) + tuple(
                self.depth_pretrained.get_intermediate_layers(prompt_depth.repeat(1, 3, 1, 1), layers, return_class_token=True))
        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: RGB uint8 image + inverse depth (numpy) -> inverse depth at the image size (numpy)
    @torch.no_grad()
    def infer_image(self, raw_image, depth_low_res, input_size: int = 518):
        # body of RGBDDepth.infer_image, calling the official forward (self.forward returns a DepthEstimatorOutput)
        inputs, (h, w) = self.image2tensor(raw_image, depth_low_res, input_size)
        pred_depth = RGBDDepth.forward(self, inputs)
        pred_depth = F.interpolate(pred_depth[:, None], (h, w), mode="nearest")[0, 0]
        return pred_depth.cpu().numpy()

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


__all__ = ["CameraDepthModelForDepthEstimation", "CameraDepthModelPreTrainedModel"]