# -*- coding: utf-8 -*-
import pymel.core as pm
import maya.cmds as cmds
from mtoa.ui.ae.shaderTemplate import ShaderAETemplate
import arnold as ai

def aiUtilityCreateColorMode(attr):
    cmds.setUITemplate('attributeEditorPresetsTemplate', pushTemplate=True)
    cmds.attrEnumOptionMenuGrp('AIUtilityColorMode', attribute=attr, label="色彩模式", 
                               enumeratedItem=[(0, 'Color'), (3, 'Normal'), (1, 'Geometric Normal'), (2, 'Un-bumped Normal'), (22, 'Bump Difference'),
                                                (4, 'Barycentric Coords'), (5, 'UV Coords'), (6, 'U Coords'), (7, 'V Coords'),
                                                (8, 'U Surface Derivative (dPdu)'), (9, 'V Surface Derivative (dPdv)'),
                                                (10, 'Shading Point (Relative to BBox)'), (11, 'Primitive ID'), (12, 'Uniform ID'),
                                                (13, 'Triangle Wireframe'), (14, 'Polygon Wireframe'), (15, 'Object'),
                                                (16, 'Edge Length'), (17, 'Floatgrid'), (18, 'Reflection Lines'),
                                                (19, 'Bad UVs'), (20, 'Number of Lights'), (21, 'Object ID')])    
    cmds.setUITemplate(popTemplate=True)

def aiUtilitySetColorMode(attr):
    cmds.attrEnumOptionMenuGrp('AIUtilityColorMode', edit=True, attribute=attr)

class AEaiUtilityTemplate(ShaderAETemplate):
    def checkShadeMode(self, nodeName):
        fullAttr = '%s.%s' % (nodeName, 'shade_mode')
        shadeModeValue = pm.getAttr(fullAttr)
        if shadeModeValue == 3:
            pm.editorTemplate(dimControl=(nodeName, 'aoDistance', False))
        else:
            pm.editorTemplate(dimControl=(nodeName, 'aoDistance', True))
        
        #roughness for plastic and metal
        if shadeModeValue == 4 or shadeModeValue == 5:
            pm.editorTemplate(dimControl=(nodeName, 'roughness', False))
        else:
            pm.editorTemplate(dimControl=(nodeName, 'roughness', True))

    def setup(self):
        self.addSwatch()
        self.beginScrollLayout()
        
        self.addCustom('message', 'AEshaderTypeNew', 'AEshaderTypeReplace')

        self.beginLayout('公用属性', collapse=False)
        self.addControl('shadeMode', changeCommand=self.checkShadeMode, label='阴影模式')
        self.addCustom('colorMode', aiUtilityCreateColorMode, aiUtilitySetColorMode)
        if int(ai.AiGetVersion()[2]) > 2:
            self.addControl('overlay_mode', label='叠加模式')
        self.addControl('color', label='颜色')
        self.addControl('opacity', label='不透明度')
        self.addControl('aoDistance', label='AO距离')
        self.addControl('roughness', label='粗糙度')
        self.endLayout()

        # include/call base class/node attributes
        pm.mel.AEdependNodeTemplate(self.nodeName)
        self.addExtraControls()
        
        self.endScrollLayout()        

