mocaptool_ops

import bpy

class MocapNormalizeNames(bpy.types.Operator):
    bl_idname = "wm.normalize_names"
    bl_label = "Normalize Names"
    bl_description = "Remove 'Left' and 'Right' from names and replace \nwith conventional names using '.L' and '.R'"

    def execute(self, context):
        for bone in bpy.data.objects[context.scene.mocap.source].pose.bones:
            if 'Left' in bone.name:
                bone.name = bone.name.replace('Left', '')
                if '.L' not in bone.name:
                     bone.name += '.L'
            elif 'Right' in bone.name:
                bone.name = bone.name.replace('Right', '')
                if '.R' not in bone.name:
                     bone.name += '.R'
        return {'FINISHED'}

class MocapAutomapNames(bpy.types.Operator):
    bl_idname = "wm.automap_names"
    bl_label = "Automap Names"
    bl_description = "Fill properties with matching names if avatible"

    def execute(self, context):
        for sbone in bpy.data.objects[context.scene.mocap.source].pose.bones:
            for tbone in bpy.data.objects[context.scene.mocap.target].pose.bones:
                if tbone.bone_source == '' and tbone.name == sbone.name:
                    tbone.bone_source = sbone.name
        return {'FINISHED'}

class MocapToggleIK(bpy.types.Operator):
    bl_idname = "wm.toggle_ikconstraints"
    bl_label = "Toggle IK"
    bl_description = "Toggle IK constraints on selected armature's bones"

    def execute(self, context):
        target_arm = bpy.data.objects[context.scene.mocap.target].pose
        if bpy.context.object.type == 'ARMATURE':
            target_arm = bpy.context.object.pose
        bool = False
        for tbone in target_arm.bones:
            if "IK" in tbone.constraints.keys():
                bool = not tbone.constraints['IK'].mute
                break
        for tbone in target_arm.bones:
            if "IK" in tbone.constraints.keys():
                tbone.constraints['IK'].mute = bool
        return {'FINISHED'}

class convert():

    def get_or_create_fcurve(self, t_action, t_data_path, array_index=-1, group=None):
        for fc in t_action.fcurves:
            if fc.data_path == t_data_path and (array_index<0 or fc.array_index == array_index):
                return fc

        fc = t_action.fcurves.new(t_data_path, array_index)
        fc.group = group
        return fc

    def add_keyframe_quat(self, t_action, quat, frame, t_data_path):
        for i in range(len(quat)):
            fc = self.get_or_create_fcurve(t_action, t_data_path, i)
            pos = len(fc.keyframe_points)
            fc.keyframe_points.add(1)
            fc.keyframe_points[pos].co = [frame, quat[i]]
            fc.update()

    def add_keyframe_euler(self, t_action, euler, frame, t_data_path):
        for i in range(len(euler)):
            fc = self.get_or_create_fcurve(t_action, t_data_path, i)
            pos = len(fc.keyframe_points)
            fc.keyframe_points.add(1)
            fc.keyframe_points[pos].co = [frame, euler[i]]
            fc.update()

    #Gets the X value from kp.co and forms an array of keyframe locations
    def frames_matching(self, action, data_path):
        frames = set()
        for fc in action.fcurves:
            if fc.data_path == data_path:
                fri = [kp.co[0] for kp in fc.keyframe_points]
                frames.update(fri)
        return frames

    # Converts only one group/bone in one action - Euler to Quat
    def group_eq(self, s_action, s_bone, s_data_path, t_action, t_bone, t_data_path, order):

        frames = self.frames_matching(s_action, s_data_path)

        for fr in frames:
            euler = s_bone.rotation_euler.copy()
            for fc in s_action.fcurves:
                if fc.data_path == s_data_path:
                    euler[fc.array_index] = fc.evaluate(fr)
            quat = euler.to_quaternion()

            self.add_keyframe_quat(t_action, quat, fr, t_data_path)
            t_bone.rotation_mode = order

    # Converts only one group/bone in one action - Quat to euler
    def group_qe(self, s_action, s_bone, s_data_path, t_action, t_bone, t_data_path, order):

        frames = self.frames_matching(s_action, data_path)

        for fr in frames:
            quat = s_bone.rotation_quaternion.copy()
            for fc in s_action.fcurves:
                if fc.data_path == data_path:
                    quat[fc.array_index] = fc.evaluate(fr)
            euler = quat.to_euler(order)

            self.add_keyframe_euler(t_action, euler, fr, t_data_path)
            bone.rotation_mode = order

    ##### BONE TO BONE #####
    def bon_to_bon(self, s_action, s_bone, t_action, t_bone, order):
        do = False
        s_data_path = ''
        t_data_path = ''

        # Check weather there needs to be a conversion
        for fcurve in s_action.fcurves:
            if s_bone.name in fcurve.data_path:

                # If To-Euler conversion
                if order != 'QUATERNION':
                    if fcurve.data_path.endswith('rotation_quaternion'):
                        do = True
                        s_data_path = fcurve.data_path #[:-len('rotation_quaternion')]
                        t_data_path = 'pose.bones["' + t_bone.name + '"].rotation_quaternion'
                        break

                # If To-Quat conversion
                else:
                    if fcurve.data_path.endswith('rotation_euler'):
                        do = True
                        s_data_path = fcurve.data_path #[:-len('rotation_euler')]
                        t_data_path = 'pose.bones["' + t_bone.name + '"].rotation_euler'
                        break

        # If To-Euler conversion
        if do and order != 'QUATERNION':
            self.group_qe(s_action, s_bone, s_data_path, t_action, t_bone, t_data_path, order)

        # If To-Quat conversion
        elif do:
            self.group_eq(s_action, s_bone, s_data_path, t_action, t_bone, t_data_path, order)

        # Changes rotation mode to new one
        t_bone.rotation_mode = order

    ##### BONE TO BONE #####
    def bon_every_bon(self, s_action, s_obj, t_action, t_obj, order):
        print('#\n# New Run \n#')

        s_bones = set()
        for t_bone in t_obj.pose.bones:
            if t_bone.bone_source is None or t_bone.bone_source != '':
                s_bones.add(t_bone.bone_source)
                s_obj.pose.bones[t_bone.bone_source].bone_source = t_bone.name
                print(t_bone.bone_source + ' -> ' + t_bone.name)

        # Collects pose_bones that are in the action
        pose_bones = set()
        # Checks all fcurves
        for fcurve in s_action.fcurves:
            # Look for the ones that has rotation_euler
            if order == 'QUATERNION':
                if fcurve.data_path.endswith('rotation_euler'):
                    # If the bone from action really exists
                    bone_name = fcurve.data_path[12:-17]
                    if bone_name in s_bones:
                        if s_obj.pose.bones[bone_name] not in pose_bones:
                            pose_bones.add(s_obj.pose.bones[bone_name])
                            print('added ' + bone_name)
                    else:
                        print(bone_name, 'is not selected')

            # Look for the ones that has rotation_quaternion
            else:
                if fcurve.data_path.endswith('rotation_quaternion'):
                    # If the bone from action really exists
                    bone_name = fcurve.data_path[12:-22]
                    if bone_name in s_bones:
                        if s_obj.pose.bones[bone_name] not in pose_bones:
                            pose_bones.add(s_obj.pose.bones[bone_name])
                            print('added ' + s_obj.pose.bones[bone_name])
                    else:
                        print(bone_name, 'is not selected')

        # Convert current action and pose_bones that are in each action
        for s_bone in pose_bones:
            if s_bone.bone_source is None or s_bone.bone_source == '':
                print('bone source cannot be ' + s_bone.bone_source)
                continue
            t_bone = t_obj.pose.bones[s_bone.bone_source]
            self.bon_to_bon(s_action, s_bone, t_action, t_bone, order)

convert = convert()

class COPY_current_action_bones_bones(bpy.types.Operator):
    bl_idname = "wm.copy_mapped_animation"
    bl_label = "Copy Animation"
    bl_description = ""

    # on mouse up:
    def invoke(self, context, event):
        self.execute(context)
        return {'FINISHED'}

    def execute(op, context):
        s_action = bpy.data.objects[context.scene.mocap.source].animation_data.action
        s_obj = bpy.data.objects[context.scene.mocap.source]

        t_action = bpy.data.objects[context.scene.mocap.target].animation_data.action
        t_obj = bpy.data.objects[context.scene.mocap.target]
        order = bpy.context.scene.mocap.rotation_mode

        convert.bon_every_bon(s_action, s_obj, t_action, t_obj, order)
        return {'FINISHED'}

def register():
    bpy.utils.register_module(__name__)

def unregister():
    bpy.utils.unregister_module(__name__)

if __name__ == '__main__':
    register()