import torch
from test_model import tiny_config

from objectmodel_v1.losses import ObjectModelCriterion


def reference_dense_targets(criterion, logits, targets, level_index):
    batch, _, height, width = logits.shape
    target_logits = torch.zeros_like(logits)
    target_boxes = torch.zeros(batch, 4, height, width, device=logits.device)
    positive = torch.zeros(batch, height, width, device=logits.device)
    for batch_index, target in enumerate(targets):
        if target["labels"].numel() == 0:
            continue
        centers = target["boxes"][:, :2]
        grid = (centers * torch.tensor([width, height], device=centers.device)).long()
        grid[:, 0].clamp_(0, width - 1)
        grid[:, 1].clamp_(0, height - 1)
        areas = target["boxes"][:, 2] * target["boxes"][:, 3]
        target_levels = torch.where(areas < 0.02, 0, torch.where(areas < 0.15, 1, 2))
        for target_index in torch.where(target_levels == level_index)[0]:
            center_x, center_y = grid[target_index]
            label = target["labels"][target_index]
            candidates = []
            for dy in (-1, 0, 1):
                for dx in (-1, 0, 1):
                    x = int((center_x + dx).clamp(0, width - 1))
                    y = int((center_y + dy).clamp(0, height - 1))
                    candidates.append((dx * dx + dy * dy, x, y))
            for _, x, y in sorted(candidates)[: criterion.dense_topk]:
                target_logits[batch_index, label, y, x] = 1.0
                target_boxes[batch_index, :, y, x] = target["boxes"][target_index]
                positive[batch_index, y, x] = 1.0
    return target_logits, target_boxes, positive


def test_dense_targets_match_reference_with_boundaries_and_collisions():
    criterion = ObjectModelCriterion(tiny_config())
    logits = torch.randn(2, 5, 8, 8)
    targets = [
        {
            "boxes": torch.tensor(
                [
                    [0.01, 0.01, 0.10, 0.10],
                    [0.02, 0.02, 0.11, 0.11],
                    [0.99, 0.99, 0.50, 0.50],
                ]
            ),
            "labels": torch.tensor([1, 2, 3]),
        },
        {"boxes": torch.empty(0, 4), "labels": torch.empty(0, dtype=torch.long)},
    ]

    for level_index in range(3):
        expected = reference_dense_targets(criterion, logits, targets, level_index)
        actual = criterion._dense_targets(logits, targets, level_index)
        for expected_tensor, actual_tensor in zip(expected, actual, strict=True):
            assert torch.equal(expected_tensor, actual_tensor)


def test_dense_targets_match_reference_for_production_topk():
    config = tiny_config()
    config["loss"]["dense_topk"] = 5
    criterion = ObjectModelCriterion(config)
    logits = torch.randn(2, 5, 12, 10)
    targets = [
        {"boxes": torch.rand(9, 4), "labels": torch.randint(0, 5, (9,))},
        {"boxes": torch.rand(4, 4), "labels": torch.randint(0, 5, (4,))},
    ]

    for level_index in range(3):
        expected = reference_dense_targets(criterion, logits, targets, level_index)
        actual = criterion._dense_targets(logits, targets, level_index)
        for expected_tensor, actual_tensor in zip(expected, actual, strict=True):
            assert torch.equal(expected_tensor, actual_tensor)