from scripts.audit_yolo_label_relationships import Box, audit_boxes, classify_pair


DEFAULTS = {
    "near_duplicate_iou": 0.9,
    "containment_threshold": 0.98,
    "max_nested_area_ratio": 4.0,
}


def test_classifies_exact_near_and_possible_nested_relationships() -> None:
    exact = Box(0, 0.5, 0.5, 0.2, 0.2)
    near = Box(0, 0.501, 0.5, 0.2, 0.2)
    outer = Box(0, 0.5, 0.5, 0.18, 0.18)
    inner = Box(0, 0.5, 0.5, 0.1, 0.1)

    assert classify_pair(exact, exact, **DEFAULTS) == ("exact_duplicate", 1.0)
    assert classify_pair(exact, near, **DEFAULTS)[0] == "near_duplicate"
    assert classify_pair(outer, inner, **DEFAULTS) == ("possible_nested", 1.0)


def test_ignores_other_classes_and_non_overlapping_boxes() -> None:
    first = Box(0, 0.2, 0.2, 0.1, 0.1)
    other_class = Box(1, 0.2, 0.2, 0.1, 0.1)
    distant = Box(0, 0.8, 0.8, 0.1, 0.1)

    assert classify_pair(first, other_class, **DEFAULTS) is None
    assert classify_pair(first, distant, **DEFAULTS) is None


def test_audit_boxes_returns_pair_indices_and_coordinates() -> None:
    first = Box(0, 0.5, 0.5, 0.2, 0.2)
    second = Box(0, 0.5, 0.5, 0.2, 0.2)
    relationships = audit_boxes([first, second], **DEFAULTS)

    assert relationships == [
        {
            "relationship": "exact_duplicate",
            "score": 1.0,
            "first_index": 0,
            "second_index": 1,
            "first_box": [0, 0.5, 0.5, 0.2, 0.2],
            "second_box": [0, 0.5, 0.5, 0.2, 0.2],
        }
    ]