import numpy as np

from spelunkai_training.targets import Box, encode_targets, gaussian_radius


def test_gaussian_radius_is_positive_for_a_reasonable_box():
    assert gaussian_radius(height=20, width=15) > 0


def test_encode_targets_places_peak_at_box_center():
    box = Box(class_index=0, x=8, y=16, width=8, height=8)  # center = (12, 20)
    targets = encode_targets([box], num_classes=1, image_width=64, image_height=64, output_stride=4)

    # center in output space: (12/4, 20/4) = (3, 5) -> (x=3, y=5)
    assert targets.heatmap.shape == (1, 16, 16)
    assert targets.heatmap[0, 5, 3] == 1.0
    assert targets.heatmap.max() == 1.0


def test_encode_targets_sets_wh_and_mask_only_at_center():
    box = Box(class_index=0, x=8, y=16, width=8, height=12)
    targets = encode_targets([box], num_classes=1, image_width=64, image_height=64, output_stride=4)

    assert targets.wh[0, 5, 3] == 8
    assert targets.wh[1, 5, 3] == 12
    assert targets.mask[5, 3] == 1.0
    assert targets.mask.sum() == 1.0


def test_encode_targets_rejects_out_of_range_class_index():
    box = Box(class_index=5, x=0, y=0, width=4, height=4)
    try:
        encode_targets([box], num_classes=2, image_width=32, image_height=32, output_stride=4)
        assert False, "expected ValueError"
    except ValueError:
        pass


def test_encode_targets_with_no_boxes_is_all_zero():
    targets = encode_targets([], num_classes=3, image_width=32, image_height=32, output_stride=4)
    assert targets.heatmap.shape == (3, 8, 8)
    assert np.all(targets.heatmap == 0)
    assert np.all(targets.mask == 0)