# ---------------------------------------------------------------------
# Copyright (c) 2026 Qualcomm Technologies, Inc. and/or its subsidiaries.
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------------

from __future__ import annotations

import torch
from transformers import DepthProForDepthEstimation
from typing_extensions import Self

from qai_hub_models.datasets.nyuv2 import NYUV2Dataset
from qai_hub_models.models._shared.depth_estimation.depth_evaluator import (
    DepthEvaluator,
)
from qai_hub_models.utils.base_dataset import BaseDataset
from qai_hub_models.utils.base_evaluator import BaseEvaluator
from qai_hub_models.utils.base_model import BaseModel
from qai_hub_models.utils.input_spec import (
    ColorFormat,
    ImageMetadata,
    InputSpec,
    IoType,
    OutputSpec,
    TensorSpec,
)

MODEL_ID = "depth_pro"
MODEL_ASSET_VERSION = 1
DEFAULT_WEIGHTS = "apple/DepthPro-hf"
DEFAULT_INPUT_SIZE = 1536


class DepthProDepthEvaluator(DepthEvaluator):
    """Adapts the shared depth δ1 evaluator to DepthPro's (depth, fov) tuple.

    DepthPro's canonical inverse-depth head is what the shared evaluator
    already scale/shift-aligns against NYUv2 ground truth, so the FoV output
    is ignored here — focal-length calibration is only needed for absolute
    metric alignment, which the δ1 metric is invariant to.
    """

    def add_batch(
        self,
        output: torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor],
        gt: torch.Tensor,
    ) -> None:
        if isinstance(output, (tuple, list)):
            output = output[0]
        if output.dim() == 3:
            # depth arrives as (B, H, W); base evaluator expects (B, 1, H, W)
            output = output.unsqueeze(1)
        super().add_batch(output, gt)


class DepthPro(BaseModel):
    """Apple DepthPro monocular metric depth estimator, end-to-end.

    Exposes two on-device outputs: canonical inverse depth at the network's
    input resolution and a scalar horizontal field of view (degrees) per
    image. Off-device post-processing (see ``DepthProApp``) converts these
    into metric depth and a focal length in pixels, matching HuggingFace's
    ``DepthProImageProcessorFast.post_process_depth_estimation``.
    """

    def __init__(self, model: torch.nn.Module) -> None:
        super().__init__()
        self.model = model.eval()

    @classmethod
    def from_pretrained(cls, ckpt: str = DEFAULT_WEIGHTS) -> Self:
        net = DepthProForDepthEstimation.from_pretrained(ckpt)
        return cls(net)

    def forward(
        self, image: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Run DepthPro on `image`.

        Parameters
        ----------
        image
            Shape ``[B, 3, 1536, 1536]`` RGB in ``[0, 1]``.

        Returns
        -------
        predicted_depth
            Shape ``[B, 1536, 1536]``. Canonical inverse depth (the raw head
            output before FoV-based metric scaling and inversion).
        field_of_view
            Shape ``[B]``, horizontal field of view in degrees.
        """
        # HF's DepthProImageProcessorFast rescales to [0, 1] then normalizes
        # with mean=std=0.5, i.e. (image - 0.5) / 0.5 = 2 * image - 1.
        pixel_values = image * 2.0 - 1.0
        predicted_depth, field_of_view = self.model(
            pixel_values, return_dict=False
        )
        return predicted_depth, field_of_view

    def get_input_spec(
        self,
        batch_size: int = 1,
        height: int = DEFAULT_INPUT_SIZE,
        width: int = DEFAULT_INPUT_SIZE,
    ) -> InputSpec:
        return {
            "image": TensorSpec(
                shape=(batch_size, 3, height, width),
                dtype="float32",
                io_type=IoType.IMAGE,
                value_range=(0.0, 1.0),
                image_metadata=ImageMetadata(color_format=ColorFormat.RGB),
                apply_runtime_channel_reordering=True,
            ),
        }

    def get_output_spec(self) -> OutputSpec:
        return {
            "predicted_depth": TensorSpec(
                io_type=IoType.TENSOR,
                description=(
                    "Canonical inverse depth at the network's input "
                    "resolution. Invert and rescale by "
                    "width / focal_length_px for metric depth."
                ),
                apply_runtime_channel_reordering=True,
            ),
            "field_of_view": TensorSpec(
                io_type=IoType.TENSOR,
                description=(
                    "Horizontal field of view in degrees, one scalar per "
                    "image; used off-device to derive focal length."
                ),
            ),
        }

    def get_evaluator(self) -> BaseEvaluator:
        return DepthProDepthEvaluator()

    @classmethod
    def get_eval_dataset_classes(cls) -> list[type[BaseDataset]]:
        return [NYUV2Dataset]

    def get_calibration_dataset_cls(self) -> type[BaseDataset]:
        return NYUV2Dataset