#!/usr/bin/env python3
"""Drift-aware navigation with collision avoidance, driven from guest memory.

Three things the pursuit loop in autopilot3.py did not do:

* **See everything.** It enumerated entities by looking for things that *move*,
  so stations, hulls and parked structures were invisible — exactly the objects
  you crash into. This scans the entity heap for words that equal a known unit
  definition address and takes `position = hit - 0x130`, which finds every
  entity whether it is moving or not.

* **Know how big they are.** Each entity's definition carries `Size_Radius`
  (+0x50) and `Size_X/Y/Z` (+0x30/34/38) — fields already solved in
  docs/re/structures/unit-struct-runtime.md — so the avoidance radius is the
  game's own number, not a guess.

* **Account for drift.** The craft does not turn where it points: velocity lags
  the nose like an aircraft with sideslip. Steering the *nose* at a target
  therefore steers the *flight path* somewhere else, wide and late. This
  measures the lag online (the angle between nose and velocity, and how fast
  the velocity vector is actually swinging) and commands the nose *ahead* of
  where the flight path should go, by that lag.

Avoidance is closest-point-of-approach, not distance: what matters is whether
the two paths will intersect within a horizon, which is why a fast crossing
target is dangerous at 2 km and a station drifting away is not at 300 m.

Usage: navigator.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

# confirmed fields of the parsed unit definition (unit-struct-runtime.md)
DEF_SIZE_X, DEF_SIZE_Y, DEF_SIZE_Z, DEF_SIZE_R = 0x30, 0x34, 0x38, 0x50


def norm(v):
    n = float(np.linalg.norm(v))
    return v / n if n > 1e-9 else v * 0.0


def ang(a, b):
    return math.acos(max(-1.0, min(1.0, float(np.dot(norm(a), norm(b))))))


class World:
    def __init__(self, cfg):
        self.cfg = cfg
        self.w = gworld.World()
        self.fd, self.size = self.w.fd, self.w.size
        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.defs = {}          # def_va -> name
        self.def_word = {}      # 4-byte BE -> def_va
        self.radius = {}        # def_va -> collision radius
        self.prev = {}          # pos_off -> (t, pos)  for velocity
        self.vel = {}           # pos_off -> smoothed velocity
        self.ents = []
        self._load_defs()

    def _load_defs(self):
        for off in self.w.scan_vtable(gworld.DEF_VTABLE):
            nm = self.w.name_of(off)
            if not (nm and nm.startswith("UN_")):
                continue
            va = gmem.primary_va(off)
            if va is None:
                continue
            self.defs[va] = nm
            self.def_word[struct.pack(">I", va)] = va
            b = os.pread(self.fd, 0x60, off)
            try:
                sx, sy, sz = (struct.unpack_from(">f", b, o)[0]
                              for o in (DEF_SIZE_X, DEF_SIZE_Y, DEF_SIZE_Z))
                sr = struct.unpack_from(">f", b, DEF_SIZE_R)[0]
            except struct.error:
                sx = sy = sz = sr = 0.0
            vals = [v for v in (sr, sx, sy, sz) if math.isfinite(v) and 0 < v < 1e5]
            self.radius[va] = max(vals) if vals else 50.0

    # ------------------------------------------------------------- entities
    def scan(self):
        """Every entity in the heap — moving or not — by its definition pointer."""
        lo = gmem.va_to_off(self.cfg["va_lo"])
        hi = gmem.va_to_off(self.cfg["va_hi"])
        found = []
        # numpy, not a per-word python loop: the entity heap is 16 MB, so
        # stepping it 4 bytes at a time costs seconds per scan and the control
        # loop spends its whole budget scanning instead of flying.
        want = np.array(sorted(self.defs.keys()), dtype=np.uint32)
        for a, b in gmem.extents(self.fd, self.size):
            a, b = max(a, lo), min(b, hi)
            n = (b - a) // 4 * 4
            if n < 64:
                continue
            arr = np.frombuffer(os.pread(self.fd, n, a), dtype=">u4").astype(np.uint32)
            for k4 in np.flatnonzero(np.isin(arr, want)):
                k = int(k4) * 4
                va = int(arr[k4])
                poff = a + k - self.delta
                if poff < 0:
                    continue
                pb = os.pread(self.fd, 12, poff)
                if len(pb) < 12:
                    continue
                p = np.array(struct.unpack(">3f", pb))
                if not np.all(np.isfinite(p)) or np.max(np.abs(p)) > 1e7:
                    continue
                found.append((poff, va))
        # De-duplicate by POSITION: one entity is mirrored at several
        # addresses, so keying on the address keeps every copy and the scene
        # looks several times more crowded than it is.
        seen, uniq = {}, {}
        for poff, va in found:
            p = self.pos(poff)
            if p is None:
                continue
            key = (va, tuple(np.round(p, 0)))
            if key in seen:
                continue
            seen[key] = poff
            uniq[poff] = va
        self.ents = list(uniq.items())
        return self.ents

    def pos(self, off):
        b = os.pread(self.fd, 12, off)
        if len(b) < 12:
            return None
        p = np.array(struct.unpack(">3f", b))
        return p if np.all(np.isfinite(p)) 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)) or np.max(np.abs(M @ M.T - np.eye(3))) > 5e-3:
            return None
        return M

    def sample(self, t):
        """[(off, name, pos, vel, radius)] with velocity by finite difference."""
        out = []
        for off, va in self.ents:
            p = self.pos(off)
            if p is None:
                continue
            # A frame the emulator did not advance gives an identical position
            # and a bogus zero velocity, which then reads as "stopped" and
            # wrecks both the drift estimate and every closing-rate. Keep the
            # last good velocity instead, and smooth it.
            prev = self.prev.get(off)
            v = self.vel.get(off, np.zeros(3))
            if prev is not None:
                dt = t - prev[0]
                moved = float(np.linalg.norm(p - prev[1]))
                if dt > 0.02 and moved > 1e-4:
                    inst = (p - prev[1]) / dt
                    if float(np.linalg.norm(inst)) < 5000.0:
                        v = 0.5 * v + 0.5 * inst if np.any(v) else inst
                    self.prev[off] = (t, p)
            else:
                self.prev[off] = (t, p)
            self.vel[off] = v
            out.append((off, self.defs[va], p, v, self.radius[va]))
        return out


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


class Navigator:
    KP, KD = 2.2, 0.45
    FIRE_CONE = math.radians(9)
    FIRE_RANGE = 6000.0
    HORIZON = 6.0            # s of look-ahead for collision checks
    MARGIN_BIG = 220.0       # clearance around hulls and structures
    MARGIN_SMALL = 45.0      # ...around fighters, which manoeuvre themselves
    BIG_RADIUS = 120.0       # above this an entity counts as a structure
    SELF_MIRROR = 25.0       # the same object is mirrored at several addresses

    def __init__(self, W, pad, dry=False, log=sys.stdout):
        self.W = W
        self.pad = pad
        self.dry = dry
        self.log = log
        self.prevM = None
        self.firing = False
        self.tau = 0.8       # velocity-lag time constant, refined online

    # ------------------------------------------------------------ drift
    def update_tau(self, fwd, vel, prev_vhat, dt):
        """How long the flight path takes to catch the nose.

        The velocity vector swings toward the nose; the angle between them
        divided by the rate the velocity is actually swinging is that lag, and
        it is what the nose has to be commanded ahead by.
        """
        if prev_vhat is None or dt <= 1e-3:
            return
        vh = norm(vel)
        if np.linalg.norm(vh) < 1e-6:
            return
        swing = ang(prev_vhat, vh) / dt          # rad/s the path is turning
        lag = ang(fwd, vh)                       # rad the path is behind
        if swing > 0.02 and lag > 0.02:
            tau = lag / swing
            if 0.05 < tau < 5.0:
                self.tau = 0.9 * self.tau + 0.1 * tau

    # ------------------------------------------------- collision avoidance
    def avoidance(self, me_p, me_v, me_r, ents, me_off):
        """Sum of escape directions, weighted by how soon and how close."""
        push = np.zeros(3)
        worst = None
        for off, nm, p, v, r in ents:
            if off == me_off:
                continue
            rel_p = p - me_p
            rel_v = v - me_v
            d = float(np.linalg.norm(rel_p))
            # The same entity exists at several mirrored addresses, so our own
            # copy shows up as an obstacle at zero distance and pins the sticks
            # at full deflection forever. Anything this close is us.
            if d < self.SELF_MIRROR:
                continue
            # A wingman flying formation is metres away by design and steers
            # itself; giving it a hull-sized margin makes the loop thrash.
            big = r >= self.BIG_RADIUS
            margin = self.MARGIN_BIG if big else self.MARGIN_SMALL
            if not big and faction(nm) == "TCAF":
                margin *= 0.5
            safe = me_r + r + margin
            if d > 1e4:
                continue
            vv = float(np.dot(rel_v, rel_v))
            t_cpa = 0.0 if vv < 1e-6 else -float(np.dot(rel_p, rel_v)) / vv
            t_cpa = max(0.0, min(self.HORIZON, t_cpa))
            cpa = rel_p + rel_v * t_cpa
            miss = float(np.linalg.norm(cpa))
            if miss >= safe:
                continue
            urgency = (1.0 - miss / safe) * (1.0 - t_cpa / self.HORIZON)
            if urgency <= 0:
                continue
            esc = -norm(cpa) if miss > 1e-3 else norm(np.cross(rel_p, me_v))
            if np.linalg.norm(esc) < 1e-6:
                esc = norm(np.cross(rel_p, np.array([0.0, 1.0, 0.0])))
            push += esc * urgency
            if worst is None or urgency > worst[0]:
                worst = (urgency, nm, d, miss, t_cpa)
        return push, worst

    # -------------------------------------------------------------- target
    def pick(self, me_p, me_v, fwd, ents, me_off):
        speed = max(float(np.linalg.norm(me_v)), 1.0)
        best, bestscore = None, 1e18
        for off, nm, p, v, r in ents:
            if off == me_off or faction(nm) != "ADAN":
                continue
            rel = p - me_p
            d = float(np.linalg.norm(rel))
            if d < 1e-3:
                continue
            # lead: where it will be when a shot gets there
            lead = p + v * (d / max(speed, 200.0))
            rel_l = lead - me_p
            theta = ang(rel_l, fwd)
            score = d * (1.0 + 3.0 * (theta / math.pi) ** 2)
            if score < bestscore:
                best, bestscore = (off, nm, lead, rel_l, d), score
        return best

    # ---------------------------------------------------------------- step
    def step(self, t, dt, prev_vhat):
        ents = self.W.sample(t)
        me = None
        for e in ents:
            if "Player" in e[1]:
                me = e
                break
        if me is None:
            return "no-player", prev_vhat
        me_off, me_nm, me_p, me_v, me_r = me
        M = self.W.rot(me_off)
        if M is None:
            return "no-orientation", prev_vhat
        fwd = M[self.W.fwd_row] * self.W.fwd_sign
        right = M[(self.W.fwd_row + 1) % 3]
        up = np.cross(fwd, right)

        speed = float(np.linalg.norm(me_v))
        vhat = norm(me_v) if speed > 1.0 else fwd
        self.update_tau(fwd, me_v, prev_vhat, dt)

        # body angular velocity for the damping term
        w = np.zeros(3)
        if self.prevM is not None and dt > 1e-3:
            D = self.prevM @ M.T
            w = np.array([D[2, 1] - D[1, 2], D[0, 2] - D[2, 0], D[1, 0] - D[0, 1]]) / (2 * dt)
        self.prevM = M

        tgt = self.pick(me_p, me_v, fwd, ents, me_off)
        goal = norm(tgt[3]) if tgt else fwd

        push, worst = self.avoidance(me_p, me_v, me_r, ents, me_off)
        pn = float(np.linalg.norm(push))
        # avoidance outranks the target when it is urgent
        want = norm(goal + push * (3.0 if pn > 0.6 else 1.5)) if pn > 1e-6 else goal

        # Command the NOSE ahead of where the flight path must go, by the
        # measured lag -- steering the nose straight at the target makes the
        # path arrive wide and late.
        nose_cmd = norm(want + (want - vhat) * min(self.tau * 1.6, 2.5))

        ex = float(np.dot(nose_cmd, right))
        ey = float(np.dot(nose_cmd, up))
        ez = float(np.dot(nose_cmd, fwd))
        yaw = math.atan2(ex, ez if abs(ez) > 1e-3 else 1e-3)
        pitch = math.atan2(ey, ez if abs(ez) > 1e-3 else 1e-3)
        if ez < 0:
            yaw = math.copysign(math.pi / 2, ex if ex else 1.0)

        sx = max(-1.0, min(1.0, self.KP * yaw - self.KD * float(np.dot(w, up))))
        sy = max(-1.0, min(1.0, -(self.KP * pitch - self.KD * float(np.dot(w, right)))))

        # fire only when the *flight path* is clear and the nose is on target
        aim_ok = tgt and abs(yaw) < self.FIRE_CONE and abs(pitch) < self.FIRE_CONE
        fire = bool(aim_ok and tgt[4] < self.FIRE_RANGE and pn < 1.2)

        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

        drift = math.degrees(ang(fwd, vhat))
        msg = (f"spd={speed:6.0f} drift={drift:5.1f}d tau={self.tau:4.2f} "
               f"yaw={math.degrees(yaw):+6.1f} pit={math.degrees(pitch):+6.1f} "
               f"stick=({sx:+.2f},{sy:+.2f}) fire={int(fire)}")
        if tgt:
            msg += f" tgt={tgt[1][3:22]:<20} d={tgt[4]:7.0f}"
        if worst:
            msg += (f" | AVOID {worst[1][3:20]} miss={worst[3]:6.0f} "
                    f"t={worst[4]:4.1f}s u={worst[0]:.2f}")
        return msg, vhat

    def run(self, secs, hz=10.0):
        self.W.scan()
        t0 = time.time()
        last, last_scan, prev_vhat = t0, 0.0, None
        while time.time() - t0 < secs:
            t = time.time()
            if t - last_scan > 4.0:
                ents = self.W.scan()
                last_scan = t
                c = Counter(faction(self.W.defs[va]) for _, va in ents)
                print(f"[{t-t0:6.1f}] scan: {len(ents)} entities {dict(c)}",
                      file=self.log, flush=True)
            msg, prev_vhat = self.step(t, t - last, prev_vhat)
            last = t
            print(f"[{t-t0:6.1f}] {msg}", file=self.log, flush=True)
            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 90.0
    W = World(cfg)
    Navigator(W, Pad(), dry="--dry" in sys.argv).run(secs)


if __name__ == "__main__":
    main()