#!/usr/bin/env python3
"""Identify a LIVE grab by correlating it against committed oracle captures.

Why not whole-image statistics (screen_id.py's green/white/mean)?  Because the
class they have to reject is MOVIE FRAMES, and a movie frame can be anything.
Measured 2026-08-29: a frame of `ADV.wmv` containing a bright green laser beam
scored green=0.0018 white=0.086 mean=(53,67,76) -- numerically indistinguishable
from the title plate, and a probe built on those features tapped (A) into the
attract movie and then waited 120 s for a menu that was never coming.

So match on CONTENT instead.  Zero-normalised correlation against the committed
captures, over a small offset search, with the movie frames that fooled the
statistics kept as permanent negative controls.

A live grab is the whole 1280x720 root: xenia's title bar and menu bar occupy
the top ~45 rows, and the game surface below them is 1279x675 -- the same size
as the committed captures, which is not a coincidence.

Usage:
    screen_match.py IMAGE [IMAGE ...]     classify each
    screen_match.py --control             run the controls and exit non-zero on failure
"""
import os, sys
import numpy as np
from PIL import Image

REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
CAP = os.path.join(REPO, "docs", "re", "captures")
REFS = {
    "title": "title-builds/live-title-press-a.png",
    "menu":  "title-builds/live-main-menu.png",
}
SURFACE_TOP = 45          # rows of xenia window chrome on a 1280x720 root
SEARCH = 8                # +/- px offset search, as the corpus does elsewhere
THRESH = 0.70
FAST_DS = 4               # decimation for the live path (see below)

# 🔴 The exact path costs 1503 ms PER FRAME, measured.  A probe that ran it on
# every frame of an 8 fps x11grab drained the pipe at 0.64 fps, so the frames it
# classified were tens of seconds stale -- and the staleness GREW, which is how
# three "latencies" of 15.6 s, 20.3 s and 25.6 s were produced by a pipeline
# rather than by the game.  Ordering survives a backlog; durations do not.
# `fast=True` decimates 4x and searches +/-2 decimated px, and is controlled
# below against the same 8 captures as the exact path.


def load(p):
    return np.array(Image.open(p).convert("L"), dtype=np.float32)


def surface(a):
    """Crop a grab to the game surface. A committed capture is passed through."""
    h, w = a.shape
    if h == 720 and w == 1280:
        return a[SURFACE_TOP:, :1279]
    return a


def zncc(x, y):
    x = x - x.mean(); y = y - y.mean()
    d = np.sqrt((x * x).sum() * (y * y).sum())
    return float((x * y).sum() / d) if d else 0.0


def best_corr(img, ref, fast=False):
    """Max ZNCC over a small 2-D offset search."""
    if fast:
        img = img[::FAST_DS, ::FAST_DS]; ref = ref[::FAST_DS, ::FAST_DS]
        rng, step = 2, 1
    else:
        rng, step = SEARCH, 2
    h = min(img.shape[0], ref.shape[0]); w = min(img.shape[1], ref.shape[1])
    best = -1.0
    for dy in range(-rng, rng + 1, step):
        for dx in range(-rng, rng + 1, step):
            ys0, ys1 = max(0, dy), min(h, h + dy)
            yr0, yr1 = max(0, -dy), min(h, h - dy)
            xs0, xs1 = max(0, dx), min(w, w + dx)
            xr0, xr1 = max(0, -dx), min(w, w - dx)
            c = zncc(img[ys0:ys1, xs0:xs1], ref[yr0:yr1, xr0:xr1])
            if c > best:
                best = c
    return best


_REF_CACHE = {}


def refs():
    if not _REF_CACHE:
        for k, v in REFS.items():
            _REF_CACHE[k] = surface(load(os.path.join(CAP, v)))
    return _REF_CACHE


def classify(a_gray, fast=False):
    """Return (label, {name: corr}). label is 'title' | 'menu' | 'other'."""
    img = surface(a_gray)
    scores = {k: best_corr(img, r, fast) for k, r in refs().items()}
    k = max(scores, key=scores.get)
    return (k if scores[k] >= THRESH else "other"), scores


def classify_array(rgb, fast=False):
    g = (0.299 * rgb[:, :, 0] + 0.587 * rgb[:, :, 1] + 0.114 * rgb[:, :, 2]).astype(np.float32)
    return classify(g, fast)


CONTROLS = [
    # (path, expected)  -- positives from the committed corpus ...
    (os.path.join(CAP, "title-builds/live-title-press-a.png"), "title"),
    (os.path.join(CAP, "title-screen-oracle.png"),             "title"),
    (os.path.join(CAP, "title-builds/live-main-menu.png"),     "menu"),
    (os.path.join(CAP, "main-menu-oracle.png"),                "menu"),
    (os.path.join(CAP, "main-menu-reached.png"),               "menu"),
    # ... and the NEGATIVES.  Movie frames are the class this oracle exists to
    # reject, so they are COMMITTED fixtures, not scratch: an earlier version of
    # this list pointed at two scratch grabs and a later run of the same probe
    # overwrote one of them, turning a negative control into a title frame and
    # failing the control for the wrong reason.
    (os.path.join(CAP, "instrument-controls/movie-frame-attract-a.png"), "other"),
    (os.path.join(CAP, "instrument-controls/movie-frame-attract-b.png"), "other"),
    (os.path.join(CAP, "difficulty-screen.png"),               "other"),
]


def control():
    import time as _t
    bad = 0
    for fast in (False, True):
        print(f"--- {'FAST (live path)' if fast else 'EXACT'} ---")
        for p, exp in CONTROLS:
            if not os.path.exists(p):
                print(f"  SKIP (missing) {os.path.basename(p)}"); continue
            t = _t.time(); got, sc = classify(load(p), fast); ms = (_t.time() - t) * 1000
            ok = "ok " if got == exp else "FAIL"
            if got != exp:
                bad += 1
            print(f"  {ok} {os.path.basename(p):<34} -> {got:<6} (exp {exp:<6}) "
                  + "  ".join(f"{k}={v:+.3f}" for k, v in sc.items())
                  + f"  [{ms:.0f} ms]")
    print(f"\n{'CONTROL PASSED' if not bad else f'CONTROL FAILED ({bad})'}")
    return 1 if bad else 0


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--control":
        sys.exit(control())
    for p in sys.argv[1:]:
        got, sc = classify(load(p))
        print(f"{p}: {got}  " + "  ".join(f"{k}={v:+.3f}" for k, v in sc.items()))