#!/usr/bin/env python3
"""Angular rate as a FUNCTION of speed — fitted, not sampled at two endpoints.

The two-point probe (`rate_probe.py`) assumed the craft held the speed its throttle
selects. It does not: full pitch bleeds 1 193 -> 589 in eight seconds *with the
throttle still at maximum*, so an 8-second average is taken over a moving speed and
lands between two different caps. That is how a measured pitch rate came out 33 %
ABOVE its own cap.

Turn the bug into the instrument. Because the speed bleeds on its own, one long hold
sweeps the whole range, so sampling position *and* attitude together gives
`(speed, rate)` pairs across it — the entire curve from a single phase, with no
assumption about what speed the craft is at.

Speed comes from the position deltas (no HUD OCR needed). Both speed and rate are
per wall-second here; the run's clock ratio converts both to game units afterwards
and cancels out of the SHAPE of the curve.

⚠️ SAMPLE IN WINDOWS, NOT PER READ. Polling at 20 Hz is faster than the guest
updates these fields, so a per-read delta is either exactly zero (no update yet) or
a whole frame's worth divided by a fraction of a frame. In a first run 111 of 352
reads were zero on BOTH channels — position and attitude update on the same frame,
so the two are perfectly correlated, and dividing each by the short wall dt
manufactured a clean "rate rises with speed" curve out of pure aliasing. Summing
|delta| over a window that spans many frames is immune: the total is right however
the updates fall inside it. (This is also why the swept-total probes were never
affected — only per-sample instantaneous rates were.)

Usage: rate_curve.py <roll|pitch> <out.csv> [dwell_s]
"""
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 pos_at(fd, off):
    return struct.unpack(">3f", os.pread(fd, 12, off))


def main():
    axis, out_csv = sys.argv[1], sys.argv[2]
    dwell = float(sys.argv[3]) if len(sys.argv) > 3 else 16.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")
    f_i = cfg["fwd_row"]
    if not alive(w, off, cfg):
        sys.exit("craft is not moving")
    u_i, w_i = pin_rows(w, off, cfg)
    print(f"# locked on {nm}, sweeping {axis}")

    # Start at maximum speed and let the turn bleed it: one hold, whole range.
    pad_state(rt=255)
    time.sleep(6.0)
    shot(f"curve_{axis}_a")
    prev_r, prev_p, prev_t = rows_at(w.fd, off, cfg), pos_at(w.fd, off), time.time()
    pad_state(rt=255, **DRIVER[axis])
    t0 = time.time()
    rows = []
    WIN = 0.5                       # seconds per emitted point: many guest frames
    win_t0, win_path, win_swept = time.time(), 0.0, 0.0
    while time.time() - t0 < dwell:
        time.sleep(0.02)
        now = time.time()
        cur_r, cur_p = rows_at(w.fd, off, cfg), pos_at(w.fd, off)
        win_path += math.dist(cur_p, prev_p)
        if axis == "roll":
            d = math.atan2(dot(cur_r[u_i], prev_r[w_i]), dot(cur_r[u_i], prev_r[u_i]))
        else:
            d = math.atan2(dot(cur_r[f_i], prev_r[u_i]), dot(cur_r[f_i], prev_r[f_i]))
        win_swept += abs(math.degrees(d))
        prev_r, prev_p = cur_r, cur_p
        if now - win_t0 >= WIN:
            span = now - win_t0
            rows.append((round(now - t0, 2), round(win_path / span, 1),
                         round(win_swept / span, 2)))
            win_t0, win_path, win_swept = now, 0.0, 0.0
    pad_state()
    shot(f"curve_{axis}_b")

    with open(out_csv, "w") as f:
        f.write("t,speed_units_per_wall_s,rate_deg_per_wall_s\n")
        for r in rows:
            f.write(",".join(str(x) for x in r) + "\n")

    # Bin by speed so the shape is visible without any curve-fitting assumption.
    print(f"# {len(rows)} windows of {WIN}s; rate binned by speed (per WALL second):")
    lo = min(r[1] for r in rows)
    hi = max(r[1] for r in rows)
    print(f"#   speed swept {lo:.0f} -> {hi:.0f} units/wall-s")
    nb = 6
    for b in range(nb):
        a = lo + (hi - lo) * b / nb
        z = lo + (hi - lo) * (b + 1) / nb
        sel = [r[2] for r in rows if a <= r[1] < z]
        if sel:
            print(f"#   speed {a:7.0f}-{z:7.0f}  n={len(sel):3d}  rate {sum(sel)/len(sel):6.1f}")
    print(f"# wrote {out_csv}; read HUD TIME off curve_{axis}_[ab] for the clock")


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