"""Video Depth Anything (DepthAnything/Video-Depth-Anything @ 4f5ae231) as a transformers depth-estimation model.

``VideoDepthAnythingForDepthEstimation`` holds the submodules of the official ``VideoDepthAnything`` under the same
names (``pretrained``: DINOv2, ``head``: temporal DPT head), so the released ``*video_depth_anything_*.pth`` tensors
load by name, and computes depth with that class's own ``forward`` (``vda_video_depth.py``), returned as the standard
``DepthEstimatorOutput`` (``predicted_depth`` (B, T, H, W)). The official whole-video inference (``infer_video_depth``:
32-frame windows, key-frame overlap, scale / shift alignment) and the streaming inference (``infer_video_depth_one``,
``vda_video_depth_stream.py``) run as they are on this model.
"""

import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import DepthEstimatorOutput

from .configuration_video_depth_anything import VideoDepthAnythingConfig
from .vda_video_depth import VideoDepthAnything
from .vda_video_depth_stream import INFER_LEN, INTERP_LEN, OVERLAP
from .vda_video_depth_stream import VideoDepthAnything as StreamingVideoDepthAnything


class _OfficialView:
    """The model as an instance of an official class: that class's methods (``forward`` included) bound to this view,
    other attributes read from the model; attributes the methods set (the streaming caches) stay on the view."""

    def __init__(self, model, official):
        object.__setattr__(self, "_model", model)
        object.__setattr__(self, "_official", official)

    def __getattr__(self, name):
        if name in vars(self._official):
            return vars(self._official)[name].__get__(self)
        return getattr(self._model, name)


class VideoDepthAnythingPreTrainedModel(PreTrainedModel):
    config_class = VideoDepthAnythingConfig
    config: VideoDepthAnythingConfig
    base_model_prefix = "video_depth_anything"
    main_input_name = "pixel_values"


class VideoDepthAnythingForDepthEstimation(VideoDepthAnythingPreTrainedModel):
    def __init__(self, config: VideoDepthAnythingConfig):
        super().__init__(config)
        net = VideoDepthAnything(encoder=config.encoder, features=config.features, out_channels=list(config.out_channels),
                                 use_bn=config.use_bn, use_clstoken=config.use_clstoken, num_frames=config.num_frames,
                                 pe=config.pe, metric=config.metric)
        # 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.metric = net.metric
        self.pretrained = net.pretrained
        self.head = net.head
        self._stream = None
        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, T, 3, H, W), H and W multiples of 14, ImageNet-normalised (``VideoDepthAnythingProcessor``);
        T up to ``config.num_frames``. ``predicted_depth`` (B, T, H, W): relative inverse depth, or metric depth in
        meters for the metric checkpoints. ``output_hidden_states``: the DINOv2 features the head reads (per layer,
        (patch tokens, class token) of the B*T frames)."""
        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 = VideoDepthAnything.forward(self, pixel_values)
        hidden_states = None
        if output_hidden_states:
            hidden_states = tuple(self.pretrained.get_intermediate_layers(
                pixel_values.flatten(0, 1), self.intermediate_layer_idx[self.encoder], 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 inference: frames (T, H, W, 3) RGB uint8 -> (depths (T, H, W) numpy at the frame size, fps)
    def infer_video_depth(self, frames, target_fps, input_size: int = 518, device: str = "cuda", fp32: bool = False):
        return VideoDepthAnything.infer_video_depth(_OfficialView(self, VideoDepthAnything), frames, target_fps,
                                                    input_size=input_size, device=device, fp32=fp32)

    def infer_video_depth_one(self, frame, input_size: int = 518, device: str = "cuda", fp32: bool = False):
        """Streaming (experimental upstream): one RGB uint8 frame (H, W, 3) -> its depth; the temporal caches persist
        across calls until ``reset_streaming``."""
        if self._stream is None:
            self._stream = _OfficialView(self, StreamingVideoDepthAnything)
            # the state StreamingVideoDepthAnything.__init__ sets
            self._stream.transform = None
            self._stream.frame_id_list = []
            self._stream.frame_cache_list = []
            self._stream.gap = (INFER_LEN - OVERLAP) * 2 - 1 - (OVERLAP - INTERP_LEN)
            self._stream.id = -1
        return StreamingVideoDepthAnything.infer_video_depth_one(self._stream, frame, input_size=input_size, device=device, fp32=fp32)

    def reset_streaming(self):
        self._stream = None


__all__ = ["VideoDepthAnythingForDepthEstimation", "VideoDepthAnythingPreTrainedModel"]