# SPDX-License-Identifier: Apache-2.0
"""C15 NMS family of the Autoware LiDAR detectors: circle NMS, rotated IoU-BEV NMS with perception_utils semantics,
and the area-based class remapper.

- :func:`circle_nms` is ``circleNMS`` of ``autoware_lidar_centerpoint/lib/postprocess/circle_nms_kernel.cu:38-142``:
  boxes in descending score order; a kept box suppresses every *later* box whose squared BEV centre distance is
  strictly below ``threshold^2`` (float32: ``powf(dx, 2) + powf(dy, 2) < powf(thr, 2)``); greedy, class-agnostic.
  A threshold <= 0 means 1.5 m in Autoware's config (``centerpoint_config.hpp:83-85,139``): the caller applies that.
- :func:`iou_bev_nms` is ``perception_utils::IouBevNms::apply`` (``perception_utils/src/iou_bev_nms.cpp:132-171``,
  universe main 9ceaccf). Object ``t`` is compared with **every** earlier object ``s < t``, kept or suppressed
  (S:centerpoint:291-295); pairs with different labels where one is PEDESTRIAN are skipped; pairs whose centre
  distance squared exceeds ``search_distance_2d^2`` are skipped; the loop over ``s`` breaks at the first
  ``iou > threshold``; ``t`` is kept iff the largest IoU seen is ``<= threshold``. ``sort=False`` (the CenterPoint
  node's call) keeps the input order. IoU: the BEV rectangles of ``autoware_utils::toPolygon2d`` intersected
  exactly (Sutherland-Hodgman on two convex quadrilaterals, float64), with the area guards 1e-6 (each box and the
  intersection) and 0.01 (union).
- :class:`ClassRemapper` is ``perception_utils::DetectionClassRemapper::mapClasses``
  (``perception_utils/src/detection_class_remapper.cpp:54-70``): with ``bev_area = float(length * width)``, the first
  destination label (in label order) whose matrix entry allows it and whose ``[min, max]`` contains the area wins.

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

import math
from dataclasses import dataclass
from typing import Any, List, Mapping, Optional, Sequence

import numpy as np

__all__ = [
    "PEDESTRIAN_LABEL",
    "circle_nms",
    "bev_box_corners",
    "polygon_area",
    "clip_convex",
    "iou_bev",
    "iou_bev_nms",
    "ClassRemapper",
]

PEDESTRIAN_LABEL = 7   # autoware_perception_msgs/ObjectClassification PEDESTRIAN

_MIN_AREA = 1.0e-6        # iou_bev_nms.cpp compute2dIoU kMinArea
_MIN_UNION_AREA = 0.01    # kMinUnionArea


def circle_nms(x: Any, y: Any, dist_threshold: float) -> np.ndarray:
    """Indices (ascending) of the boxes kept by Autoware's circle NMS. The input must already be in descending
    score order (``thrust::sort``); see the module docstring."""
    xs = np.asarray(x, dtype=np.float32)
    ys = np.asarray(y, dtype=np.float32)
    n = len(xs)
    thr2 = np.float32(dist_threshold) * np.float32(dist_threshold)
    removed = np.zeros(n, dtype=bool)
    keep: List[int] = []
    for i in range(n):
        if removed[i]:
            continue
        keep.append(i)
        dx = xs[i + 1:] - xs[i]
        dy = ys[i + 1:] - ys[i]
        removed[i + 1:] |= (dx * dx + dy * dy) < thr2
    return np.asarray(keep, dtype=np.int64)


def bev_box_corners(x: float, y: float, length: float, width: float, yaw: float) -> np.ndarray:
    """(4, 2) float64 corners in ``toPolygon2d`` order: (+l/2, +w/2), (-l/2, +w/2), (-l/2, -w/2), (+l/2, -w/2) in
    the object frame, rotated by ``yaw`` about the centre (counter-clockwise for positive sizes)."""
    c, s = math.cos(yaw), math.sin(yaw)
    hl, hw = length / 2.0, width / 2.0
    local = ((hl, hw), (-hl, hw), (-hl, -hw), (hl, -hw))
    return np.array([(x + c * px - s * py, y + s * px + c * py) for px, py in local], dtype=np.float64)


def polygon_area(poly: Any) -> float:
    """Absolute shoelace area of a polygon (k, 2)."""
    p = np.asarray(poly, dtype=np.float64)
    if len(p) < 3:
        return 0.0
    x, y = p[:, 0], p[:, 1]
    return abs(float(np.dot(x, np.roll(y, -1)) - np.dot(np.roll(x, -1), y))) / 2.0


def _signed_area(p: np.ndarray) -> float:
    x, y = p[:, 0], p[:, 1]
    return float(np.dot(x, np.roll(y, -1)) - np.dot(np.roll(x, -1), y)) / 2.0


def clip_convex(subject: Any, clipper: Any) -> np.ndarray:
    """Sutherland-Hodgman: ``subject`` polygon clipped by the convex polygon ``clipper`` (either orientation) ->
    (k, 2) float64 intersection polygon (k may be 0)."""
    clip = np.asarray(clipper, dtype=np.float64)
    if _signed_area(clip) < 0:
        clip = clip[::-1]
    out = [tuple(p) for p in np.asarray(subject, dtype=np.float64)]
    for i in range(len(clip)):
        if not out:
            break
        ax, ay = clip[i]
        bx, by = clip[(i + 1) % len(clip)]
        inp, out = out, []

        def inside(p):
            return (bx - ax) * (p[1] - ay) - (by - ay) * (p[0] - ax) >= 0.0

        def cross_point(p, q):
            x1, y1, x2, y2 = p[0], p[1], q[0], q[1]
            den = (x1 - x2) * (ay - by) - (y1 - y2) * (ax - bx)
            if den == 0.0:
                return q
            t = ((x1 - ax) * (ay - by) - (y1 - ay) * (ax - bx)) / den
            return (x1 + t * (x2 - x1), y1 + t * (y2 - y1))

        prev = inp[-1]
        for cur in inp:
            if inside(cur):
                if not inside(prev):
                    out.append(cross_point(prev, cur))
                out.append(cur)
            elif inside(prev):
                out.append(cross_point(prev, cur))
            prev = cur
    return np.asarray(out, dtype=np.float64).reshape(-1, 2)


def _iou_from(pa: np.ndarray, area_a: float, env_a: np.ndarray, pb: np.ndarray, area_b: float,
              env_b: np.ndarray) -> float:
    if area_a < _MIN_AREA or area_b < _MIN_AREA:
        return 0.0
    if env_a[0] > env_b[2] or env_b[0] > env_a[2] or env_a[1] > env_b[3] or env_b[1] > env_a[3]:
        return 0.0    # disjoint envelopes (a speed guard in perception_utils; the intersection would be empty)
    inter = polygon_area(clip_convex(pa, pb))
    if inter < _MIN_AREA:
        return 0.0
    union = area_a + area_b - inter
    if union < _MIN_UNION_AREA:
        return 0.0
    return min(1.0, inter / union)


def _envelope(p: np.ndarray) -> np.ndarray:
    return np.array([p[:, 0].min(), p[:, 1].min(), p[:, 0].max(), p[:, 1].max()])


def iou_bev(a: Sequence[float], b: Sequence[float]) -> float:
    """Rotated BEV IoU of two boxes ``(x, y, length, width, yaw)`` with perception_utils' ``compute2dIoU`` guards."""
    pa, pb = bev_box_corners(*a), bev_box_corners(*b)
    return _iou_from(pa, polygon_area(pa), _envelope(pa), pb, polygon_area(pb), _envelope(pb))


def iou_bev_nms(x: Any, y: Any, length: Any, width: Any, yaw: Any, labels: Any, *, search_distance_2d: float = 10.0,
                iou_threshold: float = 0.1, scores: Optional[Any] = None, sort: bool = False,
                pedestrian_label: int = PEDESTRIAN_LABEL) -> np.ndarray:
    """Indices of the objects kept by ``IouBevNms::apply`` (module docstring), in processing order. ``labels`` are
    ObjectClassification labels; ``yaw`` is the ROS yaw. ``sort=True`` first orders by descending ``scores``
    (``std::stable_sort``)."""
    if not math.isfinite(search_distance_2d) or search_distance_2d < 0.0:
        raise ValueError("search_distance_2d must be a finite non-negative value")
    if not math.isfinite(iou_threshold) or not 0.0 <= iou_threshold <= 1.0:
        raise ValueError("iou_threshold must be a finite value between 0 and 1")
    xs = np.asarray(x, dtype=np.float64)
    n = len(xs)
    order = np.arange(n)
    if sort:
        if scores is None:
            raise ValueError("sort=True needs scores")
        order = np.argsort(-np.asarray(scores, dtype=np.float64), kind="stable")
    ys = np.asarray(y, dtype=np.float64)[order]
    xs = xs[order]
    ls = np.asarray(length, dtype=np.float64)[order]
    ws = np.asarray(width, dtype=np.float64)[order]
    yw = np.asarray(yaw, dtype=np.float64)[order]
    lab = np.asarray(labels, dtype=np.int64)[order]
    polys = [bev_box_corners(xs[i], ys[i], ls[i], ws[i], yw[i]) for i in range(n)]
    areas = [polygon_area(p) for p in polys]
    envs = [_envelope(p) for p in polys]
    sd2 = float(search_distance_2d) * float(search_distance_2d)
    keep: List[int] = []
    for t in range(n):
        max_iou = 0.0
        for s in range(t):
            if lab[t] != lab[s] and (lab[t] == pedestrian_label or lab[s] == pedestrian_label):
                continue
            dx, dy = xs[t] - xs[s], ys[t] - ys[s]
            if not dx * dx + dy * dy <= sd2:
                continue
            iou = _iou_from(polys[t], areas[t], envs[t], polys[s], areas[s], envs[s])
            max_iou = max(max_iou, iou)
            if iou > iou_threshold:
                break
        if max_iou <= iou_threshold:
            keep.append(int(order[t]))
    return np.asarray(keep, dtype=np.int64)


@dataclass(frozen=True)
class ClassRemapper:
    """``DetectionClassRemapper`` with square ``allow`` (bool), ``min_area`` and ``max_area`` matrices indexed
    ``[source_label, destination_label]`` over the ObjectClassification labels. Non-finite bounds become
    ``DBL_MAX`` (``detection_class_remapper.cpp:43-48``)."""

    allow: np.ndarray
    min_area: np.ndarray
    max_area: np.ndarray

    @classmethod
    def from_lists(cls, allow_remapping_by_area_matrix: Sequence[int], min_area_matrix: Sequence[float],
                   max_area_matrix: Sequence[float]) -> "ClassRemapper":
        """The three flat lists of ``detection_class_remapper.param.yaml`` (row = source, column = destination)."""
        if not (len(allow_remapping_by_area_matrix) == len(min_area_matrix) == len(max_area_matrix)):
            raise ValueError("Class remapping matrices must have equal sizes.")
        n = int(math.isqrt(len(min_area_matrix)))
        if n == 0 or n * n != len(min_area_matrix):
            raise ValueError("Class remapping matrices must be non-empty and square.")

        def finite(v: Sequence[float]) -> np.ndarray:
            a = np.asarray(v, dtype=np.float64)
            return np.where(np.isfinite(a), a, np.finfo(np.float64).max).reshape(n, n)

        return cls(np.asarray(allow_remapping_by_area_matrix, dtype=np.int64).reshape(n, n) != 0,
                   finite(min_area_matrix), finite(max_area_matrix))

    @classmethod
    def from_params(cls, params: Mapping[str, Any]) -> "ClassRemapper":
        """From the ``ros__parameters`` mapping of ``detection_class_remapper.param.yaml``."""
        return cls.from_lists(params["allow_remapping_by_area_matrix"], params["min_area_matrix"],
                              params["max_area_matrix"])

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

    def apply(self, labels: Any, length: Any, width: Any) -> np.ndarray:
        """New labels (uint8) for objects with ``labels`` and BEV ``length`` x ``width``."""
        lab = np.asarray(labels, dtype=np.int64).copy()
        area = (np.asarray(length, dtype=np.float64) * np.asarray(width, dtype=np.float64)).astype(np.float32)
        area = area.astype(np.float64)
        for i in range(len(lab)):
            src = int(lab[i])
            if src < 0 or src >= self.num_labels:
                continue
            for dst in range(self.num_labels):
                if self.allow[src, dst] and self.min_area[src, dst] <= area[i] <= self.max_area[src, dst]:
                    lab[i] = dst
                    break
        return lab.astype(np.uint8)