#!/usr/bin/env python3
"""Which pad input drives which rotation axis — all three measured at once.

The control mapping was left with **yaw ❔ "no input found"**: `AV_Yaw_*` exists in
the definitions (45/25) but no stick appeared to yaw. That conclusion came from
probes that measured one axis at a time using a non-forward matrix row, which is
exactly the flaw that made roll and pitch produce identical numbers — a row like
that moves under *any* rotation, so "this input yaws" and "this input pitches"
cannot be told apart by it.

So decompose properly. For each small step, with the previous frame's orthonormal
rows `(f, u, w)` = forward, up, right:

    roll  = atan2(u_new . w_old, u_new . u_old)     rotation of UP about forward
    yaw   = atan2(f_new . w_old, f_new . f_old)     rotation of FORWARD about up
    pitch = atan2(f_new . u_old, f_new . f_old)     rotation of FORWARD about right

Each drops the components the other two produce, so one held input yields three
numbers and the axis it actually drives is whichever is large.

Every candidate is held with the file pad, which writes the WHOLE state at once —
so `rx` really means right-stick-X with every other channel exactly zero, which a
quantising virtual stick could not promise.

⚠️ TWO THINGS THIS PROBE DOES NOT YET HANDLE, both learned the hard way:

1. **The craft can DIE mid-run.** Holding a stick at full deflection spins the ship
   at 150-200 deg/s in a live combat mission; two minutes of that and the first run
   ended on GAME OVER. Inputs measured after that point are meaningless, and the
   symptom is only "0 player candidates, best displacement 0.000" *afterwards* --
   the earlier rows still look fine. Check the player is alive BETWEEN inputs, or
   probe somewhere nothing shoots back.
2. **Which row is UP and which is RIGHT is not pinned.** `entities2` measures that
   row 2 is forward (against velocity), but rows 0 and 1 are assigned by the D3D
   convention, not by evidence. Roll is unaffected -- rotation of either non-forward
   row within their shared plane is roll either way -- but YAW and PITCH swap if the
   convention is wrong, so this probe's last two columns are named on an assumption.

Usage: axis_probe.py <out.csv> [dwell_s]
"""
import json
import math
import os
import struct
import sys
import time

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

PAD_FILE = os.environ.get("XENIA_PAD_FILE", "/tmp/xenia_pad.txt")

# Candidates. Sticks at full deflection, shoulders as buttons. Triggers are the
# throttle (already settled) and are left out.
# Order matters: the craft can be shot down mid-run, so the UNKNOWN inputs go
# first while it is healthy and the already-established ones (lx = roll) go last
# as controls. A row measured after death looks like a clean zero.
CANDIDATES = [
    ("rx+", {"rx": 32767}),
    ("ry+", {"ry": 32767}),
    ("LB", {"press": "LB"}),
    ("RB", {"press": "RB"}),
    ("lx+", {"lx": 32767}),
    ("ly+", {"ly": 32767}),
]


def pad_state(**kw):
    parts = [f"{k}={v}" for k, v in kw.items() if v is not None]
    tmp = PAD_FILE + ".tmp"
    with open(tmp, "w") as f:
        f.write(" ".join(parts))
    os.replace(tmp, PAD_FILE)


def rows_at(fd, off, cfg):
    base = off + cfg["rot_delta"]
    out = []
    for r in range(3):
        v = struct.unpack(">3f", os.pread(fd, 12, base + r * cfg["rot_stride"]))
        n = math.sqrt(sum(c * c for c in v)) or 1.0
        out.append(tuple(c / n for c in v))
    return out


def dot(a, b):
    return sum(a[i] * b[i] for i in range(3))


def alive(w, off, cfg, secs=1.0):
    """Is the craft still flying? A destroyed craft stops moving, and every
    subsequent input then measures a clean, meaningless zero — the failure mode
    that invalidated this probe's first run (it ended on GAME OVER and only said
    so afterwards)."""
    p0 = struct.unpack(">3f", os.pread(w.fd, 12, off))
    time.sleep(secs)
    p1 = struct.unpack(">3f", os.pread(w.fd, 12, off))
    return math.dist(p0, p1) > 1.0


def pin_rows(w, off, cfg):
    """Which non-forward row is UP and which is RIGHT?

    `entities2` measures row 2 = forward against the velocity vector, but the other
    two are assigned by the D3D convention rather than evidence — and yaw and pitch
    SWAP if that is wrong, so naming them without this test is a guess.

    Discriminator: in level flight the craft's up-vector points along world +Y and
    its right-vector lies near the horizontal plane. Sample at neutral and compare
    |world-Y| across the rows."""
    f_i = cfg["fwd_row"]
    acc = [0.0, 0.0, 0.0]
    n = 0
    for _ in range(20):
        time.sleep(0.05)
        rs = rows_at(w.fd, off, cfg)
        for r in range(3):
            acc[r] += rs[r][1]          # world Y component
        n += 1
    ys = [a / n for a in acc]
    cand = [r for r in (0, 1, 2) if r != f_i]
    up_i = max(cand, key=lambda r: abs(ys[r]))
    right_i = [r for r in cand if r != up_i][0]
    print(f"# row world-Y means: {[round(y,3) for y in ys]}  (forward = row {f_i})")
    print(f"#   -> up = row {up_i}, right = row {right_i}"
          f"   {'CONFIDENT' if abs(ys[up_i]) - abs(ys[right_i]) > 0.3 else 'WEAK — craft may not be level'}")
    return up_i, right_i


def main():
    out_csv = sys.argv[1]
    dwell = float(sys.argv[2]) if len(sys.argv) > 2 else 5.0
    cfg = json.load(open("/tmp/nav-live.json"))
    w, off, nm = speed_law.find_player()
    if not w:
        sys.exit("player entity not found")
    print(f"# locked on {nm}")
    f_i = cfg["fwd_row"]
    if not alive(w, off, cfg):
        sys.exit("craft is not moving at the start — not in flight, or already dead")
    u_i, w_i = pin_rows(w, off, cfg)

    rows = []
    print(f"# {'input':<6} {'roll':>9} {'yaw':>9} {'pitch':>9}   (deg over the dwell)")
    for label, state in CANDIDATES:
        pad_state()               # neutral
        time.sleep(2.0)
        prev = rows_at(w.fd, off, cfg)
        pad_state(**state)
        t0 = time.time()
        sw = [0.0, 0.0, 0.0]
        while time.time() - t0 < dwell:
            time.sleep(0.05)
            cur = rows_at(w.fd, off, cfg)
            roll = math.degrees(math.atan2(dot(cur[u_i], prev[w_i]), dot(cur[u_i], prev[u_i])))
            yaw = math.degrees(math.atan2(dot(cur[f_i], prev[w_i]), dot(cur[f_i], prev[f_i])))
            pitch = math.degrees(math.atan2(dot(cur[f_i], prev[u_i]), dot(cur[f_i], prev[f_i])))
            for k, v in enumerate((roll, yaw, pitch)):
                sw[k] += abs(v)
            prev = cur
        pad_state()
        wall = time.time() - t0
        if not alive(w, off, cfg):
            print(f"# CRAFT STOPPED MOVING after {label} — everything from here is "
                  f"meaningless, aborting rather than reporting zeros")
            break
        rate = [v / wall for v in sw]
        rows.append((label, *[round(v, 2) for v in rate]))
        print(f"# {label:<6} {rate[0]:9.1f} {rate[1]:9.1f} {rate[2]:9.1f}   deg/wall-s")
        time.sleep(1.5)

    with open(out_csv, "w") as f:
        f.write("input,roll_deg_s,yaw_deg_s,pitch_deg_s\n")
        for r in rows:
            f.write(",".join(str(x) for x in r) + "\n")
    print(f"# wrote {out_csv}")


if __name__ == "__main__":
    raise SystemExit(main())