"""Map-guided geometry supervision (§2.6 item 2) + map-ground coupling (§2.5).

GS-native, expressed on the alpha-composited GS depth (no NeRF ray integral):

  L_mapdepth = sum_{p in Omega^g} Gamma(p) * Huber(D_hat(p) - D^map(p))
             + lambda_fs * sum_{g: z_g < h(x_g,y_g) - delta} sigma_g

with the MapNeRF self-paced weight  Gamma(p) = exp(-|D_hat - D^map| / (2 eps(t))).
With small eps only pixels already agreeing with the map get weight (obstacles /
gross map errors auto-down-weighted); eps tempers upward to broaden support.
Gamma is used as a *weight* (detached). The second term is the GS analog of
MapNeRF's sub-ground density penalty: opacity penalty on Gaussians below ground.
"""

from __future__ import annotations

import torch

from mapgs.hdmap.ground_field import GroundField
from mapgs.render.gaussians import Gaussians, GROUP_DYNAMIC


def huber(x: torch.Tensor, delta: float) -> torch.Tensor:
    a = x.abs()
    quad = 0.5 * x ** 2
    lin = delta * (a - 0.5 * delta)
    return torch.where(a <= delta, quad, lin)


def mapdepth_loss(
    pred_depth: torch.Tensor,   # [V, H, W]
    map_depth: torch.Tensor,    # [V, H, W]
    ground_mask: torch.Tensor,  # [V, H, W] bool  (Omega^g)
    eps: float,
    delta: float = 0.5,
) -> torch.Tensor:
    resid = pred_depth - map_depth
    gamma = torch.exp(-resid.abs() / (2 * max(float(eps), 1e-6))).detach()  # self-paced weight
    l = huber(resid, delta)
    m = ground_mask.float()
    denom = m.sum().clamp_min(1.0)
    return (gamma * l * m).sum() / denom


def free_space_loss(
    gaussians: Gaussians,
    ground: GroundField,
    delta: float = 0.15,
) -> torch.Tensor:
    """Opacity penalty on Gaussians whose centers lie below the ground surface."""
    z = gaussians.means[:, 2]
    xy = gaussians.means[:, :2]
    h, valid = ground.height_at(xy)
    below = (z < (h - delta)) & valid
    bf = below.float()
    denom = bf.sum().clamp_min(1.0)
    return (gaussians.opacities * bf).sum() / denom


def ground_coupling_loss(
    gaussians: Gaussians,
    ground: GroundField,
    eps: float,
    delta: float = 0.5,
) -> torch.Tensor:
    """Soft, Gamma-weighted pull of dynamic Gaussian z toward map ground height (§2.5)."""
    dyn = gaussians.group == GROUP_DYNAMIC
    if not dyn.any():
        return gaussians.means.sum() * 0.0
    z = gaussians.means[dyn, 2]
    xy = gaussians.means[dyn, :2]
    h, valid = ground.height_at(xy)
    resid = z - h
    gamma = torch.exp(-resid.abs() / (2 * eps)).detach()
    vf = valid.float()
    denom = vf.sum().clamp_min(1.0)
    return (vf * gamma * huber(resid, delta)).sum() / denom