bl_info = {
    "name": "Blend Bone Weights",
    "author": "Markus Ewald",
    "version": (1, 1),
    "blender": (2, 7, 9),
    "location": "Object -> Tools -> Misc",
    "description": "Blends weights between two or three bones",
    #"warning": "20 Nov 2016",
    "wiki_url": "",
    "tracker_url": "",
    "category": "Mesh"
}

import os
import bpy
import bmesh

class BlendTwoBoneWeightsOperator(bpy.types.Operator):
    """Re-blends the weights of two bones for all selected vertices"""
    bl_idname = "mesh.blend_two_bone_weights"
    bl_label = "Short Name"

    weight_factor = bpy.props.FloatProperty()

    def execute(self, context):
        armature = context.scene.blend_bone_armature
        first_bone = context.scene.blend_bone_first_bone
        second_bone = context.scene.blend_bone_second_bone
        weight_factor = self.weight_factor

        ob = bpy.context.object
        assert ob is not None and ob.type == 'MESH', "active object invalid"

        # Ensure we got the latest assignments and weights,
        # switch to object mode because Blender doesn't let me add vertex groups otherwise
        ob.update_from_editmode()
        previous_mode = bpy.context.object.mode
        if previous_mode != "OBJECT":
            bpy.ops.object.mode_set(mode="OBJECT")

        # Get the indices of the vertex groups for the bones selected by the user
        first_group = None
        second_group = None
        for group in ob.vertex_groups:
            if group.name == first_bone:
                first_group = group
            if group.name == second_bone:
                second_group = group

        # If either of the bones had no vertex groups, create new vertex groups
        # so we can work with them
        if first_group == None:
            first_group = ob.vertex_groups.new(name = first_bone)

        if second_group == None:
            second_group = ob.vertex_groups.new(name = second_bone)

        # Diagnostic output
        #self.report({'INFO'}, first_bone + " is vertex group " + str(first_group.index))
        #self.report({'INFO'}, second_bone + " is vertex group " + str(second_group.index))

        # Go over all selected vertices and adjust the blend of these two groups
        selected_verts = [v for v in ob.data.vertices if v.select]
        for v in selected_verts:
            #self.report({'INFO'}, "Vertex " + str(v.index))

            # Fetch the vertex groups the vertex is assigned to, if any
            vertex_first_group = None
            vertex_second_group = None
            for g in v.groups:
                if g.group == first_group.index:
                    vertex_first_group = g
                if g.group == second_group.index:
                    vertex_second_group = g

            # Get the sum of weights for the two bones we blend between
            total_weight = 0.0
            if vertex_first_group != None:
                total_weight += vertex_first_group.weight
            if vertex_second_group != None:
                total_weight += vertex_second_group.weight

            #self.report({'INFO'}, "  Total weight: " + str(total_weight))

            # Assign the vertex to the first vertex group with the blended weight
            first_blended_weight = total_weight * (1.0 - weight_factor)
            if vertex_first_group == None:
                first_group.add([v.index], first_blended_weight, 'ADD')
            else:
                first_group.add([v.index], first_blended_weight, 'REPLACE')

            #self.report({'INFO'}, "  New first weight: " + str(first_blended_weight))

            # Assign the vertex to the second vertex group with the blended weight
            second_blended_weight = total_weight * (weight_factor)
            if vertex_second_group == None:
                second_group.add([v.index], second_blended_weight, 'ADD')
            else:
                second_group.add([v.index], second_blended_weight, 'REPLACE')

            #self.report({'INFO'}, "  New second weight: " + str(second_blended_weight))

        # Done, restore the previous edit mode
        if previous_mode != "OBJECT":
            bpy.ops.object.mode_set(mode=previous_mode)

        return {'FINISHED'}

class BlendThreeBoneWeightsOperator(bpy.types.Operator):
    """Re-blends the weights of three bones for all selected vertices"""
    bl_idname = "mesh.blend_three_bone_weights"
    bl_label = "Short Name"

    weight_factor = bpy.props.FloatProperty()

    def execute(self, context):
        armature = context.scene.blend_bone_armature
        first_bone = context.scene.blend_bone_first_bone
        second_bone = context.scene.blend_bone_second_bone
        third_bone = context.scene.blend_bone_third_bone
        weight_factor = self.weight_factor

        self.report({'INFO'}, "Weight factor is " + str(weight_factor))

        ob = bpy.context.object
        assert ob is not None and ob.type == 'MESH', "active object invalid"

        # Ensure we got the latest assignments and weights,
        # switch to object mode because Blender doesn't let me add vertex groups otherwise
        ob.update_from_editmode()
        previous_mode = bpy.context.object.mode
        if previous_mode != "OBJECT":
            bpy.ops.object.mode_set(mode="OBJECT")

        # Get the indices of the vertex groups for the bones selected by the user
        first_group = None
        second_group = None
        third_group = None
        for group in ob.vertex_groups:
            if group.name == first_bone:
                first_group = group
            if group.name == second_bone:
                second_group = group
            if group.name == third_bone:
                third_group = group

        # If either of the bones had no vertex groups, create new vertex groups
        # so we can work with them
        if first_group == None:
            first_group = ob.vertex_groups.new(name = first_bone)

        if second_group == None:
            second_group = ob.vertex_groups.new(name = second_bone)

        if third_group == None:
            third_group = ob.vertex_groups.new(name = third_bone)

        # Diagnostic output
        #self.report({'INFO'}, first_bone + " is vertex group " + str(first_group.index))
        #self.report({'INFO'}, second_bone + " is vertex group " + str(second_group.index))
        #self.report({'INFO'}, third_bone + " is vertex group " + str(third_group.index))

        # Go over all selected vertices and adjust the blend of these two groups
        selected_verts = [v for v in ob.data.vertices if v.select]
        for v in selected_verts:
            #self.report({'INFO'}, "Vertex " + str(v.index))

            # Fetch the vertex groups the vertex is assigned to, if any
            vertex_first_group = None
            vertex_second_group = None
            vertex_third_group = None
            for g in v.groups:
                if g.group == first_group.index:
                    vertex_first_group = g
                if g.group == second_group.index:
                    vertex_second_group = g
                if g.group == third_group.index:
                    vertex_third_group = g

            # Get the sum of weights for the three bones we blend between
            total_weight = 0.0
            if vertex_first_group != None:
                total_weight += vertex_first_group.weight
            if vertex_second_group != None:
                total_weight += vertex_second_group.weight
            if vertex_third_group != None:
                total_weight += vertex_third_group.weight

            #self.report({'INFO'}, "  Total weight: " + str(total_weight))

            # Assign the vertex to the first vertex group with the blended weight
            first_weight_factor = abs(0.5 - weight_factor) * 2.0
            if weight_factor >= 0.5:
                first_weight_factor = 0.0
            first_blended_weight = total_weight * first_weight_factor

            if vertex_first_group == None:
                first_group.add([v.index], first_blended_weight, 'ADD')
            else:
                first_group.add([v.index], first_blended_weight, 'REPLACE')

            #self.report({'INFO'}, "  New first weight: " + str(first_blended_weight))

            # Assign the vertex to the second vertex group with the blended weight
            second_weight_factor = max(0.0, 1.0 - abs(0.5 - weight_factor) * 2.0)
            second_blended_weight = total_weight * second_weight_factor

            if vertex_second_group == None:
                second_group.add([v.index], second_blended_weight, 'ADD')
            else:
                second_group.add([v.index], second_blended_weight, 'REPLACE')

            #self.report({'INFO'}, "  New second weight: " + str(second_blended_weight))

            # Assign the vertex to the third vertex group with the blended weight
            third_weight_factor = abs(0.5 - weight_factor) * 2.0
            if weight_factor < 0.5:
                third_weight_factor = 0.0
            third_blended_weight = total_weight * third_weight_factor

            if vertex_third_group == None:
                third_group.add([v.index], third_blended_weight, 'ADD')
            else:
                third_group.add([v.index], third_blended_weight, 'REPLACE')

            #self.report({'INFO'}, "  New third weight: " + str(third_blended_weight))

        # Done, restore the previous edit mode
        if previous_mode != "OBJECT":
            bpy.ops.object.mode_set(mode=previous_mode)

        return {'FINISHED'}

class BlendBoneWeightsPanel(bpy.types.Panel):
    """Creates a Panel in the Object properties window"""
    bl_label = "Bone Weights"
    bl_idname = "OBJECT_PT_blend_weights"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'TOOLS'
    bl_context = 'mesh_edit'

    def draw(self, context):
        layout = self.layout
        scene = context.scene

        col = layout.column()

        # Whether to blend between two or three bones
        col.prop(scene, "blend_bone_mode")
        col.separator()

        # Selection on bones between which blending should happen
        col.prop_search(scene, "blend_bone_armature", bpy.data, "armatures")
        arma = bpy.data.armatures.get(scene.blend_bone_armature)
        if arma is not None:
            if scene.blend_bone_mode == 'TWO':
                col.prop_search(scene, "blend_bone_first_bone", arma, "bones")
                col.prop_search(scene, "blend_bone_second_bone", arma, "bones")
            if scene.blend_bone_mode == 'THREE':
                col.prop_search(scene, "blend_bone_first_bone", arma, "bones")
                col.prop_search(scene, "blend_bone_second_bone", arma, "bones")
                col.prop_search(scene, "blend_bone_third_bone", arma, "bones")
        col.separator()

        # Number of vertices between which to blend
        col.prop(scene, "blend_bone_vertex_count")
        if scene.blend_bone_mode == 'TWO':
            for x in range(0, scene.blend_bone_vertex_count + 1):
                blend_factor = x / scene.blend_bone_vertex_count
                props = col.operator("mesh.blend_two_bone_weights", text="Blend " + str(blend_factor))
                props.weight_factor = blend_factor
        if scene.blend_bone_mode == 'THREE':
            for x in range(0, scene.blend_bone_vertex_count + 1):
                blend_factor = x / scene.blend_bone_vertex_count
                props = col.operator("mesh.blend_three_bone_weights", text="Blend " + str(blend_factor))
                props.weight_factor = blend_factor
        col.separator()

        col.prop(scene, "blend_bone_weight_factor")
        if scene.blend_bone_mode == 'TWO':
            props = col.operator("mesh.blend_two_bone_weights", text='Blend Bone Weights')
            props.weight_factor = scene.blend_bone_weight_factor
        if scene.blend_bone_mode == 'THREE':
            props = col.operator("mesh.blend_three_bone_weights", text='Blend Bone Weights')
            props.weight_factor = scene.blend_bone_weight_factor

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

    bpy.types.Scene.blend_bone_mode = bpy.props.EnumProperty(
        name = "Blend Bones",
        description = "Number of bones weights should be blended for",
        items = [
            ( 'TWO', 'Two', '', 2),
            ( 'THREE', 'Three', '', 3)
        ],
        default = 'TWO'
    )
    bpy.types.Scene.blend_bone_armature = bpy.props.StringProperty(
        name = "Armature",
        description = "Armature from which bones can be selected"
    )
    bpy.types.Scene.blend_bone_first_bone = bpy.props.StringProperty(
        name = "First Bone",
        description = "First bone over which weights will be blended"
    )
    bpy.types.Scene.blend_bone_second_bone = bpy.props.StringProperty(
        name = "Second Bone",
        description = "Second bone over which weights will be blended"
    )
    bpy.types.Scene.blend_bone_third_bone = bpy.props.StringProperty(
        name = "Third Bone",
        description = "Third bone over which weights will be blended"
    )
    bpy.types.Scene.blend_bone_vertex_count = bpy.props.IntProperty(
        name = "Vertex Count",
        description = "Number of vertices along which weights will be blended",
        default = 8, min = 2, max = 25
    )
    bpy.types.Scene.blend_bone_weight_factor = bpy.props.FloatProperty(
        name = "Transition",
        description = "Weight transition between first bone and second",
        default = 0.5, min = 0.0, max = 1.0
    )

def unregister():
    del bpy.types.Scene.blend_bone_weight_factor
    del bpy.types.Scene.blend_bone_vertex_count
    del bpy.types.Scene.blend_bone_third_bone
    del bpy.types.Scene.blend_bone_second_bone
    del bpy.types.Scene.blend_bone_first_bone
    del bpy.types.Scene.blend_bone_armature
    del bpy.types.Scene.blend_bone_mode

    bpy.utils.unregister_module(__name__)

# This allows you to run the script directly from blenders text editor
# to test the addon without having to install it
if __name__ == "__main__":
    register()