##############################################################################################################################
#
#    Script:          ikfkGenerator.py
#    Author:          Audrey Paransky & Simon Burdick
#    Last Updated:    06/16/2024
#    Created:         06/16/2024
#    Description:     This script generates an IK and FK skeleton rig for arms and legs. It creates a GUI window for easy use.
#
##############################################################################################################################

import maya.cmds as cmds
import functools

def createUI (pWindowTitle, pApplyCallback ):
    windowID = "myWindowID"

    if cmds.window( windowID, exists = True):
        cmds.deleteUI( windowID )
    
    cmds.window( windowID, title=pWindowTitle, sizeable = True, resizeToFitChildren=True )

    cmds.rowColumnLayout( numberOfColumns = 6, columnWidth=[ (1,60), (2,60), (3,60), (4,60), (5,60), (6,60) ] ) #row column layout means you have to define the columns but the rows can be infinite

    # FIRST ROW --------------------------------
    cmds.text( label='Joint 1:')
    joint1Field = cmds.textField( text=' ')

    cmds.text( label='Joint 2:')
    joint2Field = cmds.textField( text=' ')

    cmds.text( label='Joint 3:')
    joint3Field = cmds.textField( text=' ')

    # SECOND ROW --------------------------------
    cmds.separator( h=10, style='none' )
    cmds.button( label='Select', command=functools.partial( getSelection, 
                                                           joint1Field ) )
    cmds.separator( h=10, style='none' )
    cmds.button( label='Select', command=functools.partial( getSelection, 
                                                           joint2Field ) )
    cmds.separator( h=10, style='none' )
    cmds.button( label='Select', command=functools.partial( getSelection, 
                                                           joint3Field ) )
    # THIRD ROW --------------------------------
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    
    # FOURTH ROW --------------------------------
    cmds.text( label='Name')
    sysNameField = cmds.textField( text='(arm, leg, etc.)')

    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )

    # FIFTH ROW --------------------------------
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )

    # SIXTH ROW --------------------------------
    cmds.text( label='Position' )

    positionRadioCol = cmds.radioCollection()

    cmds.radioButton( label='Left', collection=positionRadioCol)
    cmds.radioButton( label='Center', collection=positionRadioCol, select=True)
    cmds.radioButton( label='Right', collection=positionRadioCol)

    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    
    # SEVENTH ROW --------------------------------
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    
    # EIGHTH ROW --------------------------------
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )

    # NINTH ROW --------------------------------
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )

    cmds.button( label='Apply', command=(functools.partial( pApplyCallback,
                                                           joint1Field,
                                                           joint2Field,
                                                           joint3Field,
                                                           positionRadioCol,
                                                           sysNameField ) ) )

    def cancelCallback( *pArgs ):
        if cmds.window( windowID, exists = True ):
            cmds.deleteUI( windowID)

    cmds.button( label='Close', command = cancelCallback)
    
    cmds.separator( h=10, style='none' )
    cmds.separator( h=10, style='none' )

    #display the window
    cmds.showWindow()

def getSelection( pField, *args ):
    selection = cmds.ls( selection=True, type='joint')

    if selection:
        cmds.textField( pField, edit=True, text=selection[0])
    else:
        cmds.error("No joints selected", n=True)

def applyCallback( pJoint1Field, pJoint2Field, pJoint3Field, pRadioCol, pPrefixName, *pArgs):

    jointList = []
    jointList.clear()
    jointList.append( cmds.textField( pJoint1Field, query=True, text=True ) )
    jointList.append( cmds.textField( pJoint2Field, query=True, text=True ) )
    jointList.append( cmds.textField( pJoint3Field, query=True, text=True ) )

    prefixName = ( cmds.textField( pPrefixName, query=True, text=True ) )

    radioSelection = cmds.radioCollection( pRadioCol, query=True, select=True ) #assigning whatever selected radio button (will return which radio button is selected)(which is an entire class)
    label = cmds.radioButton( radioSelection, query=True, label=True ) #will return the name of the radio button (a string)

    if label == 'Left':
        pPositionPrefix = 'l_'
    elif label == 'Right':
        pPositionPrefix = 'r_'
    elif label == 'Center':
        pPositionPrefix = 'c_'
    else:
        pPositionPrefix = '0_' #just in case something fails

    if jointList[0]!='' and jointList[1]!='' and jointList[2]!='':
        print('passed')
        generateRig( jointList, pPositionPrefix, prefixName)
    else:
        cmds.error("Not enough joints selected", n=True)






def generateRig( pJointList, pPosition, pSysName):
    
    #GET SOURCE JOINTS AND DATA
    systemName = pSysName
    jointList = pJointList
    positionPrefix = pPosition
    
    jointRotList = []
    jointTransList = []
    
    jntScaleFactor = 10
    
    for jnt in jointList:
        jointRotList.append(cmds.xform(jnt, query = True, rotation = True, worldSpace = True))
        jointTransList.append(cmds.xform(jnt, query = True, translation = True, worldSpace = True))
    
    #BUILD FK SKELETON
    fkJointList = []
    
    for n in range(3): #0,1,2
        cmds.select(deselect = True)
        fkJointList.append(cmds.joint( position = jointTransList[n], orientation = jointRotList[n], name = jointList[n]+'_fk'))
        
        if (n > 0): #if we are at anything but the first joint (we are at the second or third joint)...
            cmds.parent(fkJointList[n], fkJointList[n-1]) #parent the current joint to the previous joint
            
    #BUILD FK SYSTEM
        controlCurveList = []
        
    for jnt in fkJointList:
        ## Joint Information
        # Get parent of current joint
        jntParent = cmds.listRelatives( jnt, parent=True )
        if( jntParent ):
            jntParentName = jntParent[0]
        else:
            jntParentName = "GARBAGE" #this is for the first joint in the heirarchy since it has no parent. there needs to be something
            
        # Get Radius of current joint
        jntRadius = cmds.getAttr( jnt + '.radius' )
        
        # Get translation and rotation of current joint
        jntTranslation = cmds.xform( jnt, query=True, translation=True, worldSpace=True )
        jntRotation = cmds.xform( jnt, query=True, rotation=True, worldSpace=True )
        
        ## Controller Manipulation
        # Create new controller
        newControl = cmds.circle( c=(0, 0, 0), name=(jnt + '_ctrl') ) #c changes the shape of the cirlce, 0,0,0 ensures that it's a full circle
        controlCurveList.append( newControl[0] ) #newControl will return a dictionary containing the transform and the shape node, you just want the name, hence [0]
        cmds.move( jntTranslation[0], jntTranslation[1], jntTranslation[2], newControl )
        cmds.rotate( 0, 90, 0, newControl )
        cmds.scale( ( 5*jntRadius ), ( 5*jntRadius ), ( 5*jntRadius ), newControl )
        cmds.makeIdentity( newControl, apply=True )
        
        # Create the offset group for our new controller
        newGroup = cmds.group( empty=True, name=(jnt + '_offset') )
        cmds.move( jntTranslation[0], jntTranslation[1], jntTranslation[2], newGroup ) #the group and the cntrl are at the same place
        cmds.makeIdentity( newGroup, apply=True ) #freeze transformations
        
        # Make new control a child of new group
        cmds.parent( newControl, newGroup )
        
        # Rotate offset group to match the rotations of the joint
        cmds.rotate( jntRotation[0], jntRotation[1], jntRotation[2], newGroup )
        
        ## Place group in the right hierarchy
        if ( jntParentName + '_ctrl' ) in controlCurveList:
            cmds.parent( newGroup, (jntParentName + '_ctrl') )
            print( newGroup + ' parented under ' + jntParentName + '_ctrl' )
        else:
            print( 'No parent found' )
            print( jntParentName ) 
            print( controlCurveList )
            
        ## Parent constrain joint under control
        cmds.parentConstraint( newControl, jnt, mo = True )
        
    #BUILD IK SKELETON
    ikJointList = []
    
    for n in range(3): #0,1,2
        cmds.select(deselect = True)
        ikJointList.append(cmds.joint( position = jointTransList[n], orientation = jointRotList[n], name = jointList[n]+'_ik'))
        
        if (n > 0): #if we are at anything but the first joint (we are at the second or third joint)...
            cmds.parent(ikJointList[n], ikJointList[n-1]) #parent the current joint to the previous joint
            
    #BUILD IK SYSTEM
    ikSysHandle = cmds.ikHandle( name = systemName + 'ikHandle', sj=ikJointList[0], ee=ikJointList[2] )
    
    #first ik controller (shoulder)
    ik0Ctrl = cmds.circle(name = jointList[0] + '_ik_ctrl')
    cmds.rotate( 0, 90, 0, ik0Ctrl )
    cmds.scale( jntScaleFactor, jntScaleFactor, jntScaleFactor, ik0Ctrl)
    cmds.makeIdentity(ik0Ctrl, apply =True)
    ik0CtrlOffset = cmds.group( name = jointList[0] + '_ik_offset')
    cmds.xform( ik0CtrlOffset, translation = jointTransList[0], rotation = jointRotList[0])
    
    cmds.parentConstraint(ik0Ctrl, ikJointList[0], mo = True) #mo means maintain offset
    
    #last ik controller (wrist)
    ik1Ctrl = cmds.circle( name = jointList[2] + '_ik_ctrl')
    cmds.rotate( 0, 90, 0, ik1Ctrl )
    cmds.scale( jntScaleFactor, jntScaleFactor, jntScaleFactor, ik1Ctrl)
    cmds.makeIdentity(ik1Ctrl, apply =True)
    ik1CtrlOffset = cmds.group( name = jointList[2] + '_ik_offset')
    cmds.xform( ik1CtrlOffset, translation = jointTransList[2], rotation = jointRotList[2])
    
    cmds.parent(ikSysHandle[0], ik1Ctrl )
    cmds.orientConstraint( ik1Ctrl, ikJointList[2], mo=True )
    
    #aim controller (elbow)
    ikAimCtrl = cmds.circle( name=jointList[1] + 'Aim_ik_ctrl')
    cmds.scale( jntScaleFactor, jntScaleFactor, jntScaleFactor, ikAimCtrl)
    cmds.makeIdentity(ikAimCtrl, apply =True)
    ikAimCtrlOffset = cmds.group( name = jointList[1] + '_ik_offset')
    cmds.xform( ikAimCtrlOffset, translation = (jointTransList[1][0], jointTransList[1][1], jointTransList[1][2]-2))
    
    cmds.poleVectorConstraint( ikAimCtrl, ikSysHandle[0] )
    
    #parent everything under the shoulder control
    cmds.parent( ik1CtrlOffset, ik0Ctrl )
    cmds.parent( ikAimCtrlOffset, ik0Ctrl )
    
    #clean up into groups
    fkSystemGroup = cmds.group( empty = True, name= systemName + '_FK_system') #empty = True is so that it doesnt get a prev thing
    cmds.parent( fkJointList[0], fkSystemGroup )
    cmds.parent( fkJointList[0] + '_offset', fkSystemGroup)
    cmds.hide( fkJointList[0]) #hide the fk shoulder joint
    
    ikSystemGroup = cmds.group( empty = True, name= systemName + '_IK_system') #empty = True is so that it doesnt get a prev thing
    cmds.parent( ikJointList[0], ikSystemGroup )
    cmds.parent( ikJointList[0] + '_offset', ikSystemGroup)
    cmds.hide( ikJointList[0]) #hide the ik shoulder joint
    cmds.hide( ikSysHandle[0]) #hide the ik handle
    
    #CONNECT THE SYSTEMS TO THE JOINTS
    parentConstraints = [] #we want to collect the parent constraints so we can access the weight/influence
    
    for n in range(3):
        parentConstraints.append( cmds.parentConstraint( fkJointList[n], jointList[n] ) )
        cmds.parentConstraint( ikJointList[n], jointList[n] ) #we dont need to append them again for ik, cos it is the same?
        
    #create switcher system
    
    ikFkSwitch = cmds.circle( name = 'IkFkSwitch')
    cmds.scale( jntScaleFactor, jntScaleFactor, jntScaleFactor, ikFkSwitch)
    cmds.xform( ikFkSwitch, translation = jointTransList[2])
    cmds.move( .5,.5,0, ikFkSwitch, relative = True)
    cmds.makeIdentity(ikFkSwitch, apply = True)
    cmds.pointConstraint( jointList[2] , ikFkSwitch, mo = True) #constrain the switch to the wrist joint
    cmds.addAttr( ikFkSwitch, shortName = 'IK_FK', longName = 'IK_FK', defaultValue=1, minValue=0, maxValue=1, keyable = True)
    
    for n in range(3):
        cmds.expression( object = parentConstraints[n][0], string = f"{fkJointList[n]}W0 = {ikFkSwitch[0]}.IK_FK" )
        cmds.expression( object = parentConstraints[n][0], string = f"{ikJointList[n]}W1 = 1 - {ikFkSwitch[0]}.IK_FK" )
    
    cmds.expression( object=fkSystemGroup, string = f"visibility = {ikFkSwitch[0]}.IK_FK")
    cmds.expression( object=ikSystemGroup, string = f"visibility = 1 - {ikFkSwitch[0]}.IK_FK")
    
    def changeColor ( pSubject, pColor ):    
        if (pColor == 'red'):
            setColor = [1, 0, 0]
        elif (pColor == 'blue'):
            setColor = [0, 0, 1]
        elif (pColor == 'green'):
            setColor = [0, 1, 0]
        elif (pColor == 'magenta'):
            setColor = [1, 0, 1]
        else:
            print('Pick a valid color (red, blue, green, magenta)')
            return
            
        cmds.setAttr(pSubject + '.overrideEnabled', 1)
        cmds.setAttr(pSubject + '.overrideRGBColors',1)
        cmds.setAttr(pSubject + '.overrideColorR', setColor[0])
        cmds.setAttr(pSubject + '.overrideColorG', setColor[1])
        cmds.setAttr(pSubject + '.overrideColorB', setColor[2])
        
    # change the controller colors
    for curves in controlCurveList:
        changeColor( curves, 'red')
        
    changeColor (ikFkSwitch[0], 'magenta')
    
    changeColor( ik1Ctrl[0], 'blue')
    changeColor( ik0Ctrl[0], 'blue')
    
    
    
    
    
#MAIN CODE
createUI("name", applyCallback)