#Lower Limb Twist Joint Script
#This script is to create some joints to simulate limb twisting of an human arm.
#Select the wrist joint before activating the script
#Author: Jiuyang Wang

import maya.cmds as cmds

def generate_joints(number, sel):
    #get the parent joint of the selected joint
    pr = cmds.listRelatives(sel, parent  = True, type = 'joint')
    
    joint_name = 'forearmJnt'
    direct = ['X', 'Y', 'Z']
    length = cmds.getAttr(sel+'.translateX')
    att = ['translate', 'rotate','scale', 'jointOrient']
    
    #if parent exists:
    if pr != None:
        #for loop to go through from 1 to the number of joints+1
        for i in range(1, int(number)+1):
            #create a joint, give it a name and assign it to a var
            jt = cmds.joint()
            #parent the new joint to the parent joint
            cmds.parent(jt,pr)
            #for loop for attributes
            for a in att:
                #for loop for xyz
                for d in direct:
                    #set the attributes of the new joint to 0
                    cmds.setAttr(jt+'.'+a+d, 0)
            #calculate the position of the new joint
            position = length * (i/ (float(number)+1.0))
            #set translateX of the new joint to the position
            cmds.setAttr(jt+'.translateX', position)
            #create a new node multiplyDivide
            md = cmds.shadingNode('multiplyDivide', asUtility = True)
            #connect the translateX of the selection to the input1X of the multiDiv
            cmds.connectAttr(sel+'.rotateX', md+'.input1X')
            #set input2X of the multiDiv to the ratio
            cmds.setAttr(md+'.input2X', i/float(number))
            #connect the outputX of the multiDiv to the new joint's rotateX
            cmds.connectAttr(md+'.outputX', jt+'.rotateX')
    else:#if parent does not exist
        cmds.warning('The joint selected does not have a parent elbow joint') #print a warning msg
            

#ui function that triggers the creation of the twist joints
#promptDialog (needs 2 buttons)
#if the user presses ok (not cancel), then call a function to genetare
#the twist joints

def create_ui():    
    isJoint = False #used to jump out while loop
    wrist_string = ['Wrist', 'wrist', 'WRIST'] #possible string
    have_string = False
    
    #get the selection which is supposed to be a joint
    selected = cmds.ls(selection = True, type = 'joint')
    
    #if seleted smt
    if selected:
        #while loop: to repeatedly construct promptdialog if isJoint is false
        while isJoint == False:
            #create the promptDialog
            pdialog = cmds.promptDialog(
                        title='Lower Limb Twist Joint Creation',
                        message='Number of Joints to Create:',
                        button=['Create', 'Cancel'],
                        defaultButton='Create',
                        cancelButton='Cancel',
                        dismissString='Cancel')
            #if the result of the prompt dialog is create
            if pdialog == 'Create':
                #get the input
                num = cmds.promptDialog(query=True, text=True)
                #if the input is not empty and is number
                if num and num.isdigit():
                    #set the while loop count 0
                    st_num = 0
                    #if count<3
                    while st_num < 3:
                        #check if the selection contains keyword string
                        if wrist_string[st_num] in selected[0]:
                            #if it is true, then bool is true
                            have_string = True
                            st_num = 3 #jump out of the loop
                        st_num+=1 #if not, continue the loop
                    #if the selection contains keyword string
                    if have_string == True:
                        #call the function to create joints
                        generate_joints(num, selected[0])
                        #stop main loop
                        isJoint = True
                    else:
                        #if the selection does not contain keyword string, ask if it is ok to create joints
                        bt = cmds.confirmDialog( title='Joint Creation Confirmation', message='The selected joint does not contain the word wrist. Please confirm if you want to create twist joints based on this joint', button=['Yes','No'], defaultButton='Yes', cancelButton='No', dismissString='No' )
                        #if the user hit yes
                        if bt == 'Yes':
                            #call the function to create joints
                            generate_joints(num, selected[0])
                            #stop main loop
                        isJoint=True
                else:#if the input is not a number, show a warning and reconstruct promptdialog
                    cmds.warning('The input is not a number. Please enter the number of joints you want to created.')
                    
            else:#jump out of the loop
                 isJoint = True
                    
    else:#if selection is empty
        cmds.warning('Nothing selected or the selection is not a joint. Please select the wrist joint')

create_ui()