#!/usr/bin/env python3
"""Measure the focus ring's SPIN PERIOD from a dense live filmstrip.

No absolute angle is estimated.  The corpus's centroid estimator fails its own
control by up to 19.8 deg, and a 360-bin angular cross-correlation also FAILED
the control written for it here (a synthetic 30 deg rotation of a live frame
came back as 0 deg, peak 0.596), so neither is trusted.

What is used instead needs no angle: the annulus's 360-bin brightness profile,
correlated against frame 0.  A rotating ring's profile returns to itself once
per revolution, so the correlation trace is periodic and its first return to a
maximum IS the period.  The ring is located from the data (the peak of the
temporal-std map over the button column), not from a declared coordinate.

Usage:  ring_period.py SECONDS OUTDIR
"""
import os, subprocess, sys, time
import numpy as np
from PIL import Image

W, H, DY, DX = 1280, 720, 45, 1
R_IN, R_OUT, NB = 8.0, 18.0, 360
SECS = float(sys.argv[1]) if len(sys.argv) > 1 else 30.0
OUT = sys.argv[2] if len(sys.argv) > 2 else "/sylph-home/re/ringcap"
COL = (480, 130, 570, 530)          # x0,y0,x1,y1 in GAME coords: the button column


def grab_stream(secs):
    p = subprocess.Popen(
        ["ffmpeg", "-loglevel", "error", "-f", "x11grab", "-draw_mouse", "0",
         "-video_size", f"{W}x{H}", "-i", ":98", "-r", "15",
         "-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
        stdout=subprocess.PIPE, bufsize=W * H * 3 * 2)
    n = W * H * 3
    t0 = time.time(); frames = []; ts = []
    x0, y0, x1, y1 = COL
    while time.time() - t0 < secs:
        b = p.stdout.read(n)
        if len(b) < n:
            break
        a = np.frombuffer(b, np.uint8).reshape(H, W, 3)
        g = (0.299 * a[..., 0] + 0.587 * a[..., 1] + 0.114 * a[..., 2]).astype(np.float32)
        frames.append(g[y0 + DY:y1 + DY, x0 + DX:x1 + DX].copy())
        ts.append(time.time() - t0)
    p.kill()
    return np.array(frames), np.array(ts)


def annulus_profile(patch, cy, cx):
    h, w = patch.shape
    yy, xx = np.mgrid[0:h, 0:w]
    r = np.hypot(yy - cy, xx - cx)
    m = (r >= R_IN) & (r <= R_OUT)
    th = (np.degrees(np.arctan2(yy - cy, xx - cx)) + 360) % 360
    idx = np.clip((th[m] / 360 * NB).astype(int), 0, NB - 1)
    v = patch[m]
    prof = np.zeros(NB); cnt = np.zeros(NB)
    np.add.at(prof, idx, v); np.add.at(cnt, idx, 1.0)
    prof = np.where(cnt > 0, prof / np.maximum(cnt, 1), np.nan)
    return np.nan_to_num(prof, nan=np.nanmean(prof)), float(v.mean())


def main():
    F, T = grab_stream(SECS)
    if len(F) < 10:
        print("too few frames"); return 1
    fps = len(F) / (T[-1] - T[0])
    print(f"{len(F)} frames over {T[-1]-T[0]:.1f}s = {fps:.2f} fps", flush=True)

    std = F.std(0)
    cy, cx = np.unravel_index(np.argmax(
        np.array([[std[max(0, i-14):i+14, max(0, j-14):j+14].mean()
                   for j in range(std.shape[1])] for i in range(std.shape[0])])), std.shape)
    print(f"ring located from the data at patch({cx},{cy}) = "
          f"GAME({COL[0]+cx},{COL[1]+cy}); local std {std[cy, cx]:.2f}", flush=True)

    profs = []; means = []
    for f in F:
        p, m = annulus_profile(f, cy, cx)
        profs.append(p); means.append(m)
    P = np.array(profs); M = np.array(means)
    print(f"annulus mean brightness: {M.mean():.2f} +/- {M.std():.3f} "
          f"({100*M.std()/M.mean():.2f}% -- a PULSE would move this)", flush=True)

    a = P[0] - P[0].mean()
    corr = np.array([float(((p - p.mean()) * a).sum() /
                           np.sqrt(((p - p.mean())**2).sum() * (a * a).sum()))
                     for p in P])
    np.save(f"{OUT}/period-corr.npy", np.vstack([T, corr, M]))
    print("\n t(s)  corr-with-frame0   annulus mean")
    for t, c, m in zip(T, corr, M):
        bar = "#" * max(0, int((c + 1) * 25))
        print(f"{t:6.2f}  {c:+.3f} {bar:<50} {m:7.2f}")

    # first return to a local maximum after the trace has dipped
    dip = np.argmax(corr < 0.3) if (corr < 0.3).any() else None
    if dip:
        after = corr[dip:]
        k = dip + int(np.argmax(after))
        print(f"\nfirst return to max after the dip: t = {T[k]:.2f}s (corr {corr[k]:+.3f})")
    return 0


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