"""Video Depth Anything + MoGe extraction wrapper for SLTK.

Extracts per-keypoint relative depth values from video frames using
Video Depth Anything (VDA) for depth maps, with optional MoGe-2 camera
intrinsic calibration, and RTMPose 2D keypoint locations.  Outputs an H5
file consumed by the SMPLFX fitter for depth-correlation supervision.

When ``use_moge=True`` (default), MoGe-2 is run on a small sample of frames
to estimate camera intrinsics (focal length, principal point).  The estimated
intrinsics are stored as root-level H5 attributes for downstream use by the
SMPLFX fitter.  Video Depth Anything is always used for the actual per-frame
depth maps.
"""

from __future__ import annotations

import logging
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import h5py
import numpy as np
import torch

logger = logging.getLogger(__name__)


def _to_numpy_batch(batch: object) -> np.ndarray:
    """Convert a decord/torch batch object to a numpy array."""
    if hasattr(batch, "asnumpy"):
        return batch.asnumpy()  # type: ignore[no-any-return]
    if hasattr(batch, "detach"):
        return batch.detach().cpu().numpy()  # type: ignore[no-any-return]
    if hasattr(batch, "numpy"):
        return batch.numpy()  # type: ignore[no-any-return]
    return np.asarray(batch)


def focal2fov(focal: float, pixels: int) -> float:
    """Convert focal length to field-of-view in radians.

    Args:
        focal: Focal length in pixels.
        pixels: Image dimension (width or height).

    Returns:
        Field of view in radians.
    """
    return 2 * math.atan(pixels / (2 * focal))


@dataclass
class DepthExtractionConfig:
    """Configuration for depth extraction.

    Video Depth Anything is always used for per-frame depth maps.  When
    ``use_moge=True``, MoGe-2 additionally estimates camera intrinsics from a
    small frame sample; the intrinsics are saved to the output H5 for use by
    the SMPLFX fitter.

    Attributes:
        device: Torch device string.
        encoder: VDA encoder variant ('vits', 'vitb', or 'vitl').
        checkpoint_path: VDA checkpoint path; auto-resolved from SLTK weights if None.
        input_size: Input resolution for VDA (must be a multiple of 14).  392 fits
            comfortably on a 24 GiB GPU after VDA vitl loads (~8.5 GiB model weights).
            518 needs ~11.4 GiB for the internal 32-frame forward pass and will OOM.
        fp32: Use FP32 instead of FP16 for VDA inference.
        use_moge: Run MoGe-2 on a frame sample to estimate camera intrinsics.
        moge_intrinsic_frames: Number of frames sampled for stable intrinsic median.
    """

    device: str = "cuda:0"
    encoder: str = "vitl"  # 'vits', 'vitb', or 'vitl'
    checkpoint_path: str | None = None  # auto-resolved if None
    input_size: int = (
        518  # must be multiple of 14; VDA upsamples output back to original resolution
    )
    fp32: bool = False
    use_moge: bool = True
    moge_intrinsic_frames: int = 10


# Model architecture configs (from Video-Depth-Anything)
_MODEL_CONFIGS = {
    "vits": {"encoder": "vits", "features": 64, "out_channels": [48, 96, 192, 384]},
    "vitb": {"encoder": "vitb", "features": 128, "out_channels": [96, 192, 384, 768]},
    "vitl": {"encoder": "vitl", "features": 256, "out_channels": [256, 512, 1024, 1024]},
}


def _sample_2d_depth(
    depths: np.ndarray,
    points_2d: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Sample depth values at 2D keypoint locations.

    Args:
        depths: Depth maps of shape ``(B, H, W)``.
        points_2d: Keypoint coordinates ``(B, N, 2)`` in ``(x, y)`` order.

    Returns:
        Tuple of ``(sampled_depths, valid_mask)`` each ``(B, N)``.
    """
    b, h, w = depths.shape
    assert points_2d.shape[0] == b
    assert points_2d.ndim == 3 and points_2d.shape[2] == 2

    x = points_2d[..., 0]
    y = points_2d[..., 1]

    valid_mask = ((x >= 0) & (x < w - 0.5) & (y >= 0) & (y < h - 0.5)).astype(np.float32)

    x_clamped = np.clip(np.rint(x), 0, w - 1).astype(int)
    y_clamped = np.clip(np.rint(y), 0, h - 1).astype(int)

    sampled = depths[np.arange(b)[:, None], y_clamped, x_clamped]
    return sampled, valid_mask


class DepthExtractor:
    """Extract per-keypoint depth from video using Video Depth Anything or MoGe-2."""

    name = "depth_anything"

    def __init__(self, config: DepthExtractionConfig | None = None) -> None:
        self.config = config or DepthExtractionConfig()
        self._model: Any = None
        self._moge_model: Any = None
        self._resolve_checkpoint()  # VDA always used for depth maps

    def _resolve_checkpoint(self) -> None:
        """Auto-resolve VDA checkpoint path from SLTK weights dir."""
        if self.config.checkpoint_path is not None:
            return

        from sltk.extraction.weights import get_weight_path

        key = f"depth_anything_{self.config.encoder}"
        resolved = get_weight_path(key)
        if resolved is not None:
            self.config.checkpoint_path = str(resolved)
        else:
            raise FileNotFoundError(
                f"Depth Anything checkpoint not found for encoder={self.config.encoder}. "
                f"Download it and place at sltk/weights/smplfx/depth_anything/"
            )

    def load_model(self) -> None:
        """Load the Video Depth Anything model."""
        cfg = self.config

        # Add depth_anything_deps to sys.path so its internal imports work
        deps_dir = str(Path(__file__).parent / "depth_anything_deps")
        if deps_dir not in sys.path:
            sys.path.insert(0, deps_dir)

        from video_depth_anything.video_depth import VideoDepthAnything

        model_cfg = _MODEL_CONFIGS[cfg.encoder]
        self._model = VideoDepthAnything(**model_cfg, metric=False)
        assert cfg.checkpoint_path is not None, "No checkpoint path configured"
        state = torch.load(cfg.checkpoint_path, map_location="cpu")
        self._model.load_state_dict(state, strict=True)
        self._model = self._model.to(cfg.device).eval()
        logger.info("Loaded Video Depth Anything (%s) on %s", cfg.encoder, cfg.device)

    def _load_moge_model(self) -> None:
        """Load the MoGe-2 model for metric depth and intrinsic estimation."""
        from moge.model.v2 import MoGeModel

        cfg = self.config
        self._moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").to(cfg.device)
        self._moge_model.eval()
        logger.info("Loaded MoGe-2 (moge-2-vitl-normal) on %s", cfg.device)

    def _estimate_intrinsics_with_moge(
        self,
        vid: Any,
        frames_idx: list[int],
        original_w: int,
        original_h: int,
    ) -> dict:
        """Estimate camera intrinsics from a sample of video frames using MoGe.

        Args:
            vid: Decord VideoReader instance.
            frames_idx: Full list of frame indices in processing order.
            original_w: Video frame width in pixels.
            original_h: Video frame height in pixels.

        Returns:
            Dict with keys ``focal_x``, ``focal_y``, ``pp_x``, ``pp_y``, ``fov_x``.
        """
        cfg = self.config
        n_frames = len(frames_idx)
        sample_idxs = np.linspace(0, n_frames - 1, cfg.moge_intrinsic_frames).astype(int)
        sampled_frame_ids = [frames_idx[i] for i in sample_idxs]

        frames_np = _to_numpy_batch(vid.get_batch(sampled_frame_ids))  # (B, H, W, 3)
        frames_tensor = torch.tensor(
            frames_np / 255.0, dtype=torch.float32, device=cfg.device
        ).permute(
            0, 3, 1, 2
        )  # (B, 3, H, W)

        with torch.inference_mode():
            moge_out = self._moge_model.infer(frames_tensor)

        scale_matrix = torch.tensor(
            [[original_w, 0.0, 0.0], [0.0, original_h, 0.0], [0.0, 0.0, 1.0]],
            dtype=moge_out["intrinsics"].dtype,
            device=moge_out["intrinsics"].device,
        )
        intrinsics_abs = scale_matrix @ moge_out["intrinsics"].median(0).values

        fx = intrinsics_abs[0, 0].item()
        fy = intrinsics_abs[1, 1].item()
        cx = intrinsics_abs[0, 2].item()
        cy = intrinsics_abs[1, 2].item()
        logger.info("MoGe intrinsics: fx=%.2f fy=%.2f cx=%.2f cy=%.2f", fx, fy, cx, cy)
        return {
            "focal_x": fx,
            "focal_y": fy,
            "pp_x": cx,
            "pp_y": cy,
            "fov_x": focal2fov(fx, original_w),
        }

    def extract(
        self,
        video_path: str | Path,
        rtm_h5: str | Path,
        output_path: str | Path,
        precomputed_intrinsics: dict | None = None,
    ) -> Path:
        """Extract per-keypoint depth and write an H5 file.

        Args:
            video_path: Path to the source video.
            rtm_h5: Path to RTMPose H5 with keypoint detections.
            output_path: Where to write the depth H5.
            precomputed_intrinsics: Optional dict with keys ``focal_x``,
                ``focal_y``, ``pp_x``, ``pp_y`` from a prior MoGe pass.
                When provided, the MoGe intrinsic estimation step is skipped
                even if ``config.use_moge=True``.

        Returns:
            Path to the written H5 file.
        """
        cfg = self.config
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # -- Load RTM keypoints --
        from .smplfx_deps.h5_adapters import load_rtm_for_fitting

        data_kpts = load_rtm_for_fitting(rtm_h5, zero_weights_out_frame=False)

        # -- Read video with decord --
        from decord import VideoReader, cpu

        vid = VideoReader(str(video_path), ctx=cpu(0))
        original_h, original_w = vid.get_batch([0]).shape[1:3]

        # Re-open at native resolution
        vid = VideoReader(str(video_path), ctx=cpu(0), width=original_w, height=original_h)
        fps = vid.get_avg_fps()
        total_frames = len(vid)
        frames_idx = list(range(total_frames))
        target_fps = fps

        # -- Build per-person frame mapping --
        person_inv_map: dict[str, dict[int, int]] = {}
        for pid in data_kpts["people"]:
            sf = data_kpts["people"][pid]["successful_frames"]
            person_inv_map[pid] = {int(fi): oi for oi, fi in enumerate(sf)}

        depths_per_person: dict[str, list] = {k: [] for k in data_kpts["people"]}
        masks_per_person: dict[str, list] = {k: [] for k in data_kpts["people"]}

        # -- Step 1: MoGe intrinsic estimation (optional) --
        if precomputed_intrinsics is not None:
            intrinsics = precomputed_intrinsics
            logger.info(
                "Using precomputed intrinsics: fx=%.2f fy=%.2f",
                intrinsics["focal_x"],
                intrinsics["focal_y"],
            )
        elif cfg.use_moge:
            if self._moge_model is None:
                self._load_moge_model()
            intrinsics = self._estimate_intrinsics_with_moge(
                vid, frames_idx, original_w, original_h
            )
            # Unload MoGe before loading VDA to avoid OOM (both are ~4-10 GB).
            self._moge_model = None
            torch.cuda.empty_cache()
            logger.info("MoGe unloaded, VRAM freed for Video Depth Anything")
        else:
            intrinsics = None

        # -- Step 2: Video Depth Anything depth maps --
        if self._model is None:
            self.load_model()

        # Pass all frames in one call.  infer_video_depth handles its own
        # temporal chunking internally (INFER_LEN=32, OVERLAP=10, stride=22)
        # with correct overlap-blending across the whole video.  Our wrapper
        # chunking was redundant and broke temporal coherence at chunk boundaries.
        logger.info(
            "Extracting depth with Video Depth Anything for %d frames (input_size=%d)...",
            len(frames_idx),
            cfg.input_size,
        )
        all_frames = _to_numpy_batch(vid.get_batch(frames_idx))  # (T, H, W, 3)
        depth_maps = self._model.infer_video_depth(
            all_frames,
            target_fps,
            input_size=cfg.input_size,
            device=cfg.device,
            fp32=cfg.fp32,
        )[
            0
        ]  # (T, H, W)

        for pid in data_kpts["people"]:
            kp_idxs = []
            depth_idxs = []
            for fi in frames_idx:
                if fi in person_inv_map[pid]:
                    kp_idxs.append(person_inv_map[pid][fi])
                    depth_idxs.append(fi)

            if kp_idxs:
                kp_xy = data_kpts["people"][pid]["keypoints"][kp_idxs, :, :2]
                sampled, mask = _sample_2d_depth(depth_maps[depth_idxs], kp_xy)
                depths_per_person[pid].append(sampled)
                masks_per_person[pid].append(mask)

        # -- Concatenate and write --
        with h5py.File(str(output_path), "w") as hf:
            hf.attrs["video_path"] = str(video_path)
            hf.attrs["h5_keypoints"] = str(rtm_h5)
            hf.attrs["encoder"] = cfg.encoder

            if intrinsics is not None:
                hf.attrs["focal_x"] = intrinsics["focal_x"]
                hf.attrs["focal_y"] = intrinsics["focal_y"]
                hf.attrs["pp_x"] = intrinsics["pp_x"]
                hf.attrs["pp_y"] = intrinsics["pp_y"]

            for pid in depths_per_person:
                if not depths_per_person[pid]:
                    continue
                d = np.concatenate(depths_per_person[pid]).astype(np.float32)
                m = np.concatenate(masks_per_person[pid]).astype(np.float32)

                grp = hf.create_group(pid)
                grp.create_dataset("depths", data=d, compression="gzip", compression_opts=4)
                grp.create_dataset("depths_mask", data=m, compression="gzip", compression_opts=4)
                grp.create_dataset(
                    "successful_frames",
                    data=np.array(
                        data_kpts["people"][pid]["successful_frames"],
                        dtype=np.int32,
                    ),
                    compression="gzip",
                    compression_opts=4,
                )

        logger.info("Depth H5 saved to %s", output_path)
        return output_path

    def close(self) -> None:
        """Release GPU resources."""
        self._model = None
        self._moge_model = None
        torch.cuda.empty_cache()
