import pytest

pytest.importorskip("numpy")
pytest.importorskip("torch")

import numpy as np  # noqa: E402

from datasets.target_encoder import CenterTargetEncoder, gaussian_radius  # noqa: E402


def test_gaussian_radius_is_positive_and_scales():
    assert gaussian_radius(40, 40) > gaussian_radius(10, 10) > 0


def test_encoder_marks_single_center():
    encoder = CenterTargetEncoder(num_classes=3, stride=4)
    boxes = np.array([[40.0, 40.0, 80.0, 80.0]])
    labels = np.array([1])

    target = encoder(boxes, labels, (256, 256))

    assert target["heatmap"].shape == (3, 64, 64)
    assert target["mask"].sum() == 1
    # 中心 (60, 60) / stride 4 = (15, 15)
    assert target["mask"][15, 15] == 1
    assert target["heatmap"][1].max() == pytest.approx(1.0)
    assert target["heatmap"][0].max() == 0
    assert target["wh"][:, 15, 15].tolist() == pytest.approx([10.0, 10.0])


def test_encoder_skips_invalid_and_out_of_range_boxes():
    encoder = CenterTargetEncoder(num_classes=2, stride=4)
    boxes = np.array([[10.0, 10.0, 5.0, 5.0], [900.0, 900.0, 950.0, 950.0]])
    target = encoder(boxes, np.array([0, 1]), (256, 256))
    assert target["mask"].sum() == 0


def test_encoder_respects_max_objects():
    encoder = CenterTargetEncoder(num_classes=1, stride=4, max_objects=2)
    boxes = np.array([[i * 20.0, 0.0, i * 20.0 + 16, 16.0] for i in range(5)])
    target = encoder(boxes, np.zeros(5, dtype=np.int64), (256, 256))
    assert target["mask"].sum() == 2