'''
FK Control Generator
This script will create a control and group for the selected joint. The joint will be constrained 
to the control
Authot: Jiuyang Wang
'''

import maya.cmds as cmds

controlGroupName = 'CtrlGrp'
controlName = 'Ctrl'
def generate_control(option, check, grpText, cText):
    #1. Get the user's selection, and store the selected joints in a variable.
    sel = cmds.ls(selection = True, type = 'joint')
    grpName = cmds.textField(grpText, q=True, text = True)
    cName = cmds.textField(cText, q=True, text = True)
    check_result = cmds.checkBox(check, q = True, value = True)
    #check if the user has selected smt
    if sel:
        if option == 'selection':
            #loop starts
            for s in sel:
                # check if the control of the joint has been already created
                if cmds.objExists(s+cName):
                    # if the control has children
                    if cmds.listRelatives(s+cName, children=True):
                        #move the children to the top of the hierarchy
                        cmds.parent(cmds.listRelatives(s+cName, children=True),world = True)
                    #delete the control
                    cmds.delete(s+cName)
                # check if the control group of the joint has been already created
                if cmds.objExists(s+grpName):
                    #delete the control group
                    cmds.delete(s+grpName)
                #2. Create a curve without history, store it in a variable
                c = cmds.circle(name = s + cName, nr=(1, 0, 0), constructionHistory = False)[0]
                #3. Group the circle. Store the result in a variablee
                gp = cmds.group(c, name=s + grpName )
                #4. Match the group to the joint
                cmds.matchTransform(gp, s)
                #5. Create parentConstraint (target = control, constrained node = joint)
                cmds.parentConstraint(c, s)
                print 'Successfully created FK Control of the joint '+ s
                #6. Parent the control group to the control of the parent joint
                #6a. Get the parent of the current joint: cmds.listRelatives(    parent = True    )
                pr = cmds.listRelatives(s, parent  = True, type = 'joint')
                #6a2. Get the children of the current joint
                cl = cmds.listRelatives(s, children = True, type = 'joint')
                #6b. check to see if the parent joint exists
                if pr != None:
                    # 6c. check if the control of the parent joint exists
                    if cmds.objExists(pr[0] + cName):
                        # get the children of the parent control group, and store it in a var
                        cd = cmds.listRelatives(pr[0] + grpName, children = True)[0]
                        if cmds.nodeType(cmds.listRelatives(cd, children = True)[0]) == 'nurbsCurve':
                            #6d. parent the joint control group to the parent control: cmds.parent(controlgroup, parentControl)
                            cmds.parent(gp, cd)
                            print 'Successfully grouped the FK Control of the joint '+ s+ ' under its parent control group'
                        else: #if the tyoe of the parent control isn't a nurbcurve
                            cmds.warning('The FK Control of the parent of the joint' + s+' is not a nurbcurve. Failed to group the control group')
                    else: #if the parent control doesn't exist
                        print 'The FK Control of the parent of the joint ' + s +" does not exist."
                else: #if the parent joint doesn't exist
                    print 'The joint '+s+' has no parent joint. A new joint group has been created.'
                # if the children of the current joint exist
                if cl:
                    #go through the children to see if any of the children have a correct control and control group, and it can be parented under the crrent control group
                    for e in cl:
                        # check if the control of the children and the control group exist, and they have correct hierarchy
                        if cmds.objExists(e+cName) and cmds.objExists(e+grpName) and cmds.listRelatives(e+grpName, children = True)[0] == e+cName:
                            #check if the control group hasn't been parented
                            if cmds.listRelatives(e+grpName, parent = True)==None:
                                #parent the child's control group to the joint control
                                print 'Successfully grouped the joint '+e+' control group to the joint '+s+' control'
                                cmds.parent(e+grpName,c) 
        elif option == 'hierarchy':
            sel_hierarchy = cmds.listRelatives(sel, allDescendents = True, type = 'joint')
            for child in sel_hierarchy:
                if not child in sel:
                    has_children = cmds.listRelatives(child, children=True)
                    if check_result or has_children:
                        sel.append(child)
            #loop starts
            for s in sel:
                # check if the control of the joint has been already created
                if cmds.objExists(s+cName):
                    # if the control has children
                    if cmds.listRelatives(s+cName, children=True):
                        #move the children to the top of the hierarchy
                        cmds.parent(cmds.listRelatives(s+cName, children=True),world = True)
                    #delete the control
                    cmds.delete(s+cName)
                # check if the control group of the joint has been already created
                if cmds.objExists(s+grpName):
                    #delete the control group
                    cmds.delete(s+grpName)
                #2. Create a curve without history, store it in a variable
                c = cmds.circle(name = s + cName, nr=(1, 0, 0), constructionHistory = False)[0]
                #3. Group the circle. Store the result in a variablee
                gp = cmds.group(c, name=s + grpName )
                #4. Match the group to the joint
                cmds.matchTransform(gp, s)
                #5. Create parentConstraint (target = control, constrained node = joint)
                cmds.parentConstraint(c, s)
                print 'Successfully created FK Control of the joint '+ s
                #6. Parent the control group to the control of the parent joint
                #6a. Get the parent of the current joint: cmds.listRelatives(    parent = True    )
                pr = cmds.listRelatives(s, parent  = True, type = 'joint')
                #6a2. Get the children of the current joint
                cl = cmds.listRelatives(s, children = True, type = 'joint')
                #6b. check to see if the parent joint exists
                if pr != None:
                    # 6c. check if the control of the parent joint exists
                    if cmds.objExists(pr[0] + cName):
                        # get the children of the parent control group, and store it in a var
                        cd = cmds.listRelatives(pr[0] + grpName, children = True)[0]
                        if cmds.nodeType(cmds.listRelatives(cd, children = True)[0]) == 'nurbsCurve':
                            #6d. parent the joint control group to the parent control: cmds.parent(controlgroup, parentControl)
                            cmds.parent(gp, cd)
                            print 'Successfully grouped the FK Control of the joint '+ s+ ' under its parent control group'
                        else: #if the tyoe of the parent control isn't a nurbcurve
                            cmds.warning('The FK Control of the parent of the joint' + s+' is not a nurbcurve. Failed to group the control group')
                    else: #if the parent control doesn't exist
                        print 'The FK Control of the parent of the joint ' + s +" does not exist."
                else: #if the parent joint doesn't exist
                    print 'The joint '+s+' has no parent joint. A new joint group has been created.'
                # if the children of the current joint exist
                if cl:
                    #go through the children to see if any of the children have a correct control and control group, and it can be parented under the crrent control group
                    for e in cl:
                        # check if the control of the children and the control group exist, and they have correct hierarchy
                        if cmds.objExists(e+cName) and cmds.objExists(e+grpName) and cmds.listRelatives(e+grpName, children = True)[0] == e+cName:
                            #check if the control group hasn't been parented
                            if cmds.listRelatives(e+grpName, parent = True)==None:
                                #parent the child's control group to the joint control
                                print 'Successfully grouped the joint '+e+' control group to the joint '+s+' control'
                                cmds.parent(e+grpName,c) 
    else:
        #if the user has selected nothing or the selected object isn't a joint, give a dialog message
        cmds.confirmDialog(message="Nothing selected or the selection isn't a joint. Please select the joint that needs a FK control")
    
    
def fk_generator():
    window_name = 'fkGenerator'
    if cmds.window(window_name, exists = True):
        cmds.deleteUI(window_name)
    cmds.window(window_name)
    cmds.rowColumnLayout()
    cmds.text(label = 'Create FK Controls:')
    cmds.separator(visible = False)#separator
    cmds.separator(visible = False)#separator
    name_layout = cmds.rowColumnLayout(numberOfColumns = 2)
    #text to describe the text field
    cmds.text(label='Control Group Naming: ')
    #text field for entering the keywords. by default it takes the var DefaultSwitchKeyword
    grpName=cmds.textField(text=controlGroupName)
     #text to describe the text field
    cmds.text(label='Control Naming: ')
    #text field for entering the keywords. by default it takes the var DefaultSwitchKeyword
    cName=cmds.textField(text=controlName)
    cmds.separator(visible = False)#separator
    cmds.separator(visible = False)#separator
    #parent the layout
    cmds.setParent('..')
    cmds.rowColumnLayout(numberOfColumns = 1)
    check = cmds.checkBox(label='Include End Joints in the hierarchy')
    cmds.rowColumnLayout(numberOfColumns = 2)
    cmds.button(label = 'Create Controls for Selected', command = lambda *args: generate_control('selection', check, grpName, cName))
    cmds.button(label = 'Create Controls for Hierarchy', command = lambda *args: generate_control('hierarchy', check, grpName, cName))
    cmds.showWindow(window_name)
fk_generator()