path: root/release/scripts/modules/rigify/neck.py

# ##### 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

# <pep8 compliant>

import bpy
from rigify import RigifyError
from rigify_utils import bone_class_instance, copy_bone_simple
from rna_prop_ui import rna_idprop_ui_prop_get



def metarig_template():
    # TODO:
    ## 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('body')
    #bone.head[:] = 0.0000, -0.0276, -0.1328
    #bone.tail[:] = 0.0000, -0.0170, -0.0197
    #bone.roll = 0.0000
    #bone.connected = False
    #bone = arm.edit_bones.new('head')
    #bone.head[:] = 0.0000, -0.0170, -0.0197
    #bone.tail[:] = 0.0000, 0.0726, 0.1354
    #bone.roll = 0.0000
    #bone.connected = True
    #bone.parent = arm.edit_bones['body']
    #bone = arm.edit_bones.new('neck.01')
    #bone.head[:] = 0.0000, -0.0170, -0.0197
    #bone.tail[:] = 0.0000, -0.0099, 0.0146
    #bone.roll = 0.0000
    #bone.connected = False
    #bone.parent = arm.edit_bones['head']
    #bone = arm.edit_bones.new('neck.02')
    #bone.head[:] = 0.0000, -0.0099, 0.0146
    #bone.tail[:] = 0.0000, -0.0242, 0.0514
    #bone.roll = 0.0000
    #bone.connected = True
    #bone.parent = arm.edit_bones['neck.01']
    #bone = arm.edit_bones.new('neck.03')
    #bone.head[:] = 0.0000, -0.0242, 0.0514
    #bone.tail[:] = 0.0000, -0.0417, 0.0868
    #bone.roll = 0.0000
    #bone.connected = True
    #bone.parent = arm.edit_bones['neck.02']
    #bone = arm.edit_bones.new('neck.04')
    #bone.head[:] = 0.0000, -0.0417, 0.0868
    #bone.tail[:] = 0.0000, -0.0509, 0.1190
    #bone.roll = 0.0000
    #bone.connected = True
    #bone.parent = arm.edit_bones['neck.03']
    #bone = arm.edit_bones.new('neck.05')
    #bone.head[:] = 0.0000, -0.0509, 0.1190
    #bone.tail[:] = 0.0000, -0.0537, 0.1600
    #bone.roll = 0.0000
    #bone.connected = True
    #bone.parent = arm.edit_bones['neck.04']
    #
    #bpy.ops.object.mode_set(mode='OBJECT')
    #pbone = obj.pose.bones['head']
    #pbone['type'] = 'neck_flex'
    pass
    

def metarig_definition(obj, orig_bone_name):
    '''
    The bone given is neck_01, its parent is the body
    eg.
        body -> neck_01 -> neck_02 -> neck_03.... etc
    '''
    arm = obj.data
    neck = arm.bones[orig_bone_name]
    body = neck.parent

    bone_definition = [body.name, neck.name]
    bone_definition.extend([child.name for child in neck.children_recursive_basename])
    return bone_definition


def deform(obj, definitions, base_names, options):
    for org_bone_name in definitions[1:]:
        bpy.ops.object.mode_set(mode='EDIT')

        # Create deform bone.
        bone = copy_bone_simple(obj.data, org_bone_name, "DEF-%s" % base_names[org_bone_name], parent=True)

        # Store name before leaving edit mode
        bone_name = bone.name

        # Leave edit mode
        bpy.ops.object.mode_set(mode='OBJECT')

        # Get the pose bone
        bone = obj.pose.bones[bone_name]

        # Constrain to the original bone
        # XXX. Todo, is this needed if the bone is connected to its parent?
        con = bone.constraints.new('COPY_TRANSFORMS')
        con.name = "copy_loc"
        con.target = obj
        con.subtarget = org_bone_name


def main(obj, bone_definition, base_names, options):
    from mathutils import Vector

    arm = obj.data
    eb = obj.data.edit_bones
    bb = obj.data.bones
    pb = obj.pose.bones
    
    body = bone_definition[0]
    
    # Create the neck and head control bones
    if "head_name" in options:
        head_name = options["head_name"]
    else:
        head_name = "head"
    
    neck_name = base_names[bone_definition[1]].split(".")[0]
    
    neck_ctrl = copy_bone_simple(arm, bone_definition[1], neck_name).name
    head_ctrl = copy_bone_simple(arm, bone_definition[len(bone_definition)-1], head_name).name
    eb[head_ctrl].tail += eb[neck_ctrl].head - eb[head_ctrl].head
    eb[head_ctrl].head = eb[neck_ctrl].head
    
    # Create hinge and socket bones
    neck_hinge = copy_bone_simple(arm, bone_definition[0], "MCH-" + neck_name + "_hinge").name
    head_hinge = copy_bone_simple(arm, neck_ctrl, "MCH-" + head_name + "_hinge").name
    eb[neck_hinge].tail += eb[neck_ctrl].head - eb[neck_hinge].head
    eb[neck_hinge].head = eb[neck_ctrl].head
    eb[head_hinge].tail += eb[neck_ctrl].head - eb[head_hinge].head
    eb[head_hinge].head = eb[neck_ctrl].head
    
    neck_socket = copy_bone_simple(arm, bone_definition[1], "MCH-" + neck_name + "_socket").name
    head_socket = copy_bone_simple(arm, bone_definition[1], "MCH-" + head_name + "_socket").name
    
    # Parent-child relationships between the body, hinges, controls, and sockets
    eb[neck_ctrl].parent = eb[neck_hinge]
    eb[head_ctrl].parent = eb[head_hinge]
    
    eb[neck_socket].parent = eb[body]
    eb[head_socket].parent = eb[body]
    
    # Create neck bones
    neck = [] # neck bones
    neck_neck = [] # bones constrained to neck control
    neck_head = [] # bones constrained to head control
    for i in range(1, len(bone_definition)):
        # Create bones
        neck_bone = copy_bone_simple(arm, bone_definition[i], base_names[bone_definition[i]]).name
        neck_neck_bone = copy_bone_simple(arm, neck_ctrl, "MCH-" + base_names[bone_definition[i]] + ".neck").name
        neck_head_bone = copy_bone_simple(arm, head_ctrl, "MCH-" + base_names[bone_definition[i]] + ".head").name
        
        # Move them all to the same place
        eb[neck_neck_bone].tail += eb[neck_bone].head - eb[neck_neck_bone].head
        eb[neck_head_bone].tail += eb[neck_bone].head - eb[neck_neck_bone].head
        eb[neck_neck_bone].head = eb[neck_bone].head
        eb[neck_head_bone].head = eb[neck_bone].head
        
        # Parent/child relationships
        eb[neck_bone].parent = eb[neck_head_bone]
        eb[neck_head_bone].parent = eb[neck_neck_bone]
        
        if i > 1:
            eb[neck_neck_bone].parent = eb[neck[i-2]]
        else:
            eb[neck_neck_bone].parent = eb[body]
        
        # Add them to the lists
        neck += [neck_bone]
        neck_neck += [neck_neck_bone]
        neck_head += [neck_head_bone]

    # Create deformation rig
    deform(obj, bone_definition, base_names, options)


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

    # Axis locks
    pb[neck_ctrl].lock_location = True, True, True
    pb[head_ctrl].lock_location = True, True, True
    
    for bone in neck:
        pb[bone].lock_location = True, True, True

    # Neck hinge
    prop = rna_idprop_ui_prop_get(pb[neck_ctrl], "hinge", create=True)
    pb[neck_ctrl]["hinge"] = 0.0
    prop["soft_min"] = 0.0
    prop["soft_max"] = 1.0
    prop["hard_min"] = 0.0
    prop["hard_max"] = 1.0

    con = pb[neck_hinge].constraints.new('COPY_LOCATION')
    con.name = "socket"
    con.target = obj
    con.subtarget = neck_socket

    con = pb[neck_hinge].constraints.new('COPY_ROTATION')
    con.name = "hinge"
    con.target = obj
    con.subtarget = body

    hinge_driver_path = pb[neck_ctrl].path_from_id() + '["hinge"]'

    fcurve = con.driver_add("influence")
    driver = fcurve.driver
    var = driver.variables.new()
    driver.type = 'AVERAGE'
    var.name = "var"
    var.targets[0].id_type = 'OBJECT'
    var.targets[0].id = obj
    var.targets[0].data_path = hinge_driver_path

    mod = fcurve.modifiers[0]
    mod.poly_order = 1
    mod.coefficients[0] = 1.0
    mod.coefficients[1] = -1.0
    
    # Head hinge
    prop = rna_idprop_ui_prop_get(pb[head_ctrl], "hinge", create=True)
    pb[head_ctrl]["hinge"] = 0.0
    prop["soft_min"] = 0.0
    prop["soft_max"] = 1.0
    prop["hard_min"] = 0.0
    prop["hard_max"] = 1.0

    con = pb[head_hinge].constraints.new('COPY_LOCATION')
    con.name = "socket"
    con.target = obj
    con.subtarget = head_socket

    con = pb[head_hinge].constraints.new('COPY_ROTATION')
    con.name = "hinge"
    con.target = obj
    con.subtarget = neck_ctrl

    hinge_driver_path = pb[head_ctrl].path_from_id() + '["hinge"]'

    fcurve = con.driver_add("influence")
    driver = fcurve.driver
    var = driver.variables.new()
    driver.type = 'AVERAGE'
    var.name = "var"
    var.targets[0].id_type = 'OBJECT'
    var.targets[0].id = obj
    var.targets[0].data_path = hinge_driver_path

    mod = fcurve.modifiers[0]
    mod.poly_order = 1
    mod.coefficients[0] = 1.0
    mod.coefficients[1] = -1.0
    
    # Neck rotation constraints
    for i in range(0, len(neck_neck)):
        con = pb[neck_neck[i]].constraints.new('COPY_ROTATION')
        con.name = "neck rotation"
        con.target = obj
        con.subtarget = neck_ctrl
        con.influence = (i+1) / len(neck_neck)
        
    
    # Head rotation constraints/drivers
    prop = rna_idprop_ui_prop_get(pb[head_ctrl], "extent", create=True)
    if "extent" in options:
        pb[head_ctrl]["extent"] = options["extent"]
    else:
        pb[head_ctrl]["extent"] = 0.5
    prop["soft_min"] = 0.0
    prop["soft_max"] = 1.0
    prop["hard_min"] = 0.0
    prop["hard_max"] = 1.0
    
    extent_prop_path = pb[head_ctrl].path_from_id() + '["extent"]'
    
    for i in range(0, len(neck_head)):
        con = pb[neck_head[i]].constraints.new('COPY_ROTATION')
        con.name = "head rotation"
        con.target = obj
        con.subtarget = head_ctrl
        
        if i < (len(neck_head)-1):
            inf = (i+1) / len(neck_head)

            fcurve = con.driver_add("influence")
            driver = fcurve.driver
            var = driver.variables.new()
            var.name = "ext"
            var.targets[0].id_type = 'OBJECT'
            var.targets[0].id = obj
            var.targets[0].data_path = extent_prop_path
            
            driver.expression = "0 if ext == 0 else (((%s-1)/ext)+1)" % inf
        else:
            con.influence = 1.0
    
    # Constrain original bones to the neck bones
    for i in range(0, len(neck)):
        con = pb[bone_definition[i+1]].constraints.new('COPY_TRANSFORMS')
        con.name = "copy_transform"
        con.target = obj
        con.subtarget = neck[i]
        
    
    # Set the controls' custom shapes to use other bones for transforms
    pb[neck_ctrl].custom_shape_transform = pb[bone_definition[len(bone_definition)//2]]
    pb[head_ctrl].custom_shape_transform = pb[bone_definition[len(bone_definition)-1]]


    # last step setup layers
    if "ex_layer" in options:
        layer = [n==options["ex_layer"] for n in range(0,32)]
    else:
        layer = list(arm.bones[bone_definition[1]].layer)
    for bone in neck:
        bb[bone].layer = layer

    layer = list(arm.bones[bone_definition[1]].layer)
    bb[neck_ctrl].layer = layer
    bb[head_ctrl].layer = layer


    # no blending the result of this
    return None