From ada937ce82ccf1a310d2139ba1862a00c8b9f7bd Mon Sep 17 00:00:00 2001
From: TigerKat <karynax23@protonmail.com>
Date: Tue, 30 Mar 2021 23:58:11 +0930
Subject: [PATCH] Now gets the bone to position the model to from the model's
 name. Now also adds the armature's position to final model's position, making
 it line up better.

---
 __init__.py   |  2 +-
 import_geo.py | 36 ++++++++++++++++++++++++++++--------
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/__init__.py b/__init__.py
index f1cfb4a..cb536bc 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,7 +2,7 @@
 bl_info = {
     "name": "City of Heroes (.geo)",
     "author": "TigerKat",
-    "version": (0, 2, 3),
+    "version": (0, 2, 4),
     "blender": (2, 80, 0),
     "location": "File > Import/Export,",
     "description": "City of Heroes (.geo)",
diff --git a/import_geo.py b/import_geo.py
index fcc2789..84183c2 100644
--- a/import_geo.py
+++ b/import_geo.py
@@ -26,6 +26,17 @@ else:
         l.reverse()
         return l
 
+def extractRootBoneFromName(name):
+    if name.startswith("GEO_"):
+        name = name[4:]
+    if name.startswith("N_"):
+        name = name[2:]
+    bone = name.split("_")[0]
+    if bone in BONES_LOOKUP:
+        return bone
+    else:
+        return None
+
 def convert_model(geo_model, mesh_data, obj, scale):
     #Convert the geo_model into a GeoMesh.
     geomesh = geo_model.saveToGeoMesh()
@@ -82,9 +93,10 @@ def convert_model(geo_model, mesh_data, obj, scale):
     #todo: attempt to load textures/images
 
 def getBonePositionBody(bone):
+    #todo: fix this: Presently assumes that the head to tail will always be the same direction and length.
     if bone.parent is None:
-        p = bone.tail
-        print("root: %s: %s (%s)" % (bone, p, bone.head))
+        p = bone.head
+        print("root: %s: %s (%s)" % (bone, p, bone.tail))
         return p
     else:
         p = bone.tail + getBonePositionBody(bone.parent)
@@ -142,13 +154,21 @@ def load(operator, context, scale = 1.0, filepath = "", global_matrix = None, us
             if bone_name is None:
                 #No bones, don't attach.
                 pass
-            elif bone_name in armature.bones:
-                bone_pos = getBonePosition(armature, bone_name)
-                obj.matrix_world = mathutils.Matrix.Translation(bone_pos)
-                obj.modifiers.new(name = 'Armature', type = 'ARMATURE')
-                obj.modifiers['Armature'].object = armature_obj
             else:
-                print("Bone '%s' not found in armature '%s', skipping." % (bone_name, armature))
+                preferred_bone_name = extractRootBoneFromName(model_name)
+                arm_pos = armature_obj.matrix_world.translation
+                if preferred_bone_name is not None:
+                    bone_pos = getBonePosition(armature, preferred_bone_name)
+                    obj.matrix_world = mathutils.Matrix.Translation(arm_pos + bone_pos)
+                    obj.modifiers.new(name = 'Armature', type = 'ARMATURE')
+                    obj.modifiers['Armature'].object = armature_obj
+                elif bone_name in armature.bones:
+                    bone_pos = getBonePosition(armature, bone_name)
+                    obj.matrix_world = mathutils.Matrix.Translation(arm_pos + bone_pos)
+                    obj.modifiers.new(name = 'Armature', type = 'ARMATURE')
+                    obj.modifiers['Armature'].object = armature_obj
+                else:
+                    print("Bone '%s' not found in armature '%s', skipping." % (bone_name, armature))
         obj.select_set(True)
         pass
     pass