io_import_scene_mhx.py

    print("Snap FK Leg%s" % suffix)
    snapIk,cnsIk = getSnapBones(rig, "LegIK", suffix)
    (uplegIk, lolegIk, kneePt, ankleIk, legIk, legFk, footIk, toeIk) = snapIk
    snapFk,cnsFk = getSnapBones(rig, "LegFK", suffix)
    (uplegFk, lolegFk, kneePtFk, footFk, toeFk) = snapFk
    muteConstraints(cnsFk, True)
    
    if not rig["MhaLegIkToAnkle" + suffix]:
        matchPoseRotation(footFk, footFk, auto)
        matchPoseRotation(toeFk, toeFk, auto)
    
    matchPoseRotation(uplegFk, uplegFk, auto)
    matchPoseScale(uplegFk, uplegFk, auto)

    matchPoseRotation(lolegFk, lolegFk, auto)
    matchPoseScale(lolegFk, lolegFk, auto)
    
    muteConstraints(cnsFk, False)
    return


def ik2fkLeg(context, suffix):
    rig = context.object
    scn = context.scene
    auto = scn.tool_settings.use_keyframe_insert_auto
    print("Snap IK Leg%s" % suffix)
    snapIk,cnsIk = getSnapBones(rig, "LegIK", suffix)
    (uplegIk, lolegIk, kneePt, ankleIk, legIk, legFk, footIk, toeIk) = snapIk
    snapFk,cnsFk = getSnapBones(rig, "LegFK", suffix)
    (uplegFk, lolegFk, kneePtFk, footFk, toeFk) = snapFk
    muteConstraints(cnsIk, True)

    #rig["MhaKneeFollowsHip" + suffix] = False
    #rig["MhaKneeFollowsFoot" + suffix] = False
    
    legIkToAnkle = rig["MhaLegIkToAnkle" + suffix]
    if legIkToAnkle:
        matchPoseTranslation(ankleIk, footFk, auto)
    matchPoseTranslation(legIk, legFk, auto)
    matchPoseRotation(legIk, legFk, auto)  
    matchPoseReverse(toeIk, toeFk, auto)
    matchPoseReverse(footIk, footFk, auto)
    setInverse(rig, ankleIk)
    matchPoseTranslation(kneePt, kneePtFk, auto)
    setInverse(rig, kneePt)
    if not legIkToAnkle:
        matchPoseTranslation(ankleIk, footFk, auto)

    muteConstraints(cnsIk, False)
    return


#
#   setInverse(rig, pb):
#

def setInverse(rig, pb):
    rig.data.bones.active = pb.bone
    pb.bone.select = True
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.mode_set(mode='POSE')
    for cns in pb.constraints:
        if cns.type == 'CHILD_OF':
            bpy.ops.constraint.childof_set_inverse(constraint=cns.name, owner='BONE')
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.mode_set(mode='POSE')
    return

#
#
#

SnapBones = {
    "ArmFK" : ["UpArm", "LoArm", "ElbowPTFK", "Hand"],
    "ArmIK" : ["UpArmIK", "LoArmIK", "Elbow", "ElbowPT", "Wrist"],
    "LegFK" : ["UpLeg", "LoLeg", "KneePTFK", "Foot", "Toe"],
    "LegIK" : ["UpLegIK", "LoLegIK", "KneePT", "Ankle", "LegIK", "LegFK", "FootRev", "ToeRev"],
}


def getSnapBones(rig, key, suffix):
    names = SnapBones[key]
    pbones = []
    constraints = []
    for name in names:
        pb = rig.pose.bones[name+suffix]
        pbones.append(pb)
        for cns in pb.constraints:
            if cns.type == 'LIMIT_ROTATION' and not cns.mute:
                constraints.append(cns)
    return tuple(pbones),constraints


def muteConstraints(constraints, value):
    for cns in constraints:
        cns.mute = value


class VIEW3D_OT_MhxSnapFk2IkButton(bpy.types.Operator):
    bl_idname = "mhx.snap_fk_ik"
    bl_label = "Snap FK"
    bl_options = {'UNDO'}
    data = StringProperty()    

    def execute(self, context):
        bpy.ops.object.mode_set(mode='POSE')
        rig = context.object
        if rig.MhxSnapExact:
            rig["MhaRotationLimits"] = 0.0
        (prop, old) = setSnapProp(rig, self.data, 1.0, context, False)
        if prop[:6] == "MhaArm":
            fk2ikArm(context, prop[-2:])
        elif prop[:6] == "MhaLeg":
            fk2ikLeg(context, prop[-2:])
        restoreSnapProp(rig, prop, old, context)
        return{'FINISHED'}    


class VIEW3D_OT_MhxSnapIk2FkButton(bpy.types.Operator):
    bl_idname = "mhx.snap_ik_fk"
    bl_label = "Snap IK"
    bl_options = {'UNDO'}
    data = StringProperty()    

    def execute(self, context):
        bpy.ops.object.mode_set(mode='POSE')
        rig = context.object
        if rig.MhxSnapExact:
            rig["MhaRotationLimits"] = 0.0
        (prop, old) = setSnapProp(rig, self.data, 0.0, context, True)
        if prop[:6] == "MhaArm":
            ik2fkArm(context, prop[-2:])
        elif prop[:6] == "MhaLeg":
            ik2fkLeg(context, prop[-2:])
        restoreSnapProp(rig, prop, old, context)
        return{'FINISHED'}    


def setSnapProp(rig, data, value, context, isIk):
    words = data.split()
    prop = words[0]
    oldValue = rig[prop]
    rig[prop] = value
    ik = int(words[1])
    fk = int(words[2])
    extra = int(words[3])
    oldIk = rig.data.layers[ik]
    oldFk = rig.data.layers[fk]
    oldExtra = rig.data.layers[extra]
    rig.data.layers[ik] = True
    rig.data.layers[fk] = True
    rig.data.layers[extra] = True
    updatePose(context)
    if isIk:
        oldValue = 1.0
        oldIk = True
        oldFk = False
    else:
        oldValue = 0.0
        oldIk = False
        oldFk = True
        oldExtra = False
    return (prop, (oldValue, ik, fk, extra, oldIk, oldFk, oldExtra))

    
def restoreSnapProp(rig, prop, old, context):
    updatePose(context)
    (oldValue, ik, fk, extra, oldIk, oldFk, oldExtra) = old
    rig[prop] = oldValue
    rig.data.layers[ik] = oldIk
    rig.data.layers[fk] = oldFk
    rig.data.layers[extra] = oldExtra
    return


class VIEW3D_OT_MhxToggleFkIkButton(bpy.types.Operator):
    bl_idname = "mhx.toggle_fk_ik"
    bl_label = "FK - IK"
    bl_options = {'UNDO'}
    toggle = StringProperty()    

    def execute(self, context):
        words = self.toggle.split()
        rig = context.object
        prop = words[0]
        value = float(words[1]) 
        onLayer = int(words[2])
        offLayer = int(words[3])
        rig.data.layers[onLayer] = True
        rig.data.layers[offLayer] = False
        rig[prop] = value
        # Don't do autokey - confusing.
        #if context.tool_settings.use_keyframe_insert_auto:
        #    rig.keyframe_insert('["%s"]' % prop, frame=scn.frame_current)        
        updatePose(context)
        return{'FINISHED'}    


#
#   MHX FK/IK Switch panel
#

class MhxFKIKPanel(bpy.types.Panel):
    bl_label = "MHX FK/IK Switch"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    #bl_options = {'DEFAULT_CLOSED'}
    
    @classmethod
    def poll(cls, context):
        return (context.object and context.object.MhxRig == 'MHX')

    def draw(self, context):
        rig = context.object
        layout = self.layout        

        row = layout.row()
        row.label("")
        row.label("Left")
        row.label("Right")

        layout.label("FK/IK switch")
        row = layout.row()
        row.label("Arm")
        self.toggleButton(row, rig, "MhaArmIk_L", " 3", " 2")
        self.toggleButton(row, rig, "MhaArmIk_R", " 19", " 18")
        row = layout.row()
        row.label("Leg")
        self.toggleButton(row, rig, "MhaLegIk_L", " 5", " 4")
        self.toggleButton(row, rig, "MhaLegIk_R", " 21", " 20")
        
        layout.label("IK Influence")
        row = layout.row()
        row.label("Arm")
        row.prop(rig, '["MhaArmIk_L"]', text="")
        row.prop(rig, '["MhaArmIk_R"]', text="")
        row = layout.row()
        row.label("Leg")
        row.prop(rig, '["MhaLegIk_L"]', text="")
        row.prop(rig, '["MhaLegIk_R"]', text="")
        
        try:
            ok = (rig["MhxVersion"] >= 12)
        except:
            ok = False
        if not ok:
            layout.label("Snapping only works with MHX version 1.12 and later.")
            return
        
        layout.separator()
        layout.label("Snapping")
        row = layout.row()
        row.label("Rotation Limits")
        row.prop(rig, '["MhaRotationLimits"]', text="")
        row.prop(rig, "MhxSnapExact", text="Exact Snapping")
        
        layout.label("Snap Arm bones")
        row = layout.row()
        row.label("FK Arm")
        row.operator("mhx.snap_fk_ik", text="Snap L FK Arm").data = "MhaArmIk_L 2 3 12"
        row.operator("mhx.snap_fk_ik", text="Snap R FK Arm").data = "MhaArmIk_R 18 19 28"
        row = layout.row()
        row.label("IK Arm")
        row.operator("mhx.snap_ik_fk", text="Snap L IK Arm").data = "MhaArmIk_L 2 3 12"
        row.operator("mhx.snap_ik_fk", text="Snap R IK Arm").data = "MhaArmIk_R 18 19 28"

        layout.label("Snap Leg bones")
        row = layout.row()
        row.label("FK Leg")
        row.operator("mhx.snap_fk_ik", text="Snap L FK Leg").data = "MhaLegIk_L 4 5 12"
        row.operator("mhx.snap_fk_ik", text="Snap R FK Leg").data = "MhaLegIk_R 20 21 28"
        row = layout.row()
        row.label("IK Leg")
        row.operator("mhx.snap_ik_fk", text="Snap L IK Leg").data = "MhaLegIk_L 4 5 12"
        row.operator("mhx.snap_ik_fk", text="Snap R IK Leg").data = "MhaLegIk_R 20 21 28"


    def toggleButton(self, row, rig, prop, fk, ik):
        if rig[prop] > 0.5:
            row.operator("mhx.toggle_fk_ik", text="IK").toggle = prop + " 0" + fk + ik
        else:
            row.operator("mhx.toggle_fk_ik", text="FK").toggle = prop + " 1" + ik + fk
            

###################################################################################    
#
#    Posing panel
#
###################################################################################          
#
#    class MhxDriversPanel(bpy.types.Panel):
#

class MhxDriversPanel(bpy.types.Panel):
    bl_label = "MHX Drivers"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_options = {'DEFAULT_CLOSED'}
    
    @classmethod
    def poll(cls, context):
        return (context.object and context.object.MhxRig)

    def draw(self, context):
        lrProps = []
        props = []
        lrFaceProps = []
        faceProps = []
        plist = list(context.object.keys())
        plist.sort()
        for prop in plist:
            if prop[0:3] == 'Mha':
                if prop[-2:] == '_L':
                    lrProps.append(prop[:-2])
                elif prop[-2:] != '_R':
                    props.append(prop)
            elif prop[0:3] == 'Mhf':
                if prop[-2:] == '_L':
                    lrFaceProps.append(prop[:-2])
                elif prop[-2:] != '_R':
                    faceProps.append(prop)
                    
        ob = context.object
        layout = self.layout
        for prop in props:
            layout.prop(ob, '["%s"]' % prop, text=prop[3:])

        layout.separator()
        row = layout.row()
        row.label("Left")
        row.label("Right")
        for prop in lrProps:
            row = layout.row()
            row.prop(ob, '["%s"]' % (prop+"_L"), text=prop[3:])
            row.prop(ob, '["%s"]' % (prop+"_R"), text=prop[3:])

        if faceProps:
            layout.separator()
            layout.label("Face shapes")
            for prop in faceProps:
                layout.prop(ob, '["%s"]' % prop, text=prop[3:])

            layout.separator()
            row = layout.row()
            row.label("Left")
            row.label("Right")
            for prop in lrFaceProps:
                row = layout.row()
                row.prop(ob, '["%s"]' % (prop+"_L"), text=prop[3:])
                row.prop(ob, '["%s"]' % (prop+"_R"), text=prop[3:])
     
        return

###################################################################################    
#
#    Visibility panel
#
###################################################################################          
#
#    class MhxVisibilityPanel(bpy.types.Panel):
#

class MhxVisibilityPanel(bpy.types.Panel):
    bl_label = "MHX Visibility"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_options = {'DEFAULT_CLOSED'}
    
    @classmethod
    def poll(cls, context):
        return (context.object and context.object.MhxRig)

    def draw(self, context):
        ob = context.object
        layout = self.layout
        props = list(ob.keys())
        props.sort()
        for prop in props:
            if prop[0:3] == "Mhh": 
                layout.prop(ob, '["%s"]' % prop, text="Hide %s" % prop[3:])
        layout.separator()
        layout.operator("mhx.update_textures")
        layout.separator()
        layout.operator("mhx.add_hiders")
        layout.operator("mhx.remove_hiders")
        return

class VIEW3D_OT_MhxUpdateTexturesButton(bpy.types.Operator):
    bl_idname = "mhx.update_textures"
    bl_label = "Update"
    bl_options = {'UNDO'}

    def execute(self, context):
        scn = context.scene
        for mat in bpy.data.materials:
            if mat.animation_data:
                try:
                    mat["MhxDriven"]
                except:
                    continue
                for driver in mat.animation_data.drivers:
                    prop = mat.path_resolve(driver.data_path)
                    value = driver.evaluate(scn.frame_current)
                    prop[driver.array_index] = value
        return{'FINISHED'}    
        
class VIEW3D_OT_MhxAddHidersButton(bpy.types.Operator):
    bl_idname = "mhx.add_hiders"
    bl_label = "Add Hide Property"
    bl_options = {'UNDO'}

    def execute(self, context):
        rig = context.object
        for ob in context.scene.objects:
            if ob.select and ob != rig:
                prop = "Mhh%s" % ob.name        
                defNewProp(prop, "Bool", "default=False")
                rig[prop] = False        
                addHider(ob, "hide", rig, prop)
                addHider(ob, "hide_render", rig, prop)
        return{'FINISHED'}    
                
def addHider(ob, attr, rig, prop):
    fcu = ob.driver_add(attr)
    drv = fcu.driver
    drv.type = 'SCRIPTED'
    drv.expression = "x"
    drv.show_debug_info = True
    var = drv.variables.new()
    var.name = "x"
    targ = var.targets[0]
    targ.id = rig
    targ.data_path = '["%s"]' % prop
    return

class VIEW3D_OT_MhxRemoveHidersButton(bpy.types.Operator):
    bl_idname = "mhx.remove_hiders"
    bl_label = "Remove Hide Property"
    bl_options = {'UNDO'}

    def execute(self, context):
        rig = context.object
        for ob in context.scene.objects:
            if ob.select and ob != rig:
                ob.driver_remove("hide")
                ob.driver_remove("hide_render")
                del rig["Mhh%s" % ob.name]
        return{'FINISHED'}    
        
###################################################################################    
#
#    Layers panel
#
###################################################################################    

MhxLayers = [
    (( 0,    'Root', 'MhxRoot'),
     ( 8,    'Face', 'MhxFace')),
    (( 9,    'Tweak', 'MhxTweak'),
     (10,    'Head', 'MhxHead')),
    (( 1,    'FK Spine', 'MhxFKSpine'),
     (17,    'IK Spine', 'MhxIKSpine')),
    ((13,    'Inv FK Spine', 'MhxInvFKSpine'),
     (16,    'Clothes', 'MhxClothes')),
    ('Left', 'Right'),
    (( 2,    'IK Arm', 'MhxIKArm'),
     (18,    'IK Arm', 'MhxIKArm')),
    (( 3,    'FK Arm', 'MhxFKArm'),
     (19,    'FK Arm', 'MhxFKArm')),
    (( 4,    'IK Leg', 'MhxIKLeg'),
     (20,    'IK Leg', 'MhxIKLeg')),
    (( 5,    'FK Leg', 'MhxFKLeg'),
     (21,    'FK Leg', 'MhxFKLeg')),
    ((12,    'Extra', 'MhxExtra'),
     (28,    'Extra', 'MhxExtra')),
    (( 6,    'Fingers', 'MhxFingers'),
     (22,    'Fingers', 'MhxFingers')),
    (( 7,    'Links', 'MhxLinks'),
     (23,    'Links', 'MhxLinks')),
    ((11,    'Palm', 'MhxPalm'),
     (27,    'Palm', 'MhxPalm')),
]

#
#    class MhxLayersPanel(bpy.types.Panel):
#

class MhxLayersPanel(bpy.types.Panel):
    bl_label = "MHX Layers"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    #bl_options = {'DEFAULT_CLOSED'}
    
    @classmethod
    def poll(cls, context):
        ob = context.object
        if (ob and ob.MhxRig == 'MHX'):
            return True
        return False

    def draw(self, context):
        layout = self.layout
        layout.operator("mhx.pose_enable_all_layers")
        layout.operator("mhx.pose_disable_all_layers")
        amt = context.object.data
        for (left,right) in MhxLayers:
            row = layout.row()
            if type(left) == str:
                row.label(left)
                row.label(right)
            else:
                for (n, name, prop) in [left,right]:
                    row.prop(amt, "layers", index=n, toggle=True, text=name)
        return

class VIEW3D_OT_MhxEnableAllLayersButton(bpy.types.Operator):
    bl_idname = "mhx.pose_enable_all_layers"
    bl_label = "Enable all layers"
    bl_options = {'UNDO'}

    def execute(self, context):
        rig,mesh = getMhxRigMesh(context.object)
        for (left,right) in MhxLayers:
            if type(left) != str:
                for (n, name, prop) in [left,right]:
                    rig.data.layers[n] = True
        return{'FINISHED'}    

class VIEW3D_OT_MhxDisableAllLayersButton(bpy.types.Operator):
    bl_idname = "mhx.pose_disable_all_layers"
    bl_label = "Disable all layers"
    bl_options = {'UNDO'}

    def execute(self, context):
        rig,mesh = getMhxRigMesh(context.object)
        layers = 32*[False]
        pb = context.active_pose_bone
        if pb:
            for n in range(32):
                if pb.bone.layers[n]:
                    layers[n] = True
                    break
        else:
            layers[0] = True
        if rig:
            rig.data.layers = layers
        return{'FINISHED'}    
                
###################################################################################    
#
#    Common functions
#
###################################################################################    
#
#   getMhxRigMesh(ob):
#

def pollMhx(ob):        
    if not ob:
        return False
    elif ob.type == 'ARMATURE':
        return ob.MhxRig
    elif ob.type == 'MESH':
        par = ob.parent
        return (par and (par.type == 'ARMATURE') and par.MhxRig)
    else:
        return False

def getMhxRigMesh(ob):
    if ob.type == 'ARMATURE':
        for mesh in ob.children:
            if mesh.MhxMesh and ob.MhxRig:
                return (ob, mesh)
        return (ob, None)                
    elif ob.type == 'MESH':
        par = ob.parent
        if (par and par.type == 'ARMATURE' and par.MhxRig):
            if ob.MhxMesh:
                return (par, ob)
            else:
                return (par, None)
        else:
            return (None, None)
    return (None, None)
    
        
#
#    setInterpolation(rig):
#

def setInterpolation(rig):
    if not rig.animation_data:
        return
    act = rig.animation_data.action
    if not act:
        return
    for fcu in act.fcurves:
        for pt in fcu.keyframe_points:
            pt.interpolation = 'LINEAR'
        fcu.extrapolation = 'CONSTANT'
    return
    
###################################################################################    
#
#    initialize and register
#
###################################################################################    

def menu_func(self, context):
    self.layout.operator(ImportMhx.bl_idname, text="MakeHuman (.mhx)...")

def register():
    bpy.types.Object.MhAlpha8 = BoolProperty(default=True)
    bpy.types.Object.MhxMesh = BoolProperty(default=False)
    bpy.types.Object.MhxRig = StringProperty(default="")
    bpy.types.Object.MhxRigify = BoolProperty(default=False)
    bpy.types.Object.MhxSnapExact = BoolProperty(default=False)
    bpy.types.Object.MhxShapekeyDrivers = BoolProperty(default=True)
    bpy.types.Object.MhxStrength = FloatProperty(
        name = "Expression strength",
        description = "Multiply expression with this factor",
        default=1.0, min=-1.0, max=2.0
        )
    bpy.utils.register_module(__name__)
    bpy.types.INFO_MT_file_import.append(menu_func)

def unregister():
    try:
        bpy.utils.unregister_module(__name__)
    except:
        pass
    try:
        bpy.types.INFO_MT_file_import.remove(menu_func)
    except:
        pass

if __name__ == "__main__":
    unregister()
    register()