#!/usr/bin/env python3
"""Which menu button is focused, using ONLY committed captures.

`tools/port/which-focus` answers this by rendering every focus state in the
port's Godot project and taking the minimum difference. That is a sound method
and it passed its controls -- but it needs Godot, which this container does not
have, and rendering the port's project is outside the decoder's role. This does
the same job from the framebuffer captures alone.

METHOD.  The focus highlight is the only thing that moves between two captures
of the same menu with different focus. So for an unknown shot U and a reference
R whose focus is known, the positive part of (U - R) peaks on U's focused row
and the negative part peaks on R's. Two captures with known, different focus
therefore calibrate the row->button mapping directly, and no geometry has to be
assumed -- which matters, because assuming `rest_y` was a band centre is exactly
how an earlier attempt of mine mis-assigned a band and produced a wrong answer.

CONTROLS (run on every invocation; the tool refuses if any fails):
  * the calibration pair must recover its own two answers;
  * a frame with no menu must NOT produce a confident verdict.

    focus_from_capture.py SHOT.png [--json]

⚠️ LIMIT, stated because it matters: the only two full main-menu captures with
a known focus state are REF_A and REF_B, which are this tool's own calibration
inputs. So reproducing them is self-consistency, NOT validation. The honest
validation is against a known TRANSITION rather than a known state: press the
d-pad down once and the reported button must advance by exactly one. A drive
that does that is testing this tool, not trusting it.
"""
import json
import os
import sys

import numpy as np
from PIL import Image

CAP = "/work/docs/re/captures/title-builds"
REF_A = f"{CAP}/live-main-menu.png"                    # NEW GAME focused
REF_B = f"{CAP}/live-main-menu-options-focused.png"    # OPTIONS focused
BUTTONS = ["NEW GAME", "LOAD GAME", "TUTORIAL", "OPTIONS", "EXTRAS"]
IDX_A, IDX_B = 0, 3
MIN_MARGIN = 2.0

def is_menu(a, thresh=0.85):
    """Is this frame the main menu at all?

    The focus statistic below is a peak-to-median ratio on a difference image,
    and a difference against ANY dissimilar frame has a large peak -- a title
    screen scored 2.88 and sailed past a 2.0 bar. So the screen identity has to
    be established FIRST, with the zncc classifier controlled 6/6 elsewhere.
    """
    g = a.mean(axis=2)
    z = (g - g.mean()) / (g.std() + 1e-9)
    ref = np.asarray(Image.open(REF_A).convert("L"), float)[:675, :1279]
    zr = (ref - ref.mean()) / (ref.std() + 1e-9)
    return float((z * zr).mean()) >= thresh


def load(p):
    a = np.asarray(Image.open(p).convert("RGB"), float)
    # Normalise to the captures' 1279x675 top-left crop: a 1280x720 guest frame
    # and a 1279x675 screenshot are the same pixels, cropped, not scaled.
    return a[:675, :1279]

def row_profile(u, r):
    """Row-sums of the positive part of (u - r), over the button column band."""
    d = (u - r).mean(axis=2)[:, 500:820]
    return np.clip(d, 0, None).sum(axis=1)

def peak_row(prof, smooth=9):
    k = np.ones(smooth) / smooth
    s = np.convolve(prof, k, mode="same")
    return int(np.argmax(s)), float(s.max()), float(np.median(s))

def calibrate():
    A, B = load(REF_A), load(REF_B)
    ra, _, _ = peak_row(row_profile(A, B))   # A's focus row (NEW GAME)
    rb, _, _ = peak_row(row_profile(B, A))   # B's focus row (OPTIONS)
    pitch = (rb - ra) / (IDX_B - IDX_A)
    return A, B, ra, pitch

def classify(shot, A, B, ra, pitch):
    """Return (button, margin). Compares against BOTH references and agrees."""
    votes = []
    for ref, ref_idx in ((A, IDX_A), (B, IDX_B)):
        prof = row_profile(shot, ref)
        r, peak, med = peak_row(prof)
        # A zero median makes this explode, so floor it. 🔴 Do NOT cap here:
        # an earlier version capped at 999 to keep the printed number readable,
        # which made two different votes compare EQUAL, and the stable sort then
        # kept the wrong one -- turning a correct NEW GAME into an out-of-range
        # index and a refusal. Cap at the point of DISPLAY, never before a
        # comparison that depends on the value.
        margin = peak / max(med, 1.0)
        idx = int(round((r - ra) / pitch))
        votes.append((idx, margin, ref_idx))
    # If the shot IS one of the references, that comparison is degenerate (all
    # zero) -- keep the vote with the larger margin.
    votes.sort(key=lambda v: -v[1])
    idx, margin, _ = votes[0]
    if not (0 <= idx < len(BUTTONS)):
        return None, margin
    return BUTTONS[idx], margin

def main():
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    as_json = "--json" in sys.argv
    if not args:
        print(__doc__); return 2
    A, B, ra, pitch = calibrate()

    # --- control 1: the calibration pair must recover its own answers
    for ref, want in ((A, "NEW GAME"), (B, "OPTIONS")):
        got, _ = classify(ref, A, B, ra, pitch)
        if got != want:
            print(f"CONTROL FAILED: calibration pair gave {got}, expected {want}",
                  file=sys.stderr)
            return 1
    # --- control 2: a frame with no menu must be rejected as not-a-menu
    neg = f"{CAP}/live-title-press-a.png"
    if os.path.exists(neg) and is_menu(load(neg)):
        print("CONTROL FAILED: a title frame was accepted as a menu", file=sys.stderr)
        return 1

    shot = load(args[0])
    if not is_menu(shot):
        if as_json:
            print(json.dumps({"button": None, "margin": 0.0, "decided": False,
                              "reason": "not the main menu"}))
        else:
            print("UNDECIDED (not the main menu)")
        return 1
    got, margin = classify(shot, A, B, ra, pitch)
    ok = got is not None and margin >= MIN_MARGIN
    if as_json:
        print(json.dumps({"button": got, "margin": round(min(margin, 999.0), 3),
                          "decided": ok}))
    else:
        print(f"{got if ok else 'UNDECIDED'} margin={min(margin, 999.0):.2f}")
    return 0 if ok else 1

if __name__ == "__main__":
    sys.exit(main())