"""检测结果的匹配与绘制。"""

import pytest

pytest.importorskip("PIL")

from engine.detection import (  # noqa: E402
    COLOR_GT,
    COLOR_MATCHED,
    COLOR_UNMATCHED,
    Box,
    draw_boxes,
    from_yolo,
    iou,
    match_boxes,
    nms,
    parse_label,
)


def test_from_yolo_converts_center_to_corners():
    assert from_yolo(2, 0.5, 0.5, 0.2, 0.4) == Box(2, 0.4, 0.3, 0.6, 0.7)


def test_parse_label_reads_boxes_and_empty_files():
    boxes = parse_label("0 0.5 0.5 0.2 0.2\n3 0.1 0.1 0.1 0.1\n")
    assert [b.class_id for b in boxes] == [0, 3]
    assert parse_label("") == []          # 负样本
    assert parse_label("\n\n") == []


@pytest.mark.parametrize(
    "a,b,expected",
    [
        ((0, 0, 1, 1), (0, 0, 1, 1), 1.0),          # 完全重合
        ((0, 0, 1, 1), (2, 2, 3, 3), 0.0),          # 不相交
        ((0, 0, 1, 1), (1, 1, 2, 2), 0.0),          # 只碰到角
        ((0, 0, 2, 2), (1, 1, 3, 3), 1 / 7),        # 交 1，并 7
    ],
)
def test_iou(a, b, expected):
    assert iou(Box(0, *a), Box(0, *b)) == pytest.approx(expected)


def test_matches_by_confidence_first():
    """两个预测抢同一个真值时，置信度高的赢——与 mAP 的规则一致。"""
    truth = [from_yolo(0, 0.5, 0.5, 0.2, 0.2)]
    weak = Box(0, 0.4, 0.4, 0.6, 0.6, 0.3)
    strong = Box(0, 0.41, 0.41, 0.59, 0.59, 0.9)
    match = match_boxes([weak, strong], truth, iou_threshold=0.3)

    assert match.pairs == [(1, 0)]                  # 下标 1 = strong
    assert match.unmatched_predictions == [0]
    assert match.unmatched_ground_truth == []


def test_low_iou_counts_as_both_miss_and_false_positive():
    truth = [from_yolo(0, 0.2, 0.2, 0.1, 0.1)]
    prediction = Box(0, 0.7, 0.7, 0.9, 0.9, 0.9)
    match = match_boxes([prediction], truth)
    assert match.pairs == []
    assert match.unmatched_predictions == [0]       # 误检
    assert match.unmatched_ground_truth == [0]      # 同时也是漏检


def test_class_mismatch_is_not_a_hit_unless_allowed():
    """定位准但分类错：默认不算命中，same_class=False 时可单看位置。"""
    truth = [from_yolo(1, 0.5, 0.5, 0.2, 0.2)]
    prediction = Box(0, 0.4, 0.4, 0.6, 0.6, 0.9)

    assert match_boxes([prediction], truth).pairs == []
    assert match_boxes([prediction], truth, same_class=False).pairs == [(0, 0)]


def test_each_ground_truth_is_claimed_once():
    """一个真值只能被一个预测认领，多出来的预测算误检。"""
    truth = [from_yolo(0, 0.5, 0.5, 0.4, 0.4)]
    predictions = [
        Box(0, 0.31, 0.31, 0.69, 0.69, 0.9),
        Box(0, 0.32, 0.32, 0.68, 0.68, 0.8),
    ]
    match = match_boxes(predictions, truth)
    assert len(match.pairs) == 1
    assert len(match.unmatched_predictions) == 1


def test_no_prediction_on_background_is_clean():
    match = match_boxes([], [])
    assert match == ([], [], [])


def test_multiple_boxes_are_matched_pairwise():
    truth = [from_yolo(0, 0.2, 0.2, 0.1, 0.1), from_yolo(1, 0.8, 0.8, 0.1, 0.1)]
    predictions = [
        Box(1, 0.75, 0.75, 0.85, 0.85, 0.7),
        Box(0, 0.15, 0.15, 0.25, 0.25, 0.9),
    ]
    match = match_boxes(predictions, truth)
    assert sorted(match.pairs) == [(0, 1), (1, 0)]
    assert not match.unmatched_predictions and not match.unmatched_ground_truth


def _colors(image):
    return {pixel for pixel in image.getdata()}


def test_draw_uses_a_distinct_colour_per_role():
    """三种角色三种颜色：绿=真值，蓝=命中的预测，红=没对上的预测。"""
    from PIL import Image

    truth = [from_yolo(0, 0.25, 0.25, 0.2, 0.2), from_yolo(0, 0.75, 0.25, 0.2, 0.2)]
    predictions = [
        Box(0, 0.16, 0.16, 0.34, 0.34, 0.9),      # 命中第一个
        Box(0, 0.5, 0.8, 0.6, 0.9, 0.4),          # 没有对应真值
    ]
    match = match_boxes(predictions, truth, iou_threshold=0.5)
    assert match.pairs == [(0, 0)]

    image = Image.new("RGB", (200, 200), "black")
    draw_boxes(image, truth, predictions, ["blackdot"], match)
    present = _colors(image)
    assert COLOR_GT in present and COLOR_MATCHED in present and COLOR_UNMATCHED in present


def test_draw_handles_boxes_at_the_border():
    """框贴着上边缘时，标签文字要画到框内，不能画到图外去。"""
    from PIL import Image

    image = Image.new("RGB", (64, 64), "black")
    prediction = Box(0, 0.0, 0.0, 0.5, 0.05, 0.9)
    draw_boxes(image, [], [prediction], ["blackdot"], match_boxes([prediction], []))
    assert COLOR_UNMATCHED in _colors(image)


# ---------------- NMS ----------------


def test_nms_removes_duplicates_of_the_same_defect():
    high = Box(0, 0.30, 0.30, 0.50, 0.50, 0.9)
    duplicate = Box(0, 0.31, 0.31, 0.51, 0.51, 0.7)
    assert nms([duplicate, high]) == [1]


def test_nms_keeps_neighbours_below_the_threshold():
    """真正分开的两个缺陷，轻微相邻也要都留住。"""
    a = Box(0, 0.0, 0.0, 0.5, 0.5, 0.9)
    b = Box(0, 0.4, 0.4, 0.9, 0.9, 0.8)                 # IoU 约 0.02
    assert nms([a, b]) == [0, 1]


def test_nms_is_per_class_by_default():
    """不同类别在同一位置的分歧是有意义的，不该互相压制。"""
    a = Box(0, 0.3, 0.3, 0.6, 0.6, 0.9)
    b = Box(4, 0.31, 0.31, 0.61, 0.61, 0.8)
    assert nms([a, b]) == [0, 1]
    assert nms([a, b], per_class=False) == [0]


def test_nms_handles_empty_input():
    assert nms([]) == []


def test_nms_can_suppress_contained_fragments():
    """切片接缝上的残框：与完整框 IoU 很小，却几乎整个落在里面。"""
    full = Box(0, 0.30, 0.50, 0.40, 0.60, 0.9)
    fragment = Box(0, 0.39, 0.50, 0.40, 0.60, 0.8)      # 只剩一条边
    assert iou(full, fragment) < 0.15

    assert nms([full, fragment]) == [0, 1]                        # 只看 IoU 压不掉
    assert nms([full, fragment], contained_threshold=0.8) == [0]  # 按包含比例能压掉