atm/tests/test_layout.py

import numpy as np
import pytest

from atm.config import ROI
from atm.layout import _strips_match, detect_strips


PALETTE = {
    "turquoise": ((0, 253, 253), 60.0),
    "yellow": ((253, 253, 0), 60.0),
    "dark_green": ((0, 128, 0), 60.0),
    "dark_red": ((128, 0, 0), 60.0),
    "light_green": ((0, 255, 0), 60.0),
    "light_red": ((255, 0, 0), 60.0),
}


def _blank(h: int = 20, w: int = 200) -> np.ndarray:
    return np.zeros((h, w, 3), dtype=np.uint8)


def _paint(img: np.ndarray, x0: int, x1: int, rgb: tuple[int, int, int]) -> None:
    """Paint vivid color into BGR image (palette stores RGB)."""
    bgr = (rgb[2], rgb[1], rgb[0])
    img[:, x0:x1] = bgr


def test_single_strip():
    img = _blank(20, 200)
    _paint(img, 0, 200, (0, 253, 253))  # turquoise
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert len(out) == 1
    assert out[0].x == 0
    assert abs(out[0].w - 200) <= 1
    assert out[0].h == 20


def test_split_50_50():
    img = _blank(20, 230)
    _paint(img, 0, 100, (0, 253, 253))
    _paint(img, 130, 230, (253, 253, 0))
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert len(out) == 2
    assert out[0].x < out[1].x  # L->R
    assert out[0].x == 0
    assert abs(out[0].w - 100) <= 1
    assert out[1].x == 130
    assert abs(out[1].w - 100) <= 1


def test_split_asymmetric():
    img = _blank(20, 230)
    _paint(img, 0, 70, (0, 253, 253))    # 35% width
    _paint(img, 100, 230, (253, 253, 0))  # 65% width
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert len(out) == 2
    assert abs(out[0].w - 70) <= 1
    assert abs(out[1].w - 130) <= 1


def test_gray_only_no_strip():
    img = _blank(20, 200)
    img[:, 50:150] = (128, 128, 128)
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert out == []


def test_cooldown_gray_dots_no_detect():
    img = _blank(20, 200)
    # scattered gray dots
    for x in (20, 50, 80, 110, 140, 170):
        img[8:12, x:x + 4] = (100, 100, 100)
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert out == []


def test_vivid_palette_match():
    img = _blank(20, 200)
    _paint(img, 50, 80, (0, 255, 0))  # light_green
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert len(out) == 1


def test_blank_frame():
    img = _blank(20, 200)
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert out == []


def test_strip_too_narrow_filtered():
    img = _blank(20, 200)
    _paint(img, 50, 53, (0, 253, 253))  # only 3px wide
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=10)
    assert out == []


def test_small_gap_fuses():
    img = _blank(20, 200)
    _paint(img, 30, 70, (0, 253, 253))
    _paint(img, 75, 120, (0, 253, 253))  # 5px gap, < min_gap_px=20
    out = detect_strips(img, PALETTE, min_gap_px=20, min_strip_px=5)
    assert len(out) == 1
    assert abs(out[0].x - 30) <= 2
    assert abs((out[0].x + out[0].w) - 120) <= 2


def test_split_two_charts_with_interleaved_gray():
    # Regresie 2 ferestre TS: fiecare row de buline e mix vivid + gri, separate
    # de un gap larg de background (dividerul dintre ferestre). Înainte de fix
    # detect_strips picka doar runs vivid contiguu și rata fereastra stângă.
    palette = {**PALETTE, "gray": ((128, 128, 128), 60.0)}
    img = _blank(35, 1796)
    # Left chart: dots vivid + gray la fiecare 26px, x=0..820
    for i, x in enumerate(range(0, 820, 26)):
        rgb = (128, 128, 128) if i % 3 else (0, 128, 0)
        _paint(img, x, x + 22, rgb)
    # Window divider gap: x=820..910 rămâne background
    # Right chart: same pattern, x=910..1796
    for i, x in enumerate(range(910, 1790, 26)):
        rgb = (128, 128, 128) if i % 3 else (0, 128, 0)
        _paint(img, x, x + 22, rgb)
    out = detect_strips(img, palette, min_gap_px=28, min_strip_px=280)
    assert len(out) == 2, f"expected 2 strips, got {len(out)}: {out}"
    assert out[0].x == 0
    assert out[1].x >= 900  # right chart starts after divider
    assert out[0].x + out[0].w < out[1].x  # disjoint


def test_three_strips():
    img = _blank(20, 300)
    _paint(img, 0, 60, (0, 253, 253))
    _paint(img, 100, 160, (253, 253, 0))
    _paint(img, 220, 300, (0, 255, 0))
    out = detect_strips(img, PALETTE, min_gap_px=10, min_strip_px=5)
    assert len(out) == 3
    assert out[0].x < out[1].x < out[2].x
    assert out[0].x == 0
    assert out[1].x == 100
    assert out[2].x == 220


def test_strips_match_identical():
    a = [ROI(x=0, y=0, w=100, h=20), ROI(x=130, y=0, w=70, h=20)]
    b = [ROI(x=0, y=0, w=100, h=20), ROI(x=130, y=0, w=70, h=20)]
    assert _strips_match(a, b) is True


def test_strips_match_jitter_5px():
    a = [ROI(x=0, y=0, w=100, h=20), ROI(x=130, y=0, w=70, h=20)]
    b = [ROI(x=5, y=0, w=95, h=20), ROI(x=128, y=0, w=70, h=20)]
    assert _strips_match(a, b, tol=10) is True


def test_strips_match_drift_12px():
    a = [ROI(x=0, y=0, w=100, h=20)]
    b = [ROI(x=12, y=0, w=100, h=20)]
    assert _strips_match(a, b, tol=10) is False


def test_strips_match_count_different():
    a = [ROI(x=0, y=0, w=100, h=20)]
    b = [ROI(x=0, y=0, w=100, h=20), ROI(x=130, y=0, w=70, h=20)]
    assert _strips_match(a, b) is False