Syplheed-Reborn/tools/re-capture/selfstate.py

#!/usr/bin/env python3
"""Derive everything the autopilot needs about the live world, and write it to
a JSON config. Each step is checked against an independent property, so a wrong
answer has to survive several unrelated tests.

  1. moving position triples                     (motion, whole-RAM diff)
  2. which one is US                             (turn axis reverses with our stick)
  3. our orientation matrix                      (a row must track our velocity)
  4. how to type any entity                      (fixed offset to its definition
                                                  pointer, learned from ours)

Output: {"pos_va":…, "rot_va":…, "fwd_row":…, "def_delta":…}

Usage: selfstate.py <out.json>
"""
import json
import math
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
import gworld  # noqa: E402
from findrot_global import scan as scan_rot  # noqa: E402

FIFO = "/tmp/sylph-vgamepad.fifo"


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


def arrays(fd, size):
    for a, b in gmem.extents(fd, size):
        n = (b - a) // 4 * 4
        if n >= 64:
            yield a, np.frombuffer(os.pread(fd, n, a), dtype=">f4")


def moving_triples(fd, size, lo=120.0, hi=1600.0, dt=0.4):
    s0 = list(arrays(fd, size))
    t0 = time.time()
    time.sleep(dt)
    s1 = list(arrays(fd, size))
    t1 = time.time()
    out = []
    for (o, a), (o2, b) in zip(s0, s1):
        if o != o2 or len(a) != len(b):
            continue
        with np.errstate(invalid="ignore"):
            af = np.nan_to_num(a.astype(np.float64), nan=0, posinf=0, neginf=0)
            bf = np.nan_to_num(b.astype(np.float64), nan=0, posinf=0, neginf=0)
        d = bf - af
        idx = np.flatnonzero(np.abs(d) > 1e-3)
        cand = idx[np.isin(idx + 1, idx) & np.isin(idx + 2, idx)]
        for i in cand:
            v = np.array([d[i], d[i + 1], d[i + 2]])
            sp = float(np.linalg.norm(v)) / (t1 - t0)
            if lo < sp < hi:
                out.append(o + int(i) * 4)
    return out


def read3(fd, off):
    b = os.pread(fd, 12, off)
    return np.array(struct.unpack(">3f", b)) if len(b) == 12 else np.full(3, np.nan)


def sample_positions(fd, offs, secs, hz):
    seq = []
    n = int(secs * hz)
    for _ in range(n):
        t = time.time()
        seq.append((t, np.array([read3(fd, o) for o in offs])))
        time.sleep(max(0, 1.0 / hz - (time.time() - t)))
    return seq


def mean_turn_axis(seq):
    ts = [t for t, _ in seq]
    ps = [p for _, p in seq]
    vs = [(ps[k + 1] - ps[k]) / max(ts[k + 1] - ts[k], 1e-3) for k in range(len(ps) - 1)]
    acc = np.zeros((ps[0].shape[0], 3))
    n = 0
    for k in range(len(vs) - 1):
        c = np.cross(vs[k], vs[k + 1])
        nr = np.linalg.norm(c, axis=1, keepdims=True)
        with np.errstate(invalid="ignore", divide="ignore"):
            acc += np.where(nr > 1e-9, c / nr, 0.0)
        n += 1
    return acc / max(n, 1)


def main():
    out_path = sys.argv[1]
    fd = os.open(gmem.mem_path(), os.O_RDONLY)
    size = os.path.getsize(gmem.mem_path())

    pad("reset")
    time.sleep(1.5)
    offs = moving_triples(fd, size)
    print(f"# {len(offs)} moving triples", flush=True)
    if not offs:
        sys.exit("nothing moving")

    # ---- 2. which one is us: turn axis must reverse with the stick ----------
    axes = {}
    for name, lx in (("L", -0.95), ("R", 0.95)):
        pad(f"axis LX {lx}")
        time.sleep(0.6)
        seq = sample_positions(fd, offs, 2.5, 6.0)
        axes[name] = mean_turn_axis(seq)
        pad("axis LX 0")
        time.sleep(1.5)
        print(f"# phase {name} sampled", flush=True)
    pad("reset")
    aL, aR = axes["L"], axes["R"]
    nL, nR = np.linalg.norm(aL, axis=1), np.linalg.norm(aR, axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        cos = np.sum(aL * aR, axis=1) / (nL * nR)
    good = np.isfinite(cos) & (nL > 0.6) & (nR > 0.6)
    if not good.any():
        sys.exit("no object reversed with the stick")
    k = int(np.argmin(np.where(good, cos, 9e9)))
    pos_off = offs[k]
    print(f"# self position: va {gmem.primary_va(pos_off):#010x}  cos={cos[k]:+.3f}")

    # ---- 3. orientation: a row must track our velocity ----------------------
    time.sleep(1.0)
    rots = scan_rot(fd, size)
    print(f"# {len(rots)} orthonormal blocks; matching one to our velocity", flush=True)
    seq = sample_positions(fd, [pos_off], 2.0, 8.0)
    ps = [p[0] for _, p in seq]
    ts = [t for t, _ in seq]
    vdirs = []
    for a, b, ta, tb in zip(ps, ps[1:], ts, ts[1:]):
        v = (b - a) / max(tb - ta, 1e-3)
        n = np.linalg.norm(v)
        if n > 1e-3:
            vdirs.append(v / n)
    vdir = np.mean(vdirs, axis=0)
    vdir /= np.linalg.norm(vdir)
    speed = float(np.mean([np.linalg.norm((b - a) / max(tb - ta, 1e-3))
                           for a, b, ta, tb in zip(ps, ps[1:], ts, ts[1:])]))
    print(f"#   our heading {vdir.round(3)}  speed {speed:.1f}/s")

    # A block found by the earlier scan may have been overwritten by the time we
    # read it, so re-verify orthonormality here -- otherwise a garbage window
    # yields a meaningless (and unbounded) "cosine".
    best = None
    for ro in rots:
        b = os.pread(fd, 36, ro)
        if len(b) < 36:
            continue
        m = np.array(struct.unpack(">9f", b))
        if not np.all(np.isfinite(m)):
            continue
        M = m.reshape(3, 3)
        if np.max(np.abs(M @ M.T - np.eye(3))) > 5e-3:
            continue
        for r in range(3):
            c = float(M[r] @ vdir)
            if best is None or abs(c) > abs(best[0]):
                best = (c, ro, r)
    if best is None:
        sys.exit("no valid orientation block")
    print(f"# best orientation match: va {gmem.primary_va(best[1]):#010x} "
          f"row {best[2]}  cos={best[0]:+.4f}")

    # ---- 4. how to type an entity ------------------------------------------
    w = gworld.World()
    defs = {}
    for doff in w.scan_vtable(gworld.DEF_VTABLE):
        nm = w.name_of(doff)
        if nm and nm.startswith("UN_"):
            va = gmem.primary_va(doff)
            if va is not None:
                defs[struct.pack(">I", va)] = nm
    delta = None
    blob = os.pread(fd, 0x1000, max(0, pos_off - 0x800))
    for i in range(0, len(blob) - 3, 4):
        nm = defs.get(blob[i:i + 4])
        if nm:
            delta = (pos_off - 0x800) + i - pos_off
            print(f"# definition pointer at pos{delta:+#07x} -> {nm}")
            break
    if delta is None:
        print("# no definition pointer near our position (entity typing unavailable)")

    cfg = {
        "pos_va": gmem.primary_va(pos_off),
        "rot_va": gmem.primary_va(best[1]),
        "fwd_row": best[2],
        "fwd_sign": 1 if best[0] > 0 else -1,
        "def_delta": delta,
        "cruise_speed": speed,
    }
    json.dump(cfg, open(out_path, "w"), indent=1)
    print("\n" + json.dumps(cfg, indent=1))


if __name__ == "__main__":
    main()