"""固定边长切片：只裁剪与镜像填充，不缩放。"""

import pytest

np = pytest.importorskip("numpy")

from engine.detection import Box, from_yolo  # noqa: E402
from engine.tiling import axis_starts, cut, plan_tiles, remap_boxes, to_source  # noqa: E402

SIZE = 160


# ---------------- 切片网格 ----------------


@pytest.mark.parametrize(
    "length,overlap,expected",
    [
        (160, 0, [0]),              # 正好一片
        (128, 0, [0]),              # 不足，填充
        (320, 0, [0, 160]),         # 整两片
        (200, 0, [0, 40]),          # 最后一片贴边，与前片重叠 120
        (480, 32, [0, 128, 256, 320]),   # 步长 128，末片贴边
    ],
)
def test_axis_starts(length, overlap, expected):
    assert axis_starts(length, SIZE, overlap) == expected


def test_last_tile_always_reaches_the_edge():
    """末尾那条不足一步的区域不能被漏掉。"""
    for length in range(161, 700):
        starts = axis_starts(length, SIZE, 16)
        assert starts[-1] + SIZE == length


def test_plan_tiles_covers_a_large_image():
    tiles = plan_tiles(524, 959, SIZE, overlap=0)
    assert len(tiles) == len(axis_starts(524, SIZE, 0)) * len(axis_starts(959, SIZE, 0))
    assert all(t.pad_left == 0 and t.pad_top == 0 for t in tiles)   # 大图纯裁剪


def test_small_image_is_centred_and_padded():
    tile, = plan_tiles(116, 83, SIZE)
    assert (tile.crop_width, tile.crop_height) == (116, 83)
    assert tile.pad_left == (160 - 116) // 2 == 22
    assert tile.pad_top == (160 - 83) // 2 == 38


@pytest.mark.parametrize("bad", [0, -1])
def test_rejects_invalid_size(bad):
    with pytest.raises(ValueError, match="size"):
        plan_tiles(100, 100, bad)


@pytest.mark.parametrize("bad", [-1, 160, 200])
def test_rejects_invalid_overlap(bad):
    with pytest.raises(ValueError, match="overlap"):
        plan_tiles(300, 300, SIZE, bad)


# ---------------- 取片 ----------------


def test_cut_crops_without_resizing():
    """大图取片必须是原像素的子块，一个插值都不能有。"""
    image = np.arange(300 * 400, dtype=np.uint8).reshape(300, 400) % 251
    tile = plan_tiles(400, 300, SIZE, overlap=0)[1]      # 第二片，x 有偏移
    patch = cut(image, tile, SIZE)

    assert patch.shape == (SIZE, SIZE)
    expected = image[tile.y:tile.y + SIZE, tile.x:tile.x + SIZE]
    assert np.array_equal(patch, expected)


def test_cut_mirrors_small_images():
    image = np.arange(1, 5, dtype=np.uint8).reshape(1, 4)    # 1x4
    tile, = plan_tiles(4, 1, 4)
    patch = cut(image, tile, 4)
    assert patch.shape == (4, 4)
    # 竖直方向镜像：原始那一行在中间，上下是它的倒影（reflect 不重复边缘行）
    assert np.array_equal(patch[tile.pad_top], image[0])
    assert patch.dtype == image.dtype


def test_cut_keeps_channels():
    image = np.zeros((50, 60, 3), dtype=np.uint8)
    tile, = plan_tiles(60, 50, SIZE)
    assert cut(image, tile, SIZE).shape == (SIZE, SIZE, 3)


def test_cut_of_an_exact_size_image_is_the_image_itself():
    image = np.random.default_rng(0).integers(0, 255, (SIZE, SIZE), dtype=np.uint8)
    tile, = plan_tiles(SIZE, SIZE, SIZE)
    assert np.array_equal(cut(image, tile, SIZE), image)


# ---------------- 框的重映射 ----------------


def test_box_shifts_with_padding():
    """小图居中填充后，框只是整体平移。"""
    width, height = 80, 80
    box = from_yolo(3, 0.5, 0.5, 0.25, 0.25)            # 原图中心 20x20 的框
    tile, = plan_tiles(width, height, SIZE)
    mapped, = remap_boxes([box], tile, width, height, SIZE)

    assert mapped.class_id == 3
    # 原图中心落在片中心
    assert (mapped.x1 + mapped.x2) / 2 == pytest.approx(0.5)
    # 尺寸不变：原来 20 像素，片上还是 20 像素（占 160 的 1/8）
    assert (mapped.x2 - mapped.x1) * SIZE == pytest.approx(20)


def test_box_outside_the_tile_is_dropped():
    width = height = 400
    box = from_yolo(0, 0.9, 0.9, 0.05, 0.05)            # 右下角
    first = plan_tiles(width, height, SIZE, overlap=0)[0]   # 左上片
    assert remap_boxes([box], first, width, height, SIZE) == []


def test_partially_visible_box_is_clipped_or_dropped():
    """跨片的框：露得多就裁剪保留，露得少就丢弃。"""
    width = height = 320
    # 框中心正好在 x=160 的切割线上，左右各露一半
    box = from_yolo(1, 160 / width, 0.25, 40 / width, 40 / height)
    left, right = plan_tiles(width, height, SIZE, overlap=0)[:2]

    kept_left, = remap_boxes([box], left, width, height, SIZE, min_visibility=0.4)
    assert kept_left.x2 == pytest.approx(1.0)            # 裁到片的右边缘
    assert (kept_left.x2 - kept_left.x1) * SIZE == pytest.approx(20)

    # 阈值提到 0.6 时，只露一半就不够了
    assert remap_boxes([box], left, width, height, SIZE, min_visibility=0.6) == []
    assert remap_boxes([box], right, width, height, SIZE, min_visibility=0.6) == []


def test_mirrored_region_gets_no_box():
    """镜像出来的是复制品，不能再给它标框——否则等于告诉模型那里真有缺陷。"""
    width = height = 80
    box = from_yolo(0, 0.5, 0.5, 0.5, 0.5)
    tile, = plan_tiles(width, height, SIZE)
    mapped = remap_boxes([box], tile, width, height, SIZE)

    assert len(mapped) == 1                              # 只有一个，不是四个倒影
    # 框整体落在原图内容的范围内
    assert mapped[0].x1 * SIZE >= tile.pad_left - 1e-6
    assert mapped[0].x2 * SIZE <= tile.pad_left + tile.crop_width + 1e-6


def test_zero_area_box_is_ignored():
    tile, = plan_tiles(100, 100, SIZE)
    degenerate = Box(0, 0.5, 0.5, 0.5, 0.5)
    assert remap_boxes([degenerate], tile, 100, 100, SIZE) == []


# ---------------- 反向映射 ----------------


@pytest.mark.parametrize("size_wh", [(80, 60), (400, 300), (160, 160), (524, 959)])
def test_round_trip_back_to_source(size_wh):
    """片上的框映射回原图，必须落回原来的像素位置。

    框放在某一片裁剪区域的正中，确保它完整落在该片内——无重叠切片时，
    原图正中的框反而可能横跨切割线，不适合用来做往返校验。
    """
    width, height = size_wh
    for tile in plan_tiles(width, height, SIZE, overlap=0):
        cx = tile.x + tile.crop_width / 2
        cy = tile.y + tile.crop_height / 2
        half_w = min(tile.crop_width, 20) / 4
        half_h = min(tile.crop_height, 20) / 4
        px = [cx - half_w, cy - half_h, cx + half_w, cy + half_h]
        box = Box(2, px[0] / width, px[1] / height, px[2] / width, px[3] / height)

        mapped, = remap_boxes([box], tile, width, height, SIZE, min_visibility=0.99)
        back = to_source(mapped, tile, SIZE)
        assert [back.x1, back.y1, back.x2, back.y2] == pytest.approx(px, abs=1e-6)


def test_score_survives_both_directions():
    tile, = plan_tiles(100, 100, SIZE)
    box = Box(1, 0.4, 0.4, 0.6, 0.6, 0.83)
    mapped, = remap_boxes([box], tile, 100, 100, SIZE)
    assert mapped.score == 0.83
    assert to_source(mapped, tile, SIZE).score == 0.83


# ---------------- 切片推理的拼回 ----------------


def test_merge_deduplicates_across_the_seam():
    """重叠区里同一处缺陷会被两片各检出一次，合并后应只剩一个。"""
    from engine.tiling import merge_tiles

    width = height = 300
    tiles = plan_tiles(width, height, SIZE, overlap=64)
    # 让每片都"检出"同一个真实位置的框（像素 150..170）
    per_tile = []
    for tile in tiles:
        px1, py1, px2, py2 = 150.0, 150.0, 170.0, 170.0
        bx1 = (px1 - tile.x + tile.pad_left) / SIZE
        by1 = (py1 - tile.y + tile.pad_top) / SIZE
        bx2 = (px2 - tile.x + tile.pad_left) / SIZE
        by2 = (py2 - tile.y + tile.pad_top) / SIZE
        inside = 0 <= bx1 and bx2 <= 1 and 0 <= by1 and by2 <= 1
        per_tile.append([Box(0, bx1, by1, bx2, by2, 0.9)] if inside else [])

    assert sum(len(b) for b in per_tile) >= 2          # 确实有多片检出
    merged = merge_tiles(per_tile, tiles, SIZE, width, height)
    assert len(merged) == 1
    assert merged[0].x1 * width == pytest.approx(150, abs=1e-6)
    assert merged[0].x2 * width == pytest.approx(170, abs=1e-6)


def test_merge_keeps_genuinely_separate_defects():
    from engine.tiling import merge_tiles

    width = height = 320
    tiles = plan_tiles(width, height, SIZE, overlap=0)
    per_tile = [[] for _ in tiles]
    per_tile[0] = [Box(0, 0.1, 0.1, 0.2, 0.2, 0.9)]     # 左上片
    per_tile[3] = [Box(0, 0.8, 0.8, 0.9, 0.9, 0.8)]     # 右下片
    assert len(merge_tiles(per_tile, tiles, SIZE, width, height)) == 2


def test_merge_clips_boxes_that_reach_into_mirrored_padding():
    """小图填充后，模型可能在镜像区出框；那里不是真实区域，要裁掉。"""
    from engine.tiling import merge_tiles

    width = height = 80
    tiles = plan_tiles(width, height, SIZE)
    # 一个横跨整片的框，两端都落在镜像填充里
    merged = merge_tiles([[Box(0, 0.0, 0.4, 1.0, 0.6, 0.9)]], tiles, SIZE, width, height)

    assert len(merged) == 1
    assert merged[0].x1 == pytest.approx(0.0)
    assert merged[0].x2 == pytest.approx(1.0)           # 裁到原图边界，不超出
    assert 0.0 <= merged[0].y1 <= merged[0].y2 <= 1.0


def test_merge_drops_boxes_entirely_inside_the_mirror():
    from engine.tiling import merge_tiles

    width = height = 80
    tiles = plan_tiles(width, height, SIZE)              # pad 40，原图占 40..120
    ghost = Box(0, 0.05, 0.05, 0.2, 0.2, 0.9)           # 像素 8..32，全在镜像区
    assert merge_tiles([[ghost]], tiles, SIZE, width, height) == []


def test_tiled_predict_wires_geometry_to_the_model():
    """端到端：切片 -> 逐片推理 -> 拼回，位置要落在真实缺陷上。"""
    from engine.tiling import tiled_predict

    image = np.zeros((300, 400), dtype=np.uint8)
    image[200:220, 100:130] = 255                       # 一块亮斑

    def predict_batch(patches):
        """假模型：把片里最亮的连通区域当成检出。"""
        out = []
        for patch in patches:
            ys, xs = np.nonzero(patch == 255)
            if len(xs) == 0:
                out.append([])
                continue
            out.append([Box(0, xs.min() / SIZE, ys.min() / SIZE,
                            (xs.max() + 1) / SIZE, (ys.max() + 1) / SIZE, 0.9)])
        return out

    merged = tiled_predict(predict_batch, image, SIZE, overlap=32)
    assert len(merged) == 1
    assert merged[0].x1 * 400 == pytest.approx(100, abs=1)
    assert merged[0].y1 * 300 == pytest.approx(200, abs=1)
    assert merged[0].x2 * 400 == pytest.approx(130, abs=1)
    assert merged[0].y2 * 300 == pytest.approx(220, abs=1)


def test_tiled_predict_on_a_small_image_uses_one_tile():
    from engine.tiling import tiled_predict

    seen = []

    def predict_batch(patches):
        seen.append(len(patches))
        assert all(p.shape[:2] == (SIZE, SIZE) for p in patches)
        return [[] for _ in patches]

    tiled_predict(predict_batch, np.zeros((83, 116), dtype=np.uint8), SIZE)
    assert seen == [1]