import maya.cmds as cmds
from collections import namedtuple

PivotCtrl = namedtuple("PivotCtrl", "jnt ctrl sdk offset")

def getInheritedRotation(obj):

    rotation = [0,0,0]

    p = cmds.listRelatives(obj, allParents=True)

    while p != None:
        if cmds.nodeType(p) != "joint":
            rotation = [x + y for x,y in zip(rotation, cmds.xform(p, q=True, rotation=True))]
        else:
            rotation = [x + y for x,y in zip(rotation, cmds.joint(p, q=True, orientation=True))]
        p = cmds.listRelatives(p, allParents=True)

    return rotation

def getChildLevel(obj):
    p = cmds.listRelatives(obj, allParents=True, type="joint")
    level = 1
    while p != None:
        level += 1
        p = cmds.listRelatives(p, allParents=True, type="joint")

    return level

def generateColour(level, lenLevels, red=True, green=True, blue=False, redSub=False, greenSub=True, blueSub=False):
    div = 1.0/lenLevels
    col = [0,0,0]


    if red and not redSub:
        col[0] = div*level
    elif red:
        col[0] = 1-div*level

    if green and not greenSub:
        col[1] = div*level
    elif green:
        col[1] = 1-div*level

    if blue and not blueSub:
        col[2] = div*level
    if blue:
        col[2] = 1-div*level

    return col


def generateColourList(lenLevels, red=True, green=True, blue=False, redSub=False, greenSub=True, blueSub=False):
    colours = []

    for i in range(1,lenLevels+1):
        colours.append(generateColour(i, lenLevels, red, green, blue))

    return colours



if __name__ == "__main__":
    jnt = [cmds.ls(selection=True, type="joint")[0]]

    jntHier = cmds.listRelatives(allDescendents=True, type="joint")[1:] + jnt

    colours = generateColourList(len(jntHier))

    allJntCtrls = []


    for i,jnt in enumerate(jntHier):
        jntPos = cmds.joint(jnt, q=True, p=True)
        jntOri = cmds.joint(jnt, q=True, o=True)



        finalJntOri = [x + y for x,y in zip(jntOri, getInheritedRotation(jnt))]

        ctrl = cmds.circle(name="{}_ctrl".format(jnt))

        cmds.xform("{}.cv[0:]".format(ctrl[0]), rotation=(0,90,0))

        ctrl[1] = cmds.listRelatives(ctrl[0], shapes=True, type="nurbsCurve")[0]

        cmds.setAttr("{}.overrideEnabled".format(ctrl[1]), 1)
        cmds.setAttr("{}.overrideRGBColors".format(ctrl[1]), 1)
        cmds.setAttr("{}.overrideColorRGB".format(ctrl[1]), *colours[i], type="double3")

        sdk = cmds.group(ctrl, name="{}_sdk".format(ctrl[0]))
        off = cmds.group(sdk, name="{}_0".format(ctrl[0]))

        cmds.xform(off, translation=jntPos, rotation=finalJntOri)

        allJntCtrls.append(PivotCtrl(jnt=jnt, ctrl=ctrl, sdk=sdk, offset=off))



    for i in range(len(allJntCtrls)):
        if i < len(allJntCtrls)-1:
            cmds.parent(allJntCtrls[i].offset, allJntCtrls[i+1].ctrl[0])

        # Just setting to orient constrain atm, need to change for other options
        cmds.orientConstraint(allJntCtrls[i].ctrl[0], allJntCtrls[i].jnt, name="{}_oc".format(allJntCtrls[i].jnt))