import pytest
import torch

from spelunkai_training.losses import detection_loss, focal_loss, masked_l1_loss


def test_focal_loss_is_zero_for_a_perfect_prediction():
    target = torch.zeros(1, 1, 8, 8)
    target[0, 0, 3, 3] = 1.0
    pred = target.clone()

    loss = focal_loss(pred, target)
    assert torch.isclose(loss, torch.tensor(0.0), atol=1e-6)


def test_focal_loss_is_positive_for_a_wrong_prediction():
    target = torch.zeros(1, 1, 8, 8)
    target[0, 0, 3, 3] = 1.0
    pred = torch.full_like(target, 0.1)

    loss = focal_loss(pred, target)
    assert loss.item() > 0


def test_masked_l1_loss_only_counts_masked_pixels():
    pred = torch.zeros(1, 2, 4, 4)
    target = torch.zeros(1, 2, 4, 4)
    mask = torch.zeros(1, 4, 4)

    # mismatch outside the mask - should not affect the loss
    pred[0, :, 0, 0] = 100.0
    loss = masked_l1_loss(pred, target, mask)
    assert torch.isclose(loss, torch.tensor(0.0))

    # mismatch inside the mask - should affect the loss. Both wh channels
    # contribute (|5-3|=2 each), summed per masked pixel then divided by the
    # number of masked pixels (1) -> 4, not 2.
    mask[0, 1, 1] = 1.0
    pred[0, :, 1, 1] = 5.0
    target[0, :, 1, 1] = 3.0
    loss = masked_l1_loss(pred, target, mask)
    assert torch.isclose(loss, torch.tensor(4.0))


def test_detection_loss_combines_both_terms():
    heatmap_target = torch.zeros(1, 1, 8, 8)
    heatmap_target[0, 0, 2, 2] = 1.0
    wh_target = torch.zeros(1, 2, 8, 8)
    wh_target[0, :, 2, 2] = 10.0
    mask = torch.zeros(1, 8, 8)
    mask[0, 2, 2] = 1.0

    pred_heatmap = torch.full_like(heatmap_target, 0.3)
    pred_wh = torch.zeros_like(wh_target)

    total, parts = detection_loss(pred_heatmap, pred_wh, heatmap_target, wh_target, mask)
    assert total.item() > 0
    assert parts["total_loss"] == pytest.approx(parts["heatmap_loss"] + 0.1 * parts["wh_loss"], rel=1e-4)