aovSettings

# Set the enabled AOVs to a particular list
# Made by Toby Sanderson and Mark Bailey 2019
# Updated May 2019 to include noice aovs

import maya.cmds as cmds
import mtoa.aovs as aovs
from mtoa.core import createOptions
createOptions()



def setupAOV():

    def setFullFloat(AOVfullFloat):
        # creates new arnoldDriver and connects tech AOVs to it. 
        fullFloatAD = cmds.createNode( 'aiAOVDriver', name='fullFloatArnoldDriver')
        cmds.setAttr("defaultArnoldDriver.halfPrecision", 1)
        #cmds.setAttr( fullFloatAD+".mergeAOVs", 1 )       
        for AOV in AOVfullFloat:
            try:
                cmds.connectAttr(fullFloatAD + ".message", "aiAOV_{}.outputs[0].driver".format(AOV), force=True)
            except:
                pass
        print "fullFloatArnoldDriver created for: "+ str(AOVfullFloat)

    def setHalfFloat(AOVhalfFloat):
        # creates new arnoldDriver and connects tech AOVs to it. 
        halfFloatAD = cmds.createNode( 'aiAOVDriver', name='halfFloatArnoldDriver')
        cmds.setAttr("defaultArnoldDriver.halfPrecision", 0)
        cmds.setAttr("halfFloatArnoldDriver.halfPrecision", 1)
        #cmds.setAttr( halfFlotAD+".mergeAOVs", 1 )       
        for AOV in AOVhalfFloat:
            try:
                cmds.connectAttr(halfFloatAD + ".message", "aiAOV_{}.outputs[0].driver".format(AOV), force=True)
            except:
                pass
        print "halfFloatArnoldDriver created for: "+ str(AOVhalfFloat)

    def setAOVrgba():
        print "Setting rgba output: {}".format(AOVrgba)
        for each in AOVrgba:
            try:
                setRGBA = "aiAOV_{}.type".format(each)
                cmds.setAttr(setRGBA, 6)
            except:
                pass
            
        
    
    def setPref():
        global AOVpref
        global AOVoff
        global AOVfullFloat
        global AOVhalfFloat
        global AOVrgba
        global noiceStatus
        
        # Define list of preferred AOVs with names *EXACTLY* as they appear in the Render Settings window...
        # Below, list for no-noice 
        # AOVpref = "AO", "coat", "crypto_material", "crypto_object", "diffuse_albedo", "diffuse_direct", "diffuse_indirect", "emission", "L_hdr", "L_key", "L_rim", "L_fill", "motionvector", "N", "P", "shadow_matte", "specular_direct", "specular_indirect", "sss", "transmission", "UV", "volume", "Z"
        AOVpref = "AO", "coat", "crypto_material", "crypto_object", "diffuse_direct", "diffuse_indirect", "emission", "L_hdr", "L_key", "L_rim", "L_fill", "motionvector", "shadow_matte", "specular_direct", "specular_indirect", "sss", "transmission", "UV", "volume"
        AOVprefNoNoice = "N", "P", "Z", "diffuse_albedo" # Not used in tool (yet?)

        AOVoff = "AO", "volume", "sss", "transmission", "emission"

        #AOVfullFloat = "Z", "P", "UV", "motionvector", "N", "volume", "crypto_object", "crypto_material"
        AOVfullFloat = "Z", "P", "UV", "motionvector", "N", "volume", "crypto_object", "crypto_material"
        AOVhalfFloat = "AO", "coat", "diffuse_direct", "diffuse_indirect", "emission", "L_hdr", "L_key", "L_rim", "L_fill", "shadow_matte", "specular_direct", "specular_indirect", "sss", "transmission"

        AOVrgba = "coat", "diffuse_albedo", "diffuse_direct", "diffuse_indirect", "emission", "specular_direct", "specular_indirect", "sss", "transmission"


    def allOff(): # Turn all the current AOVs off...
        allNiceNames = []
        deactivated = []
        allCurrentAOV = cmds.ls("*aiAOV_*")

        for AOV in allCurrentAOV:
            AOV = str(AOV)
            command = AOV + ".enabled"
            cmds.setAttr(command, 0)
            niceName = AOV[6:]
            allNiceNames.append(niceName)

            if niceName not in AOVpref: # Deactivate any active but unwanted AOVs
                deactivated.append(niceName)
                aovs.AOVInterface().removeAOV(niceName)

        if not allNiceNames:
            print "AOV-Settings: All AOVs were previously deactivated!\n",
        else:
            print "AOV-Settings: These are the AOVs you originally had: ", allNiceNames, "\n",
        if deactivated:
            print "AOV-Settings: These AOVs were deactivated because they aren't on the list: ", deactivated, "\n",



    def prefOn(): # Turn the preferred set of AOVs on...
        AOVerror = []
        for AOV in AOVpref:
            command = "aiAOV_" + AOV + ".enabled"
            try:
                cmds.setAttr(command, 1)
            except:
                try:
                    aovs.AOVInterface().addAOV(AOV)
                except:
                    AOVerror.append(AOV)
                                        
        for AOV in AOVoff:
            command = "aiAOV_" + AOV + ".enabled"
            try:
                cmds.setAttr(command, 0)
            except:
                try:
                    aovs.AOVInterface().addAOV(AOV)
                except:
                    AOVerror.append(AOV)
        # Enable Noice aovs:
        cmds.setAttr("defaultArnoldRenderOptions.outputVarianceAOVs", 1)

        """
        cmds.setAttr("aiAOV_AO.enabled", 0) # Ambient Occlusion AOV off by default
        cmds.setAttr("aiAOV_volume.enabled", 0) # Volume AOV off by default
        cmds.setAttr("aiAOV_sss.enabled", 0) # sss off by default
        cmds.setAttr("aiAOV_transmission.enabled", 0) # transmission off by default
        cmds.setAttr("aiAOV_emission.enabled", 0) # emission off by default
        """
        
        if AOVerror:
            from maya import OpenMaya
            print "AOV-Settings: These AOVs couldn't be switched on: ", AOVerror, "\n",
            OpenMaya.MGlobal.displayError("Something didn't work correctly, see the script editor for details...")

        else:
            print "AOV-Settings: The list of preferred AOVs has been turned on, and everything else switched off!\n",
            
    def createUV(): # Create and connect the nodes needed for the UV AOV...
        if cmds.objExists("green_ramp") or cmds.objExists("red_ramp"):
            try:
                cmds.defaultNavigation(source="green_ramp", destination="aiAOV_UV.defaultValue", ce=True)
                print "AOV-Settings-UV: Complete! If this hasn't worked, delete the STMap nodes from the Node Editor and try again.\n",
            except:
                print "AOV-Settings-UV: Looks like there's been an error. Try deleting the STMap nodes from the Node Editor and running again.\n",
    
        else:
            # Create Green Ramp
            cmds.shadingNode('ramp', name="green_ramp", asTexture=True)
            cmds.shadingNode('place2dTexture', name="place2dTexture_Green", asUtility=True)
            cmds.setAttr("green_ramp.colorEntryList[1].color", 0, 0, 0, type="double3")
            cmds.setAttr("green_ramp.colorEntryList[2].color", 0, 1, 0, type="double3")
            cmds.setAttr("green_ramp.colorEntryList[2].position", 1)
            
            # Create Red Ramp
            cmds.shadingNode('ramp', name="red_ramp", asTexture=True)
            cmds.shadingNode('place2dTexture', name="place2dTexture_Red", asUtility=True)
            cmds.setAttr("red_ramp.colorEntryList[1].color", 0, 0, 0, type="double3")
            cmds.setAttr("red_ramp.colorEntryList[2].color", 1, 0, 0, type="double3")
            cmds.setAttr("red_ramp.colorEntryList[2].position", 1)
            cmds.setAttr("place2dTexture_Red.rotateUV", 90)
            
            # Hook it up
            cmds.connectAttr("place2dTexture_Green.outUV", "green_ramp.uvCoord")
            cmds.connectAttr("place2dTexture_Green.outUvFilterSize", "green_ramp.uvFilterSize")
            cmds.connectAttr("place2dTexture_Red.outUV", "red_ramp.uvCoord")
            cmds.connectAttr("place2dTexture_Red.outUvFilterSize", "red_ramp.uvFilterSize")
            cmds.connectAttr("red_ramp.outColor", "green_ramp.colorOffset", force=True)
            
            # Connect it to the UV AOV
            cmds.defaultNavigation(source="green_ramp", destination="aiAOV_UV.defaultValue", ce=True)

            # Sets the AOVs to full / half float based on list at top.
            
            print "AOV-Settings-UV: STmap setup complete!\n",

        #setFullFloat(AOVfullFloat)
        setHalfFloat(AOVhalfFloat)

    # Run the script
    setPref()
    allOff()
    prefOn()
    createUV()

    

    # Hook up the Cryptomatte AOVs
    if not cmds.objExists("_aov_cryptomatte"):
       cmds.shadingNode('cryptomatte', name="_aov_cryptomatte", asUtility=True)
    if not cmds.objExists("aiAOV_crypto_object"):
       cmds.shadingNode('aiAOV', name="aiAOV_crypto_object", asTexture=True)
    if not cmds.objExists("aiAOV_crypto_material"):
       cmds.shadingNode('aiAOV', name="aiAOV_crypto_material", asTexture=True)
    try:   
        cmds.connectAttr('_aov_cryptomatte.outColor', 'aiAOV_crypto_material.defaultValue')
        cmds.connectAttr('_aov_cryptomatte.outColor', 'aiAOV_crypto_object.defaultValue')
    except:
        pass
    print "AOV-Settings-Cryptomatte: Setup complete!\n",

    # Hook up customPref things if Pref in list of AOVs - created for SHR. Change the path to the pref map when pref is required again:

    if "pref" in AOVpref:
        if not cmds.objExists("prefUtil"):
            cmds.shadingNode('aiUtility', name="prefUtil", asUtility=True)
        if not cmds.objExists("prefTex"):
            cmds.shadingNode('file', name="prefTex", asTexture=True)
        try:
            cmds.connectAttr("prefTex.outColor", "prefUtil.color")
        except:
            print "AOV-Settings-Pref: Failed to connect the prefTex and prefUtil nodes, they may have been connected already!\n",
        try:
            cmds.connectAttr("prefUtil.outColor", "aiAOV_custom_pref.defaultValue")
        except:
            print "AOV-Settings-Pref: Failed to connect the prefUtil and AOV nodes!\n",

        cmds.setAttr("prefUtil.shadeMode", 2)
        cmds.setAttr("prefTex.fileTextureName", "/mnt/projects/shr/assets/Character/creature_adult/RIG/publish/textures/pref/prefBody_prefSpaceShapeSpill_<UDIM>.exr", type="string")
        cmds.setAttr("prefTex.uvTilingMode", 3) # UDIM (Mari)
        print "AOV-Settings-Pref: Setup complete!\n",



    # Hook up the lightGroup things
    cmds.setAttr("aiAOV_L_key.lightPathExpression", "C.*<L.'key'>", type="string")
    cmds.setAttr("aiAOV_L_rim.lightPathExpression", "C.*<L.'rim'>", type="string")
    cmds.setAttr("aiAOV_L_hdr.lightPathExpression", "C.*<L.'hdr'>", type="string")
    cmds.setAttr("aiAOV_L_fill.lightPathExpression", "C.*<L.'fill'>", type="string")
    #cmds.setAttr("aiAOV_shadow_hdr", "C.*<L.'hdr'>", type="string")
    print "AOV-Settings-LightGroups: Setup complete!\n",

    # Hook up the Ambient Occlusion things
    if not cmds.objExists("ambientOcclusion_AOV"):
        cmds.shadingNode('aiAmbientOcclusion', name="ambientOcclusion_AOV", asTexture=True)
    cmds.setAttr('ambientOcclusion_AOV.samples', 5)
    cmds.setAttr('ambientOcclusion_AOV.farClip', 200)
    try:
        cmds.connectAttr('ambientOcclusion_AOV.outColor', 'aiAOV_AO.defaultValue')
    except:
        pass
    

    print "AOV-Settings-AO: Setup complete!\n",

    # Set some AOVs to rgba (see list at top)
    setAOVrgba()
    # Turn on noice denoiser AOVs 
    cmds.setAttr("defaultArnoldRenderOptions.outputVarianceAOVs", 1)


def aovSettings(): 
    
    allAOVs = []
    allAOVnames = []
    AOV_SettingsExists = []
    
    wordIn = "aiAOV_"


    def createNoiceStatusLocator():
        noiceStatus = cmds.getAttr( 'defaultArnoldRenderOptions.outputVarianceAOVs' )
        cmds.spaceLocator(n="AOV_Settings_noiceStatus")
        cmds.addAttr("AOV_Settings_noiceStatus", ln= 'noiceStatus', at="bool", defaultValue= noiceStatus) 
        return noiceStatus
       

    def getNoiceStatus():
        noiceStatus = None
        try:
            cmds.select("AOV_Settings_noiceStatus")
            noiceStatus = cmds.getAttr( 'AOV_Settings_noiceStatus.noiceStatus' )
            noiceLocator = True
        except:
            noiceLocator = False
            pass
            
        return noiceLocator, noiceStatus
      

    def restoreNoiceStatus():
        storedNoiceStatus = getNoiceStatus()
        if storedNoiceStatus[0] == True:
            storedNoice = cmds.getAttr( 'AOV_Settings_noiceStatus.noiceStatus' )
            cmds.setAttr( 'defaultArnoldRenderOptions.outputVarianceAOVs', storedNoice)
        else:
            print "No AOV_Settings_noiceStatus locator found."
        
    def loadSettings():
        # Check if AOV Settings are in scene already, and create them if not
        if cmds.objExists("AOV_Settings"):
            print "AOV-Settings: Existing AOV_Settings found in scene...\n",
            AOV_SettingsExists = True
        else:
            cmds.spaceLocator(n="AOV_Settings") # "Locator"
            print "AOV-Settings: No AOV_Settings found in the scene... created some for you.\n",
            AOV_SettingsExists = False

        
    
    def saveCurrentState():
        noiceLocator = getNoiceStatus()
        sceneNoiceStatus = cmds.getAttr( 'defaultArnoldRenderOptions.outputVarianceAOVs' )
        if noiceLocator[0] == True:
            cmds.setAttr( 'AOV_Settings_noiceStatus.noiceStatus', sceneNoiceStatus )
        else:
            createNoiceStatusLocator()

        cmds.select("AOV_Settings")

       

        if cmds.listAttr(ud = True):
            print "AOV-Settings: Updating the attributes...\n",
            # The Locator alrerady has some saved settings
            settingsExist()
        else:
            print "AOV-Settings: No custom attributes found. Creating them...\n",
            # The Locator doesn't have any saved settings
            settingsNew()
    
        print "AOV-Settings: These were the AOVs previously switched on: ", aovOn, "\n",
        print "AOV-Settings: These were the AOVs previously switched off: ", aovOff, "\n",
    
    def settingsNew():
        # The Locator doesn't have any saved settings
        allAOVnames = cmds.ls("*" + wordIn + "*")

        #cmds.addAttr("AOV_Settings", ln= noiceStatus, at="bool", defaultValue= noiceStatus)

        for x in allAOVnames:
            allAOVs.append(x)
            AOVenabled = cmds.getAttr( x + ".enabled")
    
            if AOVenabled == True:
               #aovOn.append(x)
               cmds.select("AOV_Settings")
               cmds.addAttr("AOV_Settings", ln= x, at="bool", defaultValue= 1)
    
            else:
               #aovOff.append(x)
               cmds.select("AOV_Settings")
               cmds.addAttr("AOV_Settings", ln= x, at="bool", defaultValue= 0)
    
    def settingsExist():
        # The Locator already has some saved settings
        allAOVnames = cmds.ls("*aiAOV_*")
        
        #cmds.setAttr("AOV_Settings."+"noiceStatus", noiceStatus )

        for aov in allAOVnames:
            if cmds.getAttr( aov + ".enabled") == True:
                cmds.setAttr("AOV_Settings."+aov, 1) 
                        
        for aov in allAOVnames:
            if cmds.getAttr( aov + ".enabled") == False:
                cmds.setAttr("AOV_Settings."+aov, 0)
                
    def getCurrentState(allAOVnames):
        # Gets the current state of the AOVs 
        for each in allAOVnames:
            if cmds.getAttr( each + ".enabled") == True:
                aovOn.append(each)  
            
        for each in allAOVnames:
            if cmds.getAttr( each + ".enabled") == False:
                aovOff.append(each)

        

        return allAOVnames

    
    #### START HERE ########################

    aovOn = []
    aovOff = []
    aovNamespaces = []
    aovNamespaces = cmds.ls("*aiAOV_*", r=True)

    aovNamespaces = [i for i in aovNamespaces if ":" in i]
    if aovNamespaces:
        print "AOV-Settings: Warning, these AOVs have namespaces - check they're not from a referenced scene: "+ str(aovNamespaces),

    choose = cmds.confirmDialog( title='Change AOV Settings', message='Select one of the following options...', button=['Setup Defaults','ALL OFF','ALL ON','Save State','Restore','Cancel'], defaultButton='Setup Defaults', cancelButton='Cancel', dismissString='No' )

    allAOVnames = cmds.ls("*aiAOV_*", r=True)

    #### SETUP DEFAULTS ###################
    if choose == 'Setup Defaults':
        setupAOV()
        cmds.setAttr("defaultArnoldRenderOptions.outputVarianceAOVs", 1)

        print "AOV-Settings: Default AOVs created!\n",

    #### ALL OFF ###########################
    elif choose == 'ALL OFF':
        loadSettings()
        getCurrentState(allAOVnames)
        saveCurrentState()        

        for y in aovOn:
            y = str(y)
            cmds.setAttr(y + ".enabled", 0)
            
        cmds.setAttr("defaultArnoldRenderOptions.outputVarianceAOVs", 0)

        print "AOV-Settings: All AOVs have now been turned OFF!\n",

    #### ALL ON ############################
    elif choose == 'ALL ON':
        loadSettings()
        getCurrentState(allAOVnames)
        saveCurrentState()


        for x in aovOff:
            x = str(x)
            cmds.setAttr(x + ".enabled", 1)

        cmds.setAttr("defaultArnoldRenderOptions.outputVarianceAOVs", 1)    
        print "AOV-Settings: All AOVs have now been turned ON!\n",

    #### RESTORE ###########################
    # Load the last saved settings from the Locator
    elif choose == 'Restore':

        loadSettings()
        getCurrentState(allAOVnames)
        
        if aovNamespaces:
            cmds.warning( "Warning, these aovs have namespaces :"+ str(aovNamespaces) )  
            cmds.confirmDialog( title='Restore failure', message='aiAOVs with namespaces are not restorable.\nPlease configure manually.', button=['OK'], defaultButton='OK', cancelButton='No', dismissString='No' )
            
        restoreOn = []
        restoreOff = []

        cmds.select("AOV_Settings")
        allSceneAOV = (cmds.listAttr(ud= True))                 # All user-defined attributes stored in the Locator

        if allSceneAOV == False:
            print "AOV-Settings: Error - no previous set-up found in the Locator\n",# If there are no user attributes on Locator

        else:
            for each in allSceneAOV:                          # For each user-defined Locator
               each = str(each)
               attrName = "AOV_Settings." + each
               fullName = each + ".enabled"                    # Convert unicode to string and add .enabled
               getState = cmds.getAttr(attrName)                # Get state (0 or 1)
               if getState == 1:                            # If the attribute has a value of 1 (it was previously on)
                  print "AOV-Settings: This was ON: ", each, "\n",
                  restoreOn.append(each)                     # Add to restoreOn list
                  cmds.setAttr(fullName, 1)                   # Turn them on
               else:                                      # If the attribute has a value of 0 (it was previously off)
                  print "AOV-Settings: This was OFF: ", each, "\n",
                  restoreOff.append(each)                    # Add to restoreOff list
                  cmds.setAttr(fullName, 0)                   # Turn them off

        restoreNoiceStatus()
        print "AOV-Settings: All AOVs have been restored to their last saved state!\n",


    #### SAVE ################################
    elif choose == 'Save State':
        loadSettings()
        getCurrentState(allAOVnames)
        saveCurrentState()
        

        print "AOV-Settings: Current AOV state saved! Use 'Restore' to revert any subsequent changes.\n",


    #### CANCEL ##############################
    else:
        print "AOV-Settings: Cancelled: Nothing was changed.\n",

# END #