'''
*   This script was created by
*       Trevor Sommer
*           trevor@trevorsommer.com
*
*   Module containing individual appendages and other modular setups that can be used in conjunction with each other to make a full rigged setup
'''

import maya.cmds as cmds
from . import rigUtil as rUtil
from . import rigLib as rLib

def arm_basicSetupProc(baseJoint, elbowJoint, wristJoint, palmJoint, charPfx, sidePfx = "", distroUp= 1, distroDown= 1):
    '''
    Creates a simple arm setup, features:
        ik/fk and switch and controls
        scaling ik with blend-able on/off
        scalable fk controls
        IK pole vector controller
        will create upper and lower arm distribution joints

        To Do:
        if finger joints are detected will attempt to set up FK finger controls and hand offset controllers
        make naming convention more robust so you could have infinite number of arms per side without issue 

    Since this proc should only be called from other functions or from a GUI the error checking should be done there before this proc is called
    Otherwise I would end up running the same checks twice in a row, so there will is minimal error/warning handling
    '''
    if not cmds.objExists(baseJoint) or not cmds.objExists(elbowJoint) or not cmds.objExists(wristJoint) or not cmds.objExists(palmJoint):
        cmds.warning('One or more of the objects passed into arm_basicSetupProc does not exists, canceling arm setup operation')
        return False

    #get charPfx if any exists, running this check again here because its vital that this can be determined and used
    charPfx = rUtil.get_charPfx([baseJoint, elbowJoint, wristJoint, palmJoint])

    #if no charPfx is found skip setup
    if not charPfx:
        cmds.warning('No charPfx could be obtained from input, unable to designate character, canceling arm setup operation')
        return False     

    #get the side of the character     
    if not sidePfx:
        sidePfx = palmJoint.split("_")[1] 

    armPfx = charPfx + "_" + sidePfx

    '''
    Should disconnect any joints below the palm and re-parent them after duplicating the chains 
    '''

    #duplicate the joints henceforth refered to as boundJoints, into 3 joint chains ik,fk,drv
    fkJoints = rLib.duplicateJointHierarchy(baseJoint, "fk", charPfx)
    ikJoints = rLib.duplicateJointHierarchy(baseJoint, "ik", charPfx)
    drvJoints = rLib.duplicateJointHierarchy(baseJoint, "drv", charPfx)

    #create hand Controller this holds ik/fk switch variables and is always visible
    handController = cmds.circle(center= [0,0,0], normal= [0,0,1], degree= 3, sections= 8, tolerance= 0.01, constructionHistory= False, useTolerance= 0, sweep= 360, radius=1, name= (armPfx + "_Hand_CON"))[0]
    handControllerOffset = cmds.group(handController, name= (handController + "_OFFSET"))
    cmds.parentConstraint(drvJoints[3], handControllerOffset, weight= 1, name= (handController + "_OFFSET_CONST"))

    #---------rig up the ik Chain---------
    
    #pole constraint setup
    '''
    adding stuff for future space switching and ik/fk matching of the pole vector into the setup
    '''
    poleLength = vectDist(cmds.pointPosition((ikJoints[0] + '.rotatePivot'), world= True), cmds.pointPosition((ikJoints[1] + '.rotatePivot'), world= True))
    poleMidPoint = cmds.spaceLocator(position= [0,0,0], name= (armPfx + "_arm_MidPoint_LOC"))
    poleConstraintLoc = cmds.spaceLocator(position= [0,0,0], name= (armPfx + "_arm_Constraint_LOC"))
    cmds.pointConstraint(drvJoints[0], drvJoints[2], poleMidPoint, offset= [0,0,0], name= (poleMidPoint + "_arm_pointCONST"), weight= 1)
    polePosition = rUtil.polePosCalc(poleMidPoint, ikJoints[1], poleLength)

    cmds.addAttr(poleConstraintLoc, keyable= False, longName= "magnitude", attributeType= "float")
    cmds.setAttr((poleConstraintLoc + ".Magnitude"), poleLength)

    poleController = rUtil.sphereControllerObj((armPfx + "Elbow_IKPole_CON"))
    cmds.parent(poleConstraintLoc, poleController)

    poleConOffset = cmds.group(empty= True, name= (poleController + "_Main_OFFSET"))
    cmds.setAttr((poleConOffset + ".tx"), polePosition[0])
    cmds.setAttr((poleConOffset + ".ty"), polePosition[1])
    cmds.setAttr((poleConOffset + ".tz"), polePosition[2])
    cmds.makeIdentity((poleController + "_Main_OFFSET"), apply= true, translate= 1, rotate= 1, scale= 1, normal= 0)

    rUtil.attrLocker(poleController, ["rx","ry","rz","sx","sy","sz","v"])
    cmds.setAttr((poleController + ".v"), keyable= False)
    rUtil.attrLocker(poleMidPoint, ["tx","ty","tz","rx","ry","rz","sx","sy","sz","v"])
    rUtil.attrLocker(poleConstraintLoc, ["tx","ty","tz","rx","ry","rz","sx","sy","sz","v"])

    #create wrist_CON
    handLength = vectDist(cmds.pointPosition((ikJoints[2] + '.rotatePivot'), world= True), cmds.pointPosition((ikJoints[3] + '.rotatePivot'), world= True))
    armIkController = rUtil.boxControllerObj((armPfx + "_arm_IK_CON"))
    cmds.setAttr((armIkController + ".sx"), handLength)
    cmds.setAttr((armIkController + ".sy"), handLength)
    cmds.setAttr((armIkController + ".sz"), handLength)

    cmds.group(armIkController, n= (armIkController + "_SPACE_OFFSET"))
    cmds.group((armIkController + "_SPACE_OFFSET"), n= (armIkController + "_Main_OFFSET"))
    cmds.delete(cmds.parentConstraint(ikJoints[2], (armIkController + "_Main_OFFSET"), weight= 1, maintainOffset= False))
    cmds.makeIdentity(armIkController, apply= true, translate= 1, rotate= 1, scale= 1, normal= 0)

    #create ik solvers
    armIKHandle = cmds.ikHandle(startJoint= ikJoints[0], endEffector= ikJoints[2], solver= "ikRPsolver", priority= 1, name= (armPfx + "_arm_HANDLE_IK"))
    cmds.rename(cmds.ikHandle(armIKHandle, query= True, endEffector= True), (armPfx + "_arm_HANDLE_EEF"))    
    handIKHandle = cmds.ikHandle(startJoint= ikJoints[2], endEffector= ikJoints[3], solver= "ikSCsolver", priority= 1, name= (armPfx + "_hand_HANDLE_IK"))
    cmds.rename(cmds.ikHandle(handIKHandle, query= True, endEffector= True), (armPfx + "_hand_HANDLE_EEF")) 
    cmds.parent(armIKHandle, handIKHandle, armIkController)

    #add pole vector constraint
    cmds.poleVectorConstraint(poleConstraintLoc, armIKHandle, weight= 1, name= (armIKHandle + "_pole_CON"))

    #lock the attributes that are uneeded
    rUtil.attrLocker(armIKHandle, ["tx","ty","tz","rx","ry","rz","sx","sy","sz","v"])
    rUtil.attrLocker(handIKHandle, ["tx","ty","tz","rx","ry","rz","sx","sy","sz","v"])
    rUtil.attrLocker(armIkController, ["sx","sy","sz"])
    cmds.setAttr((armIkController+ ".v"), keyable= False)

    #create cosmetic connections to keep track of associated nodes for space switching and pole vector matching ik/fk switching without having to rely on naming

    cmds.addAttr(armIkController, attributeType= 'bool', keyable= True, hidden= False, minvalue=0, maxValue= 1, defaultValue= 0, longName= "space_Switching") 
    cmds.setAttr((armIkController + ".space_Switching"), 1, lock= True)
    rLib.rig_SpaceCreator(armIkController, "", "init", poleController)

    #creates dummy attributes to keep track of objects for elbow matching 
    cmds.addAttr(poleController, keyable= False, longName= "midJoint", dataType= "string") 
    cmds.connectAttr((drvJoints[1] + ".message"), (poleController + ".midJoint"), force= True) 
    cmds.addAttr(poleController, keyable= False, longName= "midPoint", dataType= "string")  
    cmds.connectAttr((poleMidPoint + ".message"), (poleController + ".midPoint"), force= True) 

    #setup simple ik stretch
    rLib.simpleStretchIKSetup(armIKHandle, handController)


    #--------- create module offsets on the driver for parenting control rig, and on bound for stable points to run corrective shape drivers --------  
    rigOffsetJoint = cmds.duplicate(drvJoints[0], parentOnly= True, name= (drvJoints[0] + "_BASE_OFFSET"))
    cmds.parent(drvJoints[0], fkJoints[0], ikJoints[0], rigOffsetJoint)

    boundOffsetJoint = cmds.duplicate(baseJoint, parentOnly= True, name= (baseJoint + "_BASE_OFFSET"))
    cmds.parent(baseJoint, boundOffsetJoint)
    
    #--------- create and connect the distribution joints apply directly to the bound rig -------- 
    if distroUp:
        rLib.distroSubRig(baseJoint, elbowJoint, distroUp, armPfx, "up", "arm")
    if distroDown:
        rLib.distroSubRig(elbowJoint, wristJoint, distroDown, charPfx, "down", "arm")

    #------- connect FK Controls  ----------------------------- 

    fkShoulderCon = fkControllerScalable(fkJoints[0])
    fkElbowCon = fkControllerScalable(fkJoints[1])
    fkWristCon = fkControllerScalable(fkJoints[2])

    rUtil.attrLocker(armIKHandle, ["rx","rz"])

    #--------- connect IK/FK Switch -----------------------------     

    rLib.createRotationSwitch(drvJoints[0], fkJoints[0], ikJoints[0], handController)
    rLib.createRotationSwitch(drvJoints[1], fkJoints[1], ikJoints[1], handController)
    rLib.createRotationSwitch(drvJoints[2], fkJoints[2], ikJoints[2], handController)
    rLib.createRotationSwitch(drvJoints[3], fkJoints[3], ikJoints[3], handController)

    #--- connect scaleX from joints

    rLib.createScaleSwitch(drvJoints[0], fkJoints[0], ikJoints[0], handController)
    rLib.createScaleSwitch(drvJoints[1], fkJoints[1], ikJoints[1], handController)
    rLib.createScaleSwitch(drvJoints[2], fkJoints[2], ikJoints[2], handController)

    #---------- hierarchy and other misc stuff to clean up and organize things -----------------------------     

    armConGRP = cmds.group(name= (armPfx + "Arm_CONTROLLER_GRP"), empty= True)
    ikConGRP = cmds.group(name= (armPfx + "Arm_IKCON_GRP"), empty= True)
    fkConGRP = cmds.group(name= (armPfx + "Arm_FKCON_GRP"), empty= True)
    miscConGRP = cmds.group(name= (armPfx + "Arm_MiscRig_GRP"), empty= True)

    cmds.parent((fkShoulderCon + "_FK_OFFSET_GRP"), (fkElbowCon + "_FK_OFFSET_GRP"), (fkWristCon + "_FK_OFFSET_GRP"), fkConGRP)
    cmds.parent((armIkController + "_SPACE_OFFSET"), (poleController + "_SPACE_OFFSET"), ikConGRP)
    cmds.parent(poleMidPoint, miscConGRP)
    cmds.parent(fkConGRP, ikConGRP, miscConGRP, (armPfx + "Arm_CONTROLLER_GRP"))
    rUtil.attrLocker(miscConGRP, ["v"])

    #connect visibilty to switches
    cmds.connectAttr((handController + ".ik_TO_fk"), (fkConGRP + ".visibility"), force= True)

    switchCondNode = cmds.shadingNode("condition", asUtility= True, name= (handController + "_Switch_Vis_COND"))
    cmds.setAttr((switchCondNode + ".colorIfTrueR"), 1)
    cmds.setAttr((switchCondNode + ".colorIfFalseR"), 0)
    cmds.setAttr((switchCondNode + ".secondTerm"), 0)
    cmds.setAttr((switchCondNode + ".operation"), 0)

    cmds.connectAttr((handController + ".ik_TO_fk"), (switchCondNode + ".firstTerm"))
    cmds.connectAttr((switchCondNode + ".outColorR"), (ikConGRP + ".visibility"))


    print "arm_basicSetupProc has completed, check you shell for function feedback and warnings"