# ##### 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 rna_prop_ui import rna_idprop_ui_prop_get
from mathutils import Vector
from rigify import RigifyError
from rigify_utils import copy_bone_simple

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

def addget_shape_key(obj, name="Key"):
    """ Fetches a shape key, or creates it if it doesn't exist
    """
    # Create a shapekey set if it doesn't already exist
    if obj.data.shape_keys is None:
        shape = obj.add_shape_key(name="Basis", from_mix=False)
        obj.active_shape_key_index = 0

    # Get the shapekey, or create it if it doesn't already exist
    if name in obj.data.shape_keys.keys:
        shape_key = obj.data.shape_keys.keys[name]
    else:
        shape_key = obj.add_shape_key(name=name, from_mix=False)

    return shape_key


def addget_shape_key_driver(obj, name="Key"):
    """ Fetches the driver for the shape key, or creates it if it doesn't
        already exist.
    """
    driver_path = 'keys["' + name + '"].value'
    fcurve = None
    driver = None
    new = False
    if obj.data.shape_keys.animation_data is not None:
        for driver_s in obj.data.shape_keys.animation_data.drivers:
            if driver_s.data_path == driver_path:
                fcurve = driver_s
    if fcurve == None:
        fcurve = obj.data.shape_keys.keys[name].driver_add("value", 0)
        fcurve.driver.type = 'AVERAGE'
        new = True

    return fcurve, new
   
    
def create_shape_and_driver(obj, bone, meshes, shape_name, var_name, var_path, expression):
    """ Creates/gets a shape key and sets up a driver for it.

        obj = armature object
        bone = driving bone name
        meshes = list of meshes to create the shapekey/driver on
        shape_name = name of the shape key
        var_name = name of the driving variable
        var_path = path to the property on the bone to drive with
        expression = python expression for the driver
    """
    pb = obj.pose.bones
    bpy.ops.object.mode_set(mode='OBJECT')

    for mesh_name in meshes:
        mesh_obj = bpy.data.objects[mesh_name]

        # Add/get the shape key
        shape = addget_shape_key(mesh_obj, name=shape_name)

        # Add/get the shape key driver
        fcurve, a = addget_shape_key_driver(mesh_obj, name=shape_name)

        # Set up the driver
        driver = fcurve.driver
        driver.type = 'SCRIPTED'
        driver.expression = expression

        # Get the variable, or create it if it doesn't already exist
        if var_name in driver.variables:
            var = driver.variables[var_name]
        else:
            var = driver.variables.new()
            var.name = var_name

        # Set up the variable
        var.type = "SINGLE_PROP"
        var.targets[0].id_type = 'OBJECT'
        var.targets[0].id = obj
        var.targets[0].data_path = 'pose.bones["' + bone + '"]' + var_path


def mark_actions():
    for action in bpy.data.actions:
        action.tag = True

def get_unmarked_action():
    for action in bpy.data.actions:
        if action.tag != True:
            return action
    return None

def add_action(name=None):
    mark_actions()
    bpy.ops.action.new()
    action = get_unmarked_action()
    if name is not None:
        action.name = name
    return action


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.name]
    except AttributeError:
        raise RigifyError("'%s' rig type requires a parent (bone: %s)" % (RIG_TYPE, orig_bone_name))

    return chain


def deform(obj, definitions, base_names, options):
    bpy.ops.object.mode_set(mode='EDIT')

    eb = obj.data.edit_bones
    pb = obj.pose.bones

    # Get list of eyes
    if "eyes" in options:
        eye_base_names = options["eyes"].replace(" ", "").split(",")
    else:
        eye_base_names = []

    # Get their ORG- names
    eyes = []
    for name in eye_base_names:
        eyes += ["ORG-"+name]

    # Duplicate the eyes to make deformation bones
    def_eyes = [] # def/org pairs
    for eye in eyes:
        def_eyes += [(copy_bone_simple(obj.data, eye, "DEF-"+base_names[eye], parent=True).name, eye)]


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

    # Constraints
    for eye in def_eyes:
        con = pb[eye[0]].constraints.new('COPY_TRANSFORMS')
        con.target = obj
        con.subtarget = eye[1]

    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 = definitions[0]
    eye_target = definitions[1]

    # Get list of pupil mesh objects
    if "mesh" in options:
        pupil_meshes = options["mesh"].replace(" ", "").split(",")
    else:
        pupil_meshes = []

    # Get list of eyes
    if "eyes" in options:
        eye_base_names = options["eyes"].replace(" ", "").split(",")
    else:
        eye_base_names = []

    # Get their ORG- names
    eyes = []
    for name in eye_base_names:
        eyes += ["ORG-"+name]

    # Get the average position of the eyes
    center = Vector(0,0,0)
    for eye in eyes:
        center += eb[eye].head
    if len(eyes) != 0:
        center /= len(eyes)

    # Get the average length of the eyes
    length = 0.0
    for eye in eyes:
        length += eb[eye].length
    if len(eyes) == 0:
        length = 1.0
    else:
        length /= len(eyes)


    # Make the mind's eye
    minds_eye = copy_bone_simple(obj.data, eye_target, "MCH-"+base_names[eye_target]+".mind", parent=True).name
    eb[minds_eye].head = center
    eb[minds_eye].tail = eb[eye_target].head
    eb[minds_eye].roll = 0.0
    eb[minds_eye].length = length

    # Create org/copy/control eye sets
    eye_sets = []
    for eye in eyes:
        copy = copy_bone_simple(obj.data, minds_eye, "MCH-"+base_names[eye]+".cpy", parent=True).name
        eb[copy].translate(eb[eye].head - eb[copy].head)
        eb[copy].parent = eb[eye].parent

        control = copy_bone_simple(obj.data, eye, base_names[eye], parent=True).name
        eb[control].parent = eb[copy]

        eye_sets += [(eye, copy, control)]

    # Bones for parent/free switch for eye target
    target_ctrl = copy_bone_simple(obj.data, eye_target, base_names[eye_target], parent=True).name
    parent = copy_bone_simple(obj.data, head, "MCH-eye_target_parent", parent=False).name

    eb[target_ctrl].parent = eb[parent]




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

    # Axis locks
    pb[target_ctrl].lock_scale = False, True, True

    # Add eye_spread action if it doesn't already exist
    action_name = "eye_spread"
    if action_name in bpy.data.actions:
        spread_action = bpy.data.actions[action_name]
    else:
        spread_action = add_action(name=action_name)

    # Add free property
    prop_name = "free"
    prop = rna_idprop_ui_prop_get(pb[target_ctrl], prop_name, create=True)
    pb[target_ctrl][prop_name] = 0.0
    prop["soft_min"] = 0.0
    prop["soft_max"] = 1.0
    prop["min"] = 0.0
    prop["max"] = 1.0

    free_driver_path = pb[target_ctrl].path_from_id() + '["free"]'

    # Constraints
    # Mind's eye tracks eye target control
    con = pb[minds_eye].constraints.new('DAMPED_TRACK')
    con.target = obj
    con.subtarget = target_ctrl

    # Parent copies transforms of head
    con = pb[parent].constraints.new('COPY_TRANSFORMS')
    con.target = obj
    con.subtarget = head

    fcurve = con.driver_add("influence", 0)
    driver = fcurve.driver
    driver.type = 'AVERAGE'
    mod = fcurve.modifiers[0]
    mod.coefficients[0] = 1.0
    mod.coefficients[1] = -1.0

    var = driver.variables.new()
    var.name = "free"
    var.targets[0].id_type = 'OBJECT'
    var.targets[0].id = obj
    var.targets[0].data_path = free_driver_path

    # Eye set's constraints
    for eye in eye_sets:
        # Org copies transforms of control
        con = pb[eye[0]].constraints.new('COPY_TRANSFORMS')
        con.target = obj
        con.subtarget = eye[2]

        # Copy copies rotation of mind's eye
        con = pb[eye[1]].constraints.new('COPY_ROTATION')
        con.target = obj
        con.subtarget = minds_eye

        # Control gets action constraint for eye spread
        con = pb[eye[2]].constraints.new('ACTION')
        con.target = obj
        con.subtarget = target_ctrl
        con.action = spread_action
        con.transform_channel = 'SCALE_X'
        con.frame_start = -20
        con.frame_end = 20
        con.minimum = 0.0
        con.maximum = 2.0
        con.target_space = 'LOCAL'
    
    
    # Get/create the shape keys and drivers for pupil dilation
    shape_names = ["PUPILS-dilate_wide", "PUPILS-dilate_narrow"]
    slider_name = "pupil_dilate"
    
    # Set up the custom property on the bone
    prop = rna_idprop_ui_prop_get(pb[target_ctrl], slider_name, create=True)
    pb[target_ctrl][slider_name] = 0.0
    prop["min"] = 0.0
    prop["max"] = 1.0
    prop["soft_min"] = 0.0
    prop["soft_max"] = 1.0
    if len(shape_names) > 1:
        prop["min"] = -1.0
        prop["soft_min"] = -1.0

    # Add the shape drivers
    # Positive
    if shape_names[0] != "":
        # Set up the variables for creating the shape key driver
        shape_name = shape_names[0]
        var_name = slider_name.replace(".", "_").replace("-", "_")
        var_path = '["' + slider_name + '"]'
        if slider_name + "_fac" in options:
            fac = options[slider_name + "_fac"]
        else:
            fac = 1.0
        expression = var_name + " * " + str(fac)
        # Create the shape key driver
        create_shape_and_driver(obj, target_ctrl, pupil_meshes, shape_name, var_name, var_path, expression)
    # Negative
    if shape_names[0] != "" and len(shape_names) > 1:
        # Set up the variables for creating the shape key driver
        shape_name = shape_names[1]
        var_name = slider_name.replace(".", "_").replace("-", "_")
        var_path = '["' + slider_name + '"]'
        if slider_name + "_fac" in options:
            fac = options[slider_name + "_fac"]
        else:
            fac = 1.0
        expression = var_name + " * " + str(fac) + " * -1"
        # Create the shape key driver
        create_shape_and_driver(obj, target_ctrl, pupil_meshes, shape_name, var_name, var_path, expression)



    # Set layers
    #layer = list(bb[definitions[2]].layer)
    #bb[lid1].layer = layer
    #bb[lid2].layer = layer
    #bb[lid3].layer = layer
    #bb[lid4].layer = layer
    #bb[lid5].layer = layer
    #bb[lid6].layer = layer
    #bb[lid7].layer = layer
    #bb[lid8].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,)