#!/usr/bin/env python3
"""Type every live entity, and locate the player, via the definition pointer.

A live entity object keeps a pointer to its parsed `.tbl` definition (vtable
0x820af844, addresses known) at a **fixed offset from its position triple**.
Finding that offset once types every moving object in the scene at a stroke:
enemies, friendlies and us, separated from the thousands of moving particles.

The offset is *derived*, not assumed: for every moving triple, every nearby word
is checked against the set of definition addresses, and the winning delta is the
one that repeats across many independent entities.

Sub-commands:
  delta            find the (position -> definition pointer) offset
  list [delta]     type every live entity and print it
  self [delta]     the player's entity, its object window, and its orientation
"""
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


def definitions(w):
    out = {}
    for off in w.scan_vtable(gworld.DEF_VTABLE):
        nm = w.name_of(off)
        if nm and nm.startswith("UN_"):
            va = gmem.primary_va(off)
            if va is not None:
                out[struct.pack(">I", va)] = nm
    return out


# Entity objects live in one heap region; scanning only it turns a 250 MB
# double-read into a few MB. That matters for more than speed: the full scan
# competes with the emulator for every core under lavapipe, and a game that
# does not advance a frame between the two samples has nothing "moving" in it.
ENT_VA_LO, ENT_VA_HI = 0xBD000000, 0xBE000000


def moving(fd, size, dt=0.5, lo=1.0, hi=4000.0, va_range=(ENT_VA_LO, ENT_VA_HI)):
    def arrays():
        if va_range:
            f0, f1 = gmem.va_to_off(va_range[0]), gmem.va_to_off(va_range[1])
        else:
            f0, f1 = 0, size
        for a, b in gmem.extents(fd, size):
            a, b = max(a, f0), min(b, f1)
            n = (b - a) // 4 * 4
            if n >= 64:
                yield a, np.frombuffer(os.pread(fd, n, a), dtype=">f4")

    s0 = dict(arrays())
    t0 = time.time()
    time.sleep(dt)
    s1 = dict(arrays())
    t1 = time.time()
    out = []
    # Match extents by their file offset. Zipping the two lists positionally is
    # wrong: the game allocates between the samples, the sparse extent layout
    # shifts, and every subsequent pair misaligns -- which silently reports
    # almost nothing as moving.
    for o, a in s0.items():
        b = s1.get(o)
        if b is None or len(a) != len(b):
            continue
        with np.errstate(invalid="ignore"):
            af = np.nan_to_num(a.astype(np.float64), nan=0, posinf=0, neginf=0)
            bf = np.nan_to_num(b.astype(np.float64), nan=0, posinf=0, neginf=0)
        d = bf - af
        idx = np.flatnonzero(np.abs(d) > 1e-4)
        cand = idx[np.isin(idx + 1, idx) & np.isin(idx + 2, idx)]
        for i in cand:
            v = np.array([d[i], d[i + 1], d[i + 2]])
            sp = float(np.linalg.norm(v)) / (t1 - t0)
            if lo < sp < hi:
                out.append((o + int(i) * 4, tuple(af[i:i + 3]), sp))
    return out


def find_delta(fd, defs, movers, radius=0x400):
    votes = Counter()
    for off, pos, sp in movers[:4000]:
        lo = max(0, off - radius)
        blob = os.pread(fd, radius * 2, lo)
        for k in range(0, len(blob) - 3, 4):
            if blob[k:k + 4] in defs:
                votes[(lo + k) - off] += 1
    return votes


def typed(fd, defs, movers, delta):
    out = []
    for off, pos, sp in movers:
        try:
            b = os.pread(fd, 4, off + delta)
        except OSError:
            continue
        nm = defs.get(b)
        if nm:
            out.append((off, nm, pos, sp))
    return out


def main():
    cmd = sys.argv[1] if len(sys.argv) > 1 else "delta"
    w = gworld.World()
    fd, size = w.fd, w.size
    defs = definitions(w)
    print(f"# {len(defs)} unit definitions", flush=True)
    movers = moving(fd, size)
    print(f"# {len(movers)} moving triples", flush=True)

    if cmd == "delta":
        votes = find_delta(fd, defs, movers)
        print("# candidate (position -> definition pointer) deltas:")
        for d, n in votes.most_common(8):
            print(f"   {d:+#08x}  seen {n}")
        return

    delta = int(sys.argv[2], 0) if len(sys.argv) > 2 else 0x130
    ents = typed(fd, defs, movers, delta)
    uniq = {}
    for off, nm, pos, sp in ents:
        uniq.setdefault((nm, tuple(round(c, 1) for c in pos)), (off, nm, pos, sp))
    ents = list(uniq.values())
    print(f"# {len(ents)} typed live entities (delta {delta:+#x})")

    if cmd == "list":
        c = Counter(nm for _, nm, _, _ in ents)
        for nm, k in c.most_common():
            print(f"   {k:4d}  {nm}")
        for off, nm, pos, sp in sorted(ents, key=lambda e: e[1])[:60]:
            print(f"  {gmem.primary_va(off):#010x} {nm:<40} "
                  f"({pos[0]:+9.1f},{pos[1]:+9.1f},{pos[2]:+9.1f}) {sp:7.1f}/s")
        return

    if cmd == "self":
        me = [e for e in ents if "Player" in e[1]]
        if not me:
            sys.exit("player entity not found")
        off, nm, pos, sp = me[0]
        print(f"# player entity: pos va {gmem.primary_va(off):#010x} {nm}")
        print(f"#   pos ({pos[0]:+.1f},{pos[1]:+.1f},{pos[2]:+.1f})  speed {sp:.1f}/s")
        # orientation inside the same object
        # The rotation is stored with a 16-byte row stride (a 4x4 transform
        # whose translation row is the position we already have), so a test for
        # nine *contiguous* floats structurally cannot find it. Test both.
        lo = max(0, off - 0x800)
        blob = os.pread(fd, 0x1000, lo)
        found = []
        for k in range(0, len(blob) - 48, 4):
            for stride, tag in ((12, "3x3"), (16, "4x4")):
                try:
                    rows = [np.array(struct.unpack_from(">3f", blob, k + stride * r))
                            for r in range(3)]
                except struct.error:
                    continue
                M = np.array(rows)
                if not np.all(np.isfinite(M)) or np.max(np.abs(M)) > 1.001:
                    continue
                if np.max(np.abs(M @ M.T - np.eye(3))) > 3e-3:
                    continue
                if abs(np.linalg.det(M) - 1.0) > 1e-2:
                    continue
                found.append(((lo + k) - off, M, stride))
                break
        print(f"# orthonormal 3x3 blocks inside the player object: {len(found)}")
        for d, M, st in found[:6]:
            print(f"   pos{d:+#07x} stride {st}:  " + "  ".join(
                "(" + ",".join(f"{v:+.3f}" for v in row) + ")" for row in M))
        if len(sys.argv) > 3 and found:
            import json
            # Which of the orthonormal blocks is the CRAFT's attitude? The one
            # with a row along the direction of travel. Taking found[0] is what
            # produced a config with rot_delta -0x764 and a nonsense forward
            # axis: several blocks inside the object are orthonormal (bone or
            # camera frames), and only the craft's own has a row that tracks
            # where the craft is going.
            p0 = np.array(pos)
            time.sleep(0.35)
            p1 = np.array(struct.unpack(">3f", os.pread(fd, 12, off)))
            step = p1 - p0
            if np.linalg.norm(step) < 1e-3:
                sys.exit("craft is not moving — cannot bind the forward axis")
            vdir = step / np.linalg.norm(step)
            best = None
            for d, M, st in found:
                cs = [float(M[r] @ vdir) for r in range(3)]
                r = int(np.argmax([abs(c) for c in cs]))
                if best is None or abs(cs[r]) > abs(best[3]):
                    best = (d, M, st, cs[r], r)
            rot_delta, M, rot_stride, cos, row = best
            sign = 1 if cos > 0 else -1
            print(f"#   attitude block pos{rot_delta:+#07x} stride {rot_stride}: "
                  f"forward = row {row} (sign {sign:+d}), cos={cos:+.3f}")
            if abs(cos) < 0.9:
                print("#   WARNING: no block tracks the flight path (|cos| < 0.9)"
                      " — the craft may be drifting hard; re-run while flying straight")
            cfg = {"def_delta": delta, "rot_delta": rot_delta,
                   "rot_stride": rot_stride,
                   "fwd_row": row, "fwd_sign": sign, "fwd_cos": round(cos, 4),
                   "va_lo": ENT_VA_LO, "va_hi": ENT_VA_HI}
            json.dump(cfg, open(sys.argv[3], "w"), indent=1)
            print("# wrote " + sys.argv[3] + ": " + json.dumps(cfg))
        return


if __name__ == "__main__":
    main()