# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
# ##### END GPL LICENSE BLOCK #####

# <pep8 compliant>

import bpy
from math import acos
from Mathutils import Vector
from rigify import get_layer_dict
from rigify_utils import bone_class_instance, copy_bone_simple

#METARIG_NAMES = ("cpy",)
RIG_TYPE = "mouth"


def metarig_template():
    # generated by rigify.write_meta_rig
    bpy.ops.object.mode_set(mode='EDIT')
    obj = bpy.context.active_object
    arm = obj.data
    bone = arm.edit_bones.new('Bone')
    bone.head[:] = 0.0000, 0.0000, 0.0000
    bone.tail[:] = 0.0000, 0.0000, 1.0000
    bone.roll = 0.0000
    bone.connected = False

    bpy.ops.object.mode_set(mode='OBJECT')
    pbone = obj.pose.bones['Bone']
    pbone['type'] = 'copy'


def metarig_definition(obj, orig_bone_name):
    bone = obj.data.bones[orig_bone_name]
    chain = []
    
    try:
        chain += [bone.parent.name, bone.parent.parent.name, bone.name]
    except AttributeError:
        raise RigifyError("'%s' rig type requires a chain of two parents (bone: %s)" % (RIG_TYPE, base_names[0]))
    
    chain += [child.name for child in bone.children_recursive_basename]
    
    if len(chain) < 10:
        raise RigifyError("'%s' rig type requires a chain of 8 bones (bone: %s)" % (RIG_TYPE, base_names[0]))
    
    return chain[:10]


def deform(obj, definitions, base_names, options):
    bpy.ops.object.mode_set(mode='EDIT')
    
    eb = obj.data.edit_bones
    pb = obj.pose.bones
    
    # Upper lip MCH
    lip1  = make_lip_stretch_bone(obj, "MCH-lip", definitions[2], definitions[3], 1.0)
    lip2  = make_lip_stretch_bone(obj, "MCH-lip", definitions[3], definitions[4], 1.0)
    lip22 = make_lip_stretch_bone(obj, "MCH-lip", definitions[4], definitions[5], 1.0)
    lip33 = make_lip_stretch_bone(obj, "MCH-lip", definitions[4], definitions[3], 1.0)
    lip3  = make_lip_stretch_bone(obj, "MCH-lip", definitions[5], definitions[4], 1.0)
    lip4  = make_lip_stretch_bone(obj, "MCH-lip", definitions[6], definitions[5], 1.0)
    
    dlip22 = copy_bone_simple(obj.data, lip22, "MCH-lip", parent=True).name
    dlip33 = copy_bone_simple(obj.data, lip33, "MCH-lip", parent=True).name
    eb[dlip22].bbone_segments = 8
    eb[dlip33].bbone_segments = 8
    
    eb[lip1].parent = eb[definitions[2]]
    eb[lip2].parent = eb[definitions[3]]
    eb[lip22].parent = eb[definitions[4]]
    eb[lip33].parent = eb[definitions[4]]
    eb[lip3].parent = eb[definitions[5]]
    eb[lip4].parent = eb[definitions[6]]
        
    # Lower lip MCH
    lip5 =  make_lip_stretch_bone(obj, "MCH-lip", definitions[6], definitions[7], 1.0)
    lip6 =  make_lip_stretch_bone(obj, "MCH-lip", definitions[7], definitions[8], 1.0)
    lip66 = make_lip_stretch_bone(obj, "MCH-lip", definitions[8], definitions[9], 1.0)
    lip77 = make_lip_stretch_bone(obj, "MCH-lip", definitions[8], definitions[7], 1.0)
    lip7 =  make_lip_stretch_bone(obj, "MCH-lip", definitions[9], definitions[8], 1.0)
    lip8 =  make_lip_stretch_bone(obj, "MCH-lip", definitions[2], definitions[9], 1.0)
    
    dlip66 = copy_bone_simple(obj.data, lip66, "MCH-lip", parent=True).name
    dlip77 = copy_bone_simple(obj.data, lip77, "MCH-lip", parent=True).name
    eb[dlip66].bbone_segments = 8
    eb[dlip77].bbone_segments = 8
    
    eb[lip5].parent = eb[definitions[6]]
    eb[lip6].parent = eb[definitions[7]]
    eb[lip66].parent = eb[definitions[8]]
    eb[lip77].parent = eb[definitions[8]]
    eb[lip7].parent = eb[definitions[9]]
    eb[lip8].parent = eb[definitions[2]]
    
    # Upper lip DEF
    dlip1 = copy_bone_simple(obj.data, lip1, "DEF-lip", parent=True).name
    dlip2 = copy_bone_simple(obj.data, lip2, "DEF-lip", parent=True).name
    dlip3 = copy_bone_simple(obj.data, lip3, "DEF-lip", parent=True).name
    dlip4 = copy_bone_simple(obj.data, lip4, "DEF-lip", parent=True).name
    
    eb[dlip2].parent  = eb[dlip1]
    eb[dlip22].parent = eb[dlip2]
    
    eb[dlip3].parent  = eb[dlip4]
    eb[dlip33].parent = eb[dlip3]
    
    eb[dlip2].connected  = True
    eb[dlip22].connected = True
    eb[dlip3].connected  = True
    eb[dlip33].connected = True
    
    eb[dlip1].bbone_segments = 8
    eb[dlip2].bbone_segments = 8
    eb[dlip3].bbone_segments = 8
    eb[dlip4].bbone_segments = 8
        
    # Lower lip DEF
    dlip5 = copy_bone_simple(obj.data, lip5, "DEF-lip", parent=True).name
    dlip6 = copy_bone_simple(obj.data, lip6, "DEF-lip", parent=True).name
    dlip7 = copy_bone_simple(obj.data, lip7, "DEF-lip", parent=True).name
    dlip8 = copy_bone_simple(obj.data, lip8, "DEF-lip", parent=True).name
    
    eb[dlip6].parent = eb[dlip5]
    eb[dlip66].parent = eb[dlip6]
    
    eb[dlip7].parent = eb[dlip8]
    eb[dlip77].parent = eb[dlip7]
    
    eb[dlip6].connected = True
    eb[dlip66].connected = True
    eb[dlip7].connected = True
    eb[dlip77].connected = True
    
    eb[dlip5].bbone_segments = 8
    eb[dlip6].bbone_segments = 8
    eb[dlip7].bbone_segments = 8
    eb[dlip8].bbone_segments = 8
    
    
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # Constraints
    con = pb[dlip1].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip1
    
    con = pb[dlip22].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip22
    
    con = pb[dlip33].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip33
    
    con = pb[dlip2].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip2
    
    con = pb[dlip3].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip3
    
    con = pb[dlip4].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip4
    
    con = pb[dlip5].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip5
    
    con = pb[dlip6].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip6
    
    con = pb[dlip66].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip66
    
    con = pb[dlip77].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip77
    
    con = pb[dlip7].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip7
    
    con = pb[dlip8].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = lip8
    
    return (None,)




def control(obj, definitions, base_names, options):
    bpy.ops.object.mode_set(mode='EDIT')
    
    eb = obj.data.edit_bones
    bb = obj.data.bones
    pb = obj.pose.bones
    
    head_e = eb[definitions[0]]
    jaw_e = eb[definitions[1]]
    
    # Head lips
    hlip1 = copy_bone_simple(obj.data, definitions[2], "MCH-"+base_names[definitions[2]]+".head").name
    hlip2 = copy_bone_simple(obj.data, definitions[3], "MCH-"+base_names[definitions[3]]+".head").name
    hlip3 = copy_bone_simple(obj.data, definitions[4], "MCH-"+base_names[definitions[4]]+".head").name
    hlip4 = copy_bone_simple(obj.data, definitions[5], "MCH-"+base_names[definitions[5]]+".head").name
    hlip5 = copy_bone_simple(obj.data, definitions[6], "MCH-"+base_names[definitions[6]]+".head").name
    hlip6 = copy_bone_simple(obj.data, definitions[7], "MCH-"+base_names[definitions[7]]+".head").name
    hlip7 = copy_bone_simple(obj.data, definitions[8], "MCH-"+base_names[definitions[8]]+".head").name
    hlip8 = copy_bone_simple(obj.data, definitions[9], "MCH-"+base_names[definitions[9]]+".head").name
    
    eb[hlip1].parent = head_e
    eb[hlip2].parent = head_e
    eb[hlip3].parent = head_e
    eb[hlip4].parent = head_e
    eb[hlip5].parent = head_e
    eb[hlip6].parent = head_e
    eb[hlip7].parent = head_e
    eb[hlip8].parent = head_e
    
    # Jaw lips
    jlip1 = copy_bone_simple(obj.data, definitions[2], "MCH-"+base_names[definitions[2]]+".jaw").name
    jlip2 = copy_bone_simple(obj.data, definitions[3], "MCH-"+base_names[definitions[3]]+".jaw").name
    jlip3 = copy_bone_simple(obj.data, definitions[4], "MCH-"+base_names[definitions[4]]+".jaw").name
    jlip4 = copy_bone_simple(obj.data, definitions[5], "MCH-"+base_names[definitions[5]]+".jaw").name
    jlip5 = copy_bone_simple(obj.data, definitions[6], "MCH-"+base_names[definitions[6]]+".jaw").name
    jlip6 = copy_bone_simple(obj.data, definitions[7], "MCH-"+base_names[definitions[7]]+".jaw").name
    jlip7 = copy_bone_simple(obj.data, definitions[8], "MCH-"+base_names[definitions[8]]+".jaw").name
    jlip8 = copy_bone_simple(obj.data, definitions[9], "MCH-"+base_names[definitions[9]]+".jaw").name
    
    eb[jlip1].parent = jaw_e
    eb[jlip2].parent = jaw_e
    eb[jlip3].parent = jaw_e
    eb[jlip4].parent = jaw_e
    eb[jlip5].parent = jaw_e
    eb[jlip6].parent = jaw_e
    eb[jlip7].parent = jaw_e
    eb[jlip8].parent = jaw_e
    
    # Control lips
    lip1 = copy_bone_simple(obj.data, definitions[2], base_names[definitions[2]]).name
    lip2 = copy_bone_simple(obj.data, definitions[3], base_names[definitions[3]]).name
    lip3 = copy_bone_simple(obj.data, definitions[4], base_names[definitions[4]]).name
    lip4 = copy_bone_simple(obj.data, definitions[5], base_names[definitions[5]]).name
    lip5 = copy_bone_simple(obj.data, definitions[6], base_names[definitions[6]]).name
    lip6 = copy_bone_simple(obj.data, definitions[7], base_names[definitions[7]]).name
    lip7 = copy_bone_simple(obj.data, definitions[8], base_names[definitions[8]]).name
    lip8 = copy_bone_simple(obj.data, definitions[9], base_names[definitions[9]]).name
    
    size = eb[lip1].length
    eb[lip1].tail = eb[lip1].head + Vector(0,size,0)
    eb[lip2].tail = eb[lip2].head + Vector(0,size,0)
    eb[lip3].tail = eb[lip3].head + Vector(0,size,0)
    eb[lip4].tail = eb[lip4].head + Vector(0,size,0)
    eb[lip5].tail = eb[lip5].head + Vector(0,size,0)
    eb[lip6].tail = eb[lip6].head + Vector(0,size,0)
    eb[lip7].tail = eb[lip7].head + Vector(0,size,0)
    eb[lip8].tail = eb[lip8].head + Vector(0,size,0)
    
    eb[lip1].roll = 0
    eb[lip2].roll = 0
    eb[lip3].roll = 0
    eb[lip4].roll = 0
    eb[lip5].roll = 0
    eb[lip6].roll = 0
    eb[lip7].roll = 0
    eb[lip8].roll = 0
    
    eb[lip1].parent = eb[jlip1]
    eb[lip2].parent = eb[jlip2]
    eb[lip3].parent = eb[jlip3]
    eb[lip4].parent = eb[jlip4]
    eb[lip5].parent = eb[jlip5]
    eb[lip6].parent = eb[jlip6]
    eb[lip7].parent = eb[jlip7]
    eb[lip8].parent = eb[jlip8]
    
    # Link lips
    llip1 = copy_bone_simple(obj.data, definitions[2], "MCH-"+base_names[definitions[2]]+".link").name
    llip2 = copy_bone_simple(obj.data, definitions[3], "MCH-"+base_names[definitions[3]]+".link").name
    llip3 = copy_bone_simple(obj.data, definitions[4], "MCH-"+base_names[definitions[4]]+".link").name
    llip4 = copy_bone_simple(obj.data, definitions[5], "MCH-"+base_names[definitions[5]]+".link").name
    llip5 = copy_bone_simple(obj.data, definitions[6], "MCH-"+base_names[definitions[6]]+".link").name
    llip6 = copy_bone_simple(obj.data, definitions[7], "MCH-"+base_names[definitions[7]]+".link").name
    llip7 = copy_bone_simple(obj.data, definitions[8], "MCH-"+base_names[definitions[8]]+".link").name
    llip8 = copy_bone_simple(obj.data, definitions[9], "MCH-"+base_names[definitions[9]]+".link").name
    
    eb[llip1].parent = eb[lip1]
    eb[llip2].parent = eb[lip2]
    eb[llip3].parent = eb[lip3]
    eb[llip4].parent = eb[lip4]
    eb[llip5].parent = eb[lip5]
    eb[llip6].parent = eb[lip6]
    eb[llip7].parent = eb[lip7]
    eb[llip8].parent = eb[lip8]
    
    
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # Constraints
    
    # Jaw lips to head lips
    influence = [0.0, 0.1, 0.5]
    
    con = pb[jlip1].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip1
    con.influence = influence[2]
    
    con = pb[jlip2].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip2
    con.influence = influence[1]
    
    con = pb[jlip3].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip3
    con.influence = influence[0]
    
    con = pb[jlip4].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip4
    con.influence = influence[1]
    
    con = pb[jlip5].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip5
    con.influence = influence[2]
    
    con = pb[jlip6].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip6
    con.influence = 1.0 - influence[1]
    
    con = pb[jlip7].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip7
    con.influence = 1.0 - influence[0]
    
    con = pb[jlip8].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = hlip8
    con.influence = 1.0 - influence[1]
    
    # ORG bones to link lips
    con = pb[definitions[2]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip1
    
    con = pb[definitions[3]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip2
    
    con = pb[definitions[4]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip3
    
    con = pb[definitions[5]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip4
    
    con = pb[definitions[6]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip5
    
    con = pb[definitions[7]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip6
    
    con = pb[definitions[8]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip7
    
    con = pb[definitions[9]].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = llip8
    
    
    # Set layers
    layer = list(bb[definitions[2]].layer)
    bb[lip1].layer = layer
    bb[lip2].layer = layer
    bb[lip3].layer = layer
    bb[lip4].layer = layer
    bb[lip5].layer = layer
    bb[lip6].layer = layer
    bb[lip7].layer = layer
    bb[lip8].layer = layer
    
    
    return (None,)




def main(obj, bone_definition, base_names, options):
    # Create control rig
    control(obj, bone_definition, base_names, options)
    # Create deform rig
    deform(obj, bone_definition, base_names, options)

    return (None,)




def make_lip_stretch_bone(obj, name, bone1, bone2, roll_alpha):
    eb = obj.data.edit_bones
    pb = obj.pose.bones
    
    # Create the bone, pointing from bone1 to bone2
    bone_e = copy_bone_simple(obj.data, bone1, name, parent=True)
    bone_e.connected = False
    bone_e.tail = eb[bone2].head
    bone = bone_e.name
    
    # Align the bone roll with the average direction of bone1 and bone2
    vec = bone_e.y_axis.cross(((1.0-roll_alpha)*eb[bone1].y_axis) + (roll_alpha*eb[bone2].y_axis)).normalize()
    
    ang = acos(vec * bone_e.x_axis)
    
    bone_e.roll += ang
    c1 = vec * bone_e.x_axis
    bone_e.roll -= (ang*2)
    c2 = vec * bone_e.x_axis
    
    if c1 > c2:
        bone_e.roll += (ang*2)
    
    bpy.ops.object.mode_set(mode='OBJECT')
    bone_p = pb[bone]
    
    # Constrains
    con = bone_p.constraints.new('COPY_LOCATION')
    con.target = obj
    con.subtarget = bone1
    
    con = bone_p.constraints.new('DAMPED_TRACK')
    con.target = obj
    con.subtarget = bone2

    con = bone_p.constraints.new('STRETCH_TO')
    con.target = obj
    con.subtarget = bone2
    con.volume = 'NO_VOLUME'

    bpy.ops.object.mode_set(mode='EDIT')

    return bone