#!/usr/bin/env python3
"""Recover each ship part's runtime WORLD transform from a Canary draw log.

The draw logger (branch capture-drawlog) now dumps, per distinct draw:
  - the first ~8 object-space vertex positions (pre-transform), and
  - the vertex-shader float constants c0..c31 (which hold the transform matrices).

The game bakes each part's WORLD matrix into the constants, so the shader's
World*View*Proj (WVP) for a part is  VP * part_World.  The body (bdy_01) is drawn
with an identity world, so body_WVP = VP.  Hence for any part drawn in the SAME
frame (the ship's parts always are):

    part_World = inv(body_WVP) * part_WVP          (View/Proj cancel)

We don't know a-priori which 4 consecutive constants form the WVP, so we try every
block c[i..i+3] and keep the one for which inv(body_WVP)*part_WVP is a RIGID
transform (orthonormal 3x3 + translation) for the parts — that uniquely identifies
the matrix block.  The translation column of each part's world matrix is the
placement we want (esp. the ±X nacelle offset the static XBG7 graph lacks).

Usage:  analyze_drawlog_wvp.py xenia_re_draws.log
Identify the body draw via --body-indices (the draw with the most indices) or by
matching object-space positions to our decoded sub-meshes.
"""
import sys, re
import numpy as np

def parse(path):
    draws = []
    cur = None
    consts = {}
    mode = None
    for line in open(path):
        if line.startswith("DRAW "):
            if cur is not None:
                cur["consts"] = consts
                draws.append(cur)
            cur = {"header": line.strip(), "positions": [], "textures": [],
                   "size_words": 0}
            consts = {}
            mode = None
            m = re.search(r"index\[base=0x([0-9A-Fa-f]+) count=(\d+)", line)
            cur["idx_base"] = int(m.group(1), 16) if m else None
            cur["idx_count"] = int(m.group(2)) if m else None
            m = re.search(r"indices=(\d+)", line)
            cur["indices"] = int(m.group(1)) if m else None
        elif cur is not None and line.strip().startswith("stream "):
            # Largest vertex buffer (size_words) = the hull (body, identity world).
            m = re.search(r"size_words=(\d+)", line)
            if m:
                cur["size_words"] = max(cur["size_words"], int(m.group(1)))
        elif cur is None:
            continue
        elif line.strip().startswith("positions:"):
            cur["positions"] = [tuple(map(float, t.split(",")))
                                for t in re.findall(r"\(([^)]+)\)", line)]
        elif line.strip().startswith("textures:"):
            cur["textures"] = re.findall(r"base=0x([0-9A-Fa-f]+)", line)
        elif line.strip().startswith("vsconst"):
            mode = "vsconst"
        elif mode == "vsconst":
            m = re.match(r"\s*c(\d+)\s+(-?[\d.]+)\s+(-?[\d.]+)\s+(-?[\d.]+)\s+(-?[\d.]+)", line)
            if m:
                consts[int(m.group(1))] = np.array([float(m.group(i)) for i in range(2, 6)])
            else:
                mode = None
    if cur is not None:
        cur["consts"] = consts
        draws.append(cur)
    return draws

def wvp_from(consts, i):
    """4x4 from constants c[i..i+3] as rows."""
    if not all(k in consts for k in (i, i+1, i+2, i+3)):
        return None
    return np.array([consts[i], consts[i+1], consts[i+2], consts[i+3]])

def is_rigid(M):
    """3x3 upper-left orthonormal (rotation/reflection), within tolerance."""
    R = M[:3, :3]
    if not np.all(np.isfinite(R)):
        return False
    G = R @ R.T
    return np.allclose(G, np.eye(3), atol=1e-2) or np.allclose(G / (np.trace(G)/3), np.eye(3), atol=1e-2)

def main():
    if len(sys.argv) < 2:
        print(__doc__); return
    draws = parse(sys.argv[1])
    print(f"parsed {len(draws)} draws")
    # Body = the draw with the largest vertex buffer (the ~10891-vert hull, drawn
    # with an identity world). idx_count is unreliable (the hull's 9 material
    # groups dedup to the first group's partial count).
    body = max(draws, key=lambda d: d.get("size_words") or 0)
    print(f"body draw: {body['header'][:90]} size_words={body['size_words']}")
    # Try every constant block; report the one giving rigid part transforms.
    for i in range(0, 29):
        bwvp = wvp_from(body["consts"], i)
        if bwvp is None:
            continue
        try:
            binv = np.linalg.inv(bwvp)
        except np.linalg.LinAlgError:
            continue
        worlds = []
        rigid = 0
        for d in draws:
            pw = wvp_from(d["consts"], i)
            if pw is None:
                continue
            W = binv @ pw           # column-vector convention; may need transpose
            worlds.append((d, W))
            if is_rigid(W) or is_rigid(W.T):
                rigid += 1
        if rigid >= max(2, len(worlds)//2):
            print(f"\n=== WVP block c{i}..c{i+3}: {rigid}/{len(worlds)} rigid ===")
            for d, W in worlds:
                # translation is the last column (or last row, convention-dependent)
                tcol = W[:3, 3]
                trow = W[3, :3]
                print(f"  idx={d['idx_count']:>6} base={d['idx_base'] and hex(d['idx_base'])}"
                      f"  t_col=({tcol[0]:+.2f},{tcol[1]:+.2f},{tcol[2]:+.2f})"
                      f"  t_row=({trow[0]:+.2f},{trow[1]:+.2f},{trow[2]:+.2f})")

if __name__ == "__main__":
    main()