# SPDX-License-Identifier: Apache-2.0
"""Host post-processing of ``autoware_lidar_frnet`` (research/frnet/SPEC.md section 5), numpy only.

``fill_cloud_kernel`` (``PKG/lib/postprocess/postprocess_kernel.cu:200-264`` @ universe ``9ceaccf``), one thread per
point over all ``N`` points (interpolated ones included):

- ``class_id``: ``best = -1e9f``, ``class_id = num_classes - 1``, then ``if (prob > best)`` over the classes in
  order -- the **first** maximum wins and a NaN never wins (:func:`autoware_argmax`); ``probability = best``;
- segmentation cloud: (x, y, z of the network input points, class_id uint8, probability float32);
- visualization cloud: (x, y, z, ``palette[class_id]``: ``(r << 16) | (g << 8) | b`` as a float32);
- filtered cloud: a point is dropped iff **any** filter class has ``prob >= threshold`` (float32 compare against
  ``static_cast<float>(0.05)``); kept raw points are the original records, kept interpolated points are synthesized
  from x, y, z, intensity with the other fields 0. The CUDA kernel appends them with ``atomicAdd`` (any order);
  here the filtered cloud is the mask :attr:`FRNetSegmentation.filter_keep` over the output points.

AWML's offline max-probability threshold (``AWML/deploy/postprocessing.py``) is not used by Autoware and not here.
The output class adds the point arrays the node publishes to ``ttaw.outputs.Segmentation3D``.
"""
from __future__ import annotations

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

import numpy as np

from ..reference.config import NODE_PARAMS, FRNetConfig, parse_filter_classes
from ..ttaw.io import InputError, encode_array
from ..ttaw.outputs import Segmentation3D

__all__ = ["RUNTIME_PARAMS", "autoware_argmax", "filter_keep_mask", "visualization_rgb", "FRNetSegmentation",
           "FRNetPostprocessor"]

# Per-request knobs (RT-host), Autoware's defaults (frnet.param.yaml:6-7). ``filter_classes``: comma-separated
# class names, "" = none; the node allows at most 16.
RUNTIME_PARAMS: Dict[str, Tuple[Any, Any, Any, Any]] = {
    "class_probability_threshold": (float, 0.0, 1.0, NODE_PARAMS["class_probability_threshold"]),
    "filter_classes": (str, None, None, ",".join(NODE_PARAMS["filter_classes"])),
}

ARGMAX_INIT = np.float32(-1e9)


def autoware_argmax(pred_probs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """``fill_cloud_kernel`` L218-228 -> ``(class_id uint8 (N,), probability float32 (N,))``."""
    p = np.asarray(pred_probs, np.float32)
    if p.ndim != 2:
        raise ValueError(f"pred_probs must be (N, C), got {p.shape}")
    n, c = p.shape
    if n == 0:
        return np.zeros(0, np.uint8), np.zeros(0, np.float32)
    q = np.where(np.isnan(p), np.float32(-np.inf), p)
    idx = np.argmax(q, axis=1)                           # first maximum (NaN treated as never greater)
    best = q[np.arange(n), idx]
    none = ~(best > ARGMAX_INIT)                         # no class beat the -1e9 start: class C - 1, prob -1e9
    if none.any():
        idx = idx.copy()
        best = best.copy()
        idx[none] = c - 1
        best[none] = ARGMAX_INIT
    return idx.astype(np.uint8), best.astype(np.float32)


def filter_keep_mask(pred_probs: np.ndarray, class_indices: Sequence[int], threshold: float) -> np.ndarray:
    """``fill_cloud_kernel`` L245-254: True = the point stays in the filtered cloud (no filter class reaches the
    float32 threshold)."""
    p = np.asarray(pred_probs, np.float32)
    thr = np.float32(threshold)
    hit = np.zeros(p.shape[0], bool)
    for c in class_indices:
        hit |= p[:, int(c)] >= thr
    return ~hit


def visualization_rgb(class_id: np.ndarray, config: FRNetConfig) -> np.ndarray:
    """The visualization cloud's packed ``rgb`` float32 per point (``const_palette[class_id]``)."""
    return config.palette_packed[np.asarray(class_id, np.int64)]


@dataclass
class FRNetSegmentation(Segmentation3D):
    """Per-point classes of one frame in the node's output order: the input points that survive the ego crop box
    (input order, after non-finite rows are dropped), then the interpolated points (``num_points_raw`` raw rows
    first). ``label_ids`` = Autoware ``class_id``, ``scores`` = its ``probability``; ``points`` (N, 4) the network
    input x, y, z, intensity (the segmentation cloud's coordinates); ``filter_keep`` (N,) the filtered cloud's
    membership; ``input_index`` (num_points_raw,) int32 the input row of every raw output point."""

    points: Optional[np.ndarray] = None
    filter_keep: Optional[np.ndarray] = None
    input_index: Optional[np.ndarray] = None
    num_points_raw: int = 0
    extras: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        super().__post_init__()
        n = len(self)
        if self.points is not None:
            self.points = np.asarray(self.points, np.float32)
            if self.points.ndim != 2 or self.points.shape[0] != n:
                raise ValueError(f"points must be ({n}, C), got {self.points.shape}")
        if self.filter_keep is not None:
            self.filter_keep = np.asarray(self.filter_keep, bool).reshape(n)
        if self.input_index is not None:
            self.input_index = np.asarray(self.input_index, np.int32).reshape(-1)   # decoded rows < 2**31
            if self.input_index.shape[0] != self.num_points_raw:
                raise ValueError("input_index must have num_points_raw entries")
        if self.num_points_raw > n:
            raise ValueError("num_points_raw exceeds the number of points")

    def to_dict(self, output_format: str = "json") -> Dict[str, Any]:
        body = super().to_dict(output_format)
        body["num_points_raw"] = int(self.num_points_raw)
        if self.points is not None:
            body["points"] = encode_array(self.points, key="points")
        if self.filter_keep is not None:
            body["filter_keep"] = encode_array(self.filter_keep, key="filter_keep")
        if self.input_index is not None:
            body["input_index"] = encode_array(self.input_index, key="input_index")
        for k, v in self.extras.items():
            body[k] = encode_array(np.asarray(v), key=k)
        return body


class FRNetPostprocessor:
    """``pred_probs`` (N, 27) of the network input points -> :class:`FRNetSegmentation`."""

    def __init__(self, config: FRNetConfig):
        self.config = config

    def filter_indices(self, filter_classes: Any) -> Tuple[int, ...]:
        """Request ``filter_classes`` (comma-separated names) -> class indices; unknown names are a client error
        (the node throws at start-up)."""
        try:
            return self.config.filter_class_indices(parse_filter_classes(filter_classes))
        except ValueError as e:
            raise InputError(f"filter_classes: {e}") from None

    def __call__(self, pred_probs: np.ndarray, *, points: np.ndarray, num_points_raw: int,
                 input_index: Optional[np.ndarray], params: Mapping[str, Any], model: str = "",
                 frame_id: str = "sensor", meta: Optional[Dict[str, Any]] = None) -> FRNetSegmentation:
        p = {k: v[3] for k, v in RUNTIME_PARAMS.items()}
        p.update(params or {})
        indices = self.filter_indices(p["filter_classes"])
        class_id, prob = autoware_argmax(pred_probs)
        keep = filter_keep_mask(pred_probs, indices, p["class_probability_threshold"])
        meta = dict(meta or {})
        meta.update(filter={"classes": [self.config.class_names[i] for i in indices],
                            "class_probability_threshold": float(p["class_probability_threshold"]),
                            "kept": int(keep.sum())})
        return FRNetSegmentation(class_id, class_names=self.config.class_names, scores=prob, model=model,
                                 frame_id=frame_id, meta=meta, points=points, filter_keep=keep,
                                 input_index=input_index, num_points_raw=int(num_points_raw))

    def from_rows(self, rows: np.ndarray, plan: Any, *, points: np.ndarray, num_points_raw: int,
                  input_index: Optional[np.ndarray], params: Mapping[str, Any], model: str = "",
                  frame_id: str = "sensor", meta: Optional[Dict[str, Any]] = None,
                  tile_width: Optional[int] = None) -> FRNetSegmentation:
        """:meth:`__call__` straight from the device's sorted rows (``(>= N_cap, >= C)`` float32, or with
        ``tile_width`` the raw TILE-layout elements of ``[N_cap, tile_width]``): the un-permute, argmax and filter in
        one compiled pass (``host/fast.postprocess_rows`` / ``postprocess_tiles``; identical output)."""
        from .fast import postprocess_rows, postprocess_tiles

        p = {k: v[3] for k, v in RUNTIME_PARAMS.items()}
        p.update(params or {})
        indices = self.filter_indices(p["filter_classes"])
        thr = p["class_probability_threshold"]
        if tile_width is not None:
            class_id, prob, keep = postprocess_tiles(rows, int(tile_width), plan, self.config.num_classes, indices,
                                                     thr)
        else:
            class_id, prob, keep = postprocess_rows(rows, plan, self.config.num_classes, indices, thr)
        meta = dict(meta or {})
        meta.update(filter={"classes": [self.config.class_names[i] for i in indices],
                            "class_probability_threshold": float(p["class_probability_threshold"]),
                            "kept": int(keep.sum())})
        return FRNetSegmentation(class_id, class_names=self.config.class_names, scores=prob, model=model,
                                 frame_id=frame_id, meta=meta, points=points, filter_keep=keep,
                                 input_index=input_index, num_points_raw=int(num_points_raw))

    def empty(self, *, model: str = "", frame_id: str = "sensor",
              meta: Optional[Dict[str, Any]] = None) -> FRNetSegmentation:
        """No output (a frame the node refuses: profile checks), with the reason in ``meta``."""
        return FRNetSegmentation(np.zeros(0, np.uint8), class_names=self.config.class_names,
                                 scores=np.zeros(0, np.float32), model=model, frame_id=frame_id,
                                 meta=dict(meta or {}), points=np.zeros((0, 4), np.float32),
                                 filter_keep=np.zeros(0, bool), input_index=np.zeros(0, np.int32), num_points_raw=0)