# -*- coding: utf-8 -*-
"""
Created on Thu Jan 18 13:43:30 2024

@author: Jared Athias
https://www.jaredathias.com/

PYTHON for Maya
built in maya 2023
----------------------------------------------------------------------------
Light_Layer_Code v1

This script will automatically create the following light layers for all selected lights
works for Arnold Renderer

Beauty
Diffuse
Specular
SSS
Direct Diffuse
Indirect Diffuse
Direct Specular
Indirect Specular
Coat
Coat Indirect
Coat Direct

Paste code in script editor, select light, and run
-----------------------------------------------------------------------------
-----------------------------------------------------------------------------
-----------------------------------------------------------------------------
-----------------------------------------------------------------------------
"""

import os
import maya.cmds as cmds
import mtoa.aovs as aovs
import maya.mel as mel

#Collect project path
filepath = cmds.file(q=True, sn=True)
path = filepath.split('/')
path.pop(-1)
path.pop(-1)
MAIN_PATH = ('/').join(path)
IMG_PATH = MAIN_PATH + "/scripts/light_layer_script/"

SELECTED_LIGHT = cmds.ls(sl=True)
LIGHT_LIST = cmds.ls(type=cmds.listNodeTypes("light"))
LAYERS = ['beauty', 'diffuse', 'specular', 'sss','diffuseDirect','diffuseIndirect','specularDirect','specularIndirect','coat','coatIndirect','coatDirect']
LIGHT_GROUPS = []
SELECTIONS = []

# Adds all selected lights into a list
def all_lights_slected():
    if SELECTED_LIGHT:
        SELECTIONS.append(SELECTED_LIGHT)

# checks to make sure lights are seelceted and assign name to light aov group
def light_selection_filter(LIGHT): 
    LIGHT_SHAPE = LIGHT + 'Shape'
    if LIGHT_SHAPE in LIGHT_LIST:
        LIGHT = LIGHT_SHAPE
#       MEL COMAND BELOW
        mel.eval("setAttr -type \"string\"" + str(LIGHT) + '.aiAov ' + str(LIGHT) + ";")        

    else:
        print('Non-light selected or please rename light in selection')
    return LIGHT 
    
    
# Returns specific layer  EXECUTES ON SPECIFIC LAYER BUTTONS  
def layer_beauty_filter():
    for aov in LAYERS:
        if aov == str(label):
            LAYER = aov
            print(LAYER)
    return LAYER

# Returns all layers
def all_layer_filter(LAYER):
    LAYER = aov
    return LAYER

# Create Node system EXECUTES WHEN CODE IS LAUNCHED
def create_nodes():
    AOVNODE = 'aiAOV_RGBA_' + str(LIGHT) + '_' + str(LAYER)
    AOVNAME = 'RGBA_' + str(LIGHT) + '_' + str(LAYER)
    if cmds.objExists(AOVNODE) or cmds.objExists(AOVNAME):
        print('AOV Already exists')
    else:
        aovs.AOVInterface().addAOV(str(AOVNAME), aovType='rgba')
        cmds.setAttr(AOVNODE + '.lightPathExpression', LAYER_DICT[LAYER], type='string')
        
# Confirmation function post execution
def create_confirm():
    result = cmds.confirmDialog(title="Light Layer Tool Results",
                                message="Layers Created",
                                button=["Okay"],
                                defaultButton="okay")
    print(result)

# CODE TO CREATE ALL LAYERS
all_lights_slected()      
for item in SELECTED_LIGHT:
    LIGHT = light_selection_filter(item)
    print(LIGHT)
    

LAYER_DICT = {'beauty':"C.*<L.'" + str(LIGHT) + "'>",
              'diffuse':"C<RD>.*<L.'" + str(LIGHT) + "'>",
              'specular': "C<RS>.*<L.'" + str(LIGHT) + "'>",
              'sss':"C<TD>.*<L.'" + str(LIGHT) + "'>",
              'diffuseDirect':"C<RD><L.'" + str(LIGHT) + "'>",
              'diffuseIndirect':"C<RD>[DSVOB].*<L.'" + str(LIGHT) + "'>",
              'specularDirect':"C<RS[^'coat']><L.'" + str(LIGHT) + "'>",
              'specularIndirect':"C<RS[^'coat']>[DSVOB].*<L.'" + str(LIGHT) + "'>",
              'coat':"C<RS'coat'>.*<L.'" + str(LIGHT) + "'>",
              'coatDirect':"C<RS'coat'><L.'" + str(LIGHT) + "'>",
              'coatIndirect':"C<RS'coat'>[DSVOB].*<L.'" + str(LIGHT) + "'>"}    
    
    
    
for aov in LAYERS:
    LAYER = all_layer_filter()
    print(LAYER)
    
    if LIGHT == 'None' or LAYER == 'None':
        print('Please select light or layers')

    else:    
        create_nodes()