#!/usr/bin/env python3
"""Scan ALL of guest RAM for 3x3 orthonormal float blocks (rotation matrices).

Vectorised with numpy: nine shifted views of each extent give every candidate
9-float window at once, so the whole 250 MB working set tests in seconds.

A rotation matrix is a very strong signature — unit rows, mutually orthogonal —
so the hits are almost entirely real orientations. Two passes taken a moment
apart, with a known stick input in between, then say which of them is *ours*.

Usage: findrot_global.py [--track seconds]
"""
import os
import struct
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gmem  # noqa: E402

FIFO = "/tmp/sylph-vgamepad.fifo"
TOL = 2e-3


def pad(line):
    with open(FIFO, "w") as f:
        f.write(line + "\n")


def scan(fd, size):
    """[(file_offset, 9 floats)] for every orthonormal 3x3 block."""
    hits = []
    for a, b in gmem.extents(fd, size):
        n = (b - a) // 4 * 4
        if n < 64:
            continue
        arr = np.frombuffer(os.pread(fd, n, a), dtype=">f4").astype(np.float32)
        if arr.size < 16:
            continue
        m = arr.size - 8
        cols = [arr[i:i + m] for i in range(9)]
        with np.errstate(invalid="ignore", over="ignore"):
            finite = np.ones(m, dtype=bool)
            for c in cols:
                finite &= np.isfinite(c) & (np.abs(c) <= 1.001)
            # row norms
            n0 = cols[0] ** 2 + cols[1] ** 2 + cols[2] ** 2
            n1 = cols[3] ** 2 + cols[4] ** 2 + cols[5] ** 2
            n2 = cols[6] ** 2 + cols[7] ** 2 + cols[8] ** 2
            ok = finite
            ok &= np.abs(n0 - 1) < TOL
            ok &= np.abs(n1 - 1) < TOL
            ok &= np.abs(n2 - 1) < TOL
            d01 = cols[0] * cols[3] + cols[1] * cols[4] + cols[2] * cols[5]
            d02 = cols[0] * cols[6] + cols[1] * cols[7] + cols[2] * cols[8]
            d12 = cols[3] * cols[6] + cols[4] * cols[7] + cols[5] * cols[8]
            ok &= np.abs(d01) < TOL
            ok &= np.abs(d02) < TOL
            ok &= np.abs(d12) < TOL
        for i in np.flatnonzero(ok):
            hits.append(a + int(i) * 4)
    return hits


def read_mat(fd, off):
    b = os.pread(fd, 36, off)
    if len(b) < 36:
        return None
    return np.array(struct.unpack(">9f", b))


def main():
    fd = os.open(gmem.mem_path(), os.O_RDONLY)
    size = os.path.getsize(gmem.mem_path())
    pad("reset")
    time.sleep(0.6)
    t0 = time.time()
    hits = scan(fd, size)
    print(f"# {len(hits)} orthonormal 3x3 blocks in RAM ({time.time()-t0:.1f}s)", flush=True)
    if not hits:
        return

    # which of them rotate when WE yaw?  measure change under left vs right
    def deltas(lx, secs=2.5, hz=6.0):
        """Mean per-step rotation vector while the stick is held.

        Sampled incrementally: at a few degrees per step the skew part of
        A·Bᵀ is the rotation vector, which it is not over a 2 s turn. Any block
        that stops being orthonormal mid-phase is memory that got reused, not an
        orientation, and is dropped.
        """
        pad(f"axis LX {lx}")
        time.sleep(0.5)
        seq = []
        for _ in range(int(secs * hz)):
            t = time.time()
            seq.append({o: read_mat(fd, o) for o in hits})
            time.sleep(max(0, 1.0 / hz - (time.time() - t)))
        pad("axis LX 0")
        time.sleep(1.0)

        def ortho(m):
            if m is None or not np.all(np.isfinite(m)):
                return False
            M = m.reshape(3, 3)
            return np.max(np.abs(M @ M.T - np.eye(3))) < 5e-3

        out = {}
        for o in hits:
            ms = [s[o] for s in seq]
            if not all(ortho(m) for m in ms):
                continue
            ws = []
            for a, b in zip(ms, ms[1:]):
                M = a.reshape(3, 3) @ b.reshape(3, 3).T
                w = np.array([M[2, 1] - M[1, 2], M[0, 2] - M[2, 0], M[1, 0] - M[0, 1]]) / 2
                if np.linalg.norm(w) < 0.5:          # small-angle regime only
                    ws.append(w)
            if len(ws) >= 4:
                out[o] = np.mean(ws, axis=0)
        return out

    dl = deltas(-0.9)
    dr = deltas(+0.9)
    rows = []
    for o in hits:
        if o not in dl or o not in dr:
            continue
        a, b = dl[o], dr[o]
        na, nb = np.linalg.norm(a), np.linalg.norm(b)
        if na < 0.01 or nb < 0.01:
            continue
        cos = float(a @ b / (na * nb))
        rows.append((cos, o, na, nb))
    rows.sort()
    print("\n# orientation blocks that rotate OPPOSITE ways for left vs right stick")
    for cos, o, na, nb in rows[:12]:
        va = gmem.primary_va(o)
        print(f"  va {va:#010x}  cos={cos:+.3f}  |wL|={na:.3f} |wR|={nb:.3f}")
    if not rows:
        print("  (none)")


if __name__ == "__main__":
    main()