Syplheed-Reborn/tools/re-capture/autopilot2.py

#!/usr/bin/env python3
"""Memory-driven autopilot: fly and fight from the game's own world state.

The earlier screen-scraping autopilot oscillated because it only ever saw a
2-D arrow on the HUD and had no rate feedback. This one reads the actual
transform of every entity out of guest RAM, so it can do proper PD control:
the derivative term uses the craft's real body angular velocity, recovered from
two consecutive rotation matrices, not a differenced pixel position.

  world state  <- /dev/shm/xenia_memory_* (see gworld.py)
  control      -> the vgamepad FIFO, written directly at loop rate

Offsets come from flight_analyze.py and are passed in / stored in offsets.json;
nothing here hard-codes a value that was not derived from a capture.
"""
import json
import math
import os
import struct
import sys
import time

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


# --------------------------------------------------------------- 3-D helpers

def dot(a, b):
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]


def sub(a, b):
    return (a[0] - b[0], a[1] - b[1], a[2] - b[2])


def norm(a):
    return math.sqrt(dot(a, a))


def clamp(v, lo=-1.0, hi=1.0):
    return max(lo, min(hi, v))


def body_rates(R_prev, R, dt):
    """Angular velocity in body axes from two rotation matrices.

    R rows are the craft's axes in world space, so R_prev @ R^T is the
    incremental rotation; its skew part is the rotation vector.
    """
    if dt <= 0:
        return (0.0, 0.0, 0.0)
    # M = R_prev * R^T   (3x3, row-major tuples)
    M = [[sum(R_prev[i * 3 + k] * R[j * 3 + k] for k in range(3)) for j in range(3)]
         for i in range(3)]
    wx = (M[2][1] - M[1][2]) / 2.0
    wy = (M[0][2] - M[2][0]) / 2.0
    wz = (M[1][0] - M[0][1]) / 2.0
    return (wx / dt, wy / dt, wz / dt)


# ------------------------------------------------------------------- reader

class Flight:
    def __init__(self, cfg):
        self.cfg = cfg
        self.w = gworld.World()
        self.pos_off = cfg["pos"]
        self.rot_off = cfg["rot"]
        self.hp_off = cfg.get("hp")
        self.player_name = cfg.get("player", "Player")
        self.entities = []
        self.last_scan = 0.0

    def rescan(self):
        self.entities = self.w.refresh()
        self.last_scan = time.time()

    def state(self, off):
        buf = self.w.read_off(off, gworld.WINDOW)
        if len(buf) < gworld.WINDOW:
            return None
        pos = struct.unpack_from(">3f", buf, self.pos_off)
        rot = struct.unpack_from(">9f", buf, self.rot_off)
        if not all(map(math.isfinite, pos + rot)):
            return None
        hp = struct.unpack_from(">f", buf, self.hp_off)[0] if self.hp_off else None
        return {"pos": pos, "rot": rot, "hp": hp}

    def player(self):
        for va, off, nm in self.entities:
            if self.player_name in nm:
                s = self.state(off)
                if s:
                    s["name"] = nm
                    return s
        return None

    def hostiles(self):
        """ADAN units are the enemy; the disc IDs encode the faction.

        `UN_e*` / `UN_be*` = ADAN, `UN_f*` / `UN_bf*` = TCAF (ours).
        """
        out = []
        for va, off, nm in self.entities:
            base = nm[3:]
            if not (base.startswith("e") or base.startswith("be")):
                continue
            s = self.state(off)
            if s:
                s["name"] = nm
                s["off"] = off
                out.append(s)
        return out


# ---------------------------------------------------------------- controller

class Autopilot:
    KP_YAW, KD_YAW = 1.6, 0.35
    KP_PITCH, KD_PITCH = 1.6, 0.35
    FIRE_CONE = math.radians(6)
    FIRE_RANGE = 4000.0

    def __init__(self, flight, pad, log=sys.stdout):
        self.f = flight
        self.pad = pad
        self.log = log
        self.prev_rot = None
        self.prev_t = None
        self.target = None
        self.firing = False

    def pick_target(self, me):
        hs = self.f.hostiles()
        if not hs:
            return None
        # nearest, mildly preferring what is already ahead of us
        def score(h):
            v = sub(h["pos"], me["pos"])
            d = norm(v) or 1.0
            fwd = me["rot"][6:9]
            ahead = dot(v, fwd) / d
            return d * (1.0 if ahead > 0 else 2.0)
        return min(hs, key=score)

    def step(self, dt):
        me = self.f.player()
        if not me:
            return "no-player"
        R = me["rot"]
        w = body_rates(self.prev_rot, R, dt) if self.prev_rot else (0, 0, 0)
        self.prev_rot = R

        tgt = self.pick_target(me)
        if not tgt:
            self.pad.axis("LX", 0.0)
            self.pad.axis("LY", 0.0)
            self.pad.trig("RT", 0.6)
            return "no-target"

        v = sub(tgt["pos"], me["pos"])
        dist = norm(v) or 1.0
        # into body axes: rows are right / up / forward
        lx = dot(R[0:3], v)
        ly = dot(R[3:6], v)
        lz = dot(R[6:9], v)
        yaw_err = math.atan2(lx, lz if lz > 1e-3 else 1e-3)
        pitch_err = math.atan2(ly, lz if lz > 1e-3 else 1e-3)
        if lz < 0:                       # behind us: turn the short way, hard
            yaw_err = math.copysign(math.pi / 2, lx if lx else 1.0)

        stick_x = clamp(self.KP_YAW * yaw_err - self.KD_YAW * w[1])
        stick_y = clamp(-(self.KP_PITCH * pitch_err - self.KD_PITCH * w[0]))
        self.pad.axis("LX", stick_x)
        self.pad.axis("LY", stick_y)
        self.pad.trig("RT", 1.0 if dist > 1500 else 0.3)

        aligned = abs(yaw_err) < self.FIRE_CONE and abs(pitch_err) < self.FIRE_CONE
        want_fire = aligned and dist < self.FIRE_RANGE
        if want_fire != self.firing:
            (self.pad.press if want_fire else self.pad.release)(self.f.cfg["fire_btn"])
            self.firing = want_fire
        return (f"tgt={tgt['name'][3:20]:<18} d={dist:8.0f} yaw={math.degrees(yaw_err):+6.1f} "
                f"pit={math.degrees(pitch_err):+6.1f} stick=({stick_x:+.2f},{stick_y:+.2f}) "
                f"fire={int(self.firing)} hp={me['hp']}")

    def run(self, seconds, hz=15.0):
        self.f.rescan()
        t0 = time.time()
        last = t0
        while time.time() - t0 < seconds:
            t = time.time()
            if t - self.f.last_scan > 3.0:
                self.f.rescan()
            msg = self.step(t - last)
            last = t
            print(f"[{t-t0:6.1f}] {msg}", file=self.log, flush=True)
            time.sleep(max(0, 1.0 / hz - (time.time() - t)))
        self.pad.reset()


def main():
    cfg = json.load(open(sys.argv[1]))
    secs = float(sys.argv[2]) if len(sys.argv) > 2 else 60.0
    ap = Autopilot(Flight(cfg), Pad())
    ap.run(secs)


if __name__ == "__main__":
    main()