atm/scripts/diag_strip_fixes.py

r"""Verify the two proposed fixes for detect_strips on the latest 2-window capture.

Fix A: include gray in VIVID_COLORS mask (cheapest change).
Fix B: use non-background mask (any pixel where diff(bg) > bg_tol).

Reuses the most recent raw capture in logs/repro/.
"""
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

import cv2  # noqa: E402
import numpy as np  # noqa: E402

from atm.config import Config, ROI  # noqa: E402
from atm.vision import crop_roi  # noqa: E402


def _detect_strips_with_palette(
    full_dot_crop, palette, color_names, min_gap_px, min_strip_px,
):
    """Same body as layout.detect_strips but with selectable color set."""
    h, w = full_dot_crop.shape[:2]
    mask = np.zeros((h, w), dtype=np.uint8)
    img_f = full_dot_crop.astype(np.float32)
    for name in color_names:
        if name not in palette:
            continue
        rgb, tol = palette[name]
        bgr = np.array([rgb[2], rgb[1], rgb[0]], dtype=np.float32)
        diff = np.linalg.norm(img_f - bgr, axis=2)
        mask |= (diff < tol).astype(np.uint8)
    kw = max(3, min_gap_px // 2)
    kernel = np.ones((1, kw), dtype=np.uint8)
    closed = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    n_labels, _labels, stats, _centroids = cv2.connectedComponentsWithStats(closed, connectivity=8)
    out = []
    for i in range(1, n_labels):
        x = int(stats[i, cv2.CC_STAT_LEFT])
        ww = int(stats[i, cv2.CC_STAT_WIDTH])
        if ww < min_strip_px:
            continue
        out.append(ROI(x=x, y=0, w=ww, h=h))
    out.sort(key=lambda r: r.x)
    return out, mask, closed


def _detect_strips_non_bg(full_dot_crop, bg_rgb, bg_tol, min_gap_px, min_strip_px):
    """Fix B: any pixel different from background → strip mask."""
    h, w = full_dot_crop.shape[:2]
    bgr_bg = np.array([bg_rgb[2], bg_rgb[1], bg_rgb[0]], dtype=np.float32)
    diff = np.linalg.norm(full_dot_crop.astype(np.float32) - bgr_bg, axis=2)
    mask = (diff > bg_tol).astype(np.uint8)
    kw = max(3, min_gap_px // 2)
    kernel = np.ones((1, kw), dtype=np.uint8)
    closed = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    n_labels, _labels, stats, _centroids = cv2.connectedComponentsWithStats(closed, connectivity=8)
    out = []
    for i in range(1, n_labels):
        x = int(stats[i, cv2.CC_STAT_LEFT])
        ww = int(stats[i, cv2.CC_STAT_WIDTH])
        if ww < min_strip_px:
            continue
        out.append(ROI(x=x, y=0, w=ww, h=h))
    out.sort(key=lambda r: r.x)
    return out, mask, closed


def _annotate(frame, strips_abs, label, out_dir):
    annotated = frame.copy()
    for r in strips_abs:
        cv2.rectangle(annotated, (r.x, r.y), (r.x + r.w, r.y + r.h), (0, 255, 255), 2)
    p = out_dir / f"diag_fix_{label}.png"
    cv2.imwrite(str(p), annotated)
    return p


def main() -> int:
    cfg = Config.load_current(ROOT / "configs")
    raws = sorted(p for p in (ROOT / "logs" / "repro").glob("*_raw.png") if p.name[0].isdigit())
    if not raws:
        raise SystemExit("Run scripts/repro_ss_resume.py first.")
    raw_path = raws[-1]
    print(f"Using: {raw_path}")
    frame = cv2.imread(str(raw_path), cv2.IMREAD_COLOR)

    palette = {n: (s.rgb, s.tolerance) for n, s in cfg.colors.items() if n != "background"}
    bg_rgb = cfg.colors["background"].rgb
    bg_tol = cfg.colors["background"].tolerance
    full_crop = crop_roi(frame, cfg.dot_roi)
    h, w = full_crop.shape[:2]
    strip_h = cfg.dot_roi.h
    min_strip_px = max(150, strip_h * 8)
    min_gap_px = max(20, int(strip_h * 0.8))
    print(f"params: min_strip_px={min_strip_px}  min_gap_px={min_gap_px}  crop={w}x{h}")

    out_dir = ROOT / "logs" / "repro"

    # Baseline (current code: vivid only, no gray)
    vivid_only = ("turquoise", "yellow", "dark_green", "dark_red", "light_green", "light_red")
    s0, m0, c0 = _detect_strips_with_palette(full_crop, palette, vivid_only, min_gap_px, min_strip_px)
    print(f"\n[BASELINE  vivid-only ] strips={len(s0)}")
    for i, r in enumerate(s0):
        print(f"  [{i}] x={r.x:>4d} w={r.w:>4d}")
    cv2.imwrite(str(out_dir / "diag_fix_BASELINE_mask_closed.png"), (c0 * 255).astype(np.uint8))
    _annotate(frame, [ROI(x=cfg.dot_roi.x + r.x, y=cfg.dot_roi.y, w=r.w, h=r.h) for r in s0],
              "BASELINE", out_dir)

    # Fix A: include gray
    fixA_palette = vivid_only + ("gray",)
    sA, mA, cA = _detect_strips_with_palette(full_crop, palette, fixA_palette, min_gap_px, min_strip_px)
    print(f"\n[FIX A    vivid+gray  ] strips={len(sA)}")
    for i, r in enumerate(sA):
        print(f"  [{i}] x={r.x:>4d} w={r.w:>4d}")
    cv2.imwrite(str(out_dir / "diag_fix_A_mask_closed.png"), (cA * 255).astype(np.uint8))
    _annotate(frame, [ROI(x=cfg.dot_roi.x + r.x, y=cfg.dot_roi.y, w=r.w, h=r.h) for r in sA],
              "A_vivid_plus_gray", out_dir)

    # Fix B: any non-background
    sB, mB, cB = _detect_strips_non_bg(full_crop, bg_rgb, bg_tol, min_gap_px, min_strip_px)
    print(f"\n[FIX B    non-bg      ] strips={len(sB)}")
    for i, r in enumerate(sB):
        print(f"  [{i}] x={r.x:>4d} w={r.w:>4d}")
    cv2.imwrite(str(out_dir / "diag_fix_B_mask_closed.png"), (cB * 255).astype(np.uint8))
    _annotate(frame, [ROI(x=cfg.dot_roi.x + r.x, y=cfg.dot_roi.y, w=r.w, h=r.h) for r in sB],
              "B_non_background", out_dir)

    # Sanity: where is the divider between the two TS windows?
    # Project non-bg mask onto x; long zero-runs reveal the gap.
    bgr_bg = np.array([bg_rgb[2], bg_rgb[1], bg_rgb[0]], dtype=np.float32)
    diff = np.linalg.norm(full_crop.astype(np.float32) - bgr_bg, axis=2)
    nonbg = (diff > bg_tol).astype(np.uint8)
    col_any = nonbg.any(axis=0).astype(np.uint8)
    # Find longest 0-run
    longest = (0, 0, 0)  # (length, x0, x1)
    i = 0
    while i < w:
        if col_any[i] == 1:
            i += 1
            continue
        j = i
        while j < w and col_any[j] == 0:
            j += 1
        run = j - i
        if run > longest[0]:
            longest = (run, i, j - 1)
        i = j
    print(f"\nLongest empty (background-only) horizontal stretch in dot_roi: "
          f"{longest[0]}px at x={longest[1]}..{longest[2]}  "
          f"(this is where the window divider sits)")
    print(f"\nWrote diag_fix_*.png to {out_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())