io_import_scene_mhx.py

        ('PJaw', (0,0.025)),
        ('PTongue', (0,0.0))], 
    'AH' : [('PMouth', (-0.1,0)),
        ('PUpLip', (0,-0.4)),
        ('PLoLip', (0,0)),
        ('PJaw', (0,0.35)),
        ('PTongue', (0,0.0))], 
    'EH' : [('PMouth', (0.1,0)),
        ('PUpLip', (0,-0.2)),
        ('PLoLip', (0,0.2)),
        ('PJaw', (0,0.2)),
        ('PTongue', (0,0.0))], 
    'TH' : [('PMouth', (0,0)),
        ('PUpLip', (0,-0.5)),
        ('PLoLip', (0,0.5)),
        ('PJaw', (-0.2,0.1)),
        ('PTongue', (0,-0.6))], 
    'L' : [('PMouth', (0,0)),
        ('PUpLip', (0,-0.2)),
        ('PLoLip', (0,0.2)),
        ('PJaw', (0.2,0.2)),
        ('PTongue', (0,-0.8))], 
    'G' : [('PMouth', (0,0)),
        ('PUpLip', (0,-0.1)),
        ('PLoLip', (0,0.1)),
        ('PJaw', (-0.3,0.1)),
        ('PTongue', (0,-0.6))], 

    'Blink' : [('PUpLid', (0,1.0)), ('PLoLid', (0,-1.0))], 
    'Unblink' : [('PUpLid', (0,0)), ('PLoLid', (0,0))], 
})

bodyLanguageVisemes = ({
    'Rest' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,-0.6)), 
        ('PUpLipMid', (0,0)), 
        ('PLoLipMid', (0,0)), 
        ('PJaw', (0,0)), 
        ('PTongue', (0,0))], 
    'Etc' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,-0.4)), 
        ('PUpLipMid', (0,0)), 
        ('PLoLipMid', (0,0)), 
        ('PJaw', (0,0)), 
        ('PTongue', (0,0))], 
    'MBP' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,0)), 
        ('PLoLipMid', (0,0)), 
        ('PJaw', (0,0)), 
        ('PTongue', (0,0))], 
    'OO' : [
        ('PMouth', (-1.0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,0)), 
        ('PLoLipMid', (0,0)), 
        ('PJaw', (0,0.4)), 
        ('PTongue', (0,0))], 
    'O' : [
        ('PMouth', (-0.9,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,0)), 
        ('PLoLipMid', (0,0)), 
        ('PJaw', (0,0.8)), 
        ('PTongue', (0,0))], 
    'R' : [
        ('PMouth', (-0.5,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-0.2)), 
        ('PLoLipMid', (0,0.2)), 
        ('PJaw', (0,0)), 
        ('PTongue', (0,0))], 
    'FV' : [
        ('PMouth', (-0.2,0)), 
        ('PMouthMid', (0,1.0)), 
        ('PUpLipMid', (0,0)), 
        ('PLoLipMid', (-0.6,-0.3)), 
        ('PJaw', (0,0)), 
        ('PTongue', (0,0))], 
    'S' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-0.5)), 
        ('PLoLipMid', (0,0.7)), 
        ('PJaw', (0,0)), 
        ('PTongue', (0,0))], 
    'SH' : [
        ('PMouth', (-0.8,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-1.0)), 
        ('PLoLipMid', (0,1.0)), 
        ('PJaw', (0,0)), 
        ('PTongue', (0,0))], 
    'EE' : [
        ('PMouth', (0.2,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-0.6)), 
        ('PLoLipMid', (0,0.6)), 
        ('PJaw', (0,0.05)), 
        ('PTongue', (0,0))], 
    'AH' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-0.4)), 
        ('PLoLipMid', (0,0)), 
        ('PJaw', (0,0.7)), 
        ('PTongue', (0,0))], 
    'EH' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-0.5)), 
        ('PLoLipMid', (0,0.6)), 
        ('PJaw', (0,0.25)), 
        ('PTongue', (0,0))], 
    'TH' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,0)), 
        ('PLoLipMid', (0,0)), 
        ('PJaw', (0,0.2)), 
        ('PTongue', (1.0,1.0))], 
    'L' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-0.5)), 
        ('PLoLipMid', (0,0.5)), 
        ('PJaw', (0,-0.2)), 
        ('PTongue', (1.0,1.0))], 
    'G' : [
        ('PMouth', (0,0)), 
        ('PMouthMid', (0,0)), 
        ('PUpLipMid', (0,-0.5)), 
        ('PLoLipMid', (0,0.5)), 
        ('PJaw', (0,-0.2)), 
        ('PTongue', (-1.0,0))], 

    'Blink' : [('PUpLid', (0,1.0)), ('PLoLid', (0,-1.0))], 
    'Unblink' : [('PUpLid', (0,0)), ('PLoLid', (0,0))], 
})

VisemeList = [
    ('Rest', 'Etc', 'AH'),
    ('MBP', 'OO', 'O'),
    ('R', 'FV', 'S'),
    ('SH', 'EE', 'EH'),
    ('TH', 'L', 'G')
]

#
#   makeVisemes(ob, scn):
#   class VIEW3D_OT_MhxMakeVisemesButton(bpy.types.Operator):
#

def makeVisemes(ob, scn):
    if ob.type != 'MESH':
        print("Active object %s is not a mesh" % ob)
        return
    if not ob.data.shape_keys:
        print("%s has no shapekeys" % ob)
        return
    adata = ob.data.shape_keys.animation_data
    if not adata:
        print("Shapekeys have not drivers")
        return
    try:
        ob.data.shape_keys.key_blocks["VIS_Rest"]
        print("Visemes already created")
        return
    except:
        pass        
    rig = ob.parent
    scale = rig.data.bones['PFace'].length*0.2

    boneShapes = {}
    for fcu in adata.drivers:
        name = fcu.data_path.split('"')[1]
        for var in fcu.driver.variables:
            if var.type == 'TRANSFORMS':
                targ = var.targets[0]  
                fmod = fcu.modifiers[0]
                try:
                    list = boneShapes[targ.bone_target]
                except:
                    list = []
                    boneShapes[targ.bone_target] = list
                list.append((targ.transform_type, fmod, ob.data.shape_keys.key_blocks[name]))
            
    verts = ob.data.vertices            
    for (vis,bones) in bodyLanguageVisemes.items():
        if vis in ['Blink', 'Unblink']:
            continue
        vkey = ob.shape_key_add(name="VIS_%s" % vis)  
        print(vkey.name)
        for n,v in enumerate(verts):
            vkey.data[n].co = v.co
        for (bone, xz) in bones:
            try:
                boneShapes[bone]
                single = True
            except:
                single = False
            if single:
                addToVisShapeKey(boneShapes[bone], vkey, verts, xz, 1, scale)
            else:
                try:
                    boneShapes[bone+"_L"]
                    double = True
                except:
                    double = False
                if double:
                    addToVisShapeKey(boneShapes[bone+"_L"], vkey, verts, xz, 1, scale)
                    #addToVisShapeKey(boneShapes[bone+"_R"], vkey, verts, xz, -1, scale)
    print("Visemes made")                    
    return
    
def addToVisShapeKey(shapes, vkey, verts, xz, sign, scale):
    for (typ, fmod, skey) in shapes:
        factor = fmod.coefficients[1]*scale
        (x,z) = xz
        if typ == 'LOC_X':
            k = factor*sign*x
        elif typ == 'LOC_Z':
            k = factor*z
        if k < skey.slider_min:
            k = skey.slider_min
        if k > skey.slider_max:
            k = skey.slider_max
        if abs(k) < 0.001:
            continue
        print("  %s %.3f %.3f %.3f %.3f" % (skey.name, factor, x, z, k))
        for n,v in enumerate(verts):
            vkey.data[n].co += k*(skey.data[n].co - v.co)
    return            
       
class VIEW3D_OT_MhxMakeVisemesButton(bpy.types.Operator):
    bl_idname = "mhx.make_visemes"
    bl_label = "Generate viseme shapekeys"

    def execute(self, context):
        makeVisemes(context.object, context.scene)
        return{'FINISHED'}    
       
#
#    mohoVisemes
#    magpieVisemes
#

mohoVisemes = dict({
    'rest' : 'Rest', 
    'etc' : 'Etc', 
    'AI' : 'AH', 
    'O' : 'O', 
    'U' : 'OO', 
    'WQ' : 'AH', 
    'L' : 'L', 
    'E' : 'EH', 
    'MBP' : 'MBP', 
    'FV' : 'FV', 
})

magpieVisemes = dict({
    "CONS" : "t,d,k,g,T,D,s,z,S,Z,h,n,N,j,r,tS", 
    "AI" : "i,&,V,aU,I,0,@,aI", 
    "E" : "eI,3,e", 
    "O" : "O,@U,oI", 
    "UW" : "U,u,w", 
    "MBP" : "m,b,p", 
    "L" : "l", 
    "FV" : "f,v", 
    "Sh" : "dZ", 
})

#
#    setViseme(context, vis, setKey, frame):
#    setBoneLocation(context, pbone, loc, mirror, setKey, frame):
#    class VIEW3D_OT_MhxVisemeButton(bpy.types.Operator):
#

def getVisemeSet(context, rig):
    try:
        visset = rig['MhxVisemeSet']
    except:
        return bodyLanguageVisemes
    if visset == 'StopStaring':
        return stopStaringVisemes
    elif visset == 'BodyLanguage':
        return bodyLanguageVisemes
    else:
        raise MhxError("Unknown viseme set %s" % visset)

def setViseme(context, vis, setKey, frame):
    rig = getMhxRig(context.object)
    pbones = rig.pose.bones
    try:
        scale = pbones['PFace'].bone.length
    except:
        return
    visemes = getVisemeSet(context, rig)
    for (b, (x, z)) in visemes[vis]:
        loc = mathutils.Vector((float(x),0,float(z)))
        try:
            pb = pbones[b]
        except:

            pb = None
            
        if pb:
            setBoneLocation(context, pb, scale, loc, False, setKey, frame)
        else:
            setBoneLocation(context, pbones[b+'_L'], scale, loc, False, setKey, frame)
            setBoneLocation(context, pbones[b+'_R'], scale, loc, True, setKey, frame)
    updatePose(context)
    return

def setBoneLocation(context, pb, scale, loc, mirror, setKey, frame):
    if mirror:
        loc[0] = -loc[0]
    pb.location = loc*scale*0.2

    if setKey or context.tool_settings.use_keyframe_insert_auto:
        for n in range(3):
            pb.keyframe_insert('location', index=n, frame=frame, group=pb.name)
    return

class VIEW3D_OT_MhxVisemeButton(bpy.types.Operator):
    bl_idname = 'mhx.pose_viseme'
    bl_label = 'Viseme'
    viseme = StringProperty()

    def invoke(self, context, event):
        setViseme(context, self.viseme, False, context.scene.frame_current)
        return{'FINISHED'}



#
#    openFile(context, filepath):
#    readMoho(context, filepath, offs):
#    readMagpie(context, filepath, offs):
#

def openFile(context, filepath):
    (path, fileName) = os.path.split(filepath)
    (name, ext) = os.path.splitext(fileName)
    return open(filepath, "rU")

def readMoho(context, filepath, offs):
    rig = getMhxRig(context.object)
    context.scene.objects.active = rig
    bpy.ops.object.mode_set(mode='POSE')    
    fp = openFile(context, filepath)        
    for line in fp:
        words= line.split()
        if len(words) < 2:
            pass
        else:
            vis = mohoVisemes[words[1]]
            setViseme(context, vis, True, int(words[0])+offs)
    fp.close()
    setInterpolation(rig)
    updatePose(context)
    print("Moho file %s loaded" % filepath)
    return

def readMagpie(context, filepath, offs):
    rig = getMhxRig(context.object)
    context.scene.objects.active = rig
    bpy.ops.object.mode_set(mode='POSE')    
    fp = openFile(context, filepath)        
    for line in fp: 
        words= line.split()
        if len(words) < 3:
            pass
        elif words[2] == 'X':
            vis = magpieVisemes[words[3]]
            setViseme(context, vis, True, int(words[0])+offs)
    fp.close()
    setInterpolation(rig)
    updatePose(context)
    print("Magpie file %s loaded" % filepath)
    return

# 
#    class VIEW3D_OT_MhxLoadMohoButton(bpy.types.Operator):
#

class VIEW3D_OT_MhxLoadMohoButton(bpy.types.Operator):
    bl_idname = "mhx.pose_load_moho"
    bl_label = "Moho (.dat)"
    filepath = StringProperty(subtype='FILE_PATH')
    startFrame = IntProperty(name="Start frame", description="First frame to import", default=1)

    def execute(self, context):
        import bpy, os, mathutils
        readMoho(context, self.properties.filepath, self.properties.startFrame-1)        
        return{'FINISHED'}    

    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}    

#
#    class VIEW3D_OT_MhxLoadMagpieButton(bpy.types.Operator):
#

class VIEW3D_OT_MhxLoadMagpieButton(bpy.types.Operator):
    bl_idname = "mhx.pose_load_magpie"
    bl_label = "Magpie (.mag)"
    filepath = StringProperty(subtype='FILE_PATH')
    startFrame = IntProperty(name="Start frame", description="First frame to import", default=1)

    def execute(self, context):
        import bpy, os, mathutils
        readMagpie(context, self.properties.filepath, self.properties.startFrame-1)        
        return{'FINISHED'}    

    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}    

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

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

    def draw(self, context):
        rig = getMhxRig(context.object)
        if not rig:
            return

        layout = self.layout        
        layout.label(text="Visemes")
        for (vis1, vis2, vis3) in VisemeList:
            row = layout.row()
            row.operator("mhx.pose_viseme", text=vis1).viseme = vis1
            row.operator("mhx.pose_viseme", text=vis2).viseme = vis2
            row.operator("mhx.pose_viseme", text=vis3).viseme = vis3
        layout.separator()
        row = layout.row()
        row.operator("mhx.pose_viseme", text="Blink").viseme = 'Blink'
        row.operator("mhx.pose_viseme", text="Unblink").viseme = 'Unblink'
        layout.label(text="Load file")
        row = layout.row()
        row.operator("mhx.pose_load_moho")
        row.operator("mhx.pose_load_magpie")
        layout.operator("mhx.update")
        layout.separator()
        layout.operator("mhx.make_visemes")
        return
        
#
#   updatePose(context):
#   class VIEW3D_OT_MhxUpdateButton(bpy.types.Operator):
#

def updatePose(context):
    scn = context.scene
    scn.frame_current = scn.frame_current
    bpy.ops.object.posemode_toggle()
    bpy.ops.object.posemode_toggle()
    return

class VIEW3D_OT_MhxUpdateButton(bpy.types.Operator):
    bl_idname = "mhx.update"
    bl_label = "Update"

    def execute(self, context):
        updatePose(context)
        return{'FINISHED'}    
        

###################################################################################    
#
#    Expression panel
#
###################################################################################    
#
#    class VIEW3D_OT_MhxResetExpressionsButton(bpy.types.Operator):
#

class VIEW3D_OT_MhxResetExpressionsButton(bpy.types.Operator):
    bl_idname = "mhx.pose_reset_expressions"
    bl_label = "Reset expressions"

    def execute(self, context):
        rig = getMhxRig(context.object)
        props = getShapeProps(rig)
        for (prop, name) in props:
            rig[prop] = 0.0
        updatePose(context)
        return{'FINISHED'}    

#
#    class VIEW3D_OT_MhxKeyExpressionButton(bpy.types.Operator):
#

class VIEW3D_OT_MhxKeyExpressionsButton(bpy.types.Operator):
    bl_idname = "mhx.pose_key_expressions"
    bl_label = "Key expressions"

    def execute(self, context):
        rig = getMhxRig(context.object)
        props = getShapeProps(rig)
        frame = context.scene.frame_current
        for (prop, name) in props:
            rig.keyframe_insert('["%s"]' % prop, frame=frame)
        updatePose(context)
        return{'FINISHED'}    
#
#    class VIEW3D_OT_MhxPinExpressionButton(bpy.types.Operator):
#

class VIEW3D_OT_MhxPinExpressionButton(bpy.types.Operator):
    bl_idname = "mhx.pose_pin_expression"
    bl_label = "Pin"
    expression = StringProperty()

    def execute(self, context):
        rig = getMhxRig(context.object)
        props = getShapeProps(rig)
        if context.tool_settings.use_keyframe_insert_auto:
            frame = context.scene.frame_current
            for (prop, name) in props:
                old = rig[prop]
                if prop == self.expression:
                    rig[prop] = 1.0
                else:
                    rig[prop] = 0.0
                if abs(rig[prop] - old) > 1e-3:
                    rig.keyframe_insert('["%s"]' % prop, frame=frame)
        else:                    
            for (prop, name) in props:
                if prop == self.expression:
                    rig[prop] = 1.0
                else:
                    rig[prop] = 0.0
        updatePose(context)
        return{'FINISHED'}    

#
#   getShapeProps(ob):        
#

def getShapeProps(rig):
    props = []        
    plist = list(rig.keys())
    plist.sort()
    for prop in plist:
        if prop[0] == '*':
            props.append((prop, prop[1:]))
    return props                

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

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

    def draw(self, context):
        rig = getMhxRig(context.object)
        if not rig:
            return
        props = getShapeProps(rig)
        if not props:
            return
        layout = self.layout
        layout.label(text="Expressions")
        layout.operator("mhx.pose_reset_expressions")
        layout.operator("mhx.pose_key_expressions")
        #layout.operator("mhx.update")
        layout.separator()
        for (prop, name) in props:
            row = layout.split(0.85)
            row.prop(rig, '["%s"]' % prop, text=name)
            row.operator("mhx.pose_pin_expression", text="", icon='UNPINNED').expression = prop
        return

#########################################
#
#   FK-IK snapping. 
#
#########################################

def getPoseMatrix(mat, pb):
    restInv = pb.bone.matrix_local.inverted()
    if pb.parent:
        parInv = pb.parent.matrix.inverted()
        parRest = pb.parent.bone.matrix_local
        return restInv * (parRest * (parInv * mat))
    else:
        return restInv * mat

        
def getGlobalMatrix(mat, pb):
    gmat = pb.bone.matrix_local * mat
    if pb.parent:
        parMat = pb.parent.matrix
        parRest = pb.parent.bone.matrix_local
        return parMat * (parRest.inverted() * gmat)
    else:
        return gmat


def matchPoseTranslation(pb, fkPb, auto):
    mat = getPoseMatrix(fkPb.matrix, pb)
    insertLocation(pb, mat, auto)
    

def insertLocation(pb, mat, auto):    
    pb.location = mat.to_translation()
    if auto:
        pb.keyframe_insert("location", group=pb.name)
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.mode_set(mode='POSE')


def matchPoseRotation(pb, fkPb, auto):
    mat = getPoseMatrix(fkPb.matrix, pb)
    insertRotation(pb, mat, auto)
    

def insertRotation(pb, mat, auto):    
    q = mat.to_quaternion()
    if pb.rotation_mode == 'QUATERNION':
        pb.rotation_quaternion = q
        if auto:
            pb.keyframe_insert("rotation_quaternion", group=pb.name)
    else:
        pb.rotation_euler = q.to_euler(pb.rotation_mode)
        if auto:
            pb.keyframe_insert("rotation_euler", group=pb.name)
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.mode_set(mode='POSE')


def matchPoseReverse(pb, fkPb, auto):
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.mode_set(mode='POSE')
    gmat = fkPb.matrix * Matrix.Rotation(math.pi, 4, 'Z')
    offs = pb.bone.length * fkPb.matrix.col[1]    
    gmat[0][3] += offs[0]
    gmat[1][3] += offs[1]
    gmat[2][3] += offs[2]    
    mat = getPoseMatrix(gmat, pb)
    pb.matrix_basis = mat
    insertLocation(pb, mat, auto)
    insertRotation(pb, mat, auto)
    


def matchPoseScale(pb, fkPb, auto):
    mat = getPoseMatrix(fkPb.matrix, pb)
    pb.scale = mat.to_scale()
    if auto:
        pb.keyframe_insert("scale", group=pb.name)
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.mode_set(mode='POSE')


def fk2ikArm(context, suffix):
    rig = context.object
    auto = context.scene.tool_settings.use_keyframe_insert_auto
    print("Snap FK Arm%s" % suffix)
    (uparmIk, loarmIk, elbow, elbowPt, wrist) = getSnapBones(rig, "ArmIK", suffix)
    (uparmFk, loarmFk, elbowPtFk, handFk) = getSnapBones(rig, "ArmFK", suffix)

    matchPoseRotation(uparmFk, uparmIk, auto)
    matchPoseScale(uparmFk, uparmIk, auto)

    matchPoseRotation(loarmFk, loarmIk, auto)
    matchPoseScale(loarmFk, loarmIk, auto)

    if rig["&HandFollowsWrist" + suffix]:
        matchPoseRotation(handFk, wrist, auto)
        matchPoseScale(handFk, wrist, auto)
    return


def ik2fkArm(context, suffix):
    rig = context.object
    scn = context.scene
    auto = scn.tool_settings.use_keyframe_insert_auto
    print("Snap IK Arm%s" % suffix)
    (uparmIk, loarmIk, elbow, elbowPt, wrist) = getSnapBones(rig, "ArmIK", suffix)
    (uparmFk, loarmFk, elbowPtFk, handFk) = getSnapBones(rig, "ArmFK", suffix)

    #rig["&ElbowFollowsShoulder" + suffix] = False
    #rig["&ElbowFollowsWrist" + suffix] = False
    
    matchPoseTranslation(wrist, handFk, auto)
    matchPoseRotation(wrist, handFk, auto)  
    matchPoseTranslation(elbow, elbowPtFk, auto)
    matchPoseTranslation(elbowPt, elbowPtFk, auto)
    setInverse(rig, elbowPt)
    return


def fk2ikLeg(context, suffix):
    rig = context.object
    auto = context.scene.tool_settings.use_keyframe_insert_auto
    print("Snap FK Leg%s" % suffix)
    (uplegIk, lolegIk, kneePt, ankleIk, legIk, legFk, footIk, toeIk) = getSnapBones(rig, "LegIK", suffix)
    (uplegFk, lolegFk, kneePtFk, footFk, toeFk) = getSnapBones(rig, "LegFK", suffix)

    matchPoseRotation(uplegFk, uplegIk, auto)
    matchPoseScale(uplegFk, uplegIk, auto)

    matchPoseRotation(lolegFk, lolegIk, auto)
    matchPoseScale(lolegFk, lolegIk, auto)
    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)
    (uplegIk, lolegIk, kneePt, ankleIk, legIk, legFk, footIk, toeIk) = getSnapBones(rig, "LegIK", suffix)
    (uplegFk, lolegFk, kneePtFk, footFk, toeFk) = getSnapBones(rig, "LegFK", suffix)

    #rig["&KneeFollowsHip" + suffix] = False
    #rig["&KneeFollowsFoot" + suffix] = False
    
    legIkToAnkle = rig["&LegIkToAnkle" + 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)
    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


def clearInverse(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_clear_inverse(constraint=cns.name, owner='BONE')
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.ops.object.mode_set(mode='POSE')
    return


def fixAnkle(rig, suffix, scn):
    layers = list(rig.data.layers)
    try:
        rig.data.layers = 32*[True]
        setInverse(rig, rig.pose.bones["Ankle" + suffix])
        scn.frame_current = scn.frame_current
    finally:
        rig.data.layers = layers
    return


def clearAnkle(rig, suffix, scn):
    layers = list(rig.data.layers)
    try:
        rig.data.layers = 32*[True]
        clearInverse(rig, rig.pose.bones["Ankle" + suffix])
        scn.frame_current = scn.frame_current
    finally:
        rig.data.layers = layers
    return


class VIEW3D_OT_FixAnkleButton(bpy.types.Operator):
    bl_idname = "mhx.fix_ankle"
    bl_label = "Fix ankle"
    bl_description = "Set inverse for ankle Child-of constraints"
    suffix = StringProperty()

    def execute(self, context):
        fixAnkle(context.object, self.suffix, context.scene)
        return{'FINISHED'}    


class VIEW3D_OT_ClearAnkleButton(bpy.types.Operator):
    bl_idname = "mhx.clear_ankle"
    bl_label = "Clear ankle"
    bl_description = "Clear inverse for ankle Child-of constraints"
    suffix = StringProperty()

    def execute(self, context):
        clearAnkle(context.object, self.suffix, context.scene)
        return{'FINISHED'}    
#
#
#

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 = []
    for name in names:
        pb = rig.pose.bones[name+suffix]
        pbones.append(pb)
    return tuple(pbones)


class VIEW3D_OT_MhxSnapFk2IkButton(bpy.types.Operator):
    bl_idname = "mhx.snap_fk_ik"
    bl_label = "Snap FK"
    data = StringProperty()    

    def execute(self, context):
        bpy.ops.object.mode_set(mode='POSE')
        rig = context.object
        (prop, old) = setSnapProp(rig, self.data, 1.0, context, False)
        if prop[:4] == "&Arm":
            fk2ikArm(context, prop[-2:])
        elif prop[:4] == "&Leg":
            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"
    data = StringProperty()    

    def execute(self, context):
        bpy.ops.object.mode_set(mode='POSE')
        rig = context.object
        (prop, old) = setSnapProp(rig, self.data, 0.0, context, True)
        if prop[:4] == "&Arm":
            ik2fkArm(context, prop[-2:])
        elif prop[:4] == "&Leg":
            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"
    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 pollMhxRig(context.object)

    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, "&ArmIk_L", " 3", " 2")
        self.toggleButton(row, rig, "&ArmIk_R", " 19", " 18")
        row = layout.row()
        row.label("Leg")
        self.toggleButton(row, rig, "&LegIk_L", " 5", " 4")
        self.toggleButton(row, rig, "&LegIk_R", " 21", " 20")
        
        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.label("Snap Arm bones")
        row = layout.row()
        row.label("FK Arm")
        row.operator("mhx.snap_fk_ik", text="Snap L FK Arm").data = "&ArmIk_L 2 3 12"
        row.operator("mhx.snap_fk_ik", text="Snap R FK Arm").data = "&ArmIk_R 18 19 28"
        row = layout.row()
        row.label("IK Arm")
        row.operator("mhx.snap_ik_fk", text="Snap L IK Arm").data = "&ArmIk_L 2 3 12"
        row.operator("mhx.snap_ik_fk", text="Snap R IK Arm").data = "&ArmIk_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 = "&LegIk_L 4 5 12"
        row.operator("mhx.snap_fk_ik", text="Snap R FK Leg").data = "&LegIk_R 20 21 28"