"""Render map geometry into camera views.

* :func:`rasterize_map_depth` casts each pixel ray against the ground height
  field (a fixed-point intersection) to get per-view ground depth ``D^map`` and
  the ground/lane mask ``Omega^g``. ``D^map`` is a supervision *target*
  (constant w.r.t. the Gaussians), so it is computed without gradients and
  fully vectorized over pixels — no triangle z-buffer needed.
* :func:`project_polylines` projects lane / boundary vertices into each view
  (used as the map side of the lane-chamfer term, §2.6 item 3, and for lane
  mIoU, §4.4).
* :func:`render_lane_mask` splats projected lane points into a soft raster.
"""

from __future__ import annotations

from typing import List

import torch

from mapgs.geometry.cameras import camera_centers
from mapgs.hdmap.ground_field import GroundField


@torch.no_grad()
def rasterize_map_depth(
    ground: GroundField,
    K: torch.Tensor,          # [V, 3, 3]
    cam2world: torch.Tensor,  # [V, 4, 4]
    H: int,
    W: int,
    iters: int = 16,
    near: float = 0.5,
    far: float = 200.0,
    min_descent: float = 1e-3,
    tol: float = 0.25,
):
    """Ray/height-field intersection -> ``depth [V,H,W]``, ``mask [V,H,W]`` (bool)."""
    device = K.device
    V = K.shape[0]
    dtype = K.dtype

    vv, uu = torch.meshgrid(
        torch.arange(H, device=device, dtype=dtype) + 0.5,
        torch.arange(W, device=device, dtype=dtype) + 0.5,
        indexing="ij",
    )
    ones = torch.ones_like(uu)
    pix = torch.stack([uu, vv, ones], dim=-1)  # [H, W, 3]
    Kinv = torch.inverse(K)                     # [V, 3, 3]
    r_cam = torch.einsum("vij,hwj->vhwi", Kinv, pix)  # camera ray, z-comp == 1
    R_c2w = cam2world[:, :3, :3]
    m = torch.einsum("vij,vhwj->vhwi", R_c2w, r_cam)  # world delta per unit cam-depth
    o = camera_centers(cam2world)[:, None, None, :]    # [V,1,1,3]

    mz = m[..., 2]
    descending = mz < -min_descent
    # initial guess from flat ground under the camera
    h0, _ = ground.height_at(o[..., :2].expand(V, H, W, 2))
    Z = (h0 - o[..., 2]) / mz.clamp(max=-min_descent)
    Z = Z.clamp(near, far)
    for _ in range(iters):
        xy = o[..., :2] + Z.unsqueeze(-1) * m[..., :2]
        h, _ = ground.height_at(xy)
        Z = ((h - o[..., 2]) / mz.clamp(max=-min_descent)).clamp(near, far)

    xy = o[..., :2] + Z.unsqueeze(-1) * m[..., :2]
    h, valid_xy = ground.height_at(xy)
    residual = (o[..., 2] + Z * mz - h).abs()
    mask = descending & valid_xy & (residual < tol) & (Z > near) & (Z < far)
    depth = torch.where(mask, Z, torch.zeros_like(Z))
    return depth, mask


@torch.no_grad()
def project_polylines(
    polylines,                # [P, 3] tensor or list of [Li, 3]
    K: torch.Tensor,          # [V, 3, 3]
    cam2world: torch.Tensor,  # [V, 4, 4]
    H: int,
    W: int,
) -> List[torch.Tensor]:
    """Project map polyline vertices into each view. Returns list (len V) of ``[Mi, 2]`` uv."""
    from mapgs.geometry.cameras import project_points
    from mapgs.geometry.transforms import se3_inverse

    if isinstance(polylines, (list, tuple)):
        pts = torch.cat([p for p in polylines if p.numel() > 0], 0) if len(polylines) else torch.zeros(0, 3)
    else:
        pts = polylines
    device = K.device
    pts = pts.to(device)
    V = K.shape[0]
    out: List[torch.Tensor] = []
    if pts.numel() == 0:
        return [torch.zeros(0, 2, device=device) for _ in range(V)]
    w2c = se3_inverse(cam2world)
    for v in range(V):
        uv, z = project_points(pts[None], K[v : v + 1], w2c[v : v + 1])
        uv = uv[0]
        z = z[0]
        inb = (z > 0.1) & (uv[:, 0] >= 0) & (uv[:, 0] < W) & (uv[:, 1] >= 0) & (uv[:, 1] < H)
        out.append(uv[inb])
    return out


def render_lane_mask(uv: torch.Tensor, H: int, W: int, radius: float = 2.0) -> torch.Tensor:
    """Splat projected lane points ``[M, 2]`` into a soft mask ``[H, W]`` in [0,1]."""
    device = uv.device
    if uv.numel() == 0:
        return torch.zeros(H, W, device=device)
    vv, uu = torch.meshgrid(
        torch.arange(H, device=device, dtype=torch.float32),
        torch.arange(W, device=device, dtype=torch.float32),
        indexing="ij",
    )
    # distance from each pixel to nearest lane point (chunked to bound memory)
    grid = torch.stack([uu, vv], dim=-1).reshape(-1, 2)  # [HW, 2]
    min_d2 = torch.full((grid.shape[0],), 1e9, device=device)
    for i in range(0, uv.shape[0], 4096):
        chunk = uv[i : i + 4096]
        d2 = (grid[:, None, :] - chunk[None, :, :]).pow(2).sum(-1).min(dim=1).values
        min_d2 = torch.minimum(min_d2, d2)
    mask = torch.exp(-min_d2 / (2 * radius ** 2))
    return mask.reshape(H, W)