"""Validated NVIDIA PointPillars candidate decoding and sample-compatible NMS."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Final

import numpy as np

POINTPILLARS_MODEL_CLASSES: Final = ("Vehicle", "Pedestrian", "Cyclist")
POINTPILLARS_OUTPUT_FIELDS: Final = (
    "x",
    "y",
    "z",
    "length",
    "width",
    "height",
    "yaw",
    "class_id",
    "score",
)
POINTPILLARS_NMS_IOU_THRESHOLD: Final = 0.01
POINTPILLARS_PRE_NMS_TOP_N: Final = 4_096
POINTPILLARS_EMBEDDED_SCORE_THRESHOLD: Final = 0.1
POINTPILLARS_MODEL_POINT_CLOUD_RANGE: Final = (
    -51.20000076293945,
    -51.20000076293945,
    -1.399999976158142,
    51.20000076293945,
    51.20000076293945,
    4.400000095367432,
)
NVIDIA_REFERENCE_COMMIT: Final = "a540badc47812a17a94e924b537d49ad3969b5a8"
_EPSILON: Final = 1e-8


class PointPillarsPostprocessError(RuntimeError):
    """PointPillars output does not satisfy the frozen NVIDIA contract."""


@dataclass(frozen=True, slots=True)
class PointPillarsBox:
    x_m: float
    y_m: float
    z_m: float
    length_m: float
    width_m: float
    height_m: float
    yaw_rad: float
    class_id: int
    model_class: str
    score: float


@dataclass(frozen=True, slots=True)
class _Geometry:
    corners: tuple[tuple[float, float], ...]
    area: float
    minimum_x: float
    maximum_x: float
    minimum_y: float
    maximum_y: float


def decode_pointpillars_output(
    output_boxes: np.ndarray,
    num_boxes: np.ndarray,
    *,
    pre_nms_top_n: int = POINTPILLARS_PRE_NMS_TOP_N,
    nms_iou_threshold: float = POINTPILLARS_NMS_IOU_THRESHOLD,
) -> tuple[PointPillarsBox, ...]:
    """Decode validated rows and reproduce the NVIDIA sample's class-agnostic NMS."""

    boxes = np.asarray(output_boxes)
    counts = np.asarray(num_boxes)
    if boxes.shape != (1, 393_216, 9) or boxes.dtype != np.float32:
        raise PointPillarsPostprocessError("PointPillars output_boxes contract changed")
    if counts.shape != (1,) or counts.dtype != np.int32:
        raise PointPillarsPostprocessError("PointPillars num_boxes contract changed")
    count = int(counts[0])
    if count < 0 or count > boxes.shape[1]:
        raise PointPillarsPostprocessError("PointPillars candidate count is invalid")
    if (
        isinstance(pre_nms_top_n, bool)
        or not isinstance(pre_nms_top_n, int)
        or pre_nms_top_n < 1
        or not math.isfinite(nms_iou_threshold)
        or not 0.0 <= nms_iou_threshold <= 1.0
    ):
        raise PointPillarsPostprocessError("PointPillars NMS parameters are invalid")
    if count == 0:
        return ()

    candidates = boxes[0, :count]
    if not np.isfinite(candidates).all():
        raise PointPillarsPostprocessError("PointPillars candidate is non-finite")
    decoded = tuple(_decode_row(row) for row in candidates)
    ordered = tuple(
        sorted(
            decoded,
            key=lambda box: box.score,
            reverse=True,
        )[:pre_nms_top_n]
    )
    return _class_agnostic_nms(ordered, nms_iou_threshold)


def oriented_bev_iou(one: PointPillarsBox, another: PointPillarsBox) -> float:
    """Return BEV IoU for two validated oriented boxes."""

    one_geometry = _geometry(one)
    another_geometry = _geometry(another)
    overlap = _intersection_area(one_geometry, another_geometry)
    denominator = one_geometry.area + another_geometry.area - overlap
    return overlap / max(denominator, _EPSILON)


def oriented_3d_iou(one: PointPillarsBox, another: PointPillarsBox) -> float:
    """Return oriented 3D IoU for two center-based LiDAR-frame boxes."""

    one_geometry = _geometry(one)
    another_geometry = _geometry(another)
    bev_overlap = _intersection_area(one_geometry, another_geometry)
    one_minimum_z = one.z_m - one.height_m / 2.0
    one_maximum_z = one.z_m + one.height_m / 2.0
    another_minimum_z = another.z_m - another.height_m / 2.0
    another_maximum_z = another.z_m + another.height_m / 2.0
    height_overlap = max(
        0.0,
        min(one_maximum_z, another_maximum_z)
        - max(one_minimum_z, another_minimum_z),
    )
    intersection = bev_overlap * height_overlap
    one_volume = one_geometry.area * one.height_m
    another_volume = another_geometry.area * another.height_m
    return intersection / max(one_volume + another_volume - intersection, _EPSILON)


def _decode_row(row: np.ndarray) -> PointPillarsBox:
    class_value = float(row[7])
    class_id = int(round(class_value))
    if (
        abs(class_value - class_id) > 1e-5
        or class_id < 0
        or class_id >= len(POINTPILLARS_MODEL_CLASSES)
    ):
        raise PointPillarsPostprocessError("PointPillars class id is invalid")
    length_m, width_m, height_m = (float(row[index]) for index in (3, 4, 5))
    score = float(row[8])
    if (
        length_m <= 0.0
        or width_m <= 0.0
        or height_m <= 0.0
        or score < POINTPILLARS_EMBEDDED_SCORE_THRESHOLD
        or score > 1.0
    ):
        raise PointPillarsPostprocessError(
            "PointPillars dimensions or score are invalid"
        )
    x_m, y_m, z_m = (float(row[index]) for index in (0, 1, 2))
    return PointPillarsBox(
        x_m=x_m,
        y_m=y_m,
        z_m=z_m,
        length_m=length_m,
        width_m=width_m,
        height_m=height_m,
        yaw_rad=float(row[6]),
        class_id=class_id,
        model_class=POINTPILLARS_MODEL_CLASSES[class_id],
        score=score,
    )


def _class_agnostic_nms(
    boxes: tuple[PointPillarsBox, ...],
    threshold: float,
) -> tuple[PointPillarsBox, ...]:
    geometries = tuple(_geometry(box) for box in boxes)
    suppressed = [False] * len(boxes)
    accepted: list[PointPillarsBox] = []
    for index, box in enumerate(boxes):
        if suppressed[index]:
            continue
        accepted.append(box)
        geometry = geometries[index]
        for candidate_index in range(index + 1, len(boxes)):
            if suppressed[candidate_index]:
                continue
            another = geometries[candidate_index]
            if not _aabbs_overlap(geometry, another):
                continue
            overlap = _intersection_area(geometry, another)
            iou = overlap / max(geometry.area + another.area - overlap, _EPSILON)
            if iou >= threshold:
                suppressed[candidate_index] = True
    return tuple(accepted)


def _geometry(box: PointPillarsBox) -> _Geometry:
    half_length = box.length_m / 2.0
    half_width = box.width_m / 2.0
    cosine = math.cos(box.yaw_rad)
    sine = math.sin(box.yaw_rad)
    corners = tuple(
        (
            box.x_m + local_x * cosine - local_y * sine,
            box.y_m + local_x * sine + local_y * cosine,
        )
        for local_x, local_y in (
            (-half_length, -half_width),
            (half_length, -half_width),
            (half_length, half_width),
            (-half_length, half_width),
        )
    )
    x_values = [point[0] for point in corners]
    y_values = [point[1] for point in corners]
    return _Geometry(
        corners=corners,
        area=box.length_m * box.width_m,
        minimum_x=min(x_values),
        maximum_x=max(x_values),
        minimum_y=min(y_values),
        maximum_y=max(y_values),
    )


def _aabbs_overlap(one: _Geometry, another: _Geometry) -> bool:
    return not (
        one.maximum_x < another.minimum_x
        or another.maximum_x < one.minimum_x
        or one.maximum_y < another.minimum_y
        or another.maximum_y < one.minimum_y
    )


def _intersection_area(one: _Geometry, another: _Geometry) -> float:
    if not _aabbs_overlap(one, another):
        return 0.0
    polygon = list(one.corners)
    clip = another.corners
    for index, edge_start in enumerate(clip):
        edge_end = clip[(index + 1) % len(clip)]
        polygon = _clip_polygon(polygon, edge_start, edge_end)
        if not polygon:
            return 0.0
    return abs(
        sum(
            one_point[0] * another_point[1]
            - another_point[0] * one_point[1]
            for one_point, another_point in zip(
                polygon,
                (*polygon[1:], polygon[0]),
                strict=True,
            )
        )
    ) / 2.0


def _clip_polygon(
    polygon: list[tuple[float, float]],
    edge_start: tuple[float, float],
    edge_end: tuple[float, float],
) -> list[tuple[float, float]]:
    if not polygon:
        return []
    result: list[tuple[float, float]] = []
    previous = polygon[-1]
    previous_inside = _inside(previous, edge_start, edge_end)
    for current in polygon:
        current_inside = _inside(current, edge_start, edge_end)
        if current_inside:
            if not previous_inside:
                result.append(_line_intersection(previous, current, edge_start, edge_end))
            result.append(current)
        elif previous_inside:
            result.append(_line_intersection(previous, current, edge_start, edge_end))
        previous = current
        previous_inside = current_inside
    return result


def _inside(
    point: tuple[float, float],
    edge_start: tuple[float, float],
    edge_end: tuple[float, float],
) -> bool:
    return _cross(
        edge_end[0] - edge_start[0],
        edge_end[1] - edge_start[1],
        point[0] - edge_start[0],
        point[1] - edge_start[1],
    ) >= -_EPSILON


def _line_intersection(
    segment_start: tuple[float, float],
    segment_end: tuple[float, float],
    edge_start: tuple[float, float],
    edge_end: tuple[float, float],
) -> tuple[float, float]:
    segment_x = segment_end[0] - segment_start[0]
    segment_y = segment_end[1] - segment_start[1]
    edge_x = edge_end[0] - edge_start[0]
    edge_y = edge_end[1] - edge_start[1]
    denominator = _cross(segment_x, segment_y, edge_x, edge_y)
    if abs(denominator) <= _EPSILON:
        return segment_end
    offset_x = edge_start[0] - segment_start[0]
    offset_y = edge_start[1] - segment_start[1]
    ratio = _cross(offset_x, offset_y, edge_x, edge_y) / denominator
    return (
        segment_start[0] + ratio * segment_x,
        segment_start[1] + ratio * segment_y,
    )


def _cross(one_x: float, one_y: float, another_x: float, another_y: float) -> float:
    return one_x * another_y - one_y * another_x