"""Visibility loss — free tokens only (§2.6 item 5).

TokenGS's penalty keeps Gaussians inside >=1 supervision frustum (clipped at
1.0). Applied **only to T^F**: map-anchored tokens get correct positions from
the map even with zero rendering gradient, so penalizing them for projecting
outside the input views would destroy the extrapolation scaffold (we *want*
road geometry beyond the FOV).
"""

from __future__ import annotations

import torch

from mapgs.geometry.cameras import project_points
from mapgs.geometry.transforms import se3_inverse


def visibility_loss(
    means_free: torch.Tensor,  # [N, 3] free-token gaussian centers
    K: torch.Tensor,           # [V, 3, 3] supervision cameras
    cam2world: torch.Tensor,   # [V, 4, 4]
    H: int,
    W: int,
    near: float = 0.1,
) -> torch.Tensor:
    if means_free.shape[0] == 0:
        return means_free.new_tensor(0.0)
    w2c = se3_inverse(cam2world)
    uv, z = project_points(means_free[None].expand(K.shape[0], -1, -1), K, w2c)  # [V,N,2],[V,N]
    du = torch.relu(uv[..., 0] - (W - 1)) + torch.relu(-uv[..., 0])
    dv = torch.relu(uv[..., 1] - (H - 1)) + torch.relu(-uv[..., 1])
    behind = torch.relu(near - z)
    out_of_view = du / W + dv / H + behind          # [V, N], 0 if inside this view
    pen = out_of_view.min(dim=0).values              # inside >=1 view -> 0
    return pen.clamp(max=1.0).mean()