import bpy
import mathutils
from bpy.props import IntProperty, CollectionProperty, StringProperty, FloatProperty, FloatVectorProperty
from bpy.types import Panel, UIList


class OBJECT_UL_muscle_anims(UIList):
    def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
        split = layout.split()
        split.prop(item, "name", text="", emboss=False, translate=False, icon='BORDER_RECT')


class SaveMuscleAnimOperator(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "object.save_muscles_animation"
    bl_label = "Save muscles animation"

    @classmethod
    def poll(cls, context):
        return context.active_object is not None and context.active_object.type == 'MESH' and context.active_object.muscle_animation_index >= 0

    def execute(self, context):
        obj = context.active_object
        muscles = [mod for mod in obj.modifiers if mod.type == 'MUSCLE']
        #there should be only one muscle modifier
        muscle = muscles[0]
        anim = obj.muscle_animations[obj.muscle_animation_index]
        anim.contraction_values.clear()
        for muscle in muscle.muscles:
            item = anim.contraction_values.add()
            item.muscle = muscle.name
            item.value = muscle.contraction
           
        if obj.muscle_animation_jaw:
            bone = obj.parent.pose.bones[obj.muscle_animation_jaw]
            oldmode = bone.rotation_mode
            bone.rotation_mode = 'XYZ'
            anim.jaw_rotation = mathutils.Vector((bone.rotation_euler.x, bone.rotation_euler.y, bone.rotation_euler.z));
            bone.rotation_mode = oldmode
        return {'FINISHED'}


class InsertMuscleAnimOperator(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "object.insert_muscles_animation"
    bl_label = "Insert muscles animation to current frame"

    @classmethod
    def poll(cls, context):
        return context.active_object is not None and context.active_object.type == 'MESH' and context.active_object.muscle_animation_index >= 0

    def execute(self, context):
        obj = context.active_object
        anim = obj.muscle_animations[obj.muscle_animation_index]
        #should exist...
        mod = next(i for i in obj.modifiers if i.type == 'MUSCLE')
        for item in anim.contraction_values:
            muscles = [i for i in mod.muscles if i.muscle and item.muscle in i.muscle.name]
            if len(muscles) > 0:
                muscle = muscles[0]
                muscle.contraction = item.value*obj.muscle_animation_factor
                muscle.keyframe_insert(data_path="contraction")
            
        if obj.muscle_animation_jaw:
            bone = obj.parent.pose.bones[obj.muscle_animation_jaw]
            oldmode = bone.rotation_mode
            bone.rotation_mode = 'XYZ'
            bone.rotation_euler.x = anim.jaw_rotation[0]*obj.muscle_animation_factor
            bone.rotation_euler.y = anim.jaw_rotation[1]*obj.muscle_animation_factor
            bone.rotation_euler.z = anim.jaw_rotation[2]*obj.muscle_animation_factor
            bone.rotation_mode = oldmode
            if bone.rotation_mode == 'QUATERNION':
                bone.keyframe_insert(data_path="rotation_quaternion")
            else:
                bone.keyframe_insert(data_path="rotation_euler")
        return {'FINISHED'}
        
class AddMuscleAnimOperator(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "object.muscle_anim_add"
    bl_label = "Add new muscle animation"

    @classmethod
    def poll(cls, context):
        return context.active_object is not None and context.active_object.type == 'MESH'

    def execute(self, context):
        ma = context.active_object.muscle_animations.add()
        ma.name = "Muscle animation"
        return {'FINISHED'}


class RemoveMuscleAnimOperator(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "object.muscle_anim_remove"
    bl_label = "Delete currently selected muscle animation"

    @classmethod
    def poll(cls, context):
        return context.active_object is not None and context.active_object.type == 'MESH' and context.active_object.muscle_animation_index >= 0

    def execute(self, context):
        context.active_object.muscle_animations.remove(context.active_object.muscle_animation_index)
        return {'FINISHED'}

        
class SaveMusclesAnimOperator(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "object.save_muscles_anim_file"
    bl_label = "Save muscles animations"
    filepath = bpy.props.StringProperty(subtype="FILE_PATH")

    @classmethod
    def poll(cls, context):
        return context.active_object is not None and context.active_object.type == 'MESH' and len(context.active_object.muscle_animations)> 0

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

    def execute(self, context):
        obj = context.active_object
        anims = obj.muscle_animations
        f = open(self.filepath, 'w')
        f.write("%d\n" % len(anims))
        for anim in anims:
            f.write("%s\n" % (anim.name))
            f.write("%f %f %f\n" % (anim.jaw_rotation[0], anim.jaw_rotation[1], anim.jaw_rotation[2]))
            f.write("%d\n" % (len(anim.contraction_values)))
            for item in anim.contraction_values:
                f.write("%s\n" % (item.muscle))
                f.write("%f\n" % (item.value))
        f.close()
        return {'FINISHED'}
        
class LoadMusclesAnimOperator(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "object.load_muscles_anim_file"
    bl_label = "Load muscles animations"
    filepath = bpy.props.StringProperty(subtype="FILE_PATH")

    @classmethod
    def poll(cls, context):
        return context.active_object is not None and context.active_object.type == 'MESH'

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

    def execute(self, context):
        obj = context.active_object
        anims = obj.muscle_animations
        f = open(self.filepath, 'r')
        n = int(f.readline().strip())
        created = {}
        for i in range(0, n):
            name = f.readline().strip()
            rx, ry, rz = f.readline().strip().split(' ')
            cnt = int(f.readline().strip())
            anim = anims.add()
            anim.name = name
            anim.jaw_rotation = mathutils.Vector((float(rx), float(ry), float(rz)))
            for j in range(0, cnt):
                muscle = f.readline().strip()
                value = float(f.readline().strip())
                item = anim.contraction_values.add()
                item.muscle = muscle;
                item.value = value;
        f.close()
        return {'FINISHED'}   
 
 
def bone_items(self, context):
    arma = context.active_object.parent
    if arma is None or arma.type != 'ARMATURE':
        return []
    return [(bone.name, bone.name, "") for bone in arma.data.bones]

    
class PanelMuscleAnimations(Panel):
    """Creates a Panel in the Object properties window"""
    bl_label = "Muscle animation Panel"
    bl_idname = "OBJECT_PT_ui_list_muscle_animation"
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "object"   
    
    def draw(self, context):
        layout = self.layout
        ob = context.object
        row = layout.row()
        if ob.parent is not None and ob.parent.type == 'ARMATURE':
            row.prop_search(ob, "muscle_animation_jaw", ob.parent.data, "bones")            
            
        row = layout.row()
        row.template_list("OBJECT_UL_muscle_anims", "", ob, "muscle_animations", ob, "muscle_animation_index")
        col = row.column(align=True)
        col.operator("object.muscle_anim_add", icon='ZOOMIN', text="")
        col.operator("object.muscle_anim_remove", icon='ZOOMOUT', text="")

        row = layout.row()
        row.prop(ob, 'muscle_animation_factor', text='Factor');
        row = layout.row()
        sub = row.row(align=True)
        sub.operator("object.save_muscles_animation", text="Save current")
        sub.operator("object.insert_muscles_animation", text="Insert current")

        row = layout.row()
        sub = row.row(align=True)
        sub.operator("object.save_muscles_anim_file", text="Save to file")
        sub.operator("object.load_muscles_anim_file", text="Load from file")

        
class MuscleContraction(bpy.types.PropertyGroup):
    muscle = StringProperty()
    value = FloatProperty()

class MuscleAnimation(bpy.types.PropertyGroup):
    # name = StringProperty()
    contraction_values = CollectionProperty(type=MuscleContraction)
    id = IntProperty()
    jaw_rotation = FloatVectorProperty()

if __name__ == "__main__":
    register()