import pymel.core as pm


def orientJoint(joint, aimObj, upObj, 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)
    if parent:
        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 orientJointTriangle(firstJoint, secondJoint, thirdJoint, aimAxis=pm.dt.Vector(1, 0, 0), upAxis=pm.dt.Vector(0, 1, 0)):
    orientJoint(firstJoint, secondJoint, thirdJoint, aimAxis, -upAxis)
    orientJoint(secondJoint, thirdJoint, firstJoint, aimAxis, -upAxis)
    if not thirdJoint.getChildren():
        thirdJoint.r.set(0, 0, 0)
        thirdJoint.jo.set(0, 0, 0)
        
def orientJointTriangleUI():
    
    win = pm.window(title="Orient Joint Triangle")
    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 = pm.dt.Vector(aimAxisField.getValue())
        upAxis = pm.dt.Vector(upAxisField.getValue())
        if len(pm.ls(sl=1, type="joint")) < 3:
            pm.displayError("Could not orient joint triangle. Must select three joints.")
            return
        firstJoint, secondJoint, thirdJoint = pm.ls(sl=1)[:3]
        orientJointTriangle(firstJoint, secondJoint, thirdJoint, aimAxis, upAxis)
        pm.deleteUI(win)
    
    button = pm.button(label="Orient Joint Triangle", command=command)
    win.show()
    
orientJointTriangleUI()