# SPDX-License-Identifier: Apache-2.0
"""C12 voxelization: deterministic pillar builders (vectorised numpy), Autoware's pillar feature decoration, and the
two canvas forms of the PointPillars scatter (dense scatter for references, gather index for the device).

Autoware's CUDA voxelizer (``autoware_lidar_centerpoint/lib/preprocess/preprocess_kernel.cu``) is
non-deterministic: slots inside a pillar come from ``atomicAdd`` over a randomly shuffled point buffer and pillar
ids from ``atomicAdd`` over a 2-D thread grid (S:centerpoint:177-186). The ports replace it by a deterministic
policy with the same statistics (PLAN.md D1, S:centerpoint:381-389):

1. optional seeded permutation of the input rows (:func:`ttaw.pointcloud.seeded_permutation`);
2. the half-open range test and the cell index ``floor((v - min) / voxel)`` (a float32 *division*, clamped; the
   kernel's ``:135-145``);
3. the first ``max_points_per_pillar`` points of every cell in (permuted) input order;
4. pillar ids in ascending cell order, keeping the first ``max_pillars``. The cell order is ``"flipped_x"``
   (CenterPoint: ``cell = (grid_x - 1 - ix) * grid_y + iy``, front rows first, ``:144``) or ``"raster"``
   (PointPainting: ``cell = iy * grid_x + ix``, S:pointpainting:229-232).

The result (:class:`PillarSet`) holds the raw rows of every kept slot (zero-padded), the point counts and the
``(z, y, x)`` = ``(0, iy, ix)`` coordinates of mmdet3d / the CUDA ``uint3 {0, idy, idx}`` (``:189-190``).

:func:`decorate_pillar_features` is ``generateFeatures_kernel`` (``:211-343``) for 9 / 10 / 11 encoder features
(the cluster mean is summed slot by slot in float32, a fixed order standing in for the kernel's float atomics; the
pillar centre is ``voxel / 2 + coord * voxel + min`` in float32; padded slots are all zero).

Canvas: :func:`scatter_canvas` is ``scatterFeatures_kernel`` (``scatter_kernel.cu:25-48``: ``canvas[c, iy, ix]``).
The device builds the same canvas as a *gather* (C19, probe P4): :func:`canvas_gather_index` gives, for every cell
``iy * grid_x + ix`` of the NHWC canvas, the row of its pillar or a sentinel row (an explicit zero row of the
feature table), so the canvas is ``table[index]`` with no scatter, no clear and a static shape.

numpy only; no side effects on import.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Optional, Sequence, Tuple

import numpy as np

from .pointcloud import seeded_permutation

__all__ = [
    "PillarGridSpec",
    "PillarSet",
    "assign_pillars",
    "decorate_pillar_features",
    "scatter_canvas",
    "canvas_gather_index",
    "pillar_cell_index",
]

_ORDERS = ("flipped_x", "raster")


@dataclass(frozen=True)
class PillarGridSpec:
    """A pillar grid and its capacities. Ranges and sizes are float32, as the CUDA config stores them
    (``centerpoint_config.hpp:41-56``); ``grid_x = int((max_x - min_x) / voxel_x)`` in float32 (``:95-97``)."""

    range_min: Tuple[float, float, float]
    range_max: Tuple[float, float, float]
    voxel_size: Tuple[float, float, float]
    max_points_per_pillar: int = 32
    max_pillars: int = 40000
    order: str = "flipped_x"

    def __post_init__(self) -> None:
        if self.order not in _ORDERS:
            raise ValueError(f"order must be one of {_ORDERS}, got {self.order!r}")
        if self.max_points_per_pillar < 1 or self.max_pillars < 1:
            raise ValueError("capacities must be >= 1")

    @property
    def min_f32(self) -> np.ndarray:
        return np.asarray(self.range_min, dtype=np.float32)

    @property
    def max_f32(self) -> np.ndarray:
        return np.asarray(self.range_max, dtype=np.float32)

    @property
    def voxel_f32(self) -> np.ndarray:
        return np.asarray(self.voxel_size, dtype=np.float32)

    def _grid(self, axis: int) -> int:
        return int(np.float32(self.max_f32[axis] - self.min_f32[axis]) / self.voxel_f32[axis])

    @property
    def grid_x(self) -> int:
        return self._grid(0)

    @property
    def grid_y(self) -> int:
        return self._grid(1)

    @property
    def grid_z(self) -> int:
        return self._grid(2)

    @property
    def num_cells(self) -> int:
        return self.grid_x * self.grid_y


@dataclass
class PillarSet:
    """The deterministic voxelization of one cloud (fixed capacity ``max_pillars``; rows past ``num_pillars`` are
    zero)."""

    points: np.ndarray            # (max_pillars, K, F) float32: raw rows of the kept slots, zero-padded
    num_points: np.ndarray        # (max_pillars,) int32: min(count, K), 0 for unused rows
    coords: np.ndarray            # (max_pillars, 3) int32: (0, iy, ix)
    num_pillars: int              # kept pillars (<= max_pillars)
    num_pillars_total: int        # non-empty cells before the cap (Autoware's num_voxels counter)
    num_points_in_range: int
    max_points_in_pillar: int     # largest per-cell count before the slot cap
    pillars_over_cap: int         # cells holding more than K points
    num_points_input: int

    @property
    def overflow(self) -> bool:
        """More non-empty cells than ``max_pillars``: the farthest-back pillars were dropped."""
        return self.num_pillars_total > self.points.shape[0]

    def stats(self) -> Dict[str, Any]:
        return {"num_points": self.num_points_input, "num_points_in_range": self.num_points_in_range,
                "num_pillars": self.num_pillars, "num_pillars_total": self.num_pillars_total,
                "max_points_in_pillar": self.max_points_in_pillar, "pillars_over_cap": self.pillars_over_cap,
                "pillar_overflow": self.overflow}


def pillar_cell_index(points: np.ndarray, spec: PillarGridSpec) -> Tuple[np.ndarray, np.ndarray]:
    """``(ix, iy)`` int64 of each row: ``floor((v - min) / voxel)`` evaluated in float32, clamped to the grid
    (``preprocess_kernel.cu:140-143``). The caller has already applied the range test."""
    p = np.asarray(points, dtype=np.float32)
    mn, vs = spec.min_f32, spec.voxel_f32
    ix = np.floor((p[:, 0] - mn[0]) / vs[0]).astype(np.int64)
    iy = np.floor((p[:, 1] - mn[1]) / vs[1]).astype(np.int64)
    return np.clip(ix, 0, spec.grid_x - 1), np.clip(iy, 0, spec.grid_y - 1)


def _cell_key(ix: np.ndarray, iy: np.ndarray, spec: PillarGridSpec) -> np.ndarray:
    if spec.order == "flipped_x":
        return (spec.grid_x - 1 - ix) * spec.grid_y + iy
    return iy * spec.grid_x + ix


def _key_to_ixy(key: np.ndarray, spec: PillarGridSpec) -> Tuple[np.ndarray, np.ndarray]:
    if spec.order == "flipped_x":
        return (spec.grid_x - 1) - key // spec.grid_y, key % spec.grid_y
    return key % spec.grid_x, key // spec.grid_x


def assign_pillars(points: np.ndarray, spec: PillarGridSpec, *, shuffle_seed: Optional[int] = None) -> PillarSet:
    """Deterministic voxelization (module docstring). ``points``: (N, F) float32 in the model frame, x, y, z in
    columns 0-2 (the remaining columns, e.g. time_lag, are copied into the slots). Rows with a non-finite x, y or z
    never enter a pillar."""
    pts = np.ascontiguousarray(np.asarray(points, dtype=np.float32))
    if pts.ndim != 2 or pts.shape[1] < 3:
        raise ValueError(f"points must be (N, >=3), got {pts.shape}")
    n_in = len(pts)
    perm = seeded_permutation(n_in, shuffle_seed)
    if perm is not None:
        pts = pts[perm]
    mn, mx = spec.min_f32, spec.max_f32
    keep = np.ones(len(pts), dtype=bool)
    for a in range(3):  # Autoware: reject x < min || x >= max (NaN would pass there; here it fails)
        keep &= (pts[:, a] >= mn[a]) & (pts[:, a] < mx[a])
    pts = pts[keep]
    ix, iy = pillar_cell_index(pts, spec)
    key = _cell_key(ix, iy, spec)
    order = np.argsort(key, kind="stable")              # stable: input order inside each cell
    key_sorted = key[order]
    uniq, first, counts = np.unique(key_sorted, return_index=True, return_counts=True)
    slot = np.arange(len(key_sorted)) - np.repeat(first, counts)
    K, P = spec.max_points_per_pillar, spec.max_pillars
    n_total = int(len(uniq))
    n_keep = min(n_total, P)
    pid = np.repeat(np.arange(n_total), counts)          # pillar id of every sorted row (ascending cell order)
    sel = (slot < K) & (pid < n_keep)
    F = pts.shape[1]
    raw = np.zeros((P, K, F), dtype=np.float32)
    raw[pid[sel], slot[sel]] = pts[order][sel]
    num = np.zeros(P, dtype=np.int32)
    num[:n_keep] = np.minimum(counts[:n_keep], K)
    cx, cy = _key_to_ixy(uniq[:n_keep], spec)
    coords = np.zeros((P, 3), dtype=np.int32)
    coords[:n_keep, 1] = cy
    coords[:n_keep, 2] = cx
    return PillarSet(points=raw, num_points=num, coords=coords, num_pillars=n_keep, num_pillars_total=n_total,
                     num_points_in_range=int(keep.sum()), max_points_in_pillar=int(counts.max()) if n_total else 0,
                     pillars_over_cap=int((counts > K).sum()), num_points_input=n_in)


def decorate_pillar_features(pillars: PillarSet, spec: PillarGridSpec, *, encoder_in_feature_size: int = 9,
                             point_dim: Optional[int] = None) -> np.ndarray:
    """``generateFeatures_kernel`` (``preprocess_kernel.cu:211-343``) -> (max_pillars, K, encoder_in_feature_size)
    float32. ``point_dim`` (columns copied from the raw rows) is 5 (x, y, z, intensity, time_lag) for 11 features
    and 4 (x, y, z, time_lag) otherwise, as the kernel's ``point_dim``. Layouts:

    - 9:  x, y, z, t, x-x̄, y-ȳ, z-z̄, x-x_c, y-y_c
    - 10: x, y, z, t, x-x̄, y-ȳ, z-z̄, x-x_c, y-y_c, z-z_c
    - 11: x, y, z, i, t, x-x̄, y-ȳ, z-z̄, x-x_c, y-y_c, z-z_c

    x̄ is the mean of the pillar's kept points (summed slot by slot in float32, divided by the float count);
    ``x_c = voxel_x / 2 + ix * voxel_x + min_x`` in float32 (``:269-276``). Padded slots and unused pillars are 0."""
    F_enc = int(encoder_in_feature_size)
    if F_enc not in (9, 10, 11):
        raise ValueError("encoder_in_feature_size must be 9, 10 or 11")
    pd = (5 if F_enc >= 11 else 4) if point_dim is None else int(point_dim)
    raw = pillars.points
    P, K, F = raw.shape
    if F < pd:
        raise ValueError(f"pillar rows have {F} columns, decoration needs {pd}")
    n = pillars.num_pillars
    out = np.zeros((P, K, F_enc), dtype=np.float32)
    if n == 0:
        return out
    r = raw[:n]
    cnt = pillars.num_points[:n].astype(np.float32)
    s = np.zeros((n, 3), dtype=np.float32)
    for k in range(K):  # fixed slot order (the kernel's float atomicAdd order is unspecified)
        s += r[:, k, :3]
    mean = s / cnt[:, None]
    vs, mn = spec.voxel_f32, spec.min_f32
    half = vs / np.float32(2)
    xo = (half[0] + pillars.coords[:n, 2].astype(np.float32) * vs[0]) + mn[0]
    yo = (half[1] + pillars.coords[:n, 1].astype(np.float32) * vs[1]) + mn[1]
    zo = (half[2] + pillars.coords[:n, 0].astype(np.float32) * vs[2]) + mn[2]
    f = out[:n]
    f[:, :, :pd] = r[:, :, :pd]
    f[:, :, pd:pd + 3] = r[:, :, :3] - mean[:, None, :]
    f[:, :, pd + 3] = r[:, :, 0] - xo[:, None]
    f[:, :, pd + 4] = r[:, :, 1] - yo[:, None]
    if F_enc in (10, 11):
        f[:, :, pd + 5] = r[:, :, 2] - zo[:, None]
    valid = np.arange(K)[None, :] < pillars.num_points[:n, None]
    f[~valid] = 0.0
    return out


def scatter_canvas(features: np.ndarray, coords: np.ndarray, num_pillars: int, spec: PillarGridSpec) -> np.ndarray:
    """``scatterFeatures_kernel``: ``canvas[c, iy, ix] = features[p, c]`` for the first ``num_pillars`` rows; the
    canvas is (C, grid_y, grid_x) float32 and zero elsewhere (NCHW without the batch axis)."""
    f = np.asarray(features, dtype=np.float32).reshape(np.shape(features)[0], -1)
    C = f.shape[1]
    canvas = np.zeros((C, spec.grid_y, spec.grid_x), dtype=np.float32)
    c = np.asarray(coords)[:num_pillars]
    canvas[:, c[:, 1], c[:, 2]] = f[:num_pillars].T
    return canvas


def canvas_gather_index(coords: np.ndarray, num_pillars: int, spec: PillarGridSpec, *, sentinel: int,
                        dtype: Any = np.uint32) -> np.ndarray:
    """Gather form of the scatter (C19): for every cell ``iy * grid_x + ix`` of the row-major (NHWC) canvas, the row
    of its pillar, or ``sentinel`` (an all-zero row of the feature table) for an empty cell. Then
    ``canvas_nhwc = table[index]`` equals :func:`scatter_canvas` transposed to (grid_y, grid_x, C)."""
    idx = np.full(spec.num_cells, sentinel, dtype=np.int64)
    c = np.asarray(coords)[:num_pillars]
    idx[c[:, 1].astype(np.int64) * spec.grid_x + c[:, 2].astype(np.int64)] = np.arange(num_pillars)
    return idx.astype(dtype)