#!/usr/bin/env python3
"""Fly and fight from guest memory.

World state comes from entities2.py's typing rule: a live entity's position
triple sits at a fixed offset before its definition pointer, so one scan of the
entity heap yields every craft in the scene *with its unit type* — which is what
separates enemies from wingmen, capital ships and the thousands of moving
particles.

Control is PD on the aiming error, with the derivative taken from the craft's
own body angular velocity (recovered from two consecutive orientation matrices)
rather than from the differenced error — that is what the old screen-scraping
autopilot lacked, and why it oscillated.

Usage: autopilot3.py <config.json> [seconds] [--dry]
"""
import json
import math
import os
import struct
import sys
import time
from collections import Counter

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

HOSTILE = ("e", "be")          # UN_e* / UN_be* are ADAN; UN_f*/UN_bf* are ours


def unit_faction(nm):
    base = nm[3:]
    return "ADAN" if base.startswith("be") or base.startswith("e") else "TCAF"


class World:
    def __init__(self, cfg):
        self.cfg = cfg
        self.w = gworld.World()
        self.fd, self.size = self.w.fd, self.w.size
        self.defs = entities2.definitions(self.w)
        self.delta = cfg["def_delta"]
        self.rot_delta = cfg["rot_delta"]
        self.rot_stride = cfg.get("rot_stride", 12)
        self.fwd_row = cfg["fwd_row"]
        self.fwd_sign = cfg["fwd_sign"]
        self.ents = []

    def rescan(self):
        movers = entities2.moving(self.fd, self.size, dt=0.35,
                                  va_range=(self.cfg["va_lo"], self.cfg["va_hi"]))
        ents = entities2.typed(self.fd, self.defs, movers, self.delta)
        uniq = {}
        for off, nm, pos, sp in ents:
            uniq.setdefault(off, (off, nm, pos, sp))
        self.ents = list(uniq.values())
        return self.ents

    def pos(self, off):
        b = os.pread(self.fd, 12, off)
        return np.array(struct.unpack(">3f", b)) if len(b) == 12 else None

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


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


def body_rate(Mprev, M, dt):
    if Mprev is None or M is None or dt <= 0:
        return np.zeros(3)
    D = Mprev @ 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
    return w / dt


class Autopilot:
    KP, KD = 2.2, 0.45
    FIRE_CONE = math.radians(10)
    FIRE_RANGE = 6000.0

    def __init__(self, world, pad, dry=False):
        self.W = world
        self.pad = pad
        self.dry = dry
        self.prevM = None
        self.firing = False

    def me(self):
        for off, nm, pos, sp in self.W.ents:
            if "Player" in nm:
                return off, nm
        return None, None

    def step(self, dt):
        off, nm = self.me()
        if off is None:
            return "no-player"
        p = self.W.pos(off)
        M = self.W.rot(off)
        if p is None or M is None:
            return "no-state"
        fwd = M[self.W.fwd_row] * self.W.fwd_sign
        rows = [M[i] for i in range(3)]
        right = rows[(self.W.fwd_row + 1) % 3]
        up = np.cross(fwd, right)

        w = body_rate(self.prevM, M, dt)
        self.prevM = M

        # nearest hostile, preferring what is already in front
        best, bestscore = None, 1e18
        for eoff, enm, epos, esp in self.W.ents:
            if unit_faction(enm) != "ADAN":
                continue
            q = self.W.pos(eoff)
            if q is None:
                continue
            v = q - p
            d = float(np.linalg.norm(v))
            if d < 1e-3:
                continue
            # Prefer targets we can actually bring the nose onto. Nearest-first
            # picks whatever is closest even at 90 deg off the nose, and closing
            # on an off-boresight target only raises the bearing rate -- which is
            # exactly the lag that kept the first run outside its firing cone.
            ahead = float(v @ fwd) / d
            ang = math.acos(max(-1.0, min(1.0, ahead)))
            score = d * (1.0 + 3.0 * (ang / math.pi) ** 2)
            if score < bestscore:
                best, bestscore = (eoff, enm, q, v, d), score
        if best is None:
            if not self.dry:
                self.pad.axis("LX", 0.0)
                self.pad.axis("LY", 0.0)
            return "no-hostiles"

        eoff, enm, q, v, d = best
        lx = float(v @ right)
        ly = float(v @ up)
        lz = float(v @ fwd)
        yaw = math.atan2(lx, lz if abs(lz) > 1e-3 else 1e-3)
        pitch = math.atan2(ly, lz if abs(lz) > 1e-3 else 1e-3)
        if lz < 0:
            yaw = math.copysign(math.pi / 2, lx if lx else 1.0)

        sx = clamp(self.KP * yaw - self.KD * float(w @ up))
        sy = clamp(-(self.KP * pitch - self.KD * float(w @ right)))
        aligned = abs(yaw) < self.FIRE_CONE and abs(pitch) < self.FIRE_CONE
        fire = aligned and d < self.FIRE_RANGE

        if not self.dry:
            self.pad.axis("LX", sx)
            self.pad.axis("LY", sy)
            if fire != self.firing:
                (self.pad.press if fire else self.pad.release)("RB")
                self.firing = fire
        self.shots = getattr(self, "shots", 0) + (1 if fire else 0)
        return (f"tgt={enm[3:24]:<22} d={d:8.0f} yaw={math.degrees(yaw):+6.1f} "
                f"pit={math.degrees(pitch):+6.1f} stick=({sx:+.2f},{sy:+.2f}) fire={int(fire)}")

    def run(self, secs, hz=10.0):
        t0 = time.time()
        last = t0
        last_scan = 0.0
        while time.time() - t0 < secs:
            t = time.time()
            if t - last_scan > 2.5:
                ents = self.W.rescan()
                last_scan = t
                c = Counter(unit_faction(e[1]) for e in ents)
                print(f"[{t-t0:6.1f}] rescan: {len(ents)} entities {dict(c)}", flush=True)
            print(f"[{t-t0:6.1f}] {self.step(t - last)}", flush=True)
            last = t
            time.sleep(max(0, 1.0 / hz - (time.time() - t)))
        if not self.dry:
            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
    dry = "--dry" in sys.argv
    W = World(cfg)
    W.rescan()
    ap = Autopilot(W, Pad(), dry=dry)
    ap.run(secs)


if __name__ == "__main__":
    main()