'''
    File name: jointOrient.py
    Author: Bastian Kramer
    Date created: 21/10/2020
    Date last modified: 21/10/2020
    Python Version: 2.7

    This is a simple script containing a function to orient joints based on an aim and an up target object,
    storing the evaluated rotation in the joint's jointOrient attribute.

    An additional UI can get used to adjust values for aim and up axis used to evaluate the orientation.
'''

import pymel.core as pm


def orientJoint(aimObj, upObj, joint, aimAxis=pm.dt.Vector(1, 0, 0), upAxis=pm.dt.Vector(0, 1, 0)):
    sel = pm.ls(sl=1)
    parent = joint.getParent()
    children = joint.getChildren()
    
    if children:
        pm.parent(children, world=True)
    pm.parent(joint, world=True)
    joint.jointOrient.set(0, 0, 0)
    
    con = pm.aimConstraint(aimObj, joint, mo=False,
        aimVector=aimAxis, upVector=upAxis, worldUpType="object",
        worldUpObject=upObj)
    rotate = joint.rotate.get()

    pm.delete(con)
    joint.rotate.set(0, 0, 0)
    joint.jointOrient.set(rotate)
    
    if parent:
        pm.parent(joint, parent)
    if children:
        pm.parent(children, joint)
    pm.select(sel)
    
    
def orientJointUI():
    
    win = pm.window(title="Orient Joint")
    lay = pm.columnLayout(adjustableColumn=True)
    aimAxisField = pm.floatFieldGrp(numberOfFields=3, label="Aim Axis", value1=1, value2=0, value3=0)
    upAxisField = pm.floatFieldGrp(numberOfFields=3, label="Up Axis", value1=0, value2=1, value3=0)
    
    def command(*args):
        aimAxis = aimAxisField.getValue()
        upAxis = upAxisField.getValue()
        if len(pm.ls(sl=1, type="transform")) < 3:
            pm.displayError("Could not orient joint. Must select an aim transform, an up transform and a joint to orient.")
            return
        if (pm.ls(sl=1, type="transform")[2].nodeType() != "joint"):
            pm.displayError("Could not orient joint. Third selected transform must be a joint.")
            return
        aimObj, upObj, joint = pm.ls(sl=1)
        orientJoint(aimObj, upObj, joint, aimAxis, upAxis)
    
    button = pm.button(label="Orient Joint", command=command)
    win.show()
    
orientJointUI()