#!/usr/bin/env python3
"""Ask questions of a ctrl_probe.py capture: which words answer to which input?

The capture is a window of the player entity object sampled every tick, tagged
with the pad state that produced it. That makes the interesting question
mechanical: for each 4-byte offset, does its value during phase X differ from
its value during the rest phases either side? A word that only moves while `RT`
is held is that input's state — a throttle setting, an afterburner tank, a heat
gauge — and one that moves in *every* phase is just live physics.

Sub-commands
  phases                 per-phase mean of every offset that moves at all
  respond <phase>        offsets that move during <phase> and not at rest
  near <value> [tol]     offsets whose first sample is ~= value (HUD anchor)
  trace <off> [off...]   full time series of specific offsets (pos-relative hex)
"""
import os
import struct
import sys

import numpy as np

HDR = b"SYLPHCTR"
REC = struct.Struct("<d32sff3f")


def load(path):
    with open(path, "rb") as f:
        blob = f.read()
    assert blob[:8] == HDR, "not a ctrl_probe capture"
    n, win, back = struct.unpack_from("<III", blob, 8)
    stride = REC.size + win
    base = 8 + 12
    ts, phase, sp, dodge, pos, wins = [], [], [], [], [], []
    for i in range(n):
        o = base + i * stride
        t, ph, s, dg, x, y, z = REC.unpack_from(blob, o)
        ts.append(t)
        phase.append(ph.split(b"\0")[0].decode())
        sp.append(s)
        dodge.append(dg)
        pos.append((x, y, z))
        wins.append(blob[o + REC.size:o + REC.size + win])
    # A run that ended in GAME OVER keeps sampling a dead object, and those
    # frames dominate every statistic. $BINQ_TMAX truncates the capture to the
    # part that was still flying.
    tmax = float(os.environ.get("BINQ_TMAX", "inf"))
    if tmax < float("inf"):
        keep = [i for i, t in enumerate(ts) if t <= tmax]
        n = len(keep)
        ts = [ts[i] for i in keep]
        phase = [phase[i] for i in keep]
        sp = [sp[i] for i in keep]
        dodge = [dodge[i] for i in keep]
        pos = [pos[i] for i in keep]
        wins = [wins[i] for i in keep]
    A = np.frombuffer(b"".join(wins), dtype=">f4").reshape(n, win // 4).astype(np.float64)
    U = np.frombuffer(b"".join(wins), dtype=">u4").reshape(n, win // 4)
    return dict(n=n, win=win, back=back, t=np.array(ts), phase=phase,
                speed=np.array(sp), dodge=np.array(dodge),
                pos=np.array(pos), F=A, U=U)


def label(d, i):
    return f"pos{i * 4 - d['back']:+#07x}"


def finite(d):
    F = d["F"]
    return np.all(np.isfinite(F), axis=0) & (np.max(np.abs(F), axis=0) < 1e12)


def cmd_phases(d, args):
    ok = finite(d)
    phases = []
    for p in d["phase"]:
        if p not in phases:
            phases.append(p)
    F = d["F"]
    mv = np.zeros(F.shape[1])
    means = {}
    for p in phases:
        m = np.array([x == p for x in d["phase"]])
        means[p] = F[m].mean(axis=0)
    for p in phases:
        mv = np.maximum(mv, np.abs(means[p] - means[phases[0]]))
    idx = np.flatnonzero(ok & (mv > 1e-3))
    order = idx[np.argsort(-mv[idx])][:int(args[0]) if args else 25]
    print("offset     " + "".join(f"{p[:8]:>10}" for p in phases))
    for i in order:
        print(f"{label(d, i):<10}" + "".join(f"{means[p][i]:10.2f}" for p in phases))


def cmd_respond(d, args):
    want = args[0]
    F, ok = d["F"], finite(d)
    inp = np.array([p == want for p in d["phase"]])
    rest = np.array([p.startswith("rest") or p == "base" for p in d["phase"]])
    if not inp.any():
        sys.exit(f"no phase {want!r}")
    # A word answering to this input must move *while it is held* and be quiet
    # at rest; a word that also moves at rest is live physics, not the input.
    a = F[inp]
    r = F[rest]
    d_in = a.max(axis=0) - a.min(axis=0)
    d_rest = r.max(axis=0) - r.min(axis=0)
    score = d_in - 2.0 * d_rest
    idx = np.flatnonzero(ok & (d_in > 1e-3) & (score > 0))
    for i in idx[np.argsort(-score[idx])][:20]:
        print(f"{label(d, i):<10} in-phase {a[:, i].min():12.3f}..{a[:, i].max():12.3f}"
              f"   at-rest {r[:, i].min():12.3f}..{r[:, i].max():12.3f}")
    if not len(idx):
        print("(nothing moves under this input that is quiet at rest)")


def cmd_near(d, args):
    v = float(args[0])
    tol = float(args[1]) if len(args) > 1 else max(1e-3, abs(v) * 1e-3)
    F, U = d["F"], d["U"]
    for i in np.flatnonzero(np.abs(F[0] - v) <= tol):
        print(f"{label(d, i):<10} f32 {F[0, i]:.4f} -> {F[-1, i]:.4f}")
    for i in np.flatnonzero(np.abs(U[0].astype(np.float64) - v) <= tol):
        print(f"{label(d, i):<10} u32 {U[0, i]} -> {U[-1, i]}")


def cmd_trace(d, args):
    offs = [int(a, 0) for a in args]
    idx = [(o + d["back"]) // 4 for o in offs]
    print("t      phase        speed  " + "  ".join(f"{o:+#07x}" for o in offs))
    for k in range(d["n"]):
        print(f"{d['t'][k]:6.2f} {d['phase'][k]:<12} {d['speed'][k]:6.0f}  "
              + "  ".join(f"{d['F'][k, i]:8.2f}" for i in idx))


def main():
    d = load(sys.argv[1])
    print(f"# {d['n']} ticks, window {d['win']:#x} bytes, back {d['back']:#x}")
    cmd = sys.argv[2] if len(sys.argv) > 2 else "phases"
    {"phases": cmd_phases, "respond": cmd_respond, "near": cmd_near,
     "trace": cmd_trace}[cmd](d, sys.argv[3:])


if __name__ == "__main__":
    main()