tests/test_segmentation.py

"""
tests/test_segmentation.py
----------------------------
Tests for the multi-signal sprite sheet segmentation pipeline.
Run with: pytest tests/test_segmentation.py -v
"""

import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from segmentation import (
    GridLayout,
    detect_separator_grid,
    detect_tile_grid,
    segment_grid,
    segment_sheet,
)


def test_single_sprite():
    """A single solid block should be detected as exactly one sprite."""
    mask = np.zeros((20, 20), dtype=bool)
    mask[5:15, 5:15] = True
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1
    assert detected[0].bbox.width == 10
    assert detected[0].bbox.height == 10


def test_disconnected_parts_of_one_sprite_merge():
    """
    A sprite made of two small blobs a couple of pixels apart (e.g. a
    plant's leaf and stem) must be treated as ONE sprite, not two.
    """
    mask = np.zeros((30, 30), dtype=bool)
    mask[5:8, 5:8] = True       # leaf
    mask[10:12, 6:9] = True     # stem, 2px gap
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1


def test_two_far_sprites_stay_separate():
    mask = np.zeros((30, 60), dtype=bool)
    mask[5:10, 5:10] = True
    mask[5:10, 40:45] = True
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 2


def test_two_close_larger_sprites_stay_separate():
    """
    Two side-by-side sprites, each a solid 10x10 block, with only a 3px
    gap must NOT be merged into one -- this is the "yakın iki sprite"
    case from the spec that a naive dilation-based merge would get wrong.
    """
    mask = np.zeros((20, 40), dtype=bool)
    mask[5:15, 5:15] = True
    mask[5:15, 18:28] = True
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 2


def test_hundred_small_sprites_grid():
    """100 small sprites tightly packed in a grid must all be found separately."""
    mask = np.zeros((110, 110), dtype=bool)
    count = 0
    for row in range(10):
        for col in range(10):
            y, x = row * 11, col * 11
            mask[y:y + 8, x:x + 8] = True
            count += 1
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == count


def test_multi_part_plants_in_grid_stay_separate_from_neighbors():
    """
    Six plants, each made of two disconnected blobs, arranged in a grid.
    Each plant's own parts must merge, but neighboring plants must not
    merge with each other.
    """
    mask = np.zeros((60, 60), dtype=bool)
    positions = [(5, 5), (5, 20), (5, 35), (20, 5), (20, 20), (20, 35)]
    for (y, x) in positions:
        mask[y:y + 3, x + 2:x + 5] = True
        mask[y + 4:y + 7, x:x + 3] = True
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 6


def test_empty_mask_returns_no_sprites():
    mask = np.zeros((20, 20), dtype=bool)
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 0


def test_single_pixel_noise_is_dropped():
    """A lone 1px speck should not be reported as a sprite."""
    mask = np.zeros((20, 20), dtype=bool)
    mask[10, 10] = True  # single isolated pixel
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 0


def test_long_horizontal_fence_stays_one_sprite():
    """A long thin horizontal object (e.g. a fence) is a valid single sprite."""
    mask = np.zeros((10, 100), dtype=bool)
    mask[4:6, 5:95] = True
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1
    assert detected[0].bbox.width == 90


def test_long_vertical_object_stays_one_sprite():
    mask = np.zeros((100, 10), dtype=bool)
    mask[5:95, 4:6] = True
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1
    assert detected[0].bbox.height == 90


def test_touching_sprites_merge_since_bboxes_overlap():
    """Two sprites whose pixels are directly touching (0 gap) are one region."""
    mask = np.zeros((20, 20), dtype=bool)
    mask[5:10, 5:10] = True
    mask[5:10, 10:15] = True  # directly adjacent, no gap
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1


def test_sprites_are_numbered_in_reading_order():
    """Detected sprites should be ordered top-to-bottom, left-to-right."""
    mask = np.zeros((40, 40), dtype=bool)
    mask[25:30, 5:10] = True    # bottom-left
    mask[5:10, 25:30] = True    # top-right
    mask[5:10, 5:10] = True     # top-left
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 3
    xs_of_top_row = [d.bbox.x for d in detected if d.bbox.y < 15]
    assert xs_of_top_row == sorted(xs_of_top_row)


def test_low_confidence_flagged_for_oversized_region():
    """A region covering most of the sheet should get a low confidence score."""
    mask = np.ones((20, 20), dtype=bool)
    mask[0, 0] = False  # keep it from being literally 100% to avoid degenerate case
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1
    assert detected[0].confidence < 100


# ---------------------------------------------------------------------------
# Grid-corrected segmentation: animation frames and edge-to-edge tilesets
# ---------------------------------------------------------------------------

def _particle_frame(mask, x0, y0, w, h, rng):
    """A frame drawn as scattered 1px particles, like an explosion's debris."""
    for _ in range(24):
        mask[y0 + rng.integers(0, h), x0 + rng.integers(0, w)] = True
    # anchor the frame's extent so its bbox is predictable
    mask[y0, x0] = True
    mask[y0 + h - 1, x0 + w - 1] = True


def test_particle_animation_frames_are_not_lost_as_noise():
    """
    An animation strip whose frames are clouds of loose 1-2px particles.
    Every one of those particles is individually below MIN_SPRITE_SIZE_PX,
    so the gap pipeline drops them and whole frames vanish. The frame
    gutters must be used instead, giving exactly one sprite per frame with
    no ink thrown away.
    """
    rng = np.random.default_rng(11)
    mask = np.zeros((40, 200), dtype=bool)
    for i in range(5):
        _particle_frame(mask, 5 + i * 40, 5, 30, 30, rng)

    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 5

    covered = np.zeros_like(mask)
    for d in detected:
        covered[d.bbox.y:d.bbox.y2, d.bbox.x:d.bbox.x2] = True
    assert (mask & ~covered).sum() == 0, "particles were dropped from the result"


def test_separator_grid_finds_phase_shifted_pitch():
    """
    A lattice whose gutters sit at 8, 19, 30, ... -- period 11, phase 8.
    Searching phase 0 only would report a much coarser (wrong) period.
    """
    mask = np.zeros((110, 110), dtype=bool)
    for row in range(10):
        for col in range(10):
            mask[row * 11:row * 11 + 8, col * 11:col * 11 + 8] = True
    layout = detect_separator_grid(mask)
    assert layout is not None
    assert layout.cell_width == 11 and layout.cell_height == 11
    assert layout.occupied_cells == 100


def test_frame_with_internal_gap_is_not_halved_by_the_grid():
    """
    Frames separated by a wide gutter, each frame containing a 1px internal
    gap. That gap offers a cut line which is technically legal (the row
    really is empty) but slices every frame in half. The merge-only
    regrouping must keep each frame whole regardless.
    """
    mask = np.zeros((30, 120), dtype=bool)
    for i in range(4):
        x = i * 30
        mask[5:15, x:x + 10] = True       # upper half of the frame
        mask[16:25, x:x + 10] = True      # lower half, 1px gap at row 15
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 4
    for d in detected:
        assert d.bbox.y == 5 and d.bbox.y2 == 25


def test_shattered_frames_regroup_but_never_split_a_merged_sprite():
    """
    The regrouping is merge-only: a sprite the gap stage already joined
    across a background gap keeps its identity even when a grid line falls
    between its parts.
    """
    mask = np.zeros((30, 30), dtype=bool)
    mask[5:8, 5:8] = True       # leaf
    mask[10:12, 6:9] = True     # stem, 2px gap -- one sprite
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1
    assert detected[0].bbox.y == 5
    assert detected[0].bbox.y2 == 12


def _tileset_rgba(tile: int, tiles_across: int, rng):
    """Edge-to-edge tileset: no background anywhere, tiles differ sharply."""
    size = tile * tiles_across
    rgba = np.zeros((size, size, 4), dtype=np.uint8)
    rgba[:, :, 3] = 255
    for r in range(tiles_across):
        for c in range(tiles_across):
            colour = rng.integers(0, 256, size=3, dtype=np.uint16).astype(np.uint8)
            rgba[r * tile:(r + 1) * tile, c * tile:(c + 1) * tile, :3] = colour
    return rgba


def test_tile_grid_recovers_pitch_of_edge_to_edge_tileset():
    """
    A tileset drawn corner to corner has no background gutter at all, so
    the gap pipeline sees one sheet-sized blob. The tile pitch has to come
    from the sheet's own edge periodicity -- and must be the FUNDAMENTAL
    pitch, not a multiple of it (16, never 32 or 80).
    """
    rng = np.random.default_rng(5)
    rgba = _tileset_rgba(16, 10, rng)
    mask = np.ones(rgba.shape[:2], dtype=bool)

    layout = detect_tile_grid(rgba, mask)
    assert layout is not None
    assert layout.cell_width == 16 and layout.cell_height == 16

    detected, _, _ = segment_sheet(mask, rgba=rgba)
    assert len(detected) == 100


def test_tile_grid_needs_pixels_and_is_skipped_without_them():
    """Mask-only callers keep the original behaviour, never a tile grid."""
    rng = np.random.default_rng(5)
    rgba = _tileset_rgba(16, 10, rng)
    mask = np.ones(rgba.shape[:2], dtype=bool)
    assert detect_tile_grid(None, mask) is None
    detected, _, _ = segment_sheet(mask)
    assert len(detected) == 1


def test_animation_sheet_does_not_get_a_tile_grid():
    """
    A sheet with a strong pitch on ONE axis only (a horizontal strip) is
    not a tileset; requiring both axes to agree is what keeps it out.
    """
    rng = np.random.default_rng(9)
    rgba = np.zeros((64, 512, 4), dtype=np.uint8)
    for i in range(8):
        colour = rng.integers(0, 256, size=3, dtype=np.uint16).astype(np.uint8)
        rgba[10:54, i * 64 + 8:i * 64 + 56, :3] = colour
        rgba[10:54, i * 64 + 8:i * 64 + 56, 3] = 255
    mask = rgba[:, :, 3] > 0
    assert detect_tile_grid(rgba, mask) is None
    detected, _, _ = segment_sheet(mask, rgba=rgba)
    assert len(detected) == 8


def test_two_close_sprites_are_not_fused_by_a_grid():
    """
    The grid paths must never merge a couple of independent sprites that
    the gap pipeline correctly kept apart.
    """
    mask = np.zeros((20, 40), dtype=bool)
    mask[5:15, 5:15] = True
    mask[5:15, 18:28] = True
    rgba = np.zeros((20, 40, 4), dtype=np.uint8)
    rgba[mask] = (200, 120, 60, 255)
    detected, _, _ = segment_sheet(mask, rgba=rgba)
    assert len(detected) == 2


def test_grid_backed_sprites_are_not_flagged_low_confidence_for_being_sparse():
    """
    A validated grid cut is more trustworthy than the gap heuristic, so a
    frame that is legitimately a sparse cloud of particles must not be
    pushed below the manual-review threshold just for being sparse.
    """
    rng = np.random.default_rng(11)
    mask = np.zeros((40, 200), dtype=bool)
    for i in range(5):
        _particle_frame(mask, 5 + i * 40, 5, 30, 30, rng)
    detected, _, _ = segment_sheet(mask)
    assert detected and all(d.confidence >= 70 for d in detected)


# ---------------------------------------------------------------------------
# Composite objects spanning several grid cells
# ---------------------------------------------------------------------------

def test_object_spanning_several_cells_becomes_one_sprite():
    """
    A ladder drawn down three grid cells is ONE object. Cut per cell it
    yields three fragments that mean nothing on their own.
    """
    mask = np.zeros((64, 64), dtype=bool)
    mask[8:56, 20:28] = True                       # ladder, 3 cells tall
    for (r, c) in [(0, 0), (0, 3), (3, 3)]:        # unrelated single tiles
        mask[r * 16 + 3:r * 16 + 13, c * 16 + 3:c * 16 + 13] = True

    layout = GridLayout(16, 16, 4, 4, 0, 0, 90.0, "tile")
    detected = segment_grid(mask, layout)
    tall = [d for d in detected if d.bbox.height > 20]
    assert len(tall) == 1
    assert tall[0].bbox.height == 48
    assert tall[0].component_count == 4
    assert any("Composite" in r for r in tall[0].confidence_reasons)
    # The unrelated tiles are untouched.
    assert len(detected) == 4


def test_tileable_terrain_is_not_fused_into_one_object():
    """
    A large connected mass of terrain is exactly what the grid cut exists
    to break up. It is connected the same way a composite object is, so
    only its extent can tell them apart -- it must stay per-cell.
    """
    mask = np.zeros((80, 80), dtype=bool)
    mask[0:80, 0:80] = True     # 5x5 cells of solid, connected floor
    layout = GridLayout(16, 16, 5, 5, 0, 0, 90.0, "tile")
    detected = segment_grid(mask, layout)
    assert len(detected) == 25
    assert all(d.component_count == 1 for d in detected)


def test_adjacent_flat_tiles_are_not_merged_by_touching_alone():
    """
    Two flat filler tiles touch along their whole shared edge, exactly as a
    composite object's halves do. Keying the merge on connected-component
    identity plus extent is what keeps them apart from a real object --
    here they are part of one big terrain mass.
    """
    mask = np.zeros((48, 96), dtype=bool)
    mask[:, :] = True
    layout = GridLayout(16, 16, 6, 3, 0, 0, 90.0, "tile")
    detected = segment_grid(mask, layout)
    assert len(detected) == 18


def test_composite_merge_can_be_switched_off():
    mask = np.zeros((48, 16), dtype=bool)
    mask[4:44, 4:12] = True
    layout = GridLayout(16, 16, 1, 3, 0, 0, 90.0, "tile")
    assert len(segment_grid(mask, layout)) == 1
    assert len(segment_grid(mask, layout, merge_composites=False)) == 3


if __name__ == "__main__":
    # Minimal manual runner for environments without pytest installed.
    import traceback
    tests = [(name, fn) for name, fn in list(globals().items()) if name.startswith("test_") and callable(fn)]
    passed, failed = 0, 0
    for name, fn in tests:
        try:
            fn()
            print(f"PASS  {name}")
            passed += 1
        except Exception:
            print(f"FAIL  {name}")
            traceback.print_exc()
            failed += 1
    print(f"\n{passed} passed, {failed} failed")
    sys.exit(1 if failed else 0)