#!/usr/bin/env python3
"""Check the `MCOL` decode in `docs/re/structures/mcol-collision.md`.

`MCOL` (`hidden/MiscBin.pak`, 11 objects) shares `REGN`'s container, so the
`POF0` reader is imported from `regn_decode.py` rather than duplicated.

    ./mcol_probe.py stride  <MiscBin.pak>       # the 0x5C block is stride 16, not 12
    ./mcol_probe.py cells   <MiscBin.pak>       # the u16s name spheres in their cell
    ./mcol_probe.py mesh    <MiscBin.pak>       # vertices, triangles, bounding spheres
    ./mcol_probe.py obj     <MiscBin.pak> <hash> <out.obj>
    ./mcol_probe.py verify  <MiscBin.pak>       # all checks, with pass/fail

The `cells` check is the powered one: a `u16` is reached *through* a specific
grid cell, so the sphere it names must reach that cell.  A bound-check ("is it
a valid index?") cannot fail here and is deliberately not used -- see the
hazard section of the doc.
"""

import math
import os
import random
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from regn_decode import pak_entries, pof0_pointer_slots, FIXUP_BASE

SPHERE = 16  # stride of the 0x5C record: centre f32[3] then radius f32
VERTEX = 12  # stride of the 0x54 record: f32[3]
TRI = 6      # stride of the 0x58 record: u16[3]


def align16(n):
    return (n + 15) // 16 * 16


def u32(b, o): return struct.unpack_from(">I", b, o)[0]
def u16(b, o): return struct.unpack_from(">H", b, o)[0]
def f32(b, o): return struct.unpack_from(">f", b, o)[0]
def ptr(b, o): return u32(b, o) + FIXUP_BASE


def objects(pak):
    return [(h, b) for h, b in pak_entries(pak) if b[:4] == b"MCOL"]


class Mcol:
    def __init__(self, blob):
        self.b = blob
        self.p54, self.p58, self.p5c, self.p74 = (ptr(blob, o)
                                                  for o in (0x54, 0x58, 0x5C, 0x74))
        self.sphere_bytes = self.p54 - self.p5c
        self.n = self.sphere_bytes // SPHERE
        self.nv = u16(blob, 0x50)   # vertex count
        self.nt = u16(blob, 0x52)   # triangle count == sphere count
        self.bbox_min = [f32(blob, 0x10 + 4 * i) for i in range(3)]
        self.cell_size = [f32(blob, 0x40 + 4 * i) for i in range(3)]

    def vertex(self, i):
        o = self.p54 + VERTEX * i
        return [f32(self.b, o + 4 * k) for k in range(3)]

    def triangle(self, i):
        o = self.p58 + TRI * i
        return [u16(self.b, o + 2 * k) for k in range(3)]

    def sphere(self, i):
        o = self.p5c + SPHERE * i
        return [f32(self.b, o + 4 * k) for k in range(3)], f32(self.b, o + 12)

    def cell_box(self, cx, cy, cz):
        lo = [self.bbox_min[k] + self.cell_size[k] * (cx, cy, cz)[k] for k in range(3)]
        return lo, [lo[k] + self.cell_size[k] for k in range(3)]

    def cells(self):
        """(cell box, [u16]) for every A record that carries a cell index."""
        b = self.b
        for slot in sorted(s for s in set(pof0_pointer_slots(b)) if s >= self.p74):
            rec = slot - 8
            if b[rec + 3] != 1:          # the other interleaved array
                continue
            pb = ptr(b, slot)
            if pb + 8 > len(b):
                continue
            count, refs = u32(b, pb), ptr(b, pb + 4)
            if count > 100_000 or refs + 2 * count > len(b):
                continue
            idx = list(struct.unpack_from(">%dH" % count, b, refs)) if count else []
            yield self.cell_box(b[rec], b[rec + 1], b[rec + 2]), idx


def dist2(c, lo, hi):
    """Squared distance from a point to an axis-aligned box (0 if inside)."""
    d = 0.0
    for k in range(3):
        if c[k] < lo[k]:
            d += (lo[k] - c[k]) ** 2
        elif c[k] > hi[k]:
            d += (c[k] - hi[k]) ** 2
    return d


def cmd_stride(pak):
    print(f"{'object':>8} {'len(0x5C)':>10} {'/12':>10} {'/16':>6} {'max u16':>8}")
    ok12 = ok16 = total = 0
    for h, blob in objects(pak):
        m = Mcol(blob)
        top = max((v for _, idx in m.cells() for v in idx if v < m.n), default=-1)
        div12 = m.sphere_bytes % 12 == 0
        print(f"{h:08x} {m.sphere_bytes:10d} "
              f"{m.sphere_bytes / 12:10.2f}{'' if div12 else ' *'} "
              f"{m.n:6d} {top:8d}"
              f"{'   max == n-1' if top == m.n - 1 else ''}")
        total += 1
        ok12 += div12
        ok16 += top == m.n - 1
    print(f"\nlength divisible by 12: {ok12}/{total}   "
          f"max u16 == len/16 - 1: {ok16}/{total}")
    return ok16 == total and ok12 < total


def cmd_cells(pak, seed=12345):
    random.seed(seed)
    hit = ctl = tot = 0
    for _, blob in objects(pak):
        m = Mcol(blob)
        spheres = [m.sphere(i) for i in range(m.n)]
        for (lo, hi), idx in m.cells():
            for v in idx:
                if v >= m.n:
                    continue
                tot += 1
                c, r = spheres[v]
                hit += dist2(c, lo, hi) <= r * r
                c2, r2 = spheres[random.randrange(m.n)]
                ctl += dist2(c2, lo, hi) <= r2 * r2
    print(f"referenced sphere reaches its own cell : {hit}/{tot} = {100*hit/tot:.2f}%")
    print(f"random sphere, same object  (control)  : {ctl}/{tot} = {100*ctl/tot:.2f}%")
    return hit / tot > 0.99 and ctl / tot < 0.25


def cmd_mesh(pak):
    """The 0x50 counts give both blocks a stride, and sphere i bounds triangle i."""
    import collections
    print(f"{'object':>8} {'verts':>6} {'tris':>6} {'len 0x54':>9} {'len 0x58':>9} {'spheres':>8}")
    n = sv = st = ss = 0
    enc = tight = ctl = tri_total = 0
    edges2 = edges = 0
    random.seed(7)
    for h, blob in objects(pak):
        m = Mcol(blob)
        l54, l58 = m.p58 - m.p54, m.p74 - m.p58
        n += 1
        sv += align16(VERTEX * m.nv) == l54
        st += align16(TRI * m.nt) == l58
        ss += m.nt == m.n
        print(f"{h:08x} {m.nv:6d} {m.nt:6d} {l54:9d} {l58:9d} {m.n:8d}")

        V = [m.vertex(i) for i in range(m.nv)]
        T = [m.triangle(i) for i in range(m.nt)]
        ec = collections.Counter()
        for i, t in enumerate(T):
            c, r = m.sphere(i)
            far = max(math.dist(V[j], c) for j in t)
            tri_total += 1
            enc += far <= r * 1.0001
            tight += abs(far / r - 1) < 0.01
            tj = T[random.randrange(m.nt)]
            ctl += max(math.dist(V[j], c) for j in tj) <= r * 1.0001
            for k in range(3):
                a, b_ = t[k], t[(k + 1) % 3]
                ec[(min(a, b_), max(a, b_))] += 1
        edges += len(ec)
        edges2 += sum(1 for v in ec.values() if v == 2)
    print(f"\nlen(0x54) == align16(12 * verts)          : {sv}/{n}")
    print(f"len(0x58) == align16(6 * tris)            : {st}/{n}")
    print(f"triangle count == sphere count            : {ss}/{n}")
    print(f"sphere i encloses triangle i              : {enc}/{tri_total} = {100*enc/tri_total:.2f}%")
    print(f"  ...and is tight to within 1%            : {tight}/{tri_total} = {100*tight/tri_total:.2f}%")
    print(f"sphere i encloses a random triangle (ctl) : {ctl}/{tri_total} = {100*ctl/tri_total:.2f}%")
    print(f"edges shared by exactly two triangles     : {edges2}/{edges} = {100*edges2/edges:.2f}%")
    return (sv == st == ss == n and enc == tri_total and edges2 == edges
            and ctl / tri_total < 0.10)


def cmd_obj(pak, want, out):
    """Export one object as a Wavefront OBJ, so the decode can be looked at."""
    want = int(want, 16)
    for h, blob in objects(pak):
        if h != want:
            continue
        m = Mcol(blob)
        with open(out, "w") as fh:
            fh.write(f"# MCOL {h:08x} -- {m.nv} vertices, {m.nt} triangles\n")
            for i in range(m.nv):
                fh.write("v %.4f %.4f %.4f\n" % tuple(m.vertex(i)))
            for i in range(m.nt):
                a, b_, c = m.triangle(i)
                fh.write(f"f {a+1} {b_+1} {c+1}\n")
        print(f"{out}: {m.nv} vertices, {m.nt} triangles")
        return True
    raise SystemExit(f"no MCOL object {want:08x}")


def main():
    if len(sys.argv) < 3:
        raise SystemExit(__doc__)
    cmd, pak = sys.argv[1], sys.argv[2]
    if cmd == "stride":
        ok = cmd_stride(pak)
    elif cmd == "cells":
        ok = cmd_cells(pak)
    elif cmd == "mesh":
        ok = cmd_mesh(pak)
    elif cmd == "obj":
        ok = cmd_obj(pak, sys.argv[3], sys.argv[4])
    elif cmd == "verify":
        a = cmd_stride(pak)
        print()
        b = cmd_cells(pak)
        print()
        c = cmd_mesh(pak)
        ok = a and b and c
    else:
        raise SystemExit(__doc__)
    print("\nPASS" if ok else "\nFAIL")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()