"""Tests for supervision.metrics.utils.object_size — mask pixel-count path."""

from __future__ import annotations

import numpy as np
import pytest

from supervision.config import AREA_DATA_FIELD
from supervision.detection.core import Detections
from supervision.metrics.core import MetricTarget
from supervision.metrics.utils.object_size import (
    SIZE_THRESHOLDS,
    ObjectSizeCategory,
    get_detection_size_category,
    get_mask_size_category,
)


def _reference_mask_size_category(mask: np.ndarray) -> np.ndarray:
    """Reference categorization using a plain ``np.sum`` area."""
    areas = np.sum(mask, axis=(1, 2))
    result = np.full(areas.shape, ObjectSizeCategory.ANY.value)
    sm, lg = SIZE_THRESHOLDS
    result[areas < sm] = ObjectSizeCategory.SMALL.value
    result[(areas >= sm) & (areas < lg)] = ObjectSizeCategory.MEDIUM.value
    result[areas >= lg] = ObjectSizeCategory.LARGE.value
    return result


class TestGetMaskSizeCategory:
    """Verify the count_nonzero-based mask area categorization."""

    @pytest.mark.parametrize(
        ("n", "h", "w", "seed"),
        [
            pytest.param(0, 8, 8, 1, id="empty"),
            pytest.param(1, 40, 40, 2, id="single"),
            pytest.param(6, 64, 64, 3, id="batch"),
            pytest.param(4, 200, 200, 4, id="across-thresholds"),
        ],
    )
    def test_matches_sum_based_reference(
        self, n: int, h: int, w: int, seed: int
    ) -> None:
        """Categories match a plain ``np.sum(axis=(1, 2))`` reference."""
        rng = np.random.default_rng(seed)
        mask = rng.random((n, h, w)) < rng.random()
        np.testing.assert_array_equal(
            get_mask_size_category(mask), _reference_mask_size_category(mask)
        )

    def test_threshold_boundaries(self) -> None:
        """Pixel counts on each SMALL/MEDIUM/LARGE boundary land in the right bin."""
        sm, lg = SIZE_THRESHOLDS
        mask = np.zeros((3, 100, 100), dtype=bool)
        mask[0].flat[: sm - 1] = True  # below SMALL threshold -> SMALL
        mask[1].flat[:sm] = True  # exactly SMALL threshold -> MEDIUM
        mask[2].flat[:lg] = True  # exactly LARGE threshold -> LARGE
        np.testing.assert_array_equal(
            get_mask_size_category(mask),
            [
                ObjectSizeCategory.SMALL.value,
                ObjectSizeCategory.MEDIUM.value,
                ObjectSizeCategory.LARGE.value,
            ],
        )


def test_detection_size_category_prefers_explicit_area_metadata() -> None:
    """Explicit area metadata overrides geometry for size categorization."""
    detections = Detections(
        xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
        data={AREA_DATA_FIELD: np.array([2500.0], dtype=np.float32)},
    )

    result = get_detection_size_category(detections, MetricTarget.BOXES)

    np.testing.assert_array_equal(result, [ObjectSizeCategory.MEDIUM.value])


def test_detection_size_category_falls_back_without_area_metadata() -> None:
    """Geometry remains the fallback when explicit area metadata is absent."""
    detections = Detections(xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32))

    result = get_detection_size_category(detections, MetricTarget.BOXES)

    np.testing.assert_array_equal(result, [ObjectSizeCategory.SMALL.value])


@pytest.mark.parametrize(
    "area_data",
    [
        pytest.param(np.array([[2500.0]], dtype=np.float32), id="two-dimensional"),
    ],
)
def test_detection_size_category_rejects_invalid_area_metadata(
    area_data: np.ndarray,
) -> None:
    """Invalid area metadata is rejected instead of misaligning detections."""
    detections = Detections(
        xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
        data={AREA_DATA_FIELD: area_data},
    )

    with pytest.raises(ValueError, match="area metadata"):
        get_detection_size_category(detections, MetricTarget.BOXES)
