ui.py

                break


def IktoFk(rig, window='ALL'):

    scn = bpy.context.scene
    id_store = bpy.context.window_manager

    rig_id = rig.data['rig_id']
    leg_fk2ik = eval('bpy.ops.pose.rigify_leg_fk2ik_' + rig_id)
    arm_fk2ik = eval('bpy.ops.pose.rigify_arm_fk2ik_' + rig_id)
    limb_generated_names = get_limb_generated_names(rig)

    if window == 'ALL':
        frames = get_keyed_frames_in_range(bpy.context, rig)
    elif window == 'CURRENT':
        frames = [scn.frame_current]
    else:
        frames = [scn.frame_current]

    if not id_store.rigify_transfer_only_selected:
        bpy.ops.pose.select_all(action='DESELECT')
        pbones = rig.pose.bones
    else:
        pbones = bpy.context.selected_pose_bones
        bpy.ops.pose.select_all(action='DESELECT')

    for b in pbones:
        for group in limb_generated_names:
            if b.name in limb_generated_names[group].values() or b.name in limb_generated_names[group]['controls']\
                    or b.name in limb_generated_names[group]['ik_ctrl']:
                names = limb_generated_names[group]
                if names['limb_type'] == 'arm':
                    func = arm_fk2ik
                    controls = names['controls']
                    ik_ctrl = names['ik_ctrl']
                    fk_ctrl = names['fk_ctrl']
                    parent = names['parent']
                    pole = names['pole']
                    rig.pose.bones[controls[1]].bone.select = True
                    rig.pose.bones[controls[2]].bone.select = True
                    rig.pose.bones[controls[3]].bone.select = True
                    kwargs = {'uarm_fk': controls[1], 'farm_fk': controls[2], 'hand_fk': controls[3],
                              'uarm_ik': controls[0], 'farm_ik': ik_ctrl[1],
                              'hand_ik': controls[4]}
                    args = (controls[0], controls[1], controls[2], controls[3],
                            controls[4], pole, parent)
                else:
                    func = leg_fk2ik
                    controls = names['controls']
                    ik_ctrl = names['ik_ctrl']
                    fk_ctrl = names['fk_ctrl']
                    parent = names['parent']
                    pole = names['pole']
                    rig.pose.bones[controls[1]].bone.select = True
                    rig.pose.bones[controls[2]].bone.select = True
                    rig.pose.bones[controls[3]].bone.select = True
                    kwargs = {'thigh_fk': controls[1], 'shin_fk': controls[2], 'foot_fk': controls[3],
                              'mfoot_fk': controls[7], 'thigh_ik': controls[0], 'shin_ik': ik_ctrl[1],
                              'foot_ik': ik_ctrl[2], 'mfoot_ik': ik_ctrl[2]}
                    args = (controls[0], controls[1], controls[2], controls[3],
                            controls[6], controls[5], pole, parent)

                for f in frames:
                    if not bones_in_frame(f, rig, *args):
                        continue
                    scn.frame_set(f)
                    func(**kwargs)
                    bpy.ops.anim.keyframe_insert_menu(type='BUILTIN_KSI_VisualLocRot')
                    bpy.ops.anim.keyframe_insert_menu(type='Scaling')

                bpy.ops.pose.select_all(action='DESELECT')
                limb_generated_names.pop(group)
                break


def clearAnimation(act, anim_type, names):

    bones = []
    for group in names:
        if names[group]['limb_type'] == 'arm':
            if anim_type == 'IK':
                bones.extend([names[group]['controls'][0], names[group]['controls'][4]])
            elif anim_type == 'FK':
                bones.extend([names[group]['controls'][1], names[group]['controls'][2], names[group]['controls'][3]])
        else:
            if anim_type == 'IK':
                bones.extend([names[group]['controls'][0], names[group]['controls'][6], names[group]['controls'][5],
                              names[group]['controls'][4]])
            elif anim_type == 'FK':
                bones.extend([names[group]['controls'][1], names[group]['controls'][2], names[group]['controls'][3],
                              names[group]['controls'][4]])
    FCurves = []
    for fcu in act.fcurves:
        words = fcu.data_path.split('"')
        if (words[0] == "pose.bones[" and
                    words[1] in bones):
            FCurves.append(fcu)

    if FCurves == []:
        return

    for fcu in FCurves:
        act.fcurves.remove(fcu)

    # Put cleared bones back to rest pose
    bpy.ops.pose.loc_clear()
    bpy.ops.pose.rot_clear()
    bpy.ops.pose.scale_clear()

    # updateView3D()


def rotPoleToggle(rig, window='ALL', value=False, toggle=False, bake=False):

    scn = bpy.context.scene
    id_store = bpy.context.window_manager

    rig_id = rig.data['rig_id']
    leg_fk2ik = eval('bpy.ops.pose.rigify_leg_fk2ik_' + rig_id)
    arm_fk2ik = eval('bpy.ops.pose.rigify_arm_fk2ik_' + rig_id)
    leg_ik2fk = eval('bpy.ops.pose.rigify_leg_ik2fk_' + rig_id)
    arm_ik2fk = eval('bpy.ops.pose.rigify_arm_ik2fk_' + rig_id)
    limb_generated_names = get_limb_generated_names(rig)

    if window == 'ALL':
        frames = get_keyed_frames_in_range(bpy.context, rig)
    elif window == 'CURRENT':
        frames = [scn.frame_current]
    else:
        frames = [scn.frame_current]

    if not id_store.rigify_transfer_only_selected:
        bpy.ops.pose.select_all(action='DESELECT')
        pbones = rig.pose.bones
    else:
        pbones = bpy.context.selected_pose_bones
        bpy.ops.pose.select_all(action='DESELECT')

    for b in pbones:
        for group in limb_generated_names:
            names = limb_generated_names[group]

            if toggle:
                new_pole_vector_value = not rig.pose.bones[names['parent']]['pole_vector']
            else:
                new_pole_vector_value = value

            if b.name in names.values() or b.name in names['controls'] or b.name in names['ik_ctrl']:
                if names['limb_type'] == 'arm':
                    func1 = arm_fk2ik
                    func2 = arm_ik2fk
                    controls = names['controls']
                    ik_ctrl = names['ik_ctrl']
                    fk_ctrl = names['fk_ctrl']
                    parent = names['parent']
                    pole = names['pole']
                    rig.pose.bones[controls[0]].bone.select = not new_pole_vector_value
                    rig.pose.bones[controls[4]].bone.select = not new_pole_vector_value
                    rig.pose.bones[parent].bone.select = not new_pole_vector_value
                    rig.pose.bones[pole].bone.select = new_pole_vector_value

                    kwargs1 = {'uarm_fk': controls[1], 'farm_fk': controls[2], 'hand_fk': controls[3],
                              'uarm_ik': controls[0], 'farm_ik': ik_ctrl[1],
                              'hand_ik': controls[4]}
                    kwargs2 = {'uarm_fk': controls[1], 'farm_fk': controls[2], 'hand_fk': controls[3],
                              'uarm_ik': controls[0], 'farm_ik': ik_ctrl[1], 'hand_ik': controls[4],
                              'pole': pole, 'main_parent': parent}
                    args = (controls[0], controls[4], pole, parent)
                else:
                    func1 = leg_fk2ik
                    func2 = leg_ik2fk
                    controls = names['controls']
                    ik_ctrl = names['ik_ctrl']
                    fk_ctrl = names['fk_ctrl']
                    parent = names['parent']
                    pole = names['pole']
                    rig.pose.bones[controls[0]].bone.select = not new_pole_vector_value
                    rig.pose.bones[controls[6]].bone.select = not new_pole_vector_value
                    rig.pose.bones[controls[5]].bone.select = not new_pole_vector_value
                    rig.pose.bones[parent].bone.select = not new_pole_vector_value
                    rig.pose.bones[pole].bone.select = new_pole_vector_value

                    kwargs1 = {'thigh_fk': controls[1], 'shin_fk': controls[2], 'foot_fk': controls[3],
                              'mfoot_fk': controls[7], 'thigh_ik': controls[0], 'shin_ik': ik_ctrl[1],
                              'foot_ik': ik_ctrl[2], 'mfoot_ik': ik_ctrl[2]}
                    kwargs2 = {'thigh_fk': controls[1], 'shin_fk': controls[2], 'foot_fk': controls[3],
                              'mfoot_fk': controls[7], 'thigh_ik': controls[0], 'shin_ik': ik_ctrl[1],
                              'foot_ik': controls[6], 'pole': pole, 'footroll': controls[5], 'mfoot_ik': ik_ctrl[2],
                              'main_parent': parent}
                    args = (controls[0], controls[6], controls[5], pole, parent)

                for f in frames:
                    if bake and not bones_in_frame(f, rig, *args):
                        continue
                    scn.frame_set(f)
                    func1(**kwargs1)
                    rig.pose.bones[names['parent']]['pole_vector'] = new_pole_vector_value
                    func2(**kwargs2)
                    if bake:
                        bpy.ops.anim.keyframe_insert_menu(type='BUILTIN_KSI_VisualLocRot')
                        bpy.ops.anim.keyframe_insert_menu(type='Scaling')
                        overwrite_prop_animation(rig, rig.pose.bones[parent], 'pole_vector', new_pole_vector_value, [f])

                bpy.ops.pose.select_all(action='DESELECT')
                limb_generated_names.pop(group)
                break
    scn.frame_set(0)


class OBJECT_OT_IK2FK(bpy.types.Operator):
    """ Snaps IK limb on FK limb at current frame"""
    bl_idname = "rigify.ik2fk"
    bl_label = "IK2FK"
    bl_description = "Snaps IK limb on FK"
    bl_options = {'INTERNAL'}

    def execute(self,context):
        rig = context.object
        id_store = context.window_manager

        FktoIk(rig, window='CURRENT')

        return {'FINISHED'}


class OBJECT_OT_FK2IK(bpy.types.Operator):
    """ Snaps FK limb on IK limb at current frame"""
    bl_idname = "rigify.fk2ik"
    bl_label = "FK2IK"
    bl_description = "Snaps FK limb on IK"
    bl_options = {'INTERNAL'}

    def execute(self,context):
        rig = context.object

        IktoFk(rig, window='CURRENT')

        return {'FINISHED'}


class OBJECT_OT_TransferFKtoIK(bpy.types.Operator):
    """Transfers FK animation to IK"""
    bl_idname = "rigify.transfer_fk_to_ik"
    bl_label = "Transfer FK anim to IK"
    bl_description = "Transfer FK animation to IK bones"
    bl_options = {'INTERNAL'}

    def execute(self, context):
        rig = context.object
        id_store = context.window_manager

        FktoIk(rig)

        return {'FINISHED'}


class OBJECT_OT_TransferIKtoFK(bpy.types.Operator):
    """Transfers FK animation to IK"""
    bl_idname = "rigify.transfer_ik_to_fk"
    bl_label = "Transfer IK anim to FK"
    bl_description = "Transfer IK animation to FK bones"
    bl_options = {'INTERNAL'}

    def execute(self, context):
        rig = context.object

        IktoFk(rig)

        return {'FINISHED'}


class OBJECT_OT_ClearAnimation(bpy.types.Operator):
    bl_idname = "rigify.clear_animation"
    bl_label = "Clear Animation"
    bl_description = "Clear Animation For FK or IK Bones"
    bl_options = {'INTERNAL'}

    anim_type: StringProperty()

    def execute(self, context):
        rig = context.object
        scn = context.scene
        if not rig.animation_data:
            return {'FINISHED'}
        act = rig.animation_data.action
        if not act:
            return {'FINISHED'}

        clearAnimation(act, self.anim_type, names=get_limb_generated_names(rig))
        return {'FINISHED'}


class OBJECT_OT_Rot2Pole(bpy.types.Operator):
    bl_idname = "rigify.rotation_pole"
    bl_label = "Rotation - Pole toggle"
    bl_description = "Toggles IK chain between rotation and pole target"
    bl_options = {'INTERNAL'}

    bone_name: StringProperty(default='')
    window: StringProperty(default='ALL')
    toggle: BoolProperty(default=True)
    value: BoolProperty(default=True)
    bake: BoolProperty(default=True)

    def execute(self, context):
        rig = context.object

        if self.bone_name:
            bpy.ops.pose.select_all(action='DESELECT')
            rig.pose.bones[self.bone_name].bone.select = True

        rotPoleToggle(rig, window=self.window, toggle=self.toggle, value=self.value, bake=self.bake)
        return {'FINISHED'}


### Registering ###


classes = (
    DATA_OT_rigify_add_bone_groups,
    DATA_OT_rigify_use_standard_colors,
    DATA_OT_rigify_apply_selection_colors,
    DATA_OT_rigify_bone_group_add,
    DATA_OT_rigify_bone_group_add_theme,
    DATA_OT_rigify_bone_group_remove,
    DATA_OT_rigify_bone_group_remove_all,
    DATA_UL_rigify_bone_groups,
    DATA_MT_rigify_bone_groups_context_menu,
    DATA_PT_rigify_bone_groups,
    DATA_PT_rigify_layer_names,
    DATA_PT_rigify_buttons,
    BONE_PT_rigify_buttons,
    VIEW3D_PT_rigify_animation_tools,
    VIEW3D_PT_tools_rigify_dev,
    LayerInit,
    Generate,
    UpgradeMetarigTypes,
    SwitchToLegacy,
    Sample,
    EncodeMetarig,
    EncodeMetarigSample,
    EncodeWidget,
    OBJECT_OT_FK2IK,
    OBJECT_OT_IK2FK,
    OBJECT_OT_TransferFKtoIK,
    OBJECT_OT_TransferIKtoFK,
    OBJECT_OT_ClearAnimation,
    OBJECT_OT_Rot2Pole,
)


def register():
    from bpy.utils import register_class

    animation_register()

    # Classes.
    for cls in classes:
        register_class(cls)

    # Sub-modules.
    rot_mode.register()


def unregister():
    from bpy.utils import unregister_class

    # Sub-modules.
    rot_mode.unregister()

    # Classes.
    for cls in classes:
        unregister_class(cls)

    animation_unregister()