#====================== 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 ========================

import bpy
from mathutils import Vector
from rigify.utils import MetarigError
from rigify.utils import copy_bone
from rigify.utils import connected_children_names
from rigify.utils import strip_org, make_mechanism_name, make_deformer_name
from rigify.utils import get_layers
from rigify.utils import create_widget, create_line_widget, create_limb_widget
from rna_prop_ui import rna_idprop_ui_prop_get
import re


class Rig:
    """ A finger rig.  It takes a single chain of bones.
        This is a control and deformation rig.
    """
    def __init__(self, obj, bone, params):
        """ Gather and validate data about the rig.
        """
        self.obj = obj
        self.org_bones = [bone] + connected_children_names(obj, bone)
        self.params = params

        if len(self.org_bones) <= 1:
            raise MetarigError("RIGIFY ERROR: Bone '%s': input to rig type must be a chain of 2 or more bones." % (strip_org(bone)))

        # Get user-specified layers, if they exist
        if params.separate_extra_layers:
            self.ex_layers = list(params.extra_layers)
        else:
            self.ex_layers = None

        # Get other rig parameters
        self.primary_rotation_axis = params.primary_rotation_axis
        self.use_digit_twist = params.use_digit_twist

    def deform(self):
        """ Generate the deformation rig.
            Just a copy of the original bones, except the first digit which is a twist bone.
        """
        bpy.ops.object.mode_set(mode='EDIT')

        # Create the bones
        # First bone is a twist bone
        if self.use_digit_twist:
            b1a = copy_bone(self.obj, self.org_bones[0], make_deformer_name(strip_org(self.org_bones[0] + ".01")))
            b1b = copy_bone(self.obj, self.org_bones[0], make_deformer_name(strip_org(self.org_bones[0] + ".02")))
            b1tip = copy_bone(self.obj, self.org_bones[0], make_mechanism_name(strip_org(self.org_bones[0] + ".tip")))
        else:
            b1 = copy_bone(self.obj, self.org_bones[0], make_deformer_name(strip_org(self.org_bones[0])))

        # The rest are normal
        bones = []
        for bone in self.org_bones[1:]:
            bones += [copy_bone(self.obj, bone, make_deformer_name(strip_org(bone)))]

        # Position bones
        eb = self.obj.data.edit_bones
        if self.use_digit_twist:
            b1a_e = eb[b1a]
            b1b_e = eb[b1b]
            b1tip_e = eb[b1tip]

            b1tip_e.use_connect = False
            b1tip_e.tail += Vector((0.1, 0, 0))
            b1tip_e.head = b1b_e.tail
            b1tip_e.length = b1a_e.length / 4

            center = (b1a_e.head + b1a_e.tail) / 2
            b1a_e.tail = center
            b1b_e.use_connect = False
            b1b_e.head = center

        # Parenting
        if self.use_digit_twist:
            b1b_e.parent = eb[self.org_bones[0]]
            b1tip_e.parent = eb[self.org_bones[0]]
        else:
            eb[b1].use_connect = False
            eb[b1].parent = eb[self.org_bones[0]]

        for (ba, bb) in zip(bones, self.org_bones[1:]):
            eb[ba].use_connect = False
            eb[ba].parent = eb[bb]

        # Constraints
        if self.use_digit_twist:
            bpy.ops.object.mode_set(mode='OBJECT')
            pb = self.obj.pose.bones

            b1a_p = pb[b1a]

            con = b1a_p.constraints.new('COPY_LOCATION')
            con.name = "copy_location"
            con.target = self.obj
            con.subtarget = self.org_bones[0]

            con = b1a_p.constraints.new('COPY_SCALE')
            con.name = "copy_scale"
            con.target = self.obj
            con.subtarget = self.org_bones[0]

            con = b1a_p.constraints.new('DAMPED_TRACK')
            con.name = "track_to"
            con.target = self.obj
            con.subtarget = b1tip

    def control(self):
        """ Generate the control rig.
        """
        bpy.ops.object.mode_set(mode='EDIT')

        # Figure out the name for the control bone (remove the last .##)
        ctrl_name = re.sub("([0-9]+\.)", "", strip_org(self.org_bones[0])[::-1], count=1)[::-1]

        # Create the bones
        ctrl = copy_bone(self.obj, self.org_bones[0], ctrl_name)

        helpers = []
        bones = []
        for bone in self.org_bones:
            bones += [copy_bone(self.obj, bone, strip_org(bone))]
            helpers += [copy_bone(self.obj, bone, make_mechanism_name(strip_org(bone)))]

        # Position bones
        eb = self.obj.data.edit_bones

        length = 0.0
        for bone in helpers:
            length += eb[bone].length
            eb[bone].length /= 2

        eb[ctrl].length = length * 1.5

        # Parent bones
        prev = eb[self.org_bones[0]].parent
        for (b, h) in zip(bones, helpers):
            b_e = eb[b]
            h_e = eb[h]
            b_e.use_connect = False
            h_e.use_connect = False

            b_e.parent = h_e
            h_e.parent = prev

            prev = b_e

        # Transform locks and rotation mode
        bpy.ops.object.mode_set(mode='OBJECT')
        pb = self.obj.pose.bones

        for bone in bones[1:]:
            pb[bone].lock_location = True, True, True

        if pb[self.org_bones[0]].bone.use_connect == True:
            pb[bones[0]].lock_location = True, True, True

        pb[ctrl].lock_scale = True, False, True

        for bone in helpers:
            pb[bone].rotation_mode = 'XYZ'

        # Drivers
        i = 1
        val = 1.2 / (len(self.org_bones) - 1)
        for bone in helpers:
            # Add custom prop
            prop_name = "bend_%02d" % i
            prop = rna_idprop_ui_prop_get(pb[ctrl], prop_name, create=True)
            prop["min"] = 0.0
            prop["max"] = 1.0
            prop["soft_min"] = 0.0
            prop["soft_max"] = 1.0
            if i == 1:
                pb[ctrl][prop_name] = 0.0
            else:
                pb[ctrl][prop_name] = val

            # Add driver
            if 'X' in self.primary_rotation_axis:
                fcurve = pb[bone].driver_add("rotation_euler", 0)
            elif 'Y' in self.primary_rotation_axis:
                fcurve = pb[bone].driver_add("rotation_euler", 1)
            else:
                fcurve = pb[bone].driver_add("rotation_euler", 2)

            driver = fcurve.driver
            driver.type = 'SCRIPTED'

            var = driver.variables.new()
            var.name = "ctrl_y"
            var.targets[0].id_type = 'OBJECT'
            var.targets[0].id = self.obj
            var.targets[0].data_path = pb[ctrl].path_from_id() + '.scale[1]'

            var = driver.variables.new()
            var.name = "bend"
            var.targets[0].id_type = 'OBJECT'
            var.targets[0].id = self.obj
            var.targets[0].data_path = pb[ctrl].path_from_id() + '["' + prop_name + '"]'

            if '-' in self.primary_rotation_axis:
                driver.expression = "-(1.0-ctrl_y) * bend * 3.14159 * 2"
            else:
                driver.expression = "(1.0-ctrl_y) * bend * 3.14159 * 2"

            i += 1

        # Constraints
        con = pb[helpers[0]].constraints.new('COPY_LOCATION')
        con.name = "copy_location"
        con.target = self.obj
        con.subtarget = ctrl

        con = pb[helpers[0]].constraints.new('COPY_ROTATION')
        con.name = "copy_rotation"
        con.target = self.obj
        con.subtarget = ctrl

        # Constrain org bones to the control bones
        for (bone, org) in zip(bones, self.org_bones):
            con = pb[org].constraints.new('COPY_TRANSFORMS')
            con.name = "copy_transforms"
            con.target = self.obj
            con.subtarget = bone

        # Set layers for extra control bones
        if self.ex_layers:
            for bone in bones:
                pb[bone].bone.layers = self.ex_layers

        # Create control widgets
        w = create_widget(self.obj, ctrl)
        if w != None:
            mesh = w.data
            verts = [(0, 0, 0), (0, 1, 0), (0.05, 1, 0), (0.05, 1.1, 0), (-0.05, 1.1, 0), (-0.05, 1, 0)]
            if 'Z' in self.primary_rotation_axis:
                # Flip x/z coordinates
                temp = []
                for v in verts:
                    temp += [(v[2], v[1], v[0])]
                verts = temp
            edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 1)]
            mesh.from_pydata(verts, edges, [])
            mesh.update()

        for bone in bones:
            create_limb_widget(self.obj, bone)

    def generate(self):
        """ Generate the rig.
            Do NOT modify any of the original bones, except for adding constraints.
            The main armature should be selected and active before this is called.
        """
        self.deform()
        self.control()

    @classmethod
    def add_parameters(self, group):
        """ Add the parameters of this rig type to the
            RigifyParameters PropertyGroup
        """
        items = [('X', 'X', ''), ('Y', 'Y', ''), ('Z', 'Z', ''), ('-X', '-X', ''), ('-Y', '-Y', ''), ('-Z', '-Z', '')]
        group.primary_rotation_axis = bpy.props.EnumProperty(items=items, name="Primary Rotation Axis", default='X')

        group.separate_extra_layers = bpy.props.BoolProperty(name="Separate Secondary Control Layers:", default=False, description="Enable putting the secondary controls on a separate layer from the primary controls.")
        group.extra_layers = bpy.props.BoolVectorProperty(size=32, description="Layers for the secondary controls to be on.")

        group.use_digit_twist = bpy.props.BoolProperty(name="Digit Twist", default=True, description="Generate the dual-bone twist setup for the first finger digit.")

    @classmethod
    def parameters_ui(self, layout, obj, bone):
        """ Create the ui for the rig parameters.
        """
        params = obj.pose.bones[bone].rigify_parameters[0]

        r = layout.row()
        r.prop(params, "separate_extra_layers")

        r = layout.row()
        r.active = params.separate_extra_layers

        col = r.column(align=True)
        row = col.row(align=True)
        row.prop(params, "extra_layers", index=0, toggle=True, text="")
        row.prop(params, "extra_layers", index=1, toggle=True, text="")
        row.prop(params, "extra_layers", index=2, toggle=True, text="")
        row.prop(params, "extra_layers", index=3, toggle=True, text="")
        row.prop(params, "extra_layers", index=4, toggle=True, text="")
        row.prop(params, "extra_layers", index=5, toggle=True, text="")
        row.prop(params, "extra_layers", index=6, toggle=True, text="")
        row.prop(params, "extra_layers", index=7, toggle=True, text="")
        row = col.row(align=True)
        row.prop(params, "extra_layers", index=16, toggle=True, text="")
        row.prop(params, "extra_layers", index=17, toggle=True, text="")
        row.prop(params, "extra_layers", index=18, toggle=True, text="")
        row.prop(params, "extra_layers", index=19, toggle=True, text="")
        row.prop(params, "extra_layers", index=20, toggle=True, text="")
        row.prop(params, "extra_layers", index=21, toggle=True, text="")
        row.prop(params, "extra_layers", index=22, toggle=True, text="")
        row.prop(params, "extra_layers", index=23, toggle=True, text="")

        col = r.column(align=True)
        row = col.row(align=True)
        row.prop(params, "ik_layers", index=8, toggle=True, text="")
        row.prop(params, "ik_layers", index=9, toggle=True, text="")
        row.prop(params, "ik_layers", index=10, toggle=True, text="")
        row.prop(params, "ik_layers", index=11, toggle=True, text="")
        row.prop(params, "ik_layers", index=12, toggle=True, text="")
        row.prop(params, "ik_layers", index=13, toggle=True, text="")
        row.prop(params, "ik_layers", index=14, toggle=True, text="")
        row.prop(params, "ik_layers", index=15, toggle=True, text="")
        row = col.row(align=True)
        row.prop(params, "ik_layers", index=24, toggle=True, text="")
        row.prop(params, "ik_layers", index=25, toggle=True, text="")
        row.prop(params, "ik_layers", index=26, toggle=True, text="")
        row.prop(params, "ik_layers", index=27, toggle=True, text="")
        row.prop(params, "ik_layers", index=28, toggle=True, text="")
        row.prop(params, "ik_layers", index=29, toggle=True, text="")
        row.prop(params, "ik_layers", index=30, toggle=True, text="")
        row.prop(params, "ik_layers", index=31, toggle=True, text="")

        r = layout.row()
        r.label(text="Bend rotation axis:")
        r.prop(params, "primary_rotation_axis", text="")

        col = layout.column()
        col.prop(params, "use_digit_twist")

    @classmethod
    def create_sample(self, obj):
        # generated by rigify.utils.write_metarig
        bpy.ops.object.mode_set(mode='EDIT')
        arm = obj.data

        bones = {}

        bone = arm.edit_bones.new('finger.01')
        bone.head[:] = 0.0000, 0.0000, 0.0000
        bone.tail[:] = 0.2529, 0.0000, 0.0000
        bone.roll = 3.1416
        bone.use_connect = False
        bones['finger.01'] = bone.name
        bone = arm.edit_bones.new('finger.02')
        bone.head[:] = 0.2529, 0.0000, 0.0000
        bone.tail[:] = 0.4024, 0.0000, -0.0264
        bone.roll = -2.9671
        bone.use_connect = True
        bone.parent = arm.edit_bones[bones['finger.01']]
        bones['finger.02'] = bone.name
        bone = arm.edit_bones.new('finger.03')
        bone.head[:] = 0.4024, 0.0000, -0.0264
        bone.tail[:] = 0.4975, -0.0000, -0.0610
        bone.roll = -2.7925
        bone.use_connect = True
        bone.parent = arm.edit_bones[bones['finger.02']]
        bones['finger.03'] = bone.name

        bpy.ops.object.mode_set(mode='OBJECT')
        pbone = obj.pose.bones[bones['finger.01']]
        pbone.rigify_type = 'finger'
        pbone.lock_location = (True, True, True)
        pbone.lock_rotation = (False, False, False)
        pbone.lock_rotation_w = False
        pbone.lock_scale = (False, False, False)
        pbone.rotation_mode = 'YZX'
        pbone.rigify_parameters.add()
        pbone = obj.pose.bones[bones['finger.02']]
        pbone.rigify_type = ''
        pbone.lock_location = (False, False, False)
        pbone.lock_rotation = (False, False, False)
        pbone.lock_rotation_w = False
        pbone.lock_scale = (False, False, False)
        pbone.rotation_mode = 'YZX'
        pbone = obj.pose.bones[bones['finger.03']]
        pbone.rigify_type = ''
        pbone.lock_location = (False, False, False)
        pbone.lock_rotation = (False, False, False)
        pbone.lock_rotation_w = False
        pbone.lock_scale = (False, False, False)
        pbone.rotation_mode = 'YZX'

        bpy.ops.object.mode_set(mode='EDIT')
        for bone in arm.edit_bones:
            bone.select = False
            bone.select_head = False
            bone.select_tail = False
        for b in bones:
            bone = arm.edit_bones[bones[b]]
            bone.select = True
            bone.select_head = True
            bone.select_tail = True
            arm.edit_bones.active = bone