from maya import cmds
from maya import mel
from functools import partial


#Create a list of referenced top nodes (identified with idNode)
topNamespaces = cmds.ls (type="idNode")

#List declaration for characters and props
charNamespaces=list()
propNamespaces=list()

#Fixed value for button spacing
buttPadding = 28

#Creates a list of objects that are either "Creature [2]" or "Prop [6], grabs namespace"
for idName in topNamespaces:
    
    idNodeNamespace = cmds.getAttr ("%s.class" % (idName))
    if idNodeNamespace == 2:
        charNamespaces.append (idName.split(":")[0])
    if idNodeNamespace == 6:
        propNamespaces.append (idName.split(":")[0])
        
#Removes setupWrap if it exists (it's considered a prop and I don't want it on the list
if cmds.objExists ("setupWrap:setupWrap"):        
    propNamespaces.remove ('setupWrap')

#Function that runs when character ID button is pressed
def cogSelect (buttonPress,*args):
    cmds.select (buttonPress+":C_cog_CTL")
    
#Function for selecting the local (prop only)
def localSelect (buttonPress,*args):
    cmds.select (buttonPress+":C_local_CTL")
    
def delSelCon (*args):
#Function for delete all selected constraints
    items = cmds.ls(sl=1)
    for item in items:
        nodes = cmds.listRelatives(item, type='constraint')
        for node in nodes:
            if not ':' in node:
                cmds.delete(node)
    
def addToIso (buttonPress, *args):
# Resolves "Iso+" button press
    selPanel = cmds.getPanel (wf=True)
    if cmds.isolateSelect (selPanel, q=True, state=True) == False:
        isolated = cmds.isolateSelect(selPanel, state=True, query=True )
        mel.eval('enableIsolateSelect %s %d' % (selPanel,not isolated) )
        cmds.isolateSelect (selPanel, state=True)
    cmds.isolateSelect (selPanel, ado=(buttonPress+":"+buttonPress))
    if cmds.getAttr("%s.class" % (buttonPress+":"+buttonPress))  == 2:
        cmds.isolateSelect (selPanel, ado=(buttonPress+"_proxyRig:proxyRig"))
    
def removeFromIso (buttonPress, *args):
# Resolves "Iso-" button press
    selPanel = cmds.getPanel (wf=True)
    cmds.isolateSelect (selPanel, rdo=(buttonPress+":"+buttonPress))
    if cmds.getAttr("%s.class" % (buttonPress+":"+buttonPress))== 2:
        cmds.isolateSelect (selPanel, rdo=(buttonPress+"_proxyRig:proxyRig"))
    
def addIsoSet (currSet, *args):
#Adds extra button to UI if you have an atools set
    selPanel = cmds.getPanel (wf=True)
    cmds.select (currSet)
    cmds.isolateSelect (selPanel, addSelected=True)
    
def proxySwitch (buttonPress, *args):
    #Toggles proxy rig
    rigVis = cmds.getAttr (buttonPress + ":%s.visibility" % buttonPress)
    if rigVis == True:
        cmds.setAttr (buttonPress + ":%s.visibility" % buttonPress, False)
        cmds.setAttr (buttonPress + "_proxyRig:proxyRig.visibility", True)
    else:
        cmds.setAttr (buttonPress + ":%s.visibility" % buttonPress, True)
        cmds.setAttr (buttonPress + "_proxyRig:proxyRig.visibility", False)    
        
        
        
def globalSnap (buttonPress, *args):
#Performs function that snaps global control to local
    buttonPress = ":".join(cmds.ls(sl=True)[0].split(":")[:-1])
    locTrack = cmds.spaceLocator (name = buttonPress + "_snapPoint")[0]
    cmds.parentConstraint (buttonPress + ":C_local_CTL", locTrack, mo=False)
    cmds.cutKey (buttonPress + ":C_local_CTL", buttonPress + ":C_global_CTL", cl=True, at=["translateX", "translateY", "translateZ", "rotateX", "rotateY", "rotateZ"])
    cmds.delete(locTrack, cn=True )
    cmds.parentConstraint (locTrack,buttonPress + ":C_global_CTL", mo=False)
    cmds.parentConstraint (locTrack,buttonPress + ":C_local_CTL", mo=False)
    cmds.delete(locTrack)


def genNoseLoc (buttonPress, *args):
    buttonPress = ":".join(cmds.ls(sl=True)[0].split(":")[:-1])    
    
    #Checks if nose locator group already exists
    if cmds.objExists ("NoseLocatorGroup") == False:
        cmds.group (em=True, name="NoseLocatorGroup")
    #Checks to make sure you have a selection    
    if (len(cmds.ls(sl=True))) > 0:
        if cmds.objExists ("NoseLocatorGroup|%s_NoseTrackParent" % buttonPress) == False:
        #Checks to see if the object has a head (so it knows where to place locator
            try:
                cmds.select (buttonPress + ":C_head_CTL")
            except:
                cmds.error ("Selected object has no nose")
            else:
                if cmds.ls(sl=True):
                    locTrack = cmds.spaceLocator (name = buttonPress + "_Headtrack")[0]
                    locParent = cmds.spaceLocator (name = buttonPress + "_NoseTrackParent")[0]
                    cmds.parent (locTrack, locParent)
                    cmds.parentConstraint (buttonPress + ":C_head_CTL", locParent, mo=False)
                    cmds.setAttr (locTrack+'.translateY', 2.4)
                    cmds.setAttr (locTrack+'.translateZ', 12.4)
                    cmds.parent (locParent, "NoseLocatorGroup")
                    cmds.select (locTrack)
        else:
            cmds.error ("Character already has nose locator")
def massParent (*args):
    
    selObjects = cmds.ls (sl=1)
    for item in selObjects[1:]:
        cmds.parentConstraint (selObjects[0], item, mo=True)

def buttFace():
    #Function for building UI
    
    #Default size for button to check against max size
    currentMaxButtonSize = 0
    
    windowID = "nsUI"
    if cmds.window ("nsUI", exists=True):
        cmds.deleteUI ("nsUI")
    window = cmds.window("nsUI", width=25, title="Easy Asset Manager", resizeToFitChildren=1)
    cmds.rowColumnLayout("rcLayout", numberOfColumns=4, rs=[1,5])
    for chars in charNamespaces:
        #Generates button for each character
        cmds.button (bgc=(0,0,0), label=chars, align="left", command=partial (localSelect, chars), parent="rcLayout")
        cmds.button (bgc=(0,0,0.3), label="Proxy", align="left", command=partial (proxySwitch, chars), parent="rcLayout")
        cmds.button (bgc=(0.3,0,0), label="  Iso-  ", align="left", command=partial (removeFromIso, chars), parent="rcLayout")
        cmds.button (bgc=(0,0.3,0), label="  Iso+  ", align="left", command=partial (addToIso, chars), parent="rcLayout")
    for props in propNamespaces:
        #Generates button for each prop
        cmds.button (bgc=(0.9,0.9,0.9), label=props, align="left", command=partial (localSelect, props), parent="rcLayout")
        cmds.button (bgc=(1,0.9,0.6), label="Local", align="left", command=partial (localSelect, props), parent="rcLayout")
        cmds.button (bgc=(0.3,0,0), label="  Iso-  ", align="left", command=partial (removeFromIso, props), parent="rcLayout")
        cmds.button (bgc=(0,0.3,0), label="  Iso+  ", align="left", command=partial (addToIso, props), parent="rcLayout")
    totalWinHeight = 10 + (len(charNamespaces)*buttPadding) + (len(propNamespaces)*buttPadding)
   

    cmds.rowColumnLayout("extrasLayout", numberOfColumns=1, rs=[1,5])
    
    allSets = cmds.ls (type="objectSet")
    allAtoolsSets = list()
    
    for item in allSets:
        if "aToolsSet_" in item and "_Set" in item:
            allAtoolsSets.append (item)
    print (allAtoolsSets)
    
    if len(allAtoolsSets) > 0:
        print ("ifPasses")
        #Creates a set isolation button if selection set exists
        cmds.button (bgc=(0.4,0.1,0.4), label="Add Set", align="left", parent="extrasLayout", command=partial (addIsoSet, allAtoolsSets[0]))
    #Button that deletes constraints
    cmds.button ("delCon", bgc=(0.4,0.1,0.1), label="Delete Constraints", align="left", parent="extrasLayout", command=partial (delSelCon))
    cmds.button (bgc=(0.9,0.2,0.2), label="Global Snap", align="left", parent="extrasLayout", command=partial (globalSnap))
    cmds.button (bgc=(0.1,0.8,0.3), label="Nose Locator", align="left", parent="extrasLayout", command=partial (genNoseLoc))
    cmds.button (bgc=(0.8,0.3,0.6), label="Mass Parent", align="left", parent="extrasLayout", command=partial (massParent))

    
    cmds.showWindow (window)
    howManyButtons = cmds.rowColumnLayout ("extrasLayout", q=True, nch=True)
    totalWinHeight += (howManyButtons*buttPadding) - (howManyButtons/2)
    checkForButtonName = cmds.rowColumnLayout ("rcLayout", q=True, ca=True)
    print (checkForButtonName)
    for item in checkForButtonName[:-1]:
        checkButtonSize = cmds.button (item, q=True, width=True)
        print (cmds.button (item, q=True, width=True))
        if checkButtonSize > currentMaxButtonSize:
            currentMaxButtonSize = checkButtonSize
    totalWinWidth = currentMaxButtonSize + 133             # 128 is the size of the proxy and iso buttons, 5 is extra padding
    cmds.window ("nsUI", e=True, s=False, widthHeight=(totalWinWidth, totalWinHeight))


buttFace()