#!/usr/bin/env python3
"""Which control is the throttle, and where does the craft keep its own state?

Two questions, one flight. Both are answered by *consequence* rather than by
reading a field we hope is the right one:

* **Throttle.** The craft's speed is measured from its own position — a finite
  difference on the position triple in guest RAM — so no speed field has to be
  found first. The probe then holds each candidate input in turn and asks which
  one changes that measured speed. (`findspeed.py` failed the other way round:
  it assumed `RT` was the throttle and went looking for a field that rose.)

* **Own state.** Every tick also copies a window of the player entity object.
  Afterwards, offsets whose float value tracks the measured speed are candidate
  speed/throttle fields, and offsets that only ever *fall* are candidate
  hull/shield/ammo — the numbers survival needs.

Flying straight into a firefight for 90 s is how earlier runs died, so the loop
keeps the collision avoidance from navigator.py armed the whole time and marks
any sample where it had to intervene: a phase that had to dodge is not a clean
speed measurement, and says so rather than being quietly averaged in.

Usage: ctrl_probe.py <config.json> <out-prefix> [hold_s] [settle_s]
"""
import json
import math
import os
import struct
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gmem  # noqa: E402
import navigator  # noqa: E402
from flight_probe import Pad  # noqa: E402

# 0x800 each way: the shield lives at pos+0x430 (own_state.py), which a
# 0x400 window silently cropped out of the first capture.
WIN_BACK = 0x800          # bytes of the player object kept before the position
WIN_FWD = 0x800           # ...and after
WIN = WIN_BACK + WIN_FWD
HZ = 20.0

# (label, [pad commands]) — everything the pad can do that might be a throttle.
# RB is left out: it is the fire button (autopilot-memory-driven.md) and firing
# during a speed measurement only invites return fire.
CANDIDATES = [
    ("RT", [("trig", "RT", 1.0)]),
    ("LT", [("trig", "LT", 1.0)]),
    ("RT+LT", [("trig", "RT", 1.0), ("trig", "LT", 1.0)]),
    ("A", [("press", "A")]),
    ("B", [("press", "B")]),
    ("X", [("press", "X")]),
    ("Y", [("press", "Y")]),
    ("LB", [("press", "LB")]),
    ("LS", [("press", "LS")]),
    ("RS", [("press", "RS")]),
    ("RY_up", [("axis", "RY", -1.0)]),
    ("RY_down", [("axis", "RY", 1.0)]),
    ("RX_right", [("axis", "RX", 1.0)]),
    ("dpad_up", [("dpad", "up")]),
    ("dpad_down", [("dpad", "down")]),
    ("dpad_left", [("dpad", "left")]),
    ("dpad_right", [("dpad", "right")]),
]


def apply(pad, cmds):
    for c in cmds:
        if c[0] == "trig":
            pad.trig(c[1], c[2])
        elif c[0] == "axis":
            pad.axis(c[1], c[2])
        elif c[0] == "press":
            pad.press(c[1])
        elif c[0] == "dpad":
            pad.f.write(f"dpad {c[1]}\n")


def clear(pad):
    pad.reset()
    pad.f.write("dpad center\n")


class Run:
    def __init__(self, cfg, prefix, hold, settle):
        self.W = navigator.World(cfg)
        self.nav = navigator.Navigator(self.W, None, dry=True)
        self.prefix = prefix
        self.hold = hold
        self.settle = settle
        self.pad = Pad()
        self.rows = []            # (t, phase, pos, speed, dodged)
        self.win = []             # raw window bytes per tick
        ents = self.W.scan()
        me = [(off, va) for off, va in ents if "Player" in self.W.defs[va]]
        if not me:
            sys.exit("player entity not in the scan — not in flight?")
        self.me_off = me[0][0]
        self.me_name = self.W.defs[me[0][1]]
        print(f"# player {self.me_name} pos off {self.me_off:#x} "
              f"va {gmem.primary_va(self.me_off):#x} | {len(ents)} entities",
              flush=True)

    # ------------------------------------------------------------- helpers
    def pos(self):
        return self.W.pos(self.me_off)

    def sticks_for(self, vec):
        """Stick deflections that point the nose along `vec` (escape steering)."""
        M = self.W.rot(self.me_off)
        if M is None:
            return 0.0, 0.0
        fwd = M[self.W.fwd_row] * self.W.fwd_sign
        right = M[(self.W.fwd_row + 1) % 3]
        up = np.cross(fwd, right)
        ez = float(np.dot(vec, fwd))
        yaw = math.atan2(float(np.dot(vec, right)), ez if abs(ez) > 1e-3 else 1e-3)
        pitch = math.atan2(float(np.dot(vec, up)), ez if abs(ez) > 1e-3 else 1e-3)
        return (max(-1.0, min(1.0, 2.0 * yaw)), max(-1.0, min(1.0, -2.0 * pitch)))

    # ---------------------------------------------------------------- phase
    def phase(self, label, cmds, secs):
        clear(self.pad)
        apply(self.pad, cmds)
        t_end = time.time() + secs
        dodged = False
        while time.time() < t_end:
            t = time.time()
            p = self.pos()
            if p is None:
                break
            self.rows.append([t, label, p, 0.0, 0])
            self.win.append(os.pread(self.W.fd, WIN, self.me_off - WIN_BACK))

            # avoidance runs at 5 Hz; a real threat overrides the phase and the
            # samples from here on are flagged
            if len(self.rows) % 4 == 0:
                ents = self.W.sample(t)
                me = [e for e in ents if e[0] == self.me_off]
                if me:
                    _, _, mp, mv, mr = me[0]
                    push, worst = self.nav.avoidance(mp, mv, mr, ents, self.me_off)
                    if float(np.linalg.norm(push)) > 0.6:
                        sx, sy = self.sticks_for(navigator.norm(push))
                        self.pad.axis("LX", sx)
                        self.pad.axis("LY", sy)
                        dodged = True
                        self.rows[-1][4] = 1
                    elif dodged:
                        self.pad.axis("LX", 0.0)
                        self.pad.axis("LY", 0.0)
            time.sleep(max(0.0, 1.0 / HZ - (time.time() - t)))
        return dodged

    def run(self):
        t0 = time.time()
        self.phase("base", [], self.settle * 2)
        for label, cmds in CANDIDATES:
            d = self.phase(label, cmds, self.hold)
            self.phase(f"rest_{label}", [], self.settle)
            sp = self.phase_speed(label)
            print(f"[{time.time()-t0:6.1f}] {label:<10} "
                  f"v0={sp[0]:7.1f} v1={sp[1]:7.1f} d={sp[1]-sp[0]:+7.1f}"
                  f"{'  (dodged)' if d else ''}", flush=True)
        clear(self.pad)
        self.speeds()
        self.dump()

    # ------------------------------------------------------------ analysis
    def speeds(self):
        """Fill in per-tick speed by central difference on position."""
        for i, r in enumerate(self.rows):
            j, k = max(0, i - 2), min(len(self.rows) - 1, i + 2)
            dt = self.rows[k][0] - self.rows[j][0]
            if dt > 1e-3:
                r[3] = float(np.linalg.norm(self.rows[k][2] - self.rows[j][2])) / dt

    def phase_speed(self, label):
        """(speed early, speed late) within a phase — needs speeds() first."""
        self.speeds()
        v = [r[3] for r in self.rows if r[1] == label]
        if len(v) < 6:
            return (0.0, 0.0)
        n = max(2, len(v) // 4)
        return (float(np.mean(v[:n])), float(np.mean(v[-n:])))

    def dump(self):
        with open(self.prefix + ".csv", "w") as f:
            f.write("t,phase,x,y,z,speed,dodged\n")
            t0 = self.rows[0][0]
            for t, ph, p, sp, dg in self.rows:
                f.write(f"{t-t0:.3f},{ph},{p[0]:.3f},{p[1]:.3f},{p[2]:.3f},"
                        f"{sp:.3f},{dg}\n")
        with open(self.prefix + ".bin", "wb") as f:
            f.write(b"SYLPHCTR")
            f.write(struct.pack("<III", len(self.rows), WIN, WIN_BACK))
            t0 = self.rows[0][0]
            for (t, ph, p, sp, dg), w in zip(self.rows, self.win):
                f.write(struct.pack("<d32sff3f", t - t0, ph.encode()[:32], sp,
                                    float(dg), *[float(c) for c in p]))
                f.write(w.ljust(WIN, b"\0"))
        print(f"# wrote {self.prefix}.csv and {self.prefix}.bin "
              f"({len(self.rows)} ticks)", flush=True)

        # ---- per-phase summary
        print("\n# phase          n   v_early    v_late     delta  dodged")
        order, seen = [], set()
        for r in self.rows:
            if r[1] not in seen:
                seen.add(r[1])
                order.append(r[1])
        for ph in order:
            v = [r[3] for r in self.rows if r[1] == ph]
            dg = sum(r[4] for r in self.rows if r[1] == ph)
            if len(v) < 6:
                continue
            n = max(2, len(v) // 4)
            a, b = float(np.mean(v[:n])), float(np.mean(v[-n:]))
            print(f"  {ph:<12} {len(v):4d} {a:9.1f} {b:9.1f} {b-a:+9.1f} {dg:5d}")

        # ---- which words in the object track speed, and which only fall
        W_ = np.frombuffer(b"".join(x.ljust(WIN, b"\0") for x in self.win),
                           dtype=">f4").reshape(len(self.win), WIN // 4)
        sp = np.array([r[3] for r in self.rows], dtype=np.float64)
        with np.errstate(invalid="ignore", over="ignore"):
            A = W_.astype(np.float64)
        ok = np.all(np.isfinite(A), axis=0) & (np.max(np.abs(A), axis=0) < 1e9)
        var = np.std(A, axis=0)
        cand = np.flatnonzero(ok & (var > 1e-6))
        cor = []
        for i in cand:
            c = np.corrcoef(A[:, i], sp)[0, 1]
            if np.isfinite(c):
                cor.append((abs(c), c, i))
        cor.sort(reverse=True)
        print("\n# object words correlating with measured speed "
              "(offset relative to the position triple)")
        for ac, c, i in cor[:12]:
            off = i * 4 - WIN_BACK
            print(f"  pos{off:+#07x}  r={c:+.3f}  "
                  f"range {A[:, i].min():.3f} .. {A[:, i].max():.3f}")

        print("\n# words that never rise (candidate hull / shield / ammo)")
        shown = 0
        for i in cand:
            col = A[:, i]
            if col[-1] >= col[0] - 1e-6:
                continue
            if np.max(np.diff(col)) > 1e-6:
                continue
            off = i * 4 - WIN_BACK
            print(f"  pos{off:+#07x}  {col[0]:.3f} -> {col[-1]:.3f}")
            shown += 1
            if shown >= 12:
                break
        if not shown:
            print("  (none — nothing in the window decreased monotonically)")


def main():
    cfg = json.load(open(sys.argv[1]))
    prefix = sys.argv[2]
    hold = float(sys.argv[3]) if len(sys.argv) > 3 else 4.0
    settle = float(sys.argv[4]) if len(sys.argv) > 4 else 2.5
    Run(cfg, prefix, hold, settle).run()


if __name__ == "__main__":
    main()