from __future__ import annotations

import numpy as np
import numpy.typing as npt
import pytest

from supervision.config import ORIENTED_BOX_COORDINATES
from supervision.detection._geometry_dispatch import detection_area, detection_iou
from supervision.detection.compact_mask import CompactMask
from supervision.detection.core import Detections
from supervision.detection.utils.boxes import xyxyxyxy_to_xyxy
from supervision.detection.utils.iou_and_nms import (
    OverlapMetric,
    box_iou_batch,
    mask_iou_batch,
    oriented_box_iou_batch,
)
from supervision.detection.utils.masks import count_mask_pixels


def _rotated_rect(
    center_x: float, center_y: float, width: float, height: float, angle_deg: float
) -> npt.NDArray[np.float32]:
    """Return four corners of a rotated rectangle in clockwise order."""
    half_w = width / 2.0
    half_h = height / 2.0
    corners = np.array(
        [
            [-half_w, -half_h],
            [half_w, -half_h],
            [half_w, half_h],
            [-half_w, half_h],
        ],
        dtype=np.float32,
    )
    angle = np.deg2rad(angle_deg)
    rotation = np.array(
        [[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]],
        dtype=np.float32,
    )
    center = np.array([center_x, center_y], dtype=np.float32)
    return (corners @ rotation.T + center).astype(np.float32)


def _detections_from_quads(
    quads: list[npt.NDArray[np.float32]],
    xyxy: npt.NDArray[np.float32] | None = None,
) -> Detections:
    """Build Detections carrying OBB coordinates."""
    corners = np.stack(quads).astype(np.float32)
    if xyxy is None:
        xyxy = xyxyxyxy_to_xyxy(corners).astype(np.float32)
    return Detections(
        xyxy=xyxy,
        data={ORIENTED_BOX_COORDINATES: corners},
    )


def _full_image_xyxy(
    count: int, image_shape: tuple[int, int]
) -> npt.NDArray[np.float32]:
    """Return full-image xyxy boxes for CompactMask construction."""
    image_height, image_width = image_shape
    xyxy = np.array([0, 0, image_width - 1, image_height - 1], dtype=np.float32)
    return np.tile(xyxy, (count, 1))


class TestDetectionArea:
    """Tests for geometry-aware area dispatch."""

    def test_returns_mask_pixel_area(self) -> None:
        """Masks take precedence over OBBs and AABB envelopes."""
        mask = np.zeros((1, 20, 20), dtype=bool)
        mask[0, 2:5, 3:8] = True
        quad = _rotated_rect(10, 10, 12, 6, 30)
        detections = Detections(
            xyxy=np.array([[0, 0, 20, 20]], dtype=np.float32),
            mask=mask,
            data={ORIENTED_BOX_COORDINATES: quad[np.newaxis]},
        )

        area = detection_area(detections)

        np.testing.assert_array_equal(area, np.array([15], dtype=np.int64))

    def test_dense_mask_branch_reuses_count_mask_pixels(self) -> None:
        """Dense-mask area delegates to count_mask_pixels for the shared fast path."""
        mask = np.zeros((3, 10, 10), dtype=bool)
        mask[0, :2, :2] = True
        mask[1, :3, :3] = True
        detections = Detections(
            xyxy=np.tile(np.array([[0, 0, 9, 9]], dtype=np.float32), (3, 1)),
            mask=mask,
        )

        area = detection_area(detections)

        np.testing.assert_array_equal(area, count_mask_pixels(mask))
        assert area.dtype == np.int64

    def test_returns_compact_mask_area(self) -> None:
        """CompactMask inputs use CompactMask.area without dense materialisation."""
        masks = np.zeros((2, 12, 12), dtype=bool)
        masks[0, 1:4, 2:7] = True
        masks[1, 4:9, 4:10] = True
        compact_mask = CompactMask.from_dense(
            masks=masks,
            xyxy=_full_image_xyxy(len(masks), masks.shape[1:]),
            image_shape=masks.shape[1:],
        )
        detections = Detections(
            xyxy=_full_image_xyxy(len(masks), masks.shape[1:]),
            mask=compact_mask,
        )

        area = detection_area(detections)

        np.testing.assert_array_equal(area, compact_mask.area)

    def test_returns_oriented_box_area_when_present(self) -> None:
        """OBB area is used instead of the larger rotated AABB envelope."""
        quad = _rotated_rect(50, 50, 20, 10, 45)
        detections = _detections_from_quads([quad])

        area = detection_area(detections)

        np.testing.assert_allclose(area, np.array([200.0]))
        assert detections.box_area[0] > area[0]

    def test_returns_box_area_when_no_richer_geometry_is_present(self) -> None:
        """AABB area is the fallback when masks and OBB corners are absent."""
        detections = Detections(
            xyxy=np.array([[0, 0, 20, 10], [3, 4, 8, 12]], dtype=np.float32)
        )

        area = detection_area(detections)

        np.testing.assert_array_equal(area, detections.box_area)

    @pytest.mark.parametrize(
        "detections",
        [
            pytest.param(
                Detections(xyxy=np.empty((0, 4), dtype=np.float32)),
                id="empty-aabb",
            ),
            pytest.param(
                Detections(
                    xyxy=np.empty((0, 4), dtype=np.float32),
                    mask=np.empty((0, 8, 8), dtype=bool),
                ),
                id="empty-mask",
            ),
            pytest.param(
                Detections(
                    xyxy=np.empty((0, 4), dtype=np.float32),
                    data={
                        ORIENTED_BOX_COORDINATES: np.empty((0, 4, 2), dtype=np.float32)
                    },
                ),
                id="empty-obb",
            ),
            pytest.param(
                Detections(
                    xyxy=np.empty((0, 4), dtype=np.float32),
                    mask=CompactMask.from_dense(
                        masks=np.empty((0, 8, 8), dtype=bool),
                        xyxy=np.empty((0, 4), dtype=np.float32),
                        image_shape=(8, 8),
                    ),
                ),
                id="empty-compact-mask",
            ),
        ],
    )
    def test_returns_empty_array_for_empty_detections(
        self, detections: Detections
    ) -> None:
        """Empty Detections return an empty area array for every geometry branch."""
        area = detection_area(detections)

        assert area.shape == (0,)

    @pytest.mark.parametrize(
        "detections",
        [
            pytest.param(
                Detections(
                    xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
                    mask=np.ones((1, 3, 4), dtype=bool),
                ),
                id="mask",
            ),
            pytest.param(
                _detections_from_quads([_rotated_rect(50, 50, 20, 10, 30)]),
                id="obb",
            ),
            pytest.param(
                Detections(xyxy=np.array([[0, 0, 10, 5]], dtype=np.float32)),
                id="aabb",
            ),
        ],
    )
    def test_matches_detections_area_property(self, detections: Detections) -> None:
        """Detections.area delegates to the shared geometry dispatch helper."""
        area = detection_area(detections)

        np.testing.assert_array_equal(area, detections.area)

    def test_keeps_oriented_area_invariant_when_envelope_changes(self) -> None:
        """Rotating an OBB preserves exact OBB area while changing envelope area."""
        axis_aligned = _detections_from_quads([_rotated_rect(30, 30, 20, 10, 0)])
        rotated = _detections_from_quads([_rotated_rect(30, 30, 20, 10, 45)])

        areas = np.concatenate([detection_area(axis_aligned), detection_area(rotated)])

        np.testing.assert_allclose(areas, np.array([200.0, 200.0]))
        assert axis_aligned.box_area[0] != pytest.approx(rotated.box_area[0])

    def test_mask_area_handles_rotated_arrays(self) -> None:
        """Mask area counts true pixels consistently after array rotation."""
        mask = np.zeros((80, 80), dtype=bool)
        mask[30:50, 25:55] = 1
        rotated_mask = np.rot90(mask)
        detections = Detections(
            xyxy=np.array([[0, 0, 79, 79], [0, 0, 79, 79]], dtype=np.float32),
            mask=np.stack([mask, rotated_mask]),
        )

        area = detection_area(detections)

        np.testing.assert_array_equal(area, np.array([600, 600], dtype=np.int64))

    def test_degenerate_collinear_obb_has_zero_area(self) -> None:
        """A collinear (zero-area) OBB reports 0 rather than a well-formed area."""
        collinear = np.array([[0, 0], [5, 0], [10, 0], [15, 0]], dtype=np.float32)
        detections = _detections_from_quads([collinear])

        area = detection_area(detections)

        np.testing.assert_allclose(area, np.array([0.0]))


class TestDetectionIou:
    """Tests for geometry-aware IoU dispatch."""

    def test_returns_mask_iou_when_both_operands_have_masks(self) -> None:
        """Mask IoU is used when both operands carry masks."""
        masks_a = np.zeros((1, 16, 16), dtype=bool)
        masks_b = np.zeros((1, 16, 16), dtype=bool)
        masks_a[0, 2:10, 2:10] = True
        masks_b[0, 6:14, 6:14] = True
        detections_a = Detections(
            xyxy=np.array([[0, 0, 16, 16]], dtype=np.float32),
            mask=masks_a,
        )
        detections_b = Detections(
            xyxy=np.array([[0, 0, 16, 16]], dtype=np.float32),
            mask=masks_b,
        )

        iou = detection_iou(detections_a, detections_b)

        np.testing.assert_allclose(iou, mask_iou_batch(masks_a, masks_b))

    def test_returns_oriented_box_iou_when_both_operands_have_obbs(self) -> None:
        """OBB IoU is used even when AABB envelopes are identical."""
        square = np.array([[0, 0], [10, 0], [10, 10], [0, 10]], dtype=np.float32)
        diamond = np.array([[5, 0], [10, 5], [5, 10], [0, 5]], dtype=np.float32)
        shared_xyxy = np.array([[0, 0, 10, 10]], dtype=np.float32)
        detections_a = _detections_from_quads([square], xyxy=shared_xyxy)
        detections_b = _detections_from_quads([diamond], xyxy=shared_xyxy)

        iou = detection_iou(detections_a, detections_b)

        np.testing.assert_allclose(
            iou,
            oriented_box_iou_batch(square[np.newaxis], diamond[np.newaxis]),
        )
        assert iou[0, 0] < box_iou_batch(shared_xyxy, shared_xyxy)[0, 0]

    def test_degenerate_collinear_obb_yields_zero_iou_without_error(self) -> None:
        """A zero-area (collinear) OBB denominator yields 0 IoU, not NaN or a crash."""
        collinear = np.array([[0, 0], [5, 0], [10, 0], [15, 0]], dtype=np.float32)
        square = np.array([[0, 0], [10, 0], [10, 10], [0, 10]], dtype=np.float32)
        detections_a = _detections_from_quads([collinear])
        detections_b = _detections_from_quads([square])

        iou = detection_iou(detections_a, detections_b)

        assert not np.isnan(iou).any()
        np.testing.assert_allclose(iou, np.array([[0.0]]))

    def test_returns_box_iou_when_no_richer_geometry_is_present(self) -> None:
        """AABB IoU is the fallback when neither operand carries richer geometry."""
        detections_a = Detections(
            xyxy=np.array([[0, 0, 10, 10], [20, 20, 30, 30]], dtype=np.float32)
        )
        detections_b = Detections(xyxy=np.array([[5, 5, 15, 15]], dtype=np.float32))

        iou = detection_iou(detections_a, detections_b, OverlapMetric.IOS)

        np.testing.assert_allclose(
            iou,
            box_iou_batch(detections_a.xyxy, detections_b.xyxy, OverlapMetric.IOS),
        )

    def test_returns_empty_matrix_for_empty_detections(self) -> None:
        """Empty Detections return an empty pairwise IoU matrix."""
        empty = Detections(xyxy=np.empty((0, 4), dtype=np.float32))
        non_empty = Detections(xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32))

        iou = detection_iou(empty, non_empty)

        assert iou.shape == (0, 1)

    @pytest.mark.parametrize(
        ("left_kind", "right_kind", "expected_geometry"),
        [
            pytest.param("mask", "obb", "box", id="mask-vs-obb-falls-back-to-box"),
            pytest.param("mask", "aabb", "box", id="mask-vs-aabb-falls-back-to-box"),
            pytest.param("obb", "aabb", "box", id="obb-vs-aabb-falls-back-to-box"),
            pytest.param("compact", "mask", "mask", id="compact-vs-dense-uses-mask"),
            pytest.param(
                "compact", "compact", "mask", id="compact-vs-compact-uses-mask"
            ),
        ],
    )
    def test_mixed_geometry_dispatch_uses_shared_geometry_only_when_available(
        self, left_kind: str, right_kind: str, expected_geometry: str
    ) -> None:
        """Mixed geometry uses AABB fallback unless both operands carry masks."""
        image_shape = (16, 16)
        mask = np.zeros((1, *image_shape), dtype=bool)
        mask[0, 2:10, 2:10] = True
        xyxy = _full_image_xyxy(1, image_shape)
        compact_mask = CompactMask.from_dense(
            masks=mask, xyxy=xyxy, image_shape=image_shape
        )
        obb = _detections_from_quads(
            [np.array([[0, 0], [15, 0], [15, 15], [0, 15]], dtype=np.float32)],
            xyxy=xyxy,
        )
        detections_by_kind = {
            "mask": Detections(xyxy=xyxy, mask=mask),
            "compact": Detections(xyxy=xyxy, mask=compact_mask),
            "obb": obb,
            "aabb": Detections(xyxy=xyxy),
        }
        detections_left = detections_by_kind[left_kind]
        detections_right = detections_by_kind[right_kind]

        result = detection_iou(detections_left, detections_right)

        if expected_geometry == "mask":
            assert detections_left.mask is not None
            assert detections_right.mask is not None
            expected = mask_iou_batch(detections_left.mask, detections_right.mask)
        else:
            expected = box_iou_batch(detections_left.xyxy, detections_right.xyxy)
        np.testing.assert_allclose(result, expected)
