From 03f9bf8a147b6f808f1f125cdcbfb48eb0b4b65f Mon Sep 17 00:00:00 2001
From: TigerKat <karynax23@protonmail.com>
Date: Fri, 9 Apr 2021 00:48:45 +0930
Subject: [PATCH] Added support for writing .anim files, along with exporting
 both animation and skeletons.

---
 __init__.py    |  12 +--
 anim.py        | 197 +++++++++++++++++++++++++++++++++++++++++++--
 export_anim.py | 212 +++++++++++++++++++++++++++++++++++++++++++++++++
 export_skel.py |  12 +++
 util.py        |   2 +
 5 files changed, 422 insertions(+), 13 deletions(-)
 create mode 100644 export_anim.py
 create mode 100644 export_skel.py

diff --git a/__init__.py b/__init__.py
index 28084f9..8214d1d 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,7 +2,7 @@
 bl_info = {
     "name": "City of Heroes (.geo)",
     "author": "TigerKat",
-    "version": (0, 2, 7),
+    "version": (0, 2, 8),
     "blender": (2, 80, 0),
     "location": "File > Import/Export,",
     "description": "City of Heroes (.geo)",
@@ -138,7 +138,7 @@ class ExportSkel(bpy.types.Operator, ExportHelper):
         keywords = self.as_keywords(ignore=("filter_glob",
                                             "check_existing",
         ))
-        return export_skel.save(self, context, 1.0, **keywords)
+        return export_skel.save_skel(self, context, 1.0, **keywords)
 
 class ExportAnim(bpy.types.Operator, ExportHelper):
     bl_idname = "export_scene.geo_anim"
@@ -168,10 +168,10 @@ def menu_func_export(self, context):
                          text="City of Heroes (Feet) (.geo)")
     self.layout.operator(ExportGeoMetric.bl_idname,
                          text="City of Heroes (Meters) (.geo)")
-    #self.layout.operator(ExportSkel.bl_idname,
-    #                     text="City of Heroes Skeleton (skel_*.anim)")
-    #self.layout.operator(ExportAnim.bl_idname,
-    #                     text="City of Heroes Animation (.anim)")
+    self.layout.operator(ExportSkel.bl_idname,
+                         text="City of Heroes Skeleton (skel_*.anim)")
+    self.layout.operator(ExportAnim.bl_idname,
+                         text="City of Heroes Animation (.anim)")
 
 def make_annotations(cls):
     """Converts class fields to annotations if running with Blender 2.8"""
diff --git a/anim.py b/anim.py
index 1f911ac..3cdb46d 100644
--- a/anim.py
+++ b/anim.py
@@ -79,11 +79,18 @@ class BoneLink:
             return
         if next is None:
             (self.child, self.next, self.boneid) = data_or_child.decode("<iii")
+            if self.child == 0 and self.next == 0 and self.boneid == 0:
+                #All zero indicates an unused bone when written to disk.
+                #Change it to all -1 make it clear it's unused.
+                (self.child, self.next, self.boneid) = (-1, -1, -1)
         else:
             self.child = data_or_child
             self.next = next
             self.boneid = boneid
-
+    def encode(self):
+        if self.boneid == -1:
+            return struct.pack("<iii", 0, 0, 0)
+        return struct.pack("<iii", self.child, self.next, self.boneid)
     def dump(self):
         print("bone: %s (%s)" % (self.boneid, BONES_LIST[self.boneid]))
         print("  child: %s (%s)" % (self.child, BONES_LIST[self.child]))
@@ -93,6 +100,9 @@ class SkeletonHierarchy:
     def __init__(self):
         self.root = -1
         self.bones = []
+    def clear(self):
+        self.root = -1
+        self.bones = []
     def parseData(self, data, size = None):
         self.root = data.decode("<i")[0]
         self.bones = []
@@ -101,10 +111,61 @@ class SkeletonHierarchy:
         count = int(math.floor((size - 4) / 12.0))
         for i in range(count): #BONES_ON_DISK):
             self.bones.append(BoneLink(data))
+    def encode(self):
+        bone_max = len(self.bones)
+        data = struct.pack("<i", self.root)
+        blank_bone = BoneLink().encode()
+        for bl in self.bones:
+            if bl is None:
+                data += blank_bone
+            else:
+                data += bl.encode()
+        for i in range(bone_max, BONES_ON_DISK):
+            data += blank_bone
+        return data
+    def initBones(self, bone_id):
+        if bone_id == -1 or bone_id is None:
+            return
+        #Ensure the array is sized up to bone_id, and it's non-None.
+        if len(self.bones) <= bone_id:
+            #print("initBones(%s): %s: %s" % (bone_id, len(self.bones), self.bones,))
+            self.bones += [None] * (bone_id - len(self.bones) + 1)
+            #print("initBones(%s): %s: %s" % (bone_id, len(self.bones), self.bones,))
+        if self.bones[bone_id] is None or self.bones[bone_id].boneid == -1:
+            self.bones[bone_id] = BoneLink(-1, -1, bone_id)
+
+    def addBone(self, parent_id, bone_id):
+        self.initBones(bone_id)
+        self.initBones(parent_id)
+        if parent_id == -1 or parent_id is None:
+            #New bone is a root bone.
+            if self.root == -1:
+                #It's the first root bone.
+                self.root = bone_id
+            else:
+                #Find the last root bone, and add it to the end of the next chain.
+                last_root_id = self.root
+                while self.bones[last_root_id].next != -1:
+                    last_root_id = self.bones[last_root_id].next
+                self.bones[last_root_id].next = bone_id
+        else:
+            #New bone has a parent.
+            if self.bones[parent_id].child != -1:
+                #Parent already has children, add it to the end of the next chain.
+                last_sibling_id = self.bones[parent_id].child
+                while self.bones[last_sibling_id].next != -1:
+                    last_sibling_id = self.bones[last_sibling_id].next
+                self.bones[last_sibling_id].next = bone_id
+            else:
+                #It's the first child bone.
+                self.bones[parent_id].child = bone_id
     def dump(self):
         print("SkeletonHierarchy: root: %s (%s) bones: %s" % (self.root, BONES_LIST[self.root], len(self.bones)))
-        for b in self.bones:
-            b.dump()
+        for i, b in enumerate(self.bones):
+            if b is None:
+                print("No bone: %s (%s)" % (i, BONES_LIST[i]))
+            else:
+                b.dump()
 
 class BoneAnimTrack:
     def __init__(self, data = b"", srcdata = None):
@@ -183,6 +244,69 @@ class BoneAnimTrack:
             #raise Exception("bone track position type 0x%2.2x unknown" % (self.flags, ))
             pass
         self.srcdata.seek(offset)
+    def encode(self, out_offset, out_data):
+        self.encodePositions()
+        self.encodeRotations()
+        self.data = Data()
+
+        self.rotations_data_offset = len(out_data) + out_offset
+        out_data.write(self.rotation_data)
+        self.positions_data_offset = len(out_data) + out_offset
+        out_data.write(self.position_data)
+
+        self.rot_count = self.rot_fullkeycount = len(self.rotations)
+        self.pos_count = self.pos_fullkeycount = len(self.positions)
+
+        self.data.encode("<I", self.rotations_data_offset)
+        self.data.encode("<I", self.positions_data_offset)
+        self.data.encode("<H", self.rot_fullkeycount)
+        self.data.encode("<H", self.pos_fullkeycount)
+        self.data.encode("<H", self.rot_count)
+        self.data.encode("<H", self.pos_count)
+        self.data.encode("<BBBB", self.bone_id, self.flags, 0, 0)
+        self.rawdata = self.data.data
+        return self.rawdata
+
+    def encodePositions(self):
+        #Choose encoding.
+        require_full = False
+        for pos in self.positions:
+            for v in pos:
+                if abs(v) > 1.0:
+                    require_full = True
+        self.flags = self.flags & ~POSITION_MASK
+        #Set flags and endode data
+        self.position_data = b""
+        if require_full:
+            self.flags = self.flags | POSITION_UNCOMPRESS
+            for pos in self.positions:
+                self.position_data += struct.pack("<fff", pos[0], pos[1], pos[2])
+        else:
+            self.flags = self.flags | POSITION_COMPRESSED_TO_6_BYTES
+            for pos in self.positions:
+                self.position_data += compressVector3_6Byte(pos)
+        pass
+
+    def encodeRotations(self):
+        #Choose encoding.
+        require_full = False
+        #for pos in self.positions:
+        #    for v in pos:
+        #        if abs(v) > 1.0:
+        #            require_full = True
+        self.flags = self.flags & ~ROTATION_MASK
+        #Set flags and endode data
+        self.rotation_data = b""
+        if require_full:
+            self.flags = self.flags | ROTATION_UNCOMPRESS
+            for rot in self.rotattions:
+                self.rotation_data += struct.pack("<ffff", rot[0], rot[1], rot[2], rot[3])
+        else:
+            self.flags = self.flags | ROTATION_COMPRESSED_TO_5_BYTES
+            for rot in self.rotations:
+                self.rotation_data += compressQuaternion_5Byte(rot)
+        pass
+
     def dump(self):
         print("rotations_data_offset: %s" % (self.rotations_data_offset, ))
         print("positions_data_offset: %s" % (self.positions_data_offset, ))
@@ -223,7 +347,7 @@ class Anim:
         self.header_spare_room = b""
 
         self.bone_tracks = []
-        self.skeleton_hierarchy = []
+        self.skeleton_hierarchy = None
 
         pass
     def loadFromData(self, data):
@@ -238,10 +362,10 @@ class Anim:
         pass
     def saveToData(self):
         self.encode()
-        return self.data
+        return self.data.data
     def saveToFile(self, fileout):
         self.encode()
-        fileout.write(self.data)
+        fileout.write(self.data.data)
     def parseData(self):
         self.data.seek(0)
         self.header_size = self.data.decode("i")[0]
@@ -276,7 +400,66 @@ class Anim:
             self.skeleton_hierarchy.parseData(self.data, skel_size)
         pass
     def encode(self):
-        self.data = None
+        if isinstance(self.header_name, str):
+            self.header_name = bytes(self.header_name, "utf-8")
+        if isinstance(self.header_base_anim_name, str):
+            self.header_base_anim_name = bytes(self.header_base_anim_name, "utf-8")
+        header_base_size = 596
+        header_bone_tracks_size = 20 * len(self.bone_tracks)
+        if self.skeleton_hierarchy is None:
+            skel_data = b""
+            header_skeleton_hierarhy_size = 0
+        else:
+            skel_data = self.skeleton_hierarchy.encode()
+            header_skeleton_hierarhy_size = len(skel_data)
+            #todo: verify all defined bones are present in the track data
+        header_total_size = header_base_size + header_bone_tracks_size + header_skeleton_hierarhy_size
+        header_bone_track_data = b""
+        bone_track_data = Data()
+        for bt in self.bone_tracks:
+            header_bone_track_data += bt.encode(header_total_size, bone_track_data)
+        if len(header_bone_track_data) != header_bone_tracks_size:
+            #Verify the header is the right length.
+            raise Exception("Internal error: Header bone tracks size mismatch! Got: %s  Expected: %s" % (len(header_bone_tracks), header_bone_tracks_size))
+
+        self.header_size = header_total_size
+        #self.header_max_hip_displacement = 0
+        self.header_length = 0
+        for bt in self.bone_tracks:
+            self.header_length = max(self.header_length, len(bt.positions), len(bt.rotations))
+        self.header_length = max(0, self.header_length - 1)
+        self.header_bone_tracks_offset = header_base_size + header_skeleton_hierarhy_size
+        self.header_bone_track_count = len(self.bone_tracks)
+        self.header_rotation_compression_type = 0
+        self.header_position_compression_type = 0
+        if self.skeleton_hierarchy is None:
+            self.header_skeleton_hierarchy_offset = 0
+        else:
+            self.header_skeleton_hierarchy_offset = header_base_size
+        self
+        header_data = (
+            struct.pack("<i256s256sff",
+                        header_total_size,
+                        storeString(self.header_name, 256),
+                        storeString(self.header_base_anim_name, 256),
+                        self.header_max_hip_displacement,
+                        self.header_length) +
+            struct.pack("<IiiiI",
+                        self.header_bone_tracks_offset,
+                        self.header_bone_track_count,
+                        self.header_rotation_compression_type,
+                        self.header_position_compression_type,
+                        self.header_skeleton_hierarchy_offset) +
+            ZERO_BYTE * (4 * 13))
+        #Verify header lengths
+        if len(header_data) != header_base_size:
+            raise Exception("Internal error: Header base size mismatch! Got: %s  Expected: %s" % (len(header_data), header_base_size))
+        #todo: total header length
+
+        rawdata = header_data + skel_data + header_bone_track_data + bone_track_data.data
+        self.data = Data(rawdata)
+        return rawdata
+
 
     def dump(self):
         print("header_size: %s" % (self.header_size, ))
diff --git a/export_anim.py b/export_anim.py
new file mode 100644
index 0000000..1d31b39
--- /dev/null
+++ b/export_anim.py
@@ -0,0 +1,212 @@
+import bpy.path
+import bpy
+import math
+from mathutils import Vector, Quaternion
+try:
+    from .anim import *
+    from .bones import *
+except:
+    from anim import *
+    from bones import *
+
+def export_fix_coord(v):
+    return Vector((-v[0],  v[2], -v[1]))
+def export_fix_normal(v):
+    return Vector(( v[0], -v[2],  v[1]))
+def export_fix_quaternion(quat):
+    return Quaternion((-quat[1],  quat[3], -quat[2], quat[0]))
+
+
+def getBoneRotation(bone, bone_tracks, index):
+    if bone is None:
+        #rot_p = Quaternion()
+        return Quaternion()
+    else:
+        rot_p = getBoneRotation(bone.parent, bone_tracks, index)
+    if bone.name in bone_tracks:
+        trk = bone_tracks[bone.name]
+        chn = trk["rotation_quaternion"]
+        if index >= len(chn):
+            rot_s = chn[-1].copy()
+        else:
+            rot_s = chn[index].copy()
+    else:
+        rot_s = Quaternion()
+    rot_p.rotate(rot_s)
+    return rot_p
+    #rot_s.rotate(rot_p)
+    #return rot_s
+
+
+def convert_animation(context, arm_obj, arm_data, nla_track, anim, save_skel):
+    bone_tracks = {}
+    for nla_strip in nla_track.strips:
+        #todo: what's the proper way to handle multiple strips?
+        #Presently later strips will overwrite earlier strips.
+        #>>> [x.data_path for x in bpy.data.objects['fem'].animation_data.nla_tracks['run'].strips[0].action.fcurves]
+        #>>> [x.array_index for x in bpy.data.objects['fem'].animation_data.nla_tracks['run'].strips[0].action.fcurves]
+        for crv in nla_strip.action.fcurves:
+            #>>> [y.co for x in bpy.data.objects['fem'].animation_data.nla_tracks['run'].strips[0].action.fcurves for y in x.sampled_points]
+            dp = crv.data_path
+            idx = crv.array_index
+            print("crv: %s : %s" % (dp, idx))
+            #Naively convert data_path into bone, and transform type
+            parts_a = dp.split('["')
+            parts_b = parts_a[1].split('"].')
+            bone_name = parts_b[0]
+            transform = parts_b[1]
+            if bone_name not in bone_tracks:
+                bone_tracks[bone_name] = {}
+            bone_track = bone_tracks[bone_name]
+            if transform == "location":
+                data_type = Vector
+            elif transform == "rotation_quaternion":
+                data_type = Quaternion
+            else:
+                #todo:
+                raise
+                data_type = None
+            if transform not in bone_track:
+                bone_track[transform] = []
+            bone_track_channel = bone_track[transform]
+            for pnt in crv.sampled_points:
+                k = int(math.floor(pnt.co[0] + 0.5))
+                if k < 0:
+                    #ignore samples before 0
+                    continue
+                while len(bone_track_channel) <= k:
+                    bone_track_channel.append(data_type())
+                bone_track_channel[k][idx] = pnt.co[1]
+    #print("bone_tracks: %s" % bone_tracks)
+    print("bone_tracks['Head']: %s" % bone_tracks.get("Head", None))
+    #todo: convert FCurve data to track positions and rotations
+
+    #todo: Get bones required for export.
+    if save_skel:
+        #If we need to save the skeleton, ensure we have values for the T-pose loaded into bones that haven't been referenced yet.
+        for bn in arm_data.bones.keys():
+            if bn not in bone_tracks:
+                bone_tracks[bn] = {
+                    "location": [Vector()],
+                    "rotation_quaternion": [Quaternion()],
+                }
+
+    #todo: trim back tracks that have duplicates on their tail.
+    for bn, bt in bone_tracks.items():
+        #ensure missing channels have a T-pose value.
+        if "location" not in bt:
+            bt["location"] = [Vector()]
+        if "rotation_quaternion" not in bt:
+            bt["rotation_quaternion"] = [Quaternion()]
+        #Trim back duplicates at the end of a track.
+        for cn in ("location", "rotation_quaternion"):
+            chn = bt[cn]
+            while len(chn) >= 2:
+                if chn[-1] == chn[-2]:
+                    chn.pop()
+                else:
+                    break
+
+    #Convert track positions and rotations into a more convenient form.
+    for bn, bt in bone_tracks.items():
+        bone = arm_data.bones[bn]
+        #Get position of bone, relative to parent (or armature origin for root bones).
+        if bone.parent is None:
+            bone_position = bone.head
+        else:
+            bone_position = bone.head + (bone.parent.tail - bone.parent.head)
+            #print("parent[%s]: %s %s" % (bone.parent.name, bone.parent.head, bone.parent.tail))
+        print("bone_position[%s(%s)]: %s" % (bn, bone.name, bone_position))
+        bt["net_location"] = [bone_position]
+        bt["net_rotation_quaternion"] = []
+        rot_chn = bt["rotation_quaternion"]
+        for i in range(len(rot_chn)):
+            bt["net_rotation_quaternion"].append(getBoneRotation(bone, bone_tracks, i))
+        for i in range(1, len(bt["location"])):
+            if i >= len(bt["net_rotation_quaternion"]):
+                rot = bt["net_rotation_quaternion"][-1]
+            else:
+                rot = bt["net_rotation_quaternion"][i]
+            pos = bone_position + bt["location"][i]
+            print("   pos[%s]: %s" % (i, pos))
+            pos.rotate(rot)
+            bt["net_location"].append(pos)
+
+    #print("bone_tracks: %s" % bone_tracks)
+
+    #Store bone track information in the Anim
+    for bn, bt in bone_tracks.items():
+        bat = BoneAnimTrack()
+        bat.bone_id = BONES_LOOKUP[bn]
+        bat.rotations = [export_fix_quaternion(x) for x in bt["rotation_quaternion"]]
+        bat.positions = [export_fix_coord(x) for x in bt["net_location"]]
+        anim.bone_tracks.append(bat)
+
+    #Save the skeleton (if flagged).
+    if save_skel:
+        anim.skeleton_hierarchy = SkeletonHierarchy()
+        for bn in BONES_LIST:
+            if bn in arm_data.bones:
+                bone = arm_data.bones[bn]
+                if bone.parent is None:
+                    parent_id = None
+                else:
+                    parent_id = BONES_LOOKUP[bone.parent.name]
+                bone_id = BONES_LOOKUP[bn]
+                anim.skeleton_hierarchy.addBone(parent_id, bone_id)
+
+    anim.dump()
+    pass
+
+def save(operator, context, scale = 1.0, filepath = "", global_matrix = None, use_mesh_modifiers = True, save_skeleton = False):
+    #todo: prefer armature name that matches import?
+
+    #Choose the first selected armature as the armature to attach to.
+    armature = None
+    for ob in context.selected_objects:
+        if ob.type == "ARMATURE":
+            armature_obj = ob
+            armature = ob.data
+            break
+    else:
+        #If none found try again with all objects.
+        for ob in bpy.data.objects:
+            if ob.type == "ARMATURE":
+                armature_obj = ob
+                armature = ob.data
+                break
+    #todo:
+    track = None
+    skel_track = None
+    #todo: error/warning if multiple tracks are selected.
+    #todo: error/warning if multiple skel_ tracks are found
+    #todo: error if no skel_tracks are found
+    for t in armature_obj.animation_data.nla_tracks:
+        if t.select:
+            track = t
+        if t.name.lower().startswith("skel_"):
+            skel_track = t
+    if save_skeleton:
+        skel_track = track
+    arm_name = armature_obj.name
+    track_name = bpy.path.display_name_from_filepath(filepath)
+    skel_track_name = skel_track.name
+
+    #Get name and base name
+    anim_name = "%s/%s" % (arm_name, track_name)
+    anim_base_name = "%s/%s" % (arm_name, skel_track_name)
+    #todo: warning if anim_name doesn't match file path
+
+    anim = Anim()
+    anim.header_name = anim_name
+    anim.header_base_anim_name = anim_base_name
+
+    save_skel = save_skeleton or anim_name == anim_base_name
+    convert_animation(context, armature_obj, armature, track, anim, save_skel)
+
+    data = anim.saveToData()
+    fh = open(filepath, "wb")
+    fh.write(data)
+    fh.close()
+
+    return {'FINISHED'}
diff --git a/export_skel.py b/export_skel.py
new file mode 100644
index 0000000..a365c7d
--- /dev/null
+++ b/export_skel.py
@@ -0,0 +1,12 @@
+import bpy.path
+import bpy
+
+try:
+    from .export_anim import *
+except:
+    from export_anim import *
+
+
+
+def save_skel(operator, context, scale = 1.0, filepath = "", global_matrix = None, use_mesh_modifiers = True):
+    return save(operator, context, scale = scale, filepath = filepath, global_matrix = global_matrix, use_mesh_modifiers = use_mesh_modifiers, save_skeleton = True)
diff --git a/util.py b/util.py
index 43240d1..6a36384 100644
--- a/util.py
+++ b/util.py
@@ -1,6 +1,8 @@
 import struct
 import sys
 
+ZERO_BYTE = struct.pack("B", 0)
+
 if sys.version_info[0] < 3:
     byte = chr
     unbyte = ord