# SPDX-License-Identifier: Apache-2.0
"""Host pre-processing of ``autoware_lidar_frnet`` (research/frnet/SPEC.md section 3; PLAN.md 2.8), numpy only.

Two halves, both deterministic:

1. :func:`autoware_preprocess` -- the node's pre-processing, ``PKG/lib/preprocess/preprocess_kernel.cu`` +
   ``PKG/lib/lidar_frnet.cpp:196-262`` @ universe ``9ceaccf``, producing the four network inputs exactly as
   the deployed ONNX takes them (unpadded: the graph sizes its voxel tensors as ``max(inverse_map) + 1``,
   S:frnet:255-262, so the ONNX Runtime / torch goldens must never see padding):

   - **ego crop** (``projectPoints_kernel`` L211-228): with the cached static ``sensor -> base_link`` transform as 12
     float32 (row-major R, then t), a point is dropped iff it lies *strictly* inside the box on all three axes; the
     reference-frame coordinates are evaluated in float32 as the kernel writes them
     (``((r0 * x + r1 * y) + r2 * z) + t``; FMA contraction by nvcc could move a point within one ulp of a face);
   - **projection** (``project2d`` L151-163) in float32: ``depth = sqrtf(x^2+y^2+z^2)``, ``yaw = -atan2f(y, x)``,
     ``pitch = asinf(z / depth)``, ``u = 0.5 (yaw / pi + 1) W``, ``v = (1 - (pitch + |down|) / total) H``, NaN -> 0
     (``fmaxf``), clamp to ``[0, W - 1]`` / ``[0, H - 1]``, truncate. float64 math moves 0.33 % of the points to
     another cell (S:frnet:190-193); numpy's float32 ``arctan2`` / ``arcsin`` can still differ from CUDA's
     ``atan2f`` / ``asinf`` by an ulp (unavoidable, documented);
   - **interpolation grid** (L232-258): one point per pixel stored as 4 x fp16 (x, y, z, intensity); the
     **farthest** point wins, by ``depth^2`` computed in float32 from the fp16-rounded x, y, z (exact: products of
     fp16 values are exact in float32, so nvcc's FMA contraction cannot change it). The CUDA race between the
     first plain store and the CAS loop, and the order among equal depths, are fixed here: the farthest point
     wins, ties go to the later input point (the research reference's policy);
   - **interpolation** (``interpolatePoints_kernel`` L332-379): an empty pixel with ``1 <= x <= W - 2`` whose left
     and right neighbours are filled gets ``float(left + right) * 0.5`` (fp16 sum, no azimuth wrap-around), appended
     after the raw points in row-major pixel order (CUDA appends with ``atomicAdd``: non-deterministic there);
   - **frustum cells and unique** (``generateUniqueCoors`` L403-468): ``coors = (0, row, col)``, key
     ``row * W + col``; ``voxel_coors`` = the occupied cells in ascending key order, ``inverse_map`` = the rank of
     each point's cell -- here a counting sort over the ``H * W`` cells (S:frnet:460-466);
   - **profile checks** (``lidar_frnet.cpp:138-149, 227-254``): the input count, ``N`` after interpolation and ``M``
     must lie in the ``num_points`` / ``num_unique_coors`` profiles. The CUDA appends are not bounds-checked (an
     over-profile frame overflows the device buffers before the ``N`` check, S:frnet:201-205); here every count is
     checked after it is known and an out-of-profile frame is refused (:func:`profile_errors`).

2. :func:`build_plan` -- the tables of the TT device graph (PLAN.md 2.8; PORT_LOG.md "trace layout"): the points
   **sorted by cell** (stable, so a cell's points keep their input order) and padded to the capacity ``N_cap`` =
   ``num_points.max``, the per-cell segments (CSR offsets, last row) of the sorted points, the fp32 8-d point
   features with ``pre_norm`` folded in, and the gather indices that replace every scatter of the graph:
   ``point2pix[s]`` (pixel of each sorted point at stride s: pixel2point) and ``pix2vox[s]`` (cell -> voxel row
   or the zero sentinel row ``M_cap``: frustum2pixel; at s > 1 several cells fall on one pixel and the **last
   voxel in voxel order** wins, which is what ORT's sequential ScatterND does, S:frnet:285-287). The network is
   permutation-equivariant in the points, so the host un-permutes the device output with ``plan.order``.

Nothing here imports torch, onnx or ttnn.
"""
from __future__ import annotations

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

import numpy as np

from ..reference.config import STRIDES, FRNetConfig

__all__ = ["NetworkInputs", "FramePlan", "crop_transform_12", "crop_keep_mask", "project2d", "interpolation_grid",
           "interpolate_points", "unique_coors", "autoware_preprocess", "profile_errors", "decorate_points",
           "build_plan", "segmax_rounds", "segmax_shift_tables"]

PI_F32 = np.float32(np.pi)   # M_PIf


# ------------------------------------------------------------------------------------------------- the node

@dataclass
class NetworkInputs:
    """The four ONNX inputs of one frame, unpadded, plus the bookkeeping of how they were made."""

    points: np.ndarray          # (N, 4) float32: kept raw points (input order), then interpolated points
    coors: np.ndarray           # (N, 3) int64 (0, row, col) frustum cell of every point
    voxel_coors: np.ndarray     # (M, 3) int64 occupied cells in ascending key order
    inverse_map: np.ndarray     # (N,) int64 voxel row of every point
    num_points_raw: int         # rows [0, num_points_raw) of `points` are raw input points
    keep_mask: np.ndarray       # (N0,) bool: input rows that survived the ego crop box
    crop_applied: bool = True

    @property
    def num_points(self) -> int:
        return int(self.points.shape[0])

    @property
    def num_voxels(self) -> int:
        return int(self.voxel_coors.shape[0])

    @property
    def num_interpolated(self) -> int:
        return self.num_points - self.num_points_raw

    def onnx_feeds(self) -> Dict[str, np.ndarray]:
        """``{points, coors, voxel_coors, inverse_map}`` as the deployed ONNX binds them (float32 / int64)."""
        return {"points": np.ascontiguousarray(self.points, np.float32),
                "coors": np.ascontiguousarray(self.coors, np.int64),
                "voxel_coors": np.ascontiguousarray(self.voxel_coors, np.int64),
                "inverse_map": np.ascontiguousarray(self.inverse_map, np.int64)}

    def stats(self) -> Dict[str, Any]:
        counts = np.bincount(self.inverse_map) if self.num_points else np.zeros(1, np.int64)
        return {"num_points_input": int(self.keep_mask.shape[0]), "num_points_raw": self.num_points_raw,
                "num_points_interpolated": self.num_interpolated, "num_points": self.num_points,
                "num_unique_coors": self.num_voxels, "cropped_by_ego_box": int((~self.keep_mask).sum()),
                "crop_applied": self.crop_applied, "max_points_per_cell": int(counts.max()) if counts.size else 0}


def crop_transform_12(T_ref_from_sensor: Any) -> np.ndarray:
    """The node's cached crop transform (``lidar_frnet_node.cpp:160-180``): ``static_cast<float>`` of the 3x3
    rotation (row-major) then the translation of the double ``base_link <- sensor`` transform -> (12,) float32."""
    T = np.asarray(T_ref_from_sensor, np.float64)
    if T.shape not in ((4, 4), (3, 4)):
        raise ValueError(f"crop transform must be 4x4 or 3x4, got {T.shape}")
    return np.concatenate([T[:3, :3].reshape(-1), T[:3, 3]]).astype(np.float32)


def crop_keep_mask(xyz: np.ndarray, sensor_to_ref_12: np.ndarray, bounds_f32: np.ndarray) -> np.ndarray:
    """``projectPoints_kernel`` L211-228: True = keep (outside the open box ``min < p < max`` on some axis)."""
    T = np.asarray(sensor_to_ref_12, np.float32).reshape(12)
    b = np.asarray(bounds_f32, np.float32).reshape(6)
    x, y, z = (np.asarray(xyz[:, i], np.float32) for i in range(3))
    px = ((T[0] * x + T[1] * y) + T[2] * z) + T[9]
    py = ((T[3] * x + T[4] * y) + T[5] * z) + T[10]
    pz = ((T[6] * x + T[7] * y) + T[8] * z) + T[11]
    inside = (px > b[0]) & (px < b[3]) & (py > b[1]) & (py < b[4]) & (pz > b[2]) & (pz < b[5])
    return ~inside


def project2d(xyz: np.ndarray, scale_x: int, scale_y: int, fov: Tuple[np.float32, np.float32, np.float32]
              ) -> Tuple[np.ndarray, np.ndarray]:
    """``project2d`` (preprocess_kernel.cu L151-163) in float32 -> (col, row) int64, clamped then truncated."""
    _, down, total = fov
    x = np.asarray(xyz[:, 0], np.float32)
    y = np.asarray(xyz[:, 1], np.float32)
    z = np.asarray(xyz[:, 2], np.float32)
    depth = np.sqrt(x * x + y * y + z * z, dtype=np.float32)
    yaw = -np.arctan2(y, x).astype(np.float32)
    with np.errstate(invalid="ignore", divide="ignore"):
        pitch = np.arcsin(z / depth).astype(np.float32)
    proj_x = np.float32(0.5) * (yaw / PI_F32 + np.float32(1.0))
    proj_y = np.float32(1.0) - (pitch + np.float32(abs(down))) / np.float32(total)
    proj_x = proj_x * np.float32(scale_x)
    proj_y = proj_y * np.float32(scale_y)
    proj_x = np.where(np.isnan(proj_x), np.float32(0.0), proj_x)      # fmaxf(NaN, 0) == 0
    proj_y = np.where(np.isnan(proj_y), np.float32(0.0), proj_y)
    px = np.minimum(np.maximum(proj_x, np.float32(0.0)), np.float32(scale_x) - np.float32(1.0))
    py = np.minimum(np.maximum(proj_y, np.float32(0.0)), np.float32(scale_y) - np.float32(1.0))
    return px.astype(np.int64), py.astype(np.int64)


def interpolation_grid(raw: np.ndarray, config: FRNetConfig) -> Tuple[np.ndarray, np.ndarray]:
    """The interpolation grid of ``projectPoints_kernel`` (L232-258) -> ``(grid (IH * IW, 4) float16, filled
    (IH * IW,) bool)``: per pixel the farthest point by fp16 ``depth^2``, ties to the later input point."""
    iw, ih = config.interp_w, config.interp_h
    grid = np.zeros((iw * ih, 4), np.float16)
    filled = np.zeros(iw * ih, bool)
    if not len(raw):
        return grid, filled
    u, v = project2d(raw[:, :3], iw, ih, config.fov)
    pix = v * iw + u
    h = raw.astype(np.float16)                            # half4(x, y, z, intensity): __float2half (RNE)
    hf = h[:, :3].astype(np.float32)
    depth2 = hf[:, 0] * hf[:, 0] + hf[:, 1] * hf[:, 1] + hf[:, 2] * hf[:, 2]   # exact products (fp16 operands)
    order = np.lexsort((np.arange(len(raw)), depth2, pix))   # by pixel, then depth, then input index
    ps = pix[order]
    last = np.r_[ps[1:] != ps[:-1], True]                 # last of each pixel group = farthest, latest on ties
    grid[ps[last]] = h[order[last]]
    filled[ps[last]] = True
    return grid, filled


def interpolate_points(grid: np.ndarray, filled: np.ndarray, config: FRNetConfig) -> np.ndarray:
    """``interpolatePoints_kernel`` (L332-379) -> (K, 4) float32 points in row-major pixel order."""
    iw, ih = config.interp_w, config.interp_h
    f2 = filled.reshape(ih, iw)
    cand = np.zeros_like(f2)
    cand[:, 1:-1] = ~f2[:, 1:-1] & f2[:, :-2] & f2[:, 2:]
    ys, xs = np.nonzero(cand)
    g2 = grid.reshape(ih, iw, 4)
    with np.errstate(over="ignore"):                                 # an fp16 overflow is inf, as __hadd gives
        s16 = (g2[ys, xs - 1] + g2[ys, xs + 1]).astype(np.float16)  # __half operator+ (RNE)
    return s16.astype(np.float32) * np.float32(0.5)


def unique_coors(coors: np.ndarray, frustum_w: int, frustum_h: int) -> Tuple[np.ndarray, np.ndarray]:
    """``generateUniqueCoors`` (L403-468) as a counting sort over the frustum cells -> ``(voxel_coors (M, 3) int64,
    inverse_map (N,) int64)``: the occupied cells in ascending key order and each point's cell rank."""
    keys = coors[:, 1] * frustum_w + coors[:, 2]
    counts = np.bincount(keys, minlength=frustum_w * frustum_h)
    occupied = counts > 0
    cell_rank = np.cumsum(occupied) - 1
    cells = np.flatnonzero(occupied)
    voxel_coors = np.stack([np.zeros_like(cells), cells // frustum_w, cells % frustum_w], axis=1).astype(np.int64)
    return voxel_coors, cell_rank[keys].astype(np.int64)


def autoware_preprocess(points_xyzi: np.ndarray, config: FRNetConfig, *,
                        sensor_to_ref_12: Optional[np.ndarray] = None,
                        crop_bounds: Optional[Sequence[float]] = None) -> NetworkInputs:
    """The node's pre-processing of one sensor-frame cloud ``(N0, 4)`` x, y, z, intensity (raw; float32) ->
    :class:`NetworkInputs`. ``crop_bounds`` (6 numbers) overrides the config's box; the crop runs iff a bound is
    non-zero, and then needs ``sensor_to_ref_12`` (:func:`crop_transform_12`): the node skips the frame when its TF
    lookup fails, so a missing transform raises ``ValueError`` here. Profile checks are separate
    (:func:`profile_errors`)."""
    pts = np.ascontiguousarray(points_xyzi, np.float32)
    if pts.ndim != 2 or pts.shape[1] != 4:
        raise ValueError(f"points must be (N, 4) x, y, z, intensity, got {pts.shape}")
    bounds = config.crop_bounds_f32 if crop_bounds is None else np.asarray(crop_bounds, np.float64).astype(np.float32)
    crop = bool(np.any(bounds != 0.0))
    if crop:
        if sensor_to_ref_12 is None:
            raise ValueError("the ego crop box is enabled but no sensor -> base_link transform was given")
        keep = crop_keep_mask(pts[:, :3], sensor_to_ref_12, bounds)
    else:
        keep = np.ones(len(pts), bool)
    raw = pts[keep]
    grid, filled = interpolation_grid(raw, config)
    interp = interpolate_points(grid, filled, config)
    allpts = np.concatenate([raw, interp], axis=0) if len(interp) else raw.copy()
    fx, fy = project2d(allpts[:, :3], config.frustum_w, config.frustum_h, config.fov)
    coors = np.stack([np.zeros_like(fy), fy, fx], axis=1).astype(np.int64)
    voxel_coors, inverse_map = unique_coors(coors, config.frustum_w, config.frustum_h)
    return NetworkInputs(points=np.ascontiguousarray(allpts, np.float32), coors=coors, voxel_coors=voxel_coors,
                         inverse_map=inverse_map, num_points_raw=int(len(raw)), keep_mask=keep, crop_applied=crop)


def profile_errors(num_input_points: int, inputs: Optional[NetworkInputs], config: FRNetConfig) -> List[str]:
    """The node's three profile checks, in its order (``lidar_frnet.cpp:138-149`` input count before the crop,
    L227-237 ``N`` after interpolation, L243-254 ``M``); empty when the frame is accepted."""
    lo, opt, hi = config.num_points
    if not lo <= num_input_points <= hi:
        return [f"Number of points ({num_input_points}) in input pointcloud does not match the profile "
                f"[min: {lo}, opt: {opt}, max: {hi}]"]
    if inputs is None:
        return []
    if not lo <= inputs.num_points <= hi:
        return [f"Number of points ({inputs.num_points}) after interpolation does not match the profile "
                f"[min: {lo}, opt: {opt}, max: {hi}]"]
    ulo, uopt, uhi = config.num_unique
    if not ulo <= inputs.num_voxels <= uhi:
        return [f"Number of unique coors ({inputs.num_voxels}) does not match the profile "
                f"[min: {ulo}, opt: {uopt}, max: {uhi}]"]
    return []


# ------------------------------------------------------------------------------------------- the TT device plan

def decorate_points(points: np.ndarray, inverse_map: np.ndarray, num_voxels: int, *,
                    pre_norm: Optional[Tuple[np.ndarray, np.ndarray]] = None) -> np.ndarray:
    """The frustum encoder's 8-d point features ``[x, y, z, i, |xyz|, x - mx, y - my, z - mz]`` (``m`` = the
    cell mean of x, y, z; AWML ``frustum_encoder.py:79-100``), computed in float64. With ``pre_norm = (scale,
    shift)`` the ``pre_norm`` BatchNorm1d is applied (``x * scale + shift``, float64) -- the port's host fold
    (PLAN.md 2.8: xyz up to +-200 m minus a cell mean needs fp32 or better, never bf16). One rounding to float32."""
    p = np.asarray(points, np.float64)
    n = len(p)
    out = np.empty((n, 8), np.float64)
    out[:, :4] = p
    out[:, 4] = np.sqrt((p[:, :3] ** 2).sum(axis=1))
    if n:
        counts = np.bincount(inverse_map, minlength=num_voxels).astype(np.float64)
        for c in range(3):
            mean = np.bincount(inverse_map, weights=p[:, c], minlength=num_voxels) / np.maximum(counts, 1.0)
            out[:, 5 + c] = p[:, c] - mean[inverse_map]
    if pre_norm is not None:
        scale, shift = (np.asarray(a, np.float64).reshape(8) for a in pre_norm)
        out = out * scale + shift
    return out.astype(np.float32)


def segmax_rounds(max_segment_length: int) -> int:
    """Rounds of the log-step segmented max (Hillis-Steele) that cover segments up to that length:
    ``ceil(log2(L))`` (0 for L <= 1). PLAN.md 1.3: 5 for OT128 (<= 17 points per cell seen), 7 for QT128 (<= 111)."""
    return int(np.ceil(np.log2(max_segment_length))) if max_segment_length > 1 else 0


@dataclass
class FramePlan:
    """The per-frame host tables of the device graph (module docstring, part 2). Row ``i`` of every per-point table
    is the ``i``-th point in **sorted** order; rows ``N .. N_cap - 1`` are padding (one trailing sentinel segment:
    their voxel id is ``M_cap``, their pixel 0, their features 0, and nothing the output keeps reads them)."""

    n: int                                   # real points
    m: int                                   # real voxels (occupied cells)
    n_raw: int
    n_cap: int
    m_cap: int
    frustum_hw: Tuple[int, int]
    order: np.ndarray                        # (N,) int64: sorted row -> network-input point index
    seg_id: np.ndarray                       # (N_cap,) int64 voxel id of each sorted row (padding: M_cap)
    seg_offsets: np.ndarray                  # (M + 1,) int64 CSR offsets of the voxels' sorted rows
    seg_last: np.ndarray                     # (M_cap,) uint32 last sorted row of each voxel (unused voxels: N_cap - 1)
    point2pix: Dict[int, np.ndarray]         # stride -> (N_cap,) uint32 pixel of each sorted row (padding: 0)
    pix2vox: Dict[int, np.ndarray]           # stride -> ((H/s) * (W/s),) uint32 voxel row or M_cap (zero row)
    features: np.ndarray                     # (N_cap, 8) float32 decorated (+ pre_norm) features
    max_segment_length: int = 0
    strides: Tuple[int, ...] = field(default=STRIDES)

    @property
    def rank(self) -> np.ndarray:
        """(N,) int64: network-input point index -> sorted row (the inverse of ``order``)."""
        r = np.empty(self.n, np.int64)
        r[self.order] = np.arange(self.n)
        return r

    def unsort(self, sorted_rows: np.ndarray) -> np.ndarray:
        """Device rows (sorted, padded) -> the network-input point order (the first ``N`` sorted rows only)."""
        return np.asarray(sorted_rows)[: self.n][self.rank]

    def stats(self) -> Dict[str, Any]:
        return {"n": self.n, "m": self.m, "n_cap": self.n_cap, "m_cap": self.m_cap,
                "max_points_per_cell": self.max_segment_length,
                "segmax_rounds": segmax_rounds(self.max_segment_length)}


def build_plan(inputs: NetworkInputs, config: FRNetConfig, *,
               pre_norm: Optional[Tuple[np.ndarray, np.ndarray]] = None,
               n_cap: Optional[int] = None, m_cap: Optional[int] = None,
               strides: Sequence[int] = STRIDES) -> FramePlan:
    """The device tables of one accepted frame (:class:`FramePlan`). ``n_cap`` / ``m_cap`` default to the
    config's capacities (COMPILE: the profile maxima); a frame above them raises ``ValueError`` (the profile check
    refuses it first)."""
    n_cap = config.n_cap if n_cap is None else int(n_cap)
    m_cap = config.m_cap if m_cap is None else int(m_cap)
    n, m = inputs.num_points, inputs.num_voxels
    if n > n_cap or m > m_cap:
        raise ValueError(f"frame N={n} / M={m} exceeds the capacity N_cap={n_cap} / M_cap={m_cap}")
    H, W = config.frustum_h, config.frustum_w
    inv = inputs.inverse_map
    order = np.argsort(inv, kind="stable")              # voxel ids are key ranks: == the stable sort by cell key
    counts = np.bincount(inv, minlength=m).astype(np.int64)
    offsets = np.zeros(m + 1, np.int64)
    np.cumsum(counts, out=offsets[1:])
    seg_id = np.full(n_cap, m_cap, np.int64)
    seg_id[:n] = inv[order]
    seg_last = np.full(m_cap, n_cap - 1, np.uint32)
    seg_last[:m] = (offsets[1:] - 1).astype(np.uint32)
    coors_sorted = inputs.coors[order]
    vox = inputs.voxel_coors
    point2pix: Dict[int, np.ndarray] = {}
    pix2vox: Dict[int, np.ndarray] = {}
    for s in strides:
        hs, ws = H // s, W // s
        p2p = np.zeros(n_cap, np.uint32)
        p2p[:n] = ((coors_sorted[:, 1] // s) * ws + coors_sorted[:, 2] // s).astype(np.uint32)
        point2pix[s] = p2p
        last = np.full(hs * ws, -1, np.int64)
        np.maximum.at(last, (vox[:, 1] // s) * ws + vox[:, 2] // s, np.arange(m, dtype=np.int64))
        pix2vox[s] = np.where(last < 0, m_cap, last).astype(np.uint32)
    feats = np.zeros((n_cap, 8), np.float32)
    feats[:n] = decorate_points(inputs.points[order], inv[order], m, pre_norm=pre_norm)
    return FramePlan(n=n, m=m, n_raw=inputs.num_points_raw, n_cap=n_cap, m_cap=m_cap, frustum_hw=(H, W),
                     order=order, seg_id=seg_id, seg_offsets=offsets, seg_last=seg_last, point2pix=point2pix,
                     pix2vox=pix2vox, features=feats, max_segment_length=int(counts.max()) if m else 0,
                     strides=tuple(strides))


def segmax_shift_tables(seg_id: np.ndarray, rounds: int) -> List[np.ndarray]:
    """The gather tables of the exact log-step segmented max (PLAN.md 1.3, the K1 stock-op fallback): round
    ``r`` uses ``idx[i] = i - 2^r`` if row ``i - 2^r`` is in the same segment, else ``i`` (max is idempotent);
    after ``rounds`` rounds each segment's last row holds the segment max when its length is <= ``2^rounds``.
    -> ``rounds`` x (N_cap,) uint32."""
    seg = np.asarray(seg_id)
    i = np.arange(len(seg), dtype=np.int64)
    out = []
    for r in range(rounds):
        d = 1 << r
        j = i - d
        same = np.zeros(len(seg), bool)
        same[d:] = seg[d:] == seg[:-d]
        out.append(np.where(same, j, i).astype(np.uint32))
    return out