"""
test_postprocess.py — Tests for YOLOv8 post-processing and NMS.

Tests the _nms() helper and validate post-processing logic without
requiring a GPU, TRT, or running ROS2 node.
"""

import numpy as np
import pytest

from saltybot_perception.detection_utils import nms as _nms


class TestNMS:

    def test_single_box_kept(self):
        boxes = np.array([[0, 0, 10, 10]], dtype=float)
        scores = np.array([0.9])
        assert _nms(boxes, scores) == [0]

    def test_empty_input(self):
        assert _nms(np.zeros((0, 4)), np.array([])) == []

    def test_suppresses_overlapping_box(self):
        # Two heavily overlapping boxes — keep highest score
        boxes = np.array([
            [0, 0, 10, 10],   # score 0.9 — keep
            [1, 1, 11, 11],   # score 0.8 — suppress (high IoU with first)
        ], dtype=float)
        scores = np.array([0.9, 0.8])
        keep = _nms(boxes, scores, iou_threshold=0.45)
        assert keep == [0]

    def test_keeps_non_overlapping_boxes(self):
        boxes = np.array([
            [0, 0, 10, 10],
            [50, 50, 60, 60],
            [100, 100, 110, 110],
        ], dtype=float)
        scores = np.array([0.9, 0.85, 0.8])
        keep = _nms(boxes, scores, iou_threshold=0.45)
        assert sorted(keep) == [0, 1, 2]

    def test_score_ordering(self):
        # Lower score box overlaps with higher — higher should be kept
        boxes = np.array([
            [1, 1, 11, 11],   # score 0.6
            [0, 0, 10, 10],   # score 0.95 — should be kept
        ], dtype=float)
        scores = np.array([0.6, 0.95])
        keep = _nms(boxes, scores, iou_threshold=0.45)
        assert 1 in keep      # higher score (index 1) kept
        assert 0 not in keep  # lower score (index 0) suppressed

    def test_iou_threshold_controls_suppression(self):
        # Two boxes with ~0.5 IoU
        boxes = np.array([
            [0, 0, 10, 10],
            [5, 5, 15, 15],
        ], dtype=float)
        scores = np.array([0.9, 0.8])
        # High threshold — both boxes kept (IoU ~0.14 < 0.5)
        keep_high = _nms(boxes, scores, iou_threshold=0.5)
        assert sorted(keep_high) == [0, 1]
        # Low threshold — only first kept
        keep_low = _nms(boxes, scores, iou_threshold=0.0)
        assert keep_low == [0]
