#!/usr/bin/env python3
"""Type the live entities: link each moving position back to a unit definition.

findplayer.py locates positions by motion alone, but a bare (x,y,z) says nothing
about *what* is moving. Every live entity should reference its parsed `.tbl`
definition (vtable 0x820af844, one object per unit type, addresses known), so
scanning the neighbourhood of each position for a word equal to a definition
address both types the entity and reveals the live object's own layout — the
delta from that pointer to the position triple is the position offset.

Usage: liveents.py [radius_hex]
"""
import os
import struct
import sys
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):
    """{definition_va: unit_id}"""
    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[va] = nm
    return out


def main():
    radius = int(sys.argv[1], 16) if len(sys.argv) > 1 else 0x600
    w = gworld.World()
    defs = definitions(w)
    print(f"# {len(defs)} unit definitions")
    by_word = {struct.pack(">I", va): nm for va, nm in defs.items()}

    # sample twice to locate movers, exactly as findplayer.py does
    import time
    fd, size = w.fd, w.size

    def snap():
        return [(a, np.frombuffer(os.pread(fd, (b - a) // 4 * 4, a), dtype=">f4"))
                for a, b in gmem.extents(fd, size) if (b - a) >= 64]

    s0 = snap(); t0 = time.time()
    time.sleep(0.4)
    s1 = snap(); t1 = time.time()
    shapes0 = {(o, len(a)) for o, a in s0}
    pairs = [(o, a, b) for (o, a), (o2, b) in zip(s0, s1)
             if o == o2 and len(a) == len(b) and (o, len(a)) in shapes0]

    movers = []
    for base, a, b in pairs:
        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
        m = np.abs(d) > 1e-3
        idx = np.flatnonzero(m)
        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 20.0 < sp < 3000.0:
                movers.append((base + int(i) * 4, tuple(float(af[i + j]) for j in range(3)), sp))

    print(f"# {len(movers)} moving triples; typing them...")
    typed, untyped = [], 0
    seen = set()
    for off, pos, sp in movers:
        lo = max(0, off - radius)
        blob = os.pread(fd, radius * 2, lo)
        hit = None
        for k in range(0, len(blob) - 3, 4):
            nm = by_word.get(blob[k:k + 4])
            if nm:
                hit = (lo + k, nm)
                break
        if hit:
            ptr_off, nm = hit
            key = (ptr_off, nm)
            if key in seen:
                continue
            seen.add(key)
            typed.append((ptr_off, off - ptr_off, nm, pos, sp))
        else:
            untyped += 1

    print(f"# typed {len(typed)}, untyped {untyped}")
    deltas = Counter(d for _, d, _, _, _ in typed)
    print(f"# most common (def-pointer -> position) deltas: {deltas.most_common(6)}")
    for ptr_off, d, nm, pos, sp in sorted(typed, key=lambda x: x[2])[:40]:
        va = gmem.primary_va(ptr_off)
        print(f"  obj~{va:#010x} defptr+{d:#06x}  {nm:<40} "
              f"({pos[0]:+10.1f},{pos[1]:+10.1f},{pos[2]:+10.1f}) {sp:7.1f}/s")


if __name__ == "__main__":
    main()