SPRUCE-scraper/tests/test_stitch_mosaic.py

"""Tests for scripts/stitch_mosaic_from_tiles.py — grid stitch + CLI."""

from __future__ import annotations

import json
import subprocess
import sys
from pathlib import Path

from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
SCRIPT = ROOT / "scripts" / "stitch_mosaic_from_tiles.py"


def _rgb_near(
    got: tuple[int, ...], expected: tuple[int, int, int], tol: int = 8
) -> bool:
    return all(abs(a - b) <= tol for a, b in zip(got, expected))


def _write_scan(
    scan: Path,
    *,
    nx: int,
    ny: int,
    tile_size: int,
    colors: dict[tuple[int, int], tuple[int, int, int]],
) -> None:
    tiles = scan / "tiles"
    tiles.mkdir(parents=True)
    meta = {
        "nx": nx,
        "ny": ny,
        "scan_id": 999,
        "scan_time": "2024-01-01 00:00:00",
    }
    (scan / "metadata.json").write_text(json.dumps(meta), encoding="utf-8")
    for r in range(ny):
        for c in range(nx):
            rgb = colors[(r, c)]
            im = Image.new("RGB", (tile_size, tile_size), rgb)
            im.save(tiles / f"tile_r{r}_c{c}.jpg", quality=95)


def test_stitch_cli_2x2_default_layout(tmp_path: Path) -> None:
    scan = tmp_path / "scan"
    colors = {
        (0, 0): (255, 0, 0),
        (0, 1): (0, 255, 0),
        (1, 0): (0, 0, 255),
        (1, 1): (255, 255, 0),
    }
    _write_scan(scan, nx=2, ny=2, tile_size=8, colors=colors)
    out = scan / "out.jpg"
    subprocess.run(
        [
            sys.executable,
            str(SCRIPT),
            str(scan),
            "-o",
            str(out),
            "--no-flip-x",
            "--tile-gap",
            "0",
        ],
        check=True,
        cwd=str(ROOT),
    )
    im = Image.open(out).convert("RGB")
    assert im.size == (16, 16)
    assert _rgb_near(im.getpixel((4, 4)), (255, 0, 0))
    assert _rgb_near(im.getpixel((12, 4)), (0, 255, 0))
    assert _rgb_near(im.getpixel((4, 12)), (0, 0, 255))
    assert _rgb_near(im.getpixel((12, 12)), (255, 255, 0))


def test_stitch_cli_flip_y_moves_row0_to_bottom(tmp_path: Path) -> None:
    scan = tmp_path / "scan"
    colors = {
        (0, 0): (255, 0, 0),
        (0, 1): (255, 0, 0),
        (1, 0): (0, 0, 255),
        (1, 1): (0, 0, 255),
    }
    _write_scan(scan, nx=2, ny=2, tile_size=4, colors=colors)
    out = scan / "out.jpg"
    subprocess.run(
        [
            sys.executable,
            str(SCRIPT),
            str(scan),
            "-o",
            str(out),
            "--no-flip-x",
            "--flip-y",
            "--tile-gap",
            "0",
        ],
        check=True,
        cwd=str(ROOT),
    )
    im = Image.open(out).convert("RGB")
    # Without flip, top row is red; with flip-y row 0 is pasted at y=4
    assert _rgb_near(im.getpixel((2, 2)), (0, 0, 255))
    assert _rgb_near(im.getpixel((2, 6)), (255, 0, 0))


def test_stitch_missing_tile_fails(tmp_path: Path) -> None:
    scan = tmp_path / "scan"
    tiles = scan / "tiles"
    tiles.mkdir(parents=True)
    (scan / "metadata.json").write_text(
        json.dumps({"nx": 2, "ny": 2, "scan_time": "2024-01-01"}),
        encoding="utf-8",
    )
    im = Image.new("RGB", (4, 4), (128, 128, 128))
    im.save(tiles / "tile_r0_c0.jpg", quality=95)
    r = subprocess.run(
        [sys.executable, str(SCRIPT), str(scan)],
        cwd=str(ROOT),
        capture_output=True,
        text=True,
    )
    assert r.returncode != 0
    assert "Missing" in r.stderr or "Missing" in r.stdout


def test_stitch_allow_missing_black_hole(tmp_path: Path) -> None:
    scan = tmp_path / "scan"
    tiles = scan / "tiles"
    tiles.mkdir(parents=True)
    (scan / "metadata.json").write_text(
        json.dumps({"nx": 2, "ny": 2, "scan_time": "2024-01-01"}),
        encoding="utf-8",
    )
    im = Image.new("RGB", (4, 4), (255, 0, 0))
    im.save(tiles / "tile_r0_c0.jpg", quality=95)
    out = scan / "out.jpg"
    subprocess.run(
        [
            sys.executable,
            str(SCRIPT),
            str(scan),
            "-o",
            str(out),
            "--no-flip-x",
            "--allow-missing",
            "--tile-gap",
            "0",
        ],
        check=True,
        cwd=str(ROOT),
    )
    im_out = Image.open(out).convert("RGB")
    assert _rgb_near(im_out.getpixel((2, 2)), (255, 0, 0))
    # JPEG output may not keep (0,0,0) exact on chroma-subsampled black regions
    hole = im_out.getpixel((6, 6))
    assert max(hole) <= 20, hole


def test_compare_mosaic_fit(tmp_path: Path) -> None:
    scan = tmp_path / "scan"
    colors = {(0, 0): (100, 150, 200)}
    _write_scan(scan, nx=1, ny=1, tile_size=20, colors=colors)
    ref = Image.new("RGB", (10, 10), (100, 150, 200))
    ref.save(scan / "mosaic.jpg", quality=95)
    out = scan / "out.jpg"
    r = subprocess.run(
        [
            sys.executable,
            str(SCRIPT),
            str(scan),
            "-o",
            str(out),
            "--compare-mosaic",
            "--fit",
        ],
        check=True,
        cwd=str(ROOT),
        capture_output=True,
        text=True,
    )
    assert "MAE=" in r.stdout
    assert "exact_pixels=" in r.stdout


def test_stitch_tile_gap_inserts_gutter(tmp_path: Path) -> None:
    scan = tmp_path / "scan"
    colors = {(0, 0): (255, 0, 0), (0, 1): (0, 255, 0), (1, 0): (0, 0, 255), (1, 1): (128, 128, 128)}
    _write_scan(scan, nx=2, ny=2, tile_size=8, colors=colors)
    out = scan / "out.jpg"
    subprocess.run(
        [
            sys.executable,
            str(SCRIPT),
            str(scan),
            "-o",
            str(out),
            "--no-flip-x",
            "--tile-gap",
            "2",
        ],
        check=True,
        cwd=str(ROOT),
    )
    im = Image.open(out).convert("RGB")
    assert im.size == (18, 18)
    gutter = im.getpixel((8, 4))
    assert min(gutter) >= 250, gutter
    assert _rgb_near(im.getpixel((4, 4)), (255, 0, 0))


def test_stitch_default_flip_x_swaps_left_right_columns(tmp_path: Path) -> None:
    """Default matches RootView mosaic: col 0 is placed on the right."""
    scan = tmp_path / "scan"
    colors = {
        (0, 0): (255, 0, 0),
        (0, 1): (0, 255, 0),
        (1, 0): (0, 0, 255),
        (1, 1): (255, 255, 0),
    }
    _write_scan(scan, nx=2, ny=2, tile_size=8, colors=colors)
    out = scan / "out.jpg"
    subprocess.run(
        [sys.executable, str(SCRIPT), str(scan), "-o", str(out)],
        check=True,
        cwd=str(ROOT),
    )
    im = Image.open(out).convert("RGB")
    assert _rgb_near(im.getpixel((4, 4)), (0, 255, 0))
    assert _rgb_near(im.getpixel((12, 4)), (255, 0, 0))