#!/usr/bin/env python3
"""Which input is the AFTERBURNER, and what does it do?

The definition carries `AB_ConsumeShield_Begin 50`, `AB_ConsumeShield 10` and a set
of `AB_AV_*` turn caps far below the normal ones (roll 40 vs 125-200 °/s, pitch-up
30 vs 70-150) — so the burner should cost shield and cut agility. It carries no
`AB_*Velocity`, so whether it raises the speed target at all is an open question.

Method: hold `RT` (max-speed baseline), then hold each candidate button in turn and
measure speed over 1-second windows. A button that pushes speed past the RT plateau
is the burner. Buttons known to do something else are skipped: `RB` is the nose gun,
`Y` the missile mount, the d-pad the tactical map (see flight-controls-runtime.md).

Usage: ab_probe.py <out.csv> [hold_s]
"""
import math, os, subprocess, sys, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import speed_law

def pad(*a):
    subprocess.run(["vgamepad", *a], capture_output=True)

def windowed(seq, lo, hi):
    a = [s for s in seq if s[0] >= lo][0]
    b = [s for s in seq if s[0] <= hi][-1]
    return math.dist(a[1:], b[1:]) / (b[0] - a[0])

def main():
    out_csv = sys.argv[1]
    hold = float(sys.argv[2]) if len(sys.argv) > 2 else 5.0
    w, off, nm = speed_law.find_player()
    if not w:
        sys.exit("player entity not found after retries")
    print(f"# locked on {nm}")
    rows = []
    pad("trig", "RT", "1.0")          # max-speed baseline for every phase
    time.sleep(2.0)
    for label in ("baseline", "A", "B", "X", "LB", "baseline2"):
        if label.startswith("baseline"):
            pass
        else:
            pad("press", label)
        seq, t0 = [], time.time()
        while time.time() - t0 < hold:
            p = speed_law.pos_at(w.fd, off)
            seq.append((round(time.time() - t0, 3), *p))
            time.sleep(0.05)
        if not label.startswith("baseline"):
            pad("release", label)
        for s in seq:
            rows.append((label, *s))
        try:
            v = [windowed(seq, t, t + 1.0) for t in range(int(hold) - 1)]
            print(f"# {label:<10} speed per 1 s window: " + " ".join(f"{x:6.0f}" for x in v))
        except Exception as e:
            print(f"# {label:<10} failed: {e}")
    pad("trig", "RT", "0.0"); pad("reset")
    with open(out_csv, "w") as f:
        f.write("phase,t,x,y,z\n")
        for r in rows:
            f.write(",".join(str(x) for x in r) + "\n")
    print(f"# wrote {out_csv}")

if __name__ == "__main__":
    main()