#!/usr/bin/env python3
"""Angular rate cap for one axis, at minimum and maximum speed.

Generalises `roll_axis.py`. Two things it does that the probes before it did not:

* **Pins the matrix rows.** `entities2` measures row 2 = forward against velocity,
  but up-vs-right was previously taken from the D3D convention — and measuring
  world-Y in flight showed the convention is *backwards* here (up = row 1, right =
  row 0). Roll is immune to that (rotation of either non-forward row within their
  shared plane is roll either way), but **pitch is not**, so a pitch number taken
  without this pinning is naming an axis by assumption.
* **Brackets each phase with the HUD clock.** Rates are per GAME second and the
  clock ratio is a property of the moment — 1.260, 1.311, 1.383 and 1.247/1.218
  across five flights — so it cannot be assumed and must be read per phase.

    roll   driven by lx, measured as rotation of UP about forward
    pitch  driven by ly, measured as rotation of FORWARD toward up

Usage: rate_probe.py <roll|pitch> <out.csv> [dwell_s]
Then read HUD TIME off the `rate_<axis>_<phase>_[ab]` screenshots for the ratio.
"""
import json
import math
import os
import struct
import subprocess
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import speed_law  # noqa: E402
from axis_probe import alive, pad_state, pin_rows, rows_at, dot  # noqa: E402

DRIVER = {"roll": {"lx": 32767}, "pitch": {"ly": 32767}}


def shot(n):
    subprocess.run(["screenshot", f"/sylph-home/re/shots/{n}.png"], capture_output=True)


def main():
    axis = sys.argv[1]
    out_csv = sys.argv[2]
    dwell = float(sys.argv[3]) if len(sys.argv) > 3 else 8.0
    if axis not in DRIVER:
        sys.exit(f"axis must be one of {list(DRIVER)}")
    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}, measuring {axis}")
    f_i = cfg["fwd_row"]
    if not alive(w, off, cfg):
        sys.exit("craft is not moving — not in flight, or already dead")
    u_i, w_i = pin_rows(w, off, cfg)

    rows = []
    for label, trig in (("slow", "lt"), ("fast", "rt")):
        pad_state(**{trig: 255})
        time.sleep(5.0)                     # settle: 2 s was shown wrong twice
        shot(f"rate_{axis}_{label}_a")
        prev = rows_at(w.fd, off, cfg)
        pad_state(**{trig: 255}, **DRIVER[axis])
        t0, swept, seq = time.time(), 0.0, []
        while time.time() - t0 < dwell:
            time.sleep(0.05)
            cur = rows_at(w.fd, off, cfg)
            if axis == "roll":
                d = math.atan2(dot(cur[u_i], prev[w_i]), dot(cur[u_i], prev[u_i]))
            else:
                d = math.atan2(dot(cur[f_i], prev[u_i]), dot(cur[f_i], prev[f_i]))
            d = math.degrees(d)
            swept += abs(d)
            seq.append((round(time.time() - t0, 3), round(d, 4)))
            prev = cur
        pad_state(**{trig: 255})
        wall = seq[-1][0]
        shot(f"rate_{axis}_{label}_b")
        rows += [(label, *s) for s in seq]
        print(f"# {label:<5} wall {wall:5.2f}s   {axis} swept {swept:7.1f}deg   "
              f"rate {swept / wall:6.1f} deg/wall-s")
        if not alive(w, off, cfg):
            print("# CRAFT STOPPED MOVING — aborting rather than reporting zeros")
            break
    pad_state()
    with open(out_csv, "w") as f:
        f.write(f"phase,t,d{axis}_deg\n")
        for r in rows:
            f.write(",".join(str(x) for x in r) + "\n")
    print(f"# wrote {out_csv} — read HUD TIME off rate_{axis}_*_[ab] for the clock")


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