Added IK -> FK snapping in rigify arms.

Also cleaned up how some of the operators work.  They weren't playing nice
with undo.

4 files changed, 115 insertions, 22 deletions

diff --git a/rigify/generate.py b/rigify/generate.py
index c9eec71f..a439d357 100644
--- a/rigify/generate.py
+++ b/rigify/generate.py
@@ -55,8 +55,6 @@ def generate_rig(context, metarig):
     rig_id = random_string(12)  # Random so that different rigs don't collide id's
 
     # Initial configuration
-    use_global_undo = context.user_preferences.edit.use_global_undo
-    context.user_preferences.edit.use_global_undo = False
     mode_orig = context.mode
     rest_backup = metarig.data.pose_position
     metarig.data.pose_position = 'REST'
@@ -214,7 +212,6 @@ def generate_rig(context, metarig):
     except Exception as e:
         # Cleanup if something goes wrong
         print("Rigify: failed to generate rig.")
-        context.user_preferences.edit.use_global_undo = use_global_undo
         metarig.data.pose_position = rest_backup
         obj.data.pose_position = 'POSE'
         bpy.ops.object.mode_set(mode='OBJECT')
@@ -301,7 +298,6 @@ def generate_rig(context, metarig):
     bpy.ops.object.mode_set(mode='OBJECT')
     metarig.data.pose_position = rest_backup
     obj.data.pose_position = 'POSE'
-    context.user_preferences.edit.use_global_undo = use_global_undo
 
 
 def get_bone_rigs(obj, bone_name, halt_on_missing=False):
diff --git a/rigify/rig_ui_template.py b/rigify/rig_ui_template.py
index 32b0d61c..ac0e919b 100644
--- a/rigify/rig_ui_template.py
+++ b/rigify/rig_ui_template.py
@@ -18,17 +18,18 @@
 
 UI_SLIDERS = '''
 import bpy
+from mathutils import Matrix, Vector
+from math import acos
 
 rig_id = "%s"
 
 
-def get_pose_matrix_in_other_space(pose_bone, pbs):
-    """ Returns the transform matrix of pose_bone relative to pbs's transform space.
+def get_pose_matrix_in_other_space(mat, pbs):
+    """ Returns the transform matrix relative to pbs's transform space.
         In other words, you could take the matrix returned from this, slap
-        it into pbs, and it would have the same world transform as pb.
+        it into pbs, and it would have the same world transform as mat.
         This is with constraints applied.
     """
-    mat = pose_bone.matrix.copy()
     rest_inv = pbs.bone.matrix_local.inverted()
     
     if pbs.parent:
@@ -46,7 +47,7 @@ def get_local_matrix_with_constraints(pose_bone):
     """ Returns the local matrix of the given pose bone, with constraints
         applied.
     """
-    return get_pose_matrix_in_other_space(pose_bone, pose_bone)
+    return get_pose_matrix_in_other_space(pose_bone.matrix, pose_bone)
 
 
 def set_pose_rotation(pb, mat):
@@ -77,7 +78,19 @@ def set_pose_translation(pb, mat):
     """ Sets the pose bone's translation to the same translation as the given matrix.
         Matrix should be given in local space.
     """
-    pb.location = mat.to_translation()
+    if pb.bone.use_local_location == True:
+        pb.location = mat.to_translation()
+    else:
+        loc = mat.to_translation()
+        
+        rest = pb.bone.matrix_local.copy()
+        if pb.bone.parent:
+            par_rest = pb.bone.parent.matrix_local.copy()
+        else:
+            par_rest = Matrix()
+        
+        q = (par_rest.inverted() * rest).to_quaternion()
+        pb.location = loc * q
 
 
 def fk2ik(obj, fk, ik):
@@ -93,25 +106,25 @@ def fk2ik(obj, fk, ik):
     farmi = pb[ik[1]]
     handi = pb[ik[2]]
     
-    uarmmat = get_pose_matrix_in_other_space(uarmi, uarm)
+    uarmmat = get_pose_matrix_in_other_space(uarmi.matrix, uarm)
     set_pose_rotation(uarm, uarmmat)
     set_pose_scale(uarm, uarmmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
     
-    farmmat = get_pose_matrix_in_other_space(farmi, farm)
+    farmmat = get_pose_matrix_in_other_space(farmi.matrix, farm)
     set_pose_rotation(farm, farmmat)
     set_pose_scale(farm, farmmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
     
-    handmat = get_pose_matrix_in_other_space(handi, hand)
+    handmat = get_pose_matrix_in_other_space(handi.matrix, hand)
     set_pose_rotation(hand, handmat)
     set_pose_scale(hand, handmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
     
-
+    
 def ik2fk(obj, fk, ik):
     """ Matches the ik bones in an arm rig to the fk bones.
     """
@@ -126,19 +139,73 @@ def ik2fk(obj, fk, ik):
     handi = pb[ik[2]]
     pole = pb[ik[3]]
     
-    handmat = get_pose_matrix_in_other_space(hand, handi)
+    # Hand position
+    handmat = get_pose_matrix_in_other_space(hand.matrix, handi)
     set_pose_translation(handi, handmat)
     set_pose_rotation(handi, handmat)
     set_pose_scale(handi, handmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
+    
+    # Pole target position
+    a = uarm.matrix.to_translation()
+    b = farm.matrix.to_translation() + farm.vector
+    
+    # Vector from the head of the upper arm to the
+    # tip of the forearm
+    armv = b - a
+    
+    # Create a vector that is not aligned with armv.
+    # It doesn't matter what vector.  Just any vector
+    # that's guaranteed to not be pointing in the same
+    # direction.
+    if abs(armv[0]) < abs(armv[1]) and abs(armv[0]) < abs(armv[2]):
+        v = Vector((1,0,0))
+    elif abs(armv[1]) < abs(armv[2]):
+        v = Vector((0,1,0))
+    else:
+        v = Vector((0,0,1))
+    
+    # cross v with armv to get a vector perpendicular to armv
+    pv = v.cross(armv).normalized() * (uarm.length + farm.length)
+    
+    def set_pole(pvi):
+        # Translate pvi into armature space
+        ploc = a + (armv/2) + pvi
+        
+        # Set pole target to location
+        mat = get_pose_matrix_in_other_space(Matrix().Translation(ploc), pole)
+        set_pose_translation(pole, mat)
+        
+        bpy.ops.object.mode_set(mode='OBJECT')
+        bpy.ops.object.mode_set(mode='POSE')
+    
+    set_pole(pv)
+    
+    # Get the rotation difference between the ik and fk upper arms
+    q1 = uarm.matrix.to_quaternion()
+    q2 = uarmi.matrix.to_quaternion()
+    angle = acos(min(1,max(-1,q1.dot(q2)))) * 2
+    
+    # Compensate for the rotation difference
+    if angle > 0.00001:
+        pv *= Matrix.Rotation(angle, 4, armv).to_quaternion()
+        set_pole(pv)
+    
+        q1 = uarm.matrix.to_quaternion()
+        q2 = uarmi.matrix.to_quaternion()
+        angle2 = acos(min(1,max(-1,q1.dot(q2)))) * 2
+        if angle2 > 0.00001:
+            pv *= Matrix.Rotation((angle*(-2)), 4, armv).to_quaternion()
+            set_pole(pv)
 
 
 class Rigify_Arm_FK2IK(bpy.types.Operator):
     """ Snaps an FK arm to an IK arm.
     """
-    bl_idname = "pose.rigify_arm_fk2ik"
+    bl_idname = "pose.rigify_arm_fk2ik_" + rig_id
     bl_label = "Rigify Snap FK arm to IK"
+    bl_options = {'UNDO'}
     
     uarm_fk = bpy.props.StringProperty(name="Upper Arm FK Name")
     farm_fk = bpy.props.StringProperty(name="Forerm FK Name")
@@ -153,15 +220,21 @@ class Rigify_Arm_FK2IK(bpy.types.Operator):
         return (context.active_object != None and context.mode == 'POSE')
 
     def execute(self, context):
-        fk2ik(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik])
+        use_global_undo = context.user_preferences.edit.use_global_undo
+        context.user_preferences.edit.use_global_undo = False
+        try:
+            fk2ik(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik])
+        finally:
+            context.user_preferences.edit.use_global_undo = use_global_undo
         return {'FINISHED'}
 
 
 class Rigify_Arm_IK2FK(bpy.types.Operator):
     """ Snaps an IK arm to an FK arm.
     """
-    bl_idname = "pose.rigify_arm_ik2fk"
+    bl_idname = "pose.rigify_arm_ik2fk_" + rig_id
     bl_label = "Snap IK arm to FK"
+    bl_options = {'UNDO'}
     
     uarm_fk = bpy.props.StringProperty(name="Upper Arm FK Name")
     farm_fk = bpy.props.StringProperty(name="Forerm FK Name")
@@ -177,7 +250,12 @@ class Rigify_Arm_IK2FK(bpy.types.Operator):
         return (context.active_object != None and context.mode == 'POSE')
 
     def execute(self, context):
-        ik2fk(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik, self.pole])
+        use_global_undo = context.user_preferences.edit.use_global_undo
+        context.user_preferences.edit.use_global_undo = False
+        try:
+            ik2fk(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik, self.pole])
+        finally:
+            context.user_preferences.edit.use_global_undo = use_global_undo
         return {'FINISHED'}
 
 
diff --git a/rigify/rigs/biped/arm/__init__.py b/rigify/rigs/biped/arm/__init__.py
index 7d4214fd..c09b39f4 100644
--- a/rigify/rigs/biped/arm/__init__.py
+++ b/rigify/rigs/biped/arm/__init__.py
@@ -30,15 +30,29 @@ fk_arm = ["%s", "%s", "%s"]
 ik_arm = ["%s", "%s", "%s", "%s"]
 if is_selected(fk_arm+ik_arm):
     layout.prop(pose_bones[ik_arm[2]], '["ikfk_switch"]', text="FK / IK (" + ik_arm[2] + ")", slider=True)
-    p = layout.operator("pose.rigify_arm_fk2ik", text="Snap FK->IK (" + fk_arm[0] + ")")
+if is_selected(fk_arm):
+    try:
+        pose_bones[fk_arm[0]]["isolate"]
+        layout.prop(pose_bones[fk_arm[0]], '["isolate"]', text="Isolate Rotation (" + fk_arm[0] + ")", slider=True)
+    except KeyError:
+        pass
+if is_selected(fk_arm+ik_arm):
+    p = layout.operator("pose.rigify_arm_fk2ik_" + rig_id, text="Snap FK->IK (" + fk_arm[0] + ")")
     p.uarm_fk = fk_arm[0]
     p.farm_fk = fk_arm[1]
     p.hand_fk = fk_arm[2]
     p.uarm_ik = ik_arm[0]
     p.farm_ik = ik_arm[1]
     p.hand_ik = ik_arm[2]
-if is_selected(fk_arm):
-    layout.prop(pose_bones[fk_arm[0]], '["isolate"]', text="Isolate Rotation (" + fk_arm[0] + ")", slider=True)
+    p = layout.operator("pose.rigify_arm_ik2fk_" + rig_id, text="Snap IK->FK (" + fk_arm[0] + ")")
+    p.uarm_fk = fk_arm[0]
+    p.farm_fk = fk_arm[1]
+    p.hand_fk = fk_arm[2]
+    p.uarm_ik = ik_arm[0]
+    p.farm_ik = ik_arm[1]
+    p.hand_ik = ik_arm[2]
+    p.pole = ik_arm[3]
+
 """
 
 
diff --git a/rigify/ui.py b/rigify/ui.py
index 4c3353bb..c9f1e472 100644
--- a/rigify/ui.py
+++ b/rigify/ui.py
@@ -193,15 +193,20 @@ class Generate(bpy.types.Operator):
 
     bl_idname = "pose.rigify_generate"
     bl_label = "Rigify Generate Rig"
+    bl_options = {'UNDO'}
 
     def execute(self, context):
         import imp
         imp.reload(generate)
 
+        use_global_undo = context.user_preferences.edit.use_global_undo
+        context.user_preferences.edit.use_global_undo = False
         try:
             generate.generate_rig(context, context.object)
         except rigify.utils.MetarigError as rig_exception:
             rigify_report_exception(self, rig_exception)
+        finally:
+            context.user_preferences.edit.use_global_undo = use_global_undo
 
         return {'FINISHED'}