#!/usr/bin/env python3
"""Learn which body axis each stick drives, and which button fires.

The first autopilot run nulled yaw but held a steady ~27 deg pitch error, which
is the signature of a wrong pitch axis or sign — guessing `right = row0` and
`up = fwd x right` assumes a handedness the game need not share. So measure it:
hold each stick axis, read the craft's body angular velocity from consecutive
orientation matrices, and see which body axis actually responds and in which
direction.

The fire button is found the same way — by consequence, not assumption: the
nose ammo counter in the player's object must go down.

Usage: calibrate.py <config.json> [out.json]
"""
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 gworld  # noqa: E402
import entities2  # noqa: E402
from flight_probe import Pad  # noqa: E402


def player_off(w, fd, size, delta):
    defs = entities2.definitions(w)
    for _ in range(6):
        mv = entities2.moving(fd, size, dt=0.5)
        ents = entities2.typed(fd, defs, mv, delta)
        for off, nm, pos, sp in ents:
            if "Player" in nm:
                return off
        time.sleep(0.5)
    return None


def read_rot(fd, off, rot_delta, stride):
    n = stride * 2 + 12
    b = os.pread(fd, n, off + rot_delta)
    if len(b) < n:
        return None
    M = np.array([struct.unpack_from(">3f", b, stride * r) for r in range(3)])
    if not np.all(np.isfinite(M)) or np.max(np.abs(M @ M.T - np.eye(3))) > 5e-3:
        return None
    return M


def mean_body_rate(fd, off, cfg, secs=2.0, hz=8.0):
    """Mean angular velocity expressed in the craft's own axes."""
    prev, t_prev = None, None
    acc, n = np.zeros(3), 0
    end = time.time() + secs
    while time.time() < end:
        t = time.time()
        M = read_rot(fd, off, cfg["rot_delta"], cfg.get("rot_stride", 12))
        if M is not None and prev is not None:
            dt = t - t_prev
            D = prev @ M.T
            w = np.array([D[2, 1] - D[1, 2], D[0, 2] - D[2, 0], D[1, 0] - D[0, 1]]) / 2.0
            if dt > 0 and np.linalg.norm(w) < 0.5:
                # into body axes
                acc += M @ (w / dt)
                n += 1
        if M is not None:
            prev, t_prev = M, t
        time.sleep(max(0, 1.0 / hz - (time.time() - t)))
    return acc / max(n, 1)


def main():
    cfg = json.load(open(sys.argv[1]))
    out = sys.argv[2] if len(sys.argv) > 2 else sys.argv[1]
    w = gworld.World()
    fd, size = w.fd, w.size
    off = player_off(w, fd, size, cfg["def_delta"])
    if off is None:
        sys.exit("player not found")
    print(f"# player pos va {gmem.primary_va(off):#010x}")
    pad = Pad()

    res = {}
    for axis, name in (("LX", "yaw"), ("LY", "pitch")):
        pad.reset()
        time.sleep(1.0)
        base = mean_body_rate(fd, off, cfg, 1.2)
        pad.axis(axis, 0.9)
        time.sleep(0.4)
        wpos = mean_body_rate(fd, off, cfg, 2.0)
        pad.axis(axis, 0.0)
        time.sleep(1.2)
        d = wpos - base
        k = int(np.argmax(np.abs(d)))
        res[name] = {"axis": k, "sign": 1 if d[k] > 0 else -1,
                     "mag": float(abs(d[k]))}
        print(f"# {axis} (+0.9) -> body rate {d.round(3)}  => {name} axis {k} "
              f"sign {res[name]['sign']:+d}")
    pad.reset()

    # ---- fire button: the nose ammo counter must fall -----------------------
    blob = os.pread(fd, 0x1000, max(0, off - 0x800))
    ammo_offs = []
    for i in range(0, len(blob) - 3, 4):
        (u,) = struct.unpack_from(">I", blob, i)
        (f,) = struct.unpack_from(">f", blob, i)
        if u == 6000 or (math.isfinite(f) and abs(f - 6000.0) < 0.5):
            ammo_offs.append((max(0, off - 0x800) + i) - off)
    print(f"# ammo-like words (==6000) near the player: "
          f"{[hex(o) for o in ammo_offs] or 'none'}")

    def ammo():
        vals = []
        for d in ammo_offs:
            b = os.pread(fd, 4, off + d)
            if len(b) == 4:
                vals.append(struct.unpack(">I", b)[0])
        return vals

    fire_btn = None
    if ammo_offs:
        for btn in ("RB", "LB", "A", "B", "X", "Y"):
            before = ammo()
            pad.press(btn)
            time.sleep(1.2)
            pad.release(btn)
            time.sleep(0.5)
            after = ammo()
            drop = [a - b for a, b in zip(before, after)]
            print(f"#   {btn}: ammo delta {drop}")
            if any(x > 0 for x in drop):
                fire_btn = btn
                break
        for trig in ("RT", "LT"):
            if fire_btn:
                break
            before = ammo()
            pad.trig(trig, 1.0)
            time.sleep(1.2)
            pad.trig(trig, 0.0)
            time.sleep(0.5)
            after = ammo()
            drop = [a - b for a, b in zip(before, after)]
            print(f"#   {trig}: ammo delta {drop}")
            if any(x > 0 for x in drop):
                fire_btn = trig
    pad.reset()

    cfg["yaw_axis"] = res["yaw"]["axis"]
    cfg["yaw_sign"] = res["yaw"]["sign"]
    cfg["pitch_axis"] = res["pitch"]["axis"]
    cfg["pitch_sign"] = res["pitch"]["sign"]
    cfg["fire"] = fire_btn
    cfg["ammo_offs"] = ammo_offs
    json.dump(cfg, open(out, "w"), indent=1)
    print("\n" + json.dumps(cfg, indent=1))


if __name__ == "__main__":
    main()