firestorm/MW4COMPARE/tools/decompile/armature.py

#!/usr/bin/env python3
"""
armature.py - rebuild a mech's .armature source from its packed records.

    python3 armature.py <mechRecordDir> [-o out.armature]
    python3 armature.py --verify                 # check against all 65 known chassis

The packer merges <mech>.armature into <mech>.contents (via `!include=`) and
then, for every contents page that has Child= entries, emits two records
(MWMover_Tool.cpp ~78):

    <mech>.contents[<joint>]{sites}      children whose name starts 'site_'
                                         but is not 'site_eye*' - stored as
                                         YawPitchRoll + Point3D + name
    <mech>.contents[<joint>]{armature}   every other child, as a CreateMessage
                                         carrying jointName + localToParent

Between them those two records hold every page of the original .armature.
Verified on Atlas: 61 of 61 page names recovered, values exact.

Rotation notes
--------------
* {sites} stores YawPitchRoll directly, so those angles are exact.
* Joint messages only carry a matrix. Across all 65 chassis, 1348 of 1453 joint
  matrices are the identity, i.e. `Rotation=0 0 0`.
* site_lfoot / site_rfoot appear in BOTH streams - a deliberate hack in
  MWMover_Tool.cpp ("Jerry this will get deleted when you fix your foot
  problem") - and their matrix is the identity while the real angle is in the
  {sites} record. The {sites} value always wins.
* site_eye* goes only to the armature stream, so its angle comes from the
  matrix.
"""
import sys, os, glob, math, struct, re, argparse, collections

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mw4msg

OUR_MECH_SOURCE = "/home/rich/Repositories/firestorm/Gameleap/mw4/Content/Mechs"
OUR_RECORDS = "/home/rich/Repositories/FS_Ours_extracted/core/mechs"

REC_RE = re.compile(r'^(?P<chassis>.+?)\.contents\[(?P<parent>[^\]]+)\]\{(?P<kind>armature|sites)\}$',
                    re.I)


def ypr_from_matrix(r):
    """3x3 -> the (a, b, c) triple as written in Rotation=, in degrees.

    Stuff's YawPitchRoll writes the X-axis angle in the middle slot; confirmed
    against kodiak/site_eyepoint (matrix gives 0.0688 deg, source says
    0.069104).
    """
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = r
    pitch = math.asin(max(-1.0, min(1.0, m21)))
    if abs(m21) < 0.999999:
        yaw = math.atan2(-m20, m22)
        roll = math.atan2(-m01, m11)
    else:
        yaw = math.atan2(m02, m00)
        roll = 0.0
    return (math.degrees(yaw), math.degrees(pitch), math.degrees(roll))


def is_identity(r, eps=1e-5):
    ident = (1, 0, 0, 0, 1, 0, 0, 0, 1)
    return all(abs(a - b) < eps for a, b in zip(r, ident))


def rebuild(mech_dir):
    """-> (chassisName, entries, children)

    entries  list of {'name','parent','rot','trans'} - one per armature page.
             Keyed on (parent, name), not name alone: Victor legitimately has
             two different [site_lshellport] pages hanging off different joints.
    children {parentName: [childName, ...]} preserving order and multiplicity.
    """
    entries = []                       # one dict per armature page occurrence
    children = collections.defaultdict(list)
    site_children = collections.defaultdict(set)
    chassis = None
    # {sites} first: site_lfoot/site_rfoot are written to BOTH streams and only
    # the {sites} copy carries their real angle.
    records = (sorted(glob.glob(os.path.join(mech_dir, "*{sites}"))) +
               sorted(glob.glob(os.path.join(mech_dir, "*{armature}"))))
    for path in records:
        m = REC_RE.match(os.path.basename(path))
        if not m:
            continue
        chassis = chassis or m.group("chassis")
        parent, kind = m.group("parent"), m.group("kind").lower()
        with open(path, "rb") as fh:
            data = fh.read()
        children[parent]                               # ensure the parent exists

        if kind == "sites":
            for name, rot, trans in mw4msg.read_sites(data):
                entries.append({
                    "name": name, "parent": parent,
                    "rot": tuple(math.degrees(a) for a in rot), "trans": trans})
                children[parent].append(name)
                site_children[parent].add(name)
        else:
            for _off, msg in mw4msg.walk(data):
                name = mw4msg.joint_name(msg)
                if not name:
                    continue
                # site_lfoot/site_rfoot are deliberately written to both streams;
                # counting the armature copy too would double them.
                if name in site_children[parent]:
                    continue
                children[parent].append(name)
                r = mw4msg.rotation3x3(msg)
                entries.append({
                    "name": name, "parent": parent,
                    "rot": (0.0, 0.0, 0.0) if is_identity(r) else ypr_from_matrix(r),
                    "trans": mw4msg.translation(msg)})
    return chassis, entries, children


def emit(entries, children):
    """Render as .armature text, children before the joint that owns them."""
    by_name = {}
    for e in entries:
        by_name.setdefault(e["name"], []).append(e)
    out, done = [], set()

    def visit(name):
        if name in done:
            return
        done.add(name)
        for c in children.get(name, ()):
            visit(c)
        for e in by_name.get(name, [{"rot": (0.0, 0.0, 0.0), "trans": (0.0, 0.0, 0.0)}]):
            out.append(f"[{name}]")
            out.append("Rotation=%.6f %.6f %.6f" % (e["rot"] or (0.0, 0.0, 0.0)))
            out.append("Translation=%.6f %.6f %.6f" % (e["trans"] or (0.0, 0.0, 0.0)))
            for c in children.get(name, ()):
                out.append(f"Child={c}")
            out.append("")

    all_children = {c for lst in children.values() for c in lst}
    for name in list(children) + list(by_name):
        if name not in all_children:
            visit(name)
    for name in list(children) + list(by_name):
        visit(name)
    return "\r\n".join(out) + "\r\n"


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("mechdir", nargs="?", help="directory of packed records for one mech")
    ap.add_argument("-o", "--out", help="write .armature here instead of stdout")
    ap.add_argument("--verify", action="store_true",
                    help="rebuild all 65 known chassis and diff against their source")
    args = ap.parse_args()

    if args.verify:
        here = os.path.dirname(os.path.abspath(__file__))
        os.execvp("python3", ["python3", os.path.join(here, "verify_armature.py")])
    if not args.mechdir:
        ap.error("give a mech record directory, or --verify")

    chassis, entries, children = rebuild(args.mechdir)
    text = emit(entries, children)
    if args.out:
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        open(args.out, "wb").write(text.encode("latin-1"))
        print(f"{chassis}: {len(entries)} pages -> {args.out}")
    else:
        sys.stdout.write(text)


if __name__ == "__main__":
    main()