# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
# Modified 2026 by The PaGeR Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Depth Anything 3 API module.

This module provides the main API for Depth Anything 3, including model loading,
inference, and export capabilities. It supports both single and nested model architectures.
"""

from __future__ import annotations

import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin

from depth_anything_3.cfg import create_object, load_config
from depth_anything_3.registry import MODEL_REGISTRY
from depth_anything_3.utils.geometry import affine_inverse
from depth_anything_3.model.utils.valid_conv_padding import set_valid_pad_conv
torch.backends.cudnn.benchmark = False
# logger.info("CUDNN Benchmark Disabled")

SAFETENSORS_NAME = "model.safetensors"
CONFIG_NAME = "config.json"


class DepthAnything3(nn.Module, PyTorchModelHubMixin):
    """
    Depth Anything 3 main API class.

    This class provides a high-level interface for depth estimation using Depth Anything 3.
    It supports both single and nested model architectures with metric scaling capabilities.

    Features:
    - Hugging Face Hub integration via PyTorchModelHubMixin
    - Support for multiple model presets (vitb, vitg, nested variants)
    - Automatic mixed precision inference
    - Export capabilities for various formats (GLB, PLY, NPZ, etc.)
    - Camera pose estimation and metric depth scaling

    Usage:
        # Load from Hugging Face Hub
        model = DepthAnything3.from_pretrained("huggingface/model-name")

        # Or create with specific preset
        model = DepthAnything3(preset="vitg")

        # Run inference
        prediction = model.inference(images, export_dir="output", export_format="glb")
    """

    _commit_hash: str | None = None  # Set by mixin when loading from Hub

    def __init__(self, model_name: str = "da3-large", **kwargs):
        """
        Initialize DepthAnything3 with specified preset.

        Args:
        model_name: The name of the model preset to use.
                    Examples: 'da3-giant', 'da3-large', 'da3metric-large', 'da3nested-giant-large'.
        **kwargs: Additional keyword arguments (currently unused).
        """
        super().__init__()
        self.model_name = model_name

        # Build the underlying network from the bundled YAML preset.
        # Every released PaGeR checkpoint shares the same backbone topology
        # (DA3-Giant, no global PE / RoPE LoRA / seam-local-attn / fused-res
        # head ordering), so those flags are inlined and not re-read from
        # cfg.model. Only ``head_names``, ``valid_conv_padding`` and
        # ``log_depth`` differ between released variants.
        self.config = load_config(MODEL_REGISTRY[self.model_name])

        model_cfg = kwargs["model_cfg"]
        self.config.head["head_names"] = model_cfg.modalities
        self.config.head["valid_conv_padding"] = model_cfg.valid_conv_padding
        self.config.head["log_depth"] = model_cfg.log_depth
        self.config.head["with_confidence"] = True
        self.model = create_object(self.config)

        if model_cfg.valid_conv_padding:
            set_valid_pad_conv(self.model.head)

        # Device management (set by user)
        self.device = None

    def forward(
        self,
        image: torch.Tensor,
        extrinsics: torch.Tensor | None = None,
        intrinsics: torch.Tensor | None = None,
        face_ids: torch.Tensor | None = None,
        skip_heads=None,
    ) -> dict[str, torch.Tensor]:
        """
        Forward pass through the model.

        Args:
            image: Input batch with shape ``(B, N, 3, H, W)`` on the model device.
            extrinsics: Optional camera extrinsics with shape ``(B, N, 4, 4)``.
            intrinsics: Optional camera intrinsics with shape ``(B, N, 3, 3)``.
            face_ids: Optional (B, N) long tensor with values in [0, 5]; when
                the input only contains a subset of cubemap faces (N < 6),
                this tells the model which of the 6 canonical positions each
                N-slot occupies so the global PE / RoPE still line up.

        Returns:
            Dictionary containing model predictions
        """
        # Accept either (N, 4, 4) legacy input (broadcast across batch via
        # [None]) or (B, N, 4, 4) per-sample input. Same for intrinsics.
        ext_in = extrinsics if extrinsics.dim() == 4 else extrinsics[None]
        intr_in = intrinsics if intrinsics.dim() == 4 else intrinsics[None]
        ex_t_norm = self._normalize_extrinsics(ext_in.clone())
        prediction = self.model(
            image, ex_t_norm, intr_in, face_ids=face_ids, skip_heads=skip_heads,
        )
        return prediction
    

    def _normalize_extrinsics(self, ex_t: torch.Tensor | None) -> torch.Tensor | None:
        """Normalize extrinsics"""
        if ex_t is None:
            return None
        transform = affine_inverse(ex_t[:, :1])
        ex_t_norm = ex_t @ transform
        c2ws = affine_inverse(ex_t_norm)
        translations = c2ws[..., :3, 3]
        dists = translations.norm(dim=-1)
        median_dist = torch.median(dists)
        median_dist = torch.clamp(median_dist, min=1e-1)
        ex_t_norm[..., :3, 3] = ex_t_norm[..., :3, 3] / median_dist
        return ex_t_norm