atm/scripts/diag_strip_detection.py

r"""Diagnose why detect_strips finds 1 strip when there are 2 TS windows.

Reuses the most recent raw capture under logs/repro/. For each strip detected,
prints the connected-components vivid mask, gaps, and the contiguous run lengths
so we can see exactly where the threshold is killing the second strip.
"""
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  # noqa: E402
from atm.layout import VIVID_COLORS, detect_strips  # noqa: E402
from atm.vision import crop_roi, find_top_dots, classify_pixel, pixel_rgb  # noqa: E402


def _latest_raw() -> Path:
    candidates = sorted((ROOT / "logs" / "repro").glob("*_raw.png"))
    if not candidates:
        raise SystemExit("No *_raw.png in logs/repro — run scripts/repro_ss_resume.py first.")
    return candidates[-1]


def main() -> int:
    cfg = Config.load_current(ROOT / "configs")
    raw_path = _latest_raw()
    print(f"Using raw frame: {raw_path}")
    frame = cv2.imread(str(raw_path), cv2.IMREAD_COLOR)
    if frame is None:
        raise SystemExit(f"Could not read {raw_path}")

    palette = {
        name: (spec.rgb, spec.tolerance)
        for name, spec in cfg.colors.items()
        if name != "background"
    }
    full_crop = crop_roi(frame, cfg.dot_roi)
    h, w = full_crop.shape[:2]
    print(f"dot_roi crop shape: {h}x{w}  (cfg.dot_roi={cfg.dot_roi})")

    # 1) Build the vivid mask used by detect_strips
    mask = np.zeros((h, w), dtype=np.uint8)
    img_f = full_crop.astype(np.float32)
    per_color_pixels = {}
    for name in VIVID_COLORS:
        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)
        m = (diff < tol).astype(np.uint8)
        per_color_pixels[name] = int(m.sum())
        mask |= m

    print("\nVIVID_COLORS pixel counts in dot_roi crop (pre-morphology):")
    for n, c in per_color_pixels.items():
        print(f"  {n:12s} {c:>7d} px")
    print(f"  total mask:  {int(mask.sum()):>7d} px")

    # 2) Column projection: which x columns have any vivid pixel?
    col_has = mask.any(axis=0).astype(np.uint8)
    runs = []
    i = 0
    while i < w:
        if col_has[i] == 0:
            i += 1
            continue
        j = i
        while j < w and col_has[j] == 1:
            j += 1
        runs.append((i, j - 1, j - i))
        i = j
    print(f"\nContiguous vivid-column runs (raw, no morphology): {len(runs)}")
    for x0, x1, run_w in runs[:25]:
        print(f"  x={x0:>4d}..{x1:>4d}  width={run_w}")
    if len(runs) > 25:
        print(f"  ... +{len(runs) - 25} more")

    # 3) Apply same morphology + connected components as detect_strips
    strip_h = cfg.dot_roi.h
    min_strip_px = max(150, strip_h * 8)
    min_gap_px = max(20, int(strip_h * 0.8))
    kw = max(3, min_gap_px // 2)
    print(f"\ndetect_strips params: min_strip_px={min_strip_px} min_gap_px={min_gap_px} morphology kw={kw}")

    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)
    print(f"\nConnected components after CLOSE (kw={kw}):  n={n_labels - 1} (excluding bg)")
    for i in range(1, n_labels):
        x = int(stats[i, cv2.CC_STAT_LEFT])
        y = int(stats[i, cv2.CC_STAT_TOP])
        ww = int(stats[i, cv2.CC_STAT_WIDTH])
        hh = int(stats[i, cv2.CC_STAT_HEIGHT])
        passes = "PASS" if ww >= min_strip_px else "drop"
        print(f"  [{i:>2d}] x={x:>4d} y={y:>3d} w={ww:>4d} h={hh:>3d}  -> {passes}")

    strips = detect_strips(full_crop, palette, min_gap_px, min_strip_px)
    print(f"\ndetect_strips() final result: {len(strips)} strip(s)")
    for i, s in enumerate(strips):
        print(f"  strip[{i}] x={s.x} y={s.y} w={s.w} h={s.h}")

    # 4) Run find_top_dots on the FULL dot_roi (single-chart fallback path used
    # by /ss when ctx.charts < 2). This is what users see when the layout
    # detector hasn't promoted the layout to multi-chart yet.
    bg_rgb = cfg.colors["background"].rgb if "background" in cfg.colors else (18, 18, 18)
    bg_tol = cfg.colors["background"].tolerance if "background" in cfg.colors else 15.0
    full_dots = find_top_dots(full_crop, bg_rgb, bg_tol, n=3)
    print(f"\nfind_top_dots on FULL dot_roi (n=3, single-chart fallback path):")
    for i, (cx, cy) in enumerate(full_dots):
        rgb = pixel_rgb(full_crop, cx, cy)
        m = classify_pixel(rgb, palette)
        print(f"  c{i + 1}: pos=({cx + cfg.dot_roi.x},{cy + cfg.dot_roi.y})  rgb={rgb}  -> {m.name} d={m.distance:.1f}")

    # 5) Save the binary mask + closed mask for visual inspection
    out_dir = ROOT / "logs" / "repro"
    cv2.imwrite(str(out_dir / "diag_mask_raw.png"), (mask * 255).astype(np.uint8))
    cv2.imwrite(str(out_dir / "diag_mask_closed.png"), (closed * 255).astype(np.uint8))
    print(f"\nWrote diag_mask_raw.png + diag_mask_closed.png to {out_dir}")

    return 0


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