"""MapNuRec — per-pixel feed-forward 3DGS, InstantNuRec-style, warm-started from
Depth-Anything-V2. Per input view i: DA-V2 predicts relative inverse-depth (disparity)
`disp`; a learned global affine maps it to metric depth `z = 1/(a*disp + b)` (a,b>0),
so DA-V2's strong relative geometry is inherited immediately and only the metric scale
is learned (anchored by the map-depth loss on ground). A small fresh per-pixel head on
[rgb, disp] predicts opacity / log-scale-mult / rotation; color = the source pixel RGB.
Each pixel is lifted to a world-space Gaussian; the union over views is rendered by gsplat.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

IMAGENET = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])


def lift_to_world(z, K, c2w):
    """z [V,H,W] metric depth (along optical axis) -> world points [V,H,W,3] (OpenCV +z fwd)."""
    V, H, W = z.shape
    dev = z.device
    ys, xs = torch.meshgrid(torch.arange(H, device=dev), torch.arange(W, device=dev), indexing="ij")
    u, v = (xs + 0.5).float(), (ys + 0.5).float()
    fx, fy = K[:, 0, 0, None, None], K[:, 1, 1, None, None]
    cx, cy = K[:, 0, 2, None, None], K[:, 1, 2, None, None]
    x = (u - cx) / fx * z
    y = (v - cy) / fy * z
    p_cam = torch.stack([x, y, z], dim=-1)                               # [V,H,W,3]
    R, t = c2w[:, :3, :3], c2w[:, :3, 3]
    return torch.einsum("vij,vhwj->vhwi", R, p_cam) + t[:, None, None, :]


class MapNuRec(nn.Module):
    def __init__(self, depth_min=2.0, depth_max=80.0):
        super().__init__()
        from transformers import AutoModelForDepthEstimation
        self.da = AutoModelForDepthEstimation.from_pretrained(
            "depth-anything/Depth-Anything-V2-Small-hf")               # DINOv2-S + DPT (warm-start)
        # disparity -> metric depth affine z = 1/(softplus(a)*disp + softplus(b)).
        # init from DA-V2 driving-disp stats (median~2.5): a=softplus(-3.8)=0.022,
        # b=softplus(-4.6)=0.010 -> z(med)~15m, z(horizon,disp~0)~100m. Matches the map
        # scale so the gamma-weighted map-depth anchor is active from step 0 (not zeroed).
        self.aff = nn.Parameter(torch.tensor([-3.8, -4.6]))
        self.dmin, self.dmax = depth_min, depth_max
        # fresh per-pixel head on [rgb(3), disp(1)] -> opacity(1), log_scale_mult(1), rot(4)
        self.head = nn.Sequential(
            nn.Conv2d(4, 64, 3, padding=1), nn.GELU(),
            nn.Conv2d(64, 64, 3, padding=1), nn.GELU(),
            nn.Conv2d(64, 9, 1))
        nn.init.zeros_(self.head[-1].weight); nn.init.zeros_(self.head[-1].bias)
        # channel layout: [0]=opacity, [1:4]=log-scale (ANISOTROPIC), [4:8]=rot quat, [8]=spare
        self.head[-1].bias.data[0] = 2.0                               # opacity logit -> sigmoid~0.88
        self.head[-1].bias.data[4] = 1.0                               # rot quat w=1 (identity)

    def disp(self, images):
        """images [V,3,H,W] in [0,1] -> per-pixel disparity [V,H,W] (larger=closer)."""
        x = IMAGENET(images)
        out = self.da(pixel_values=x).predicted_depth
        if out.dim() == 4:
            out = out[:, 0]
        if out.shape[-2:] != images.shape[-2:]:
            out = F.interpolate(out[:, None], size=images.shape[-2:], mode="bilinear",
                                align_corners=False)[:, 0]
        return out.clamp(min=1e-3)

    def forward(self, images, K, c2w):
        """images [V,3,H,W] (0..1), K [V,3,3], c2w [V,4,4] -> gaussian dict (world frame)."""
        V, _, H, W = images.shape
        disp = self.disp(images)                                       # [V,H,W]
        a, b = F.softplus(self.aff[0]), F.softplus(self.aff[1])
        z = (1.0 / (a * disp + b)).clamp(self.dmin, self.dmax)         # metric depth
        h = self.head(torch.cat([images, disp[:, None]], dim=1))       # [V,9,H,W]
        h = h.permute(0, 2, 3, 1)                                      # [V,H,W,9]
        opacity = torch.sigmoid(h[..., 0])                            # [V,H,W]
        base = (z / K[:, 0, 0, None, None]).clamp(min=1e-4)           # pixel footprint at depth
        scale = base[..., None] * torch.exp(h[..., 1:4].clamp(-3, 3))  # [V,H,W,3] anisotropic
        quat = F.normalize(h[..., 4:8] + torch.tensor([1.0, 0, 0, 0], device=images.device), dim=-1)
        xyz = lift_to_world(z, K, c2w)                                 # [V,H,W,3]
        rgb = images.permute(0, 2, 3, 1)                              # color = source pixel
        flat = lambda t, c: t.reshape(-1, c) if c > 1 else t.reshape(-1)
        return dict(means=xyz.reshape(-1, 3), scales=scale.reshape(-1, 3),
                    quats=quat.reshape(-1, 4), opacities=opacity.reshape(-1),
                    colors=rgb.reshape(-1, 3), depth=z)