from __future__ import annotations

import math

import numpy as np
import pytest

from k1link.compute.pointpillars_postprocess import (
    PointPillarsBox,
    PointPillarsPostprocessError,
    decode_pointpillars_output,
    oriented_3d_iou,
    oriented_bev_iou,
)


def _outputs(rows: list[list[float]]) -> tuple[np.ndarray, np.ndarray]:
    boxes = np.zeros((1, 393_216, 9), dtype=np.float32)
    boxes[0, : len(rows)] = np.asarray(rows, dtype=np.float32)
    return boxes, np.asarray([len(rows)], dtype=np.int32)


def _row(
    *,
    x: float,
    y: float = 0.0,
    length: float = 4.0,
    width: float = 2.0,
    yaw: float = 0.0,
    class_id: int = 0,
    score: float = 0.9,
) -> list[float]:
    return [x, y, 0.0, length, width, 1.5, yaw, float(class_id), score]


def _box(*, yaw: float = 0.0, x: float = 0.0) -> PointPillarsBox:
    return PointPillarsBox(
        x_m=x,
        y_m=0.0,
        z_m=0.0,
        length_m=4.0,
        width_m=2.0,
        height_m=1.5,
        yaw_rad=yaw,
        class_id=0,
        model_class="Vehicle",
        score=0.9,
    )


def test_decodes_native_label_order_and_sorts_by_score() -> None:
    output_boxes, num_boxes = _outputs(
        [
            _row(x=20.0, class_id=2, score=0.6),
            _row(x=0.0, class_id=0, score=0.9),
            _row(x=10.0, class_id=1, score=0.8),
        ]
    )

    decoded = decode_pointpillars_output(output_boxes, num_boxes)

    assert [box.model_class for box in decoded] == [
        "Vehicle",
        "Pedestrian",
        "Cyclist",
    ]
    assert [box.score for box in decoded] == pytest.approx([0.9, 0.8, 0.6])


def test_nms_reproduces_nvidia_sample_class_agnostic_suppression() -> None:
    output_boxes, num_boxes = _outputs(
        [
            _row(x=0.0, class_id=0, score=0.9),
            _row(x=0.1, class_id=1, score=0.8),
            _row(x=20.0, class_id=1, score=0.7),
        ]
    )

    decoded = decode_pointpillars_output(output_boxes, num_boxes)

    assert [(box.x_m, box.model_class) for box in decoded] == [
        (0.0, "Vehicle"),
        (20.0, "Pedestrian"),
    ]


def test_pre_nms_cap_is_applied_after_stable_score_ordering() -> None:
    output_boxes, num_boxes = _outputs(
        [
            _row(x=0.0, score=0.7),
            _row(x=10.0, score=0.9),
            _row(x=20.0, score=0.8),
        ]
    )

    decoded = decode_pointpillars_output(
        output_boxes,
        num_boxes,
        pre_nms_top_n=2,
    )

    assert [box.x_m for box in decoded] == [10.0, 20.0]


def test_oriented_bev_iou_handles_rotation_and_separation() -> None:
    assert oriented_bev_iou(_box(), _box()) == pytest.approx(1.0)
    assert oriented_bev_iou(_box(), _box(yaw=math.pi / 2.0)) == pytest.approx(
        1.0 / 3.0
    )
    assert oriented_bev_iou(_box(), _box(x=20.0)) == 0.0
    assert oriented_3d_iou(_box(), _box()) == pytest.approx(1.0)
    assert oriented_3d_iou(_box(), _box(x=20.0)) == 0.0


@pytest.mark.parametrize(
    ("row", "message"),
    [
        (_row(x=0.0, class_id=3), "class id"),
        (_row(x=0.0, length=0.0), "dimensions or score"),
        (_row(x=0.0, score=0.09), "dimensions or score"),
        (_row(x=0.0, score=1.1), "dimensions or score"),
    ],
)
def test_invalid_candidate_fails_closed(row: list[float], message: str) -> None:
    output_boxes, num_boxes = _outputs([row])

    with pytest.raises(PointPillarsPostprocessError, match=message):
        decode_pointpillars_output(output_boxes, num_boxes)


def test_output_tensor_contract_is_exact() -> None:
    with pytest.raises(PointPillarsPostprocessError, match="output_boxes contract"):
        decode_pointpillars_output(
            np.zeros((1, 1, 9), dtype=np.float32),
            np.asarray([0], dtype=np.int32),
        )