data/fps.py

"""Batched GPU furthest-point sampling for the stage-2 pipeline.

Designed for the v2 (FPS) stage-2 path where the dataloader returns the full
hull-filtered point cloud padded to a fixed cap, and we subsample to the
training budget (e.g. 8k+8k) *after* collate — i.e. once per batch on the GPU
inside the training/inference loop. CPU dataloader workers cannot use CUDA, so
this is the canonical place to run FPS.

Two entry points:

* :func:`fps_batched_indices` — single-provenance FPS. Returns ``(B, k)`` long
  indices into ``(B, N, 3)`` xyz. Honors a per-point ``valid_mask`` (padding
  slots are never selected) and an optional per-sample RNG seed.

* :func:`fps_subsample_stage2_batch` — convenience wrapper that operates on a
  stage-2 batch dict produced by ``stage2_collate_fn`` with the FPS variant.
  Runs FPS once per provenance (COLMAP, depth) and gathers all the per-point
  fields (xyz, semantics, RGB, …) to the resampled order. Returns a NEW batch
  dict with the same keys but with the per-point tensors shrunk from ``N`` to
  ``n_col + n_dep`` and ``scene_xyz`` re-normalized into the post-subsample
  bbox (matching the random-sampling path).

The kernel is a vectorized version of the single-loop torch FPS in
``viz_stage2_density._fps_indices``: one ``(B, N)`` distance buffer, ``k``
iterations, each step gathers the just-picked point and updates the running
minimum-distance buffer in place. No host syncs inside the loop.
"""
from __future__ import annotations

from typing import Optional

import torch


# Sentinel for "this slot is ineligible / already picked." Using -inf so the
# argmax never selects it, while leaving real (non-negative) squared distances
# untouched by the min().
def _neg_inf(dtype: torch.dtype) -> float:
    return torch.finfo(dtype).min


@torch.no_grad()
def fps_batched_indices(
    xyz: torch.Tensor,
    valid_mask: torch.Tensor,
    k: int,
    seed_per_sample: Optional[torch.Tensor] = None,
    max_exact_iters: Optional[int] = None,
) -> torch.Tensor:
    """Batched furthest-point sampling on GPU (or CPU; CUDA strongly preferred).

    Args:
        xyz: ``(B, N, 3)`` float tensor of point coordinates.
        valid_mask: ``(B, N)`` bool tensor. ``False`` slots are padding /
            wrong-provenance / otherwise ineligible. They are never selected.
        k: number of points to pick per batch element.
        seed_per_sample: optional ``(B,)`` long tensor. Each entry is taken
            modulo the per-sample valid count to pick the seed within the
            eligible set. ``None`` → seed at the first eligible index
            (deterministic).
        max_exact_iters: optional cap for exact FPS iterations. If set and
            smaller than ``k``, exact FPS is used for the first anchors and the
            remainder is filled with deterministic random eligible points. This
            keeps broad spatial anchors while avoiding thousands of sequential
            distance-buffer updates in the stage-2 training hot path.

    Returns:
        ``(B, k)`` long tensor of indices into the ``N`` axis. For batch
        elements with fewer than ``k`` eligible points, the deficit slots are
        padded by repeating earlier picks (FPS naturally cycles when the
        eligible set is exhausted and argmax falls back to index 0; the
        downstream encoder doesn't care about duplicates).
    """
    if xyz.dim() != 3 or xyz.shape[-1] != 3:
        raise ValueError(f"xyz must be (B, N, 3), got {tuple(xyz.shape)}")
    if valid_mask.shape != xyz.shape[:2]:
        raise ValueError(
            f"valid_mask shape {tuple(valid_mask.shape)} != xyz[:2] "
            f"{tuple(xyz.shape[:2])}"
        )
    B, N, _ = xyz.shape
    device = xyz.device

    if k <= 0:
        return torch.empty((B, 0), device=device, dtype=torch.long)

    k_exact = k
    if max_exact_iters is not None:
        k_exact = max(0, min(k, int(max_exact_iters)))
    if k_exact == 0:
        return _deterministic_random_fill_indices(
            valid_mask,
            k=k,
            seed_per_sample=seed_per_sample,
            exclude_idx=None,
        )

    NEG = _neg_inf(xyz.dtype)
    dist2 = torch.where(
        valid_mask,
        torch.full_like(xyz[..., 0], float("inf")),
        torch.full_like(xyz[..., 0], NEG),
    )

    # --- Seed selection -----------------------------------------------------
    if seed_per_sample is None:
        # First eligible index (argmax over bool returns index of first True).
        seed_idx = valid_mask.to(torch.int8).argmax(dim=1)
    else:
        if seed_per_sample.shape != (B,):
            raise ValueError(
                f"seed_per_sample shape {tuple(seed_per_sample.shape)} != ({B},)"
            )
        counts = valid_mask.sum(dim=1).clamp(min=1)
        offset = (seed_per_sample.to(device=device, dtype=torch.long).abs()
                  % counts)
        # Find the (offset+1)-th True in valid_mask along dim=1.
        cum = valid_mask.to(torch.long).cumsum(dim=1)
        target = (offset + 1).unsqueeze(1)
        seed_idx = (cum >= target).to(torch.int8).argmax(dim=1)
        # Batches with zero eligible points end up with seed_idx=0; that's
        # fine — they'll just emit duplicate index 0s and the caller's
        # valid_mask handling deals with it.

    arange_b = torch.arange(B, device=device)
    out = torch.empty((B, k_exact), device=device, dtype=torch.long)
    out[:, 0] = seed_idx
    # Mark seed as picked so it can't be re-selected.
    dist2[arange_b, seed_idx] = NEG

    last_idx = seed_idx
    for i in range(1, k_exact):
        last_pts = xyz[arange_b, last_idx]                # (B, 3)
        diff = xyz - last_pts.unsqueeze(1)                # (B, N, 3)
        new_d2 = (diff * diff).sum(dim=-1)                # (B, N)
        # Only update eligible (still-in-the-game) slots. Picked/ineligible
        # slots sit at NEG and must stay there.
        eligible = dist2 > NEG
        dist2 = torch.where(eligible, torch.minimum(dist2, new_d2), dist2)
        next_idx = dist2.argmax(dim=1)
        out[:, i] = next_idx
        dist2[arange_b, next_idx] = NEG
        last_idx = next_idx

    if k_exact < k:
        fill = _deterministic_random_fill_indices(
            valid_mask,
            k=k - k_exact,
            seed_per_sample=seed_per_sample,
            exclude_idx=out,
        )
        out = torch.cat([out, fill], dim=1)

    return out


@torch.no_grad()
def _deterministic_random_fill_indices(
    valid_mask: torch.Tensor,
    k: int,
    seed_per_sample: Optional[torch.Tensor],
    exclude_idx: Optional[torch.Tensor],
) -> torch.Tensor:
    """Pick deterministic pseudo-random valid indices, excluding anchors.

    This is intentionally one vectorized top-k instead of a loop. It is used
    only for the fast hybrid FPS mode, where the expensive exact FPS anchors
    have already provided spatial coverage.
    """
    B, N = valid_mask.shape
    device = valid_mask.device
    if k <= 0:
        return torch.empty((B, 0), device=device, dtype=torch.long)

    remaining = valid_mask.clone()
    if exclude_idx is not None and exclude_idx.numel() > 0:
        remaining.scatter_(1, exclude_idx, False)

    n_top = min(k, N)
    idx_f = torch.arange(N, device=device, dtype=torch.float32).unsqueeze(0)
    batch_f = torch.arange(B, device=device, dtype=torch.float32).unsqueeze(1)
    if seed_per_sample is None:
        seed_f = torch.zeros((B, 1), device=device, dtype=torch.float32)
    else:
        seed_i = seed_per_sample.to(device=device, dtype=torch.long).remainder(1_000_003)
        seed_f = seed_i.to(dtype=torch.float32).view(B, 1)

    # A small deterministic hash in float space. Quality only needs to be good
    # enough to decorrelate the non-anchor fill; exact FPS has already handled
    # the coverage-critical prefix.
    scores = torch.sin((idx_f + 1.0) * 12.9898 + seed_f * 78.233 + batch_f * 37.719)
    scores = scores * 43758.5453
    scores = scores - torch.floor(scores)
    scores = torch.where(remaining, scores, torch.full_like(scores, -1.0))

    fill = scores.topk(n_top, dim=1).indices
    fill_ok = remaining.gather(1, fill)
    first_valid = valid_mask.to(torch.int8).argmax(dim=1)
    if exclude_idx is not None and exclude_idx.numel() > 0:
        fallback_src = exclude_idx
    else:
        fallback_src = first_valid.unsqueeze(1)
    fallback = fallback_src[:, torch.arange(n_top, device=device) % fallback_src.shape[1]]
    fill = torch.where(fill_ok, fill, fallback)

    if n_top < k:
        pad = fill[:, torch.arange(k - n_top, device=device) % n_top]
        fill = torch.cat([fill, pad], dim=1)
    return fill


@torch.no_grad()
def fps_per_provenance_indices(
    xyz: torch.Tensor,
    type_ids: torch.Tensor,
    valid_mask: torch.Tensor,
    n_col: int,
    n_dep: int,
    seed_per_sample: Optional[torch.Tensor] = None,
    max_exact_iters_per_provenance: Optional[int] = None,
) -> torch.Tensor:
    """Per-provenance FPS: COLMAP (type_id==0) first, then depth (type_id==1).

    Returns ``(B, n_col + n_dep)`` indices into ``N``. The two halves are
    concatenated so the encoder sees the same provenance ordering as the
    random-sampling path.
    """
    col_mask = valid_mask & (type_ids == 0)
    dep_mask = valid_mask & (type_ids == 1)
    # Use independent seeds so the two provenances don't pick correlated
    # starting points across the batch.
    seed_col = seed_per_sample
    seed_dep = (None if seed_per_sample is None
                else seed_per_sample ^ 0x9E3779B97F4A7C15)
    idx_col = fps_batched_indices(
        xyz, col_mask, n_col, seed_col,
        max_exact_iters=max_exact_iters_per_provenance,
    )
    idx_dep = fps_batched_indices(
        xyz, dep_mask, n_dep, seed_dep,
        max_exact_iters=max_exact_iters_per_provenance,
    )
    return torch.cat([idx_col, idx_dep], dim=1)


# ---------------------------------------------------------------------------
# Stage-2 batch convenience: gather + re-normalize
# ---------------------------------------------------------------------------

# Per-point tensors in the stage-2 batch. Matches STAGE2_POINT_KEYS in
# data.stage2_build, with the "raw_" prefix dropped by stage2_row_to_sample
# and remapped to the "scene_*" namespace.
_STAGE2_SCENE_PER_POINT_KEYS = (
    "scene_xyz",
    "scene_type_ids",
    "scene_gestalt_ids",
    "scene_gestalt_id2",
    "scene_gestalt_w1",
    "scene_ade_ids",
    "scene_geom_conf",
    "scene_sem_conf",
    "scene_rgb",
)


def _robust_norm_params_batched(
    xyz: torch.Tensor,
    valid_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Per-batch (median centre, P95 scale) computed only over valid points.

    Mirrors :func:`data.stage2_dataset._robust_norm_params` for a batched
    tensor. Implements median + P95 via :func:`torch.quantile` on padded
    rows; invalid slots are set to NaN so quantile ignores them.
    """
    B = xyz.shape[0]
    device = xyz.device
    # Set invalid slots to NaN so torch.quantile skips them (`nanmedian` /
    # `nanquantile`).
    mask3 = valid_mask.unsqueeze(-1)
    safe_xyz = torch.where(mask3, xyz, torch.full_like(xyz, float("nan")))
    center = torch.nanmedian(safe_xyz, dim=1).values            # (B, 3)
    d = torch.linalg.norm(xyz - center.unsqueeze(1), dim=-1)    # (B, N)
    d = torch.where(valid_mask, d, torch.full_like(d, float("nan")))
    scale = torch.nanquantile(d, 0.95, dim=1)                   # (B,)
    scale = scale.clamp(min=1e-3)
    return center.to(xyz.dtype), scale.to(xyz.dtype)


def _robust_norm_params_all_valid_batched(
    xyz: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fast median centre + P95 scale when every slot is valid."""
    center = xyz.median(dim=1).values
    d = torch.linalg.norm(xyz - center.unsqueeze(1), dim=-1)
    scale = torch.quantile(d, 0.95, dim=1).clamp(min=1e-3)
    return center.to(xyz.dtype), scale.to(xyz.dtype)


@torch.no_grad()
def fps_subsample_stage2_batch(
    batch: dict,
    n_col: int,
    n_dep: int,
    seed_per_sample: Optional[torch.Tensor] = None,
    renormalize_bbox: bool = True,
    max_exact_iters_per_provenance: Optional[int] = None,
) -> dict:
    """Apply per-provenance FPS to a stage-2 batch, in place where possible.

    The batch is expected to come from ``stage2_collate_fn`` with the FPS
    variant: per-point tensors have shape ``(B, N_max, …)`` and a
    ``scene_valid_mask`` ``(B, N_max)`` marks real points vs. padding.
    Coordinates are in WORLD space (no bbox normalisation applied) — i.e.
    produced by ``stage2_row_to_sample(..., pre_subsample=False)``.

    Args:
        batch: dict with at least ``scene_xyz``, ``scene_type_ids``,
            ``scene_valid_mask``, ``bbox_center`` (filler placeholder, will be
            overwritten), ``bbox_scale``, ``bbox_R``. Optional per-point keys
            (gestalt/ADE/conf/RGB) are gathered if present.
        n_col / n_dep: per-provenance budgets.
        seed_per_sample: optional ``(B,)`` long for reproducible seeding.
        renormalize_bbox: if True, recompute ``bbox_center`` / ``bbox_scale``
            from the FPS-subsampled cloud (matches the random path which also
            recomputes per augmented sample). If False, leaves the existing
            ``bbox_center`` / ``bbox_scale`` alone.
        max_exact_iters_per_provenance: optional cap for exact FPS anchors in
            each provenance. ``None`` preserves exact FPS. Smaller values enable
            the fast hybrid sampler.

    Returns:
        A new dict with the per-point tensors shrunk to ``(B, n_col + n_dep,
        …)`` and ``scene_xyz`` normalized into the post-subsample bbox.
        Non-point tensors (init_verts, verts_gt, etc.) are passed through
        un-changed. ``scene_valid_mask`` is set to all-True at the new size.
    """
    xyz_world = batch["scene_xyz"]                       # (B, N, 3) WORLD
    type_ids = batch["scene_type_ids"]                   # (B, N)
    valid_mask = batch["scene_valid_mask"]               # (B, N)
    B, N, _ = xyz_world.shape
    device = xyz_world.device

    idx = fps_per_provenance_indices(
        xyz_world, type_ids, valid_mask,
        n_col=n_col, n_dep=n_dep,
        seed_per_sample=seed_per_sample,
        max_exact_iters_per_provenance=max_exact_iters_per_provenance,
    )                                                    # (B, K)
    K = idx.shape[1]
    arange_b = torch.arange(B, device=device).unsqueeze(1).expand(B, K)

    out = dict(batch)

    # Gather all per-point tensors. Index supports any trailing shape.
    for key in _STAGE2_SCENE_PER_POINT_KEYS:
        t = batch.get(key)
        if t is None:
            continue
        if t.dim() == 2:
            out[key] = t[arange_b, idx]
        else:
            # (B, N, C) — gather along dim=1
            out[key] = t[arange_b, idx]

    # After FPS the valid mask is trivially all-True (we picked real or padded
    # slots; padded slots only sneak in when n_col / n_dep exceeds the
    # available count, in which case duplicate real picks are fine).
    out["scene_valid_mask"] = torch.ones((B, K), dtype=torch.bool, device=device)

    if renormalize_bbox:
        sub_xyz_world = out["scene_xyz"]                  # (B, K, 3) WORLD
        center, scale = _robust_norm_params_all_valid_batched(sub_xyz_world)
        out["scene_xyz"] = ((sub_xyz_world - center.unsqueeze(1))
                            / scale.view(B, 1, 1)).to(sub_xyz_world.dtype)
        out["bbox_center"] = center
        out["bbox_scale"] = scale
        # bbox_R is identity for the FPS path (no augmentation rotation
        # baked into the cache; matches the v1 random path during augment=False).
        out["bbox_R"] = torch.eye(3, device=device, dtype=sub_xyz_world.dtype) \
            .unsqueeze(0).expand(B, 3, 3).contiguous()

        # Re-normalize init_verts / verts_gt into the new bbox so the model's
        # init and supervision land in the same frame as scene_xyz.
        init_world = batch.get("init_verts_world")
        if isinstance(init_world, torch.Tensor):
            out["init_verts"] = (
                (init_world - center.unsqueeze(1)) / scale.view(B, 1, 1)
            ).to(init_world.dtype)
        verts_gt_world = batch.get("verts_gt_world")
        if isinstance(verts_gt_world, torch.Tensor):
            # verts_gt is (B, K_v, 4) — xyz + validity flag.
            xyz_g = verts_gt_world[..., :3]
            flag = verts_gt_world[..., 3:]
            xyz_g_n = (xyz_g - center.unsqueeze(1)) / scale.view(B, 1, 1)
            out["verts_gt"] = torch.cat([xyz_g_n.to(verts_gt_world.dtype), flag], dim=-1)

    return out