cad-agent/tests/test_heatmap.py

"""`dwg_review.heatmap` 单元测试（纯函数，不依赖 CAD / LLM / Pillow）。"""

from dwg_review import heatmap as hm

WIN = {"min_x": 0.0, "min_y": 0.0, "max_x": 100.0, "max_y": 100.0}


def _grid(boxes, cell=50.0):
    return hm.build_grid(WIN, boxes, cell)


# --- norm_box / entity_boxes ------------------------------------------------


def test_norm_box_valid():
    assert hm.norm_box([0, 0, 10, 20]) == (0.0, 0.0, 10.0, 20.0)


def test_norm_box_swaps_reversed():
    assert hm.norm_box([10, 20, 0, 0]) == (0.0, 0.0, 10.0, 20.0)


def test_norm_box_invalid():
    assert hm.norm_box([1, 2, 3]) is None
    assert hm.norm_box("x") is None
    assert hm.norm_box(None) is None
    assert hm.norm_box([1, 2, 3, "a"]) is None
    assert hm.norm_box([0, 0, float("nan"), 1]) is None


def test_entity_boxes_accepts_single_envelope():
    env = {"items": [{"bbox": [0, 0, 1, 1]}, {"bbox": [2, 2, 3, 3]}], "total": 2}
    assert hm.entity_boxes(env) == [(0.0, 0.0, 1.0, 1.0), (2.0, 2.0, 3.0, 3.0)]


def test_entity_boxes_accepts_iterable_and_skips_invalid():
    envs = [
        {"items": [{"bbox": [0, 0, 1, 1]}, {"nope": 1}, "junk"]},
        {"items": "not-a-list"},
        None,
        {"items": [{"bbox": [5, 5, 6, 6]}]},
    ]
    assert hm.entity_boxes(envs) == [(0.0, 0.0, 1.0, 1.0), (5.0, 5.0, 6.0, 6.0)]


def test_entity_boxes_empty():
    assert hm.entity_boxes(None) == []
    assert hm.entity_boxes({}) == []
    assert hm.entity_boxes({"items": []}) == []


# --- box helpers ------------------------------------------------------------


def test_box_area_and_union():
    assert hm.box_area((0, 0, 10, 5)) == 50.0
    assert hm.box_area((10, 0, 0, 5)) == 0.0  # 退化（未归一）
    assert hm.union_box([(0, 0, 10, 10), (5, 5, 20, 30)]) == (0.0, 0.0, 20.0, 30.0)
    assert hm.union_box([]) is None


def test_div_to_cell_uses_long_edge():
    assert hm.div_to_cell(WIN, 4) == 25.0


# --- build_grid -------------------------------------------------------------


def test_build_grid_shape_and_left_half():
    g = _grid([(0, 0, 50, 100)])
    assert (g["cols"], g["rows"]) == (2, 2)
    assert (g["cell_w"], g["cell_h"]) == (50.0, 50.0)
    # 左半格两行各命中一次，右半格不命中（恰好落在 x=50 边界上不算越界）
    assert [row[0] for row in g["counts"]] == [1, 1]
    assert [row[1] for row in g["counts"]] == [0, 0]
    assert g["coverage"][0][0] == 1.0
    assert g["coverage"][1][0] == 1.0


def test_build_grid_row0_is_top():
    g = _grid([(0, 60, 10, 70)])
    assert g["counts"][0][0] == 1
    assert g["counts"][1][0] == 0


def test_build_grid_degenerate_horizontal_line_counted():
    # 高度为 0 的线（bbox 退化）也要能被计入，且横跨两列
    g = _grid([(0, 25, 100, 25)])
    assert g["counts"][1] == [1, 1]
    assert g["counts"][0] == [0, 0]


def test_build_grid_ignores_boxes_outside_window():
    g = _grid([(200, 200, 300, 300)])
    assert sum(sum(row) for row in g["counts"]) == 0


def test_build_grid_invalid_window():
    g = hm.build_grid({"min_x": 0, "min_y": 0, "max_x": 0, "max_y": 10}, [], 5)
    assert g["cols"] == 0 and g["rows"] == 0 and g["counts"] == []


# --- occupied / cluster -----------------------------------------------------


def test_occupied_cells_row_major():
    g = _grid([(0, 0, 50, 100)])
    assert hm.occupied_cells(g) == [(0, 0), (1, 0)]


def test_cluster_cells_separate_diagonal():
    g = _grid([(5, 5, 15, 15), (85, 85, 95, 95)])
    assert sorted(hm.occupied_cells(g)) == [(0, 1), (1, 0)]
    clusters = hm.cluster_cells(g)
    assert len(clusters) == 2  # 对角相邻不算连通
    boxes = sorted(cl["bbox"] for cl in clusters)
    assert boxes == [(0.0, 0.0, 50.0, 50.0), (50.0, 50.0, 100.0, 100.0)]


def test_cluster_cells_dilate_merges_nearby():
    g = _grid([(5, 5, 15, 15), (85, 85, 95, 95)])
    clusters = hm.cluster_cells(g, dilate=1)
    assert len(clusters) == 1
    # 两格合并后 bbox 为两真实占用格的并集（膨胀出来的空格不参与）
    assert clusters[0]["bbox"] == (0.0, 0.0, 100.0, 100.0)


def test_cluster_cells_sorted_by_count_desc():
    # 左下格被 3 个实体命中，右上格 1 个
    g = _grid([(0, 0, 10, 10), (0, 20, 10, 30), (0, 40, 10, 50), (90, 90, 100, 100)])
    clusters = hm.cluster_cells(g)
    assert clusters[0]["count"] == 3
    assert clusters[0]["bbox"] == (0.0, 0.0, 50.0, 50.0)


def test_cluster_cells_empty_grid():
    assert hm.cluster_cells(_grid([])) == []


# --- windows / stats --------------------------------------------------------


def test_grid_to_windows_applies_margin():
    clusters = [{"bbox": (0.0, 0.0, 100.0, 50.0), "cells": [], "count": 1}]
    (win,) = hm.grid_to_windows(clusters, margin=0.1)
    assert win == {"min_x": -10.0, "min_y": -5.0, "max_x": 110.0, "max_y": 55.0}


def test_grid_to_windows_preserves_order():
    clusters = [
        {"bbox": (0.0, 0.0, 1.0, 1.0), "cells": [], "count": 9},
        {"bbox": (2.0, 2.0, 3.0, 3.0), "cells": [], "count": 1},
    ]
    wins = hm.grid_to_windows(clusters)
    assert [w["min_x"] for w in wins] == [0.0, 2.0]


def test_grid_stats():
    g = _grid([(0, 0, 50, 100)])
    clusters = hm.cluster_cells(g)
    st = hm.grid_stats(g, clusters)
    assert st["rows"] == 2 and st["cols"] == 2
    assert st["occupied"] == 2
    assert st["occupied_ratio"] == 0.5
    assert st["max_count"] == 1
    assert st["clusters"] == 1
    assert st["max_cluster_cells"] == 2


def test_grid_stats_without_clusters():
    st = hm.grid_stats(_grid([]))
    assert st["occupied"] == 0
    assert "clusters" not in st