import pymel.core as pm
import sys as sys

def create_substance_shader(mesh_textureSet_name, extension):

# $mesh_$textureSet_Diffuse
# $mesh_$textureSet_Specular
# $mesh_$textureSet_Roughness
# $mesh_$textureSet_f0
# $mesh_$textureSet_Normal
  diffuse_filename = 'sourceImages\\' + mesh_textureSet_name + '_Diffuse.'+extension
  specular_filename = 'sourceImages\\' + mesh_textureSet_name + '_Specular.'+extension
  roughness_filename = 'sourceImages\\' + mesh_textureSet_name + '_Roughness.'+extension
  f0_file_nanme = 'sourceImages\\' + mesh_textureSet_name + '_f0.'+extension
  normal_filename = 'sourceImages\\' + mesh_textureSet_name + '_Normal.'+extension
  # Height map
  # Emissive map

  shader_node = pm.shadingNode('aiStandard', asShader=True)
  gammaCorrect_roughness = pm.shadingNode('gammaCorrect', asShader=True)
  gammaCorrect_ksn = pm.shadingNode('gammaCorrect', asShader=True)
  roughness_uv_node = pm.shadingNode('place2dTexture', asUtility=True)
  bump2d_node = pm.shadingNode('bump2d', asShader=True)
  pm.setAttr('%s.bumpInterp' % bump2d_node, 1)
  f0_uv_node = pm.shadingNode('place2dTexture', asUtility=True)
  bump_uv_node = pm.shadingNode('place2dTexture', asUtility=True)
  f0_file_node = pm.shadingNode("file",asTexture=True)
  roughess_file_node = pm.shadingNode("file",asTexture=True)
  diffuse_file_node = pm.shadingNode("file",asTexture=True)
  sepcular_file_node = pm.shadingNode("file",asTexture=True)
  normal_file_node = pm.shadingNode("file",asTexture=True)
  diffuse_uv_node = pm.shadingNode('place2dTexture', asUtility=True)
  specular_uv_node = pm.shadingNode('place2dTexture', asUtility=True)

  pm.setAttr(gammaCorrect_roughness + ".gamma", 2.2, 2.2, 2.2)
  pm.setAttr(gammaCorrect_ksn + ".gamma", 2.2, 2.2, 2.2)

  pm.setAttr(shader_node + ".Kd", 1)
  pm.setAttr(shader_node + ".Ks", 1)
  # pm.setAttr(shader_node + ".Fresnel", 1)
  pm.setAttr(shader_node + ".FresnelAffectDiff", 1)
  pm.setAttr(shader_node + ".specularFresnel", 1)
  # print('{0}.specularRoughness'.format(shader_node))
  pm.connectAttr( '{0}.outValueX'.format(gammaCorrect_roughness), '{0}.specularRoughness'.format(shader_node))
  pm.connectAttr( '{0}.outValueX'.format(gammaCorrect_ksn), '{0}.Ksn'.format(shader_node))
  pm.connectAttr( '{0}.outNormal'.format(bump2d_node), '{0}.normalCamera'.format(shader_node))
  pm.connectAttr( '{0}.outColor'.format(roughess_file_node), '{0}.value'.format(gammaCorrect_roughness))
  pm.connectAttr( '{0}.outColor'.format(f0_file_node), '{0}.value'.format(gammaCorrect_ksn))
  pm.connectAttr( '{0}.outColor'.format(diffuse_file_node), '{0}.color'.format(shader_node))
  pm.connectAttr( '{0}.outColor'.format(sepcular_file_node), '{0}.KsColor'.format(shader_node))
  # pm.connectAttr( '{0}.outColorR'.format(normal_file_node), '{0}.bumpValue'.format(bump2d_node))
  pm.connectAttr( '{0}.outAlpha'.format(normal_file_node), '{0}.bumpValue'.format(bump2d_node))

  pm.connectAttr( '{0}.outUV'.format(roughness_uv_node), '{0}.uvCoord'.format(roughess_file_node))
  pm.connectAttr( '{0}.outUV'.format(f0_uv_node), '{0}.uvCoord'.format(f0_file_node))
  pm.connectAttr( '{0}.outUV'.format(bump_uv_node), '{0}.uvCoord'.format(normal_file_node))
  pm.connectAttr( '{0}.outUV'.format(diffuse_uv_node), '{0}.uvCoord'.format(diffuse_file_node))
  pm.connectAttr( '{0}.outUV'.format(specular_uv_node), '{0}.uvCoord'.format(sepcular_file_node))

  pm.connectAttr( '{0}.outUvFilterSize'.format(roughness_uv_node), '{0}.uvFilterSize'.format(roughess_file_node))
  pm.connectAttr( '{0}.outUvFilterSize'.format(f0_uv_node), '{0}.uvFilterSize'.format(f0_file_node))
  pm.connectAttr( '{0}.outUvFilterSize'.format(bump_uv_node), '{0}.uvFilterSize'.format(normal_file_node))
  pm.connectAttr( '{0}.outUvFilterSize'.format(diffuse_uv_node), '{0}.uvFilterSize'.format(diffuse_file_node))
  pm.connectAttr( '{0}.outUvFilterSize'.format(specular_uv_node), '{0}.uvFilterSize'.format(sepcular_file_node))

  pm.setAttr('%s.fileTextureName' %diffuse_file_node, diffuse_filename , type="string")
  pm.setAttr('%s.fileTextureName' %sepcular_file_node, specular_filename , type="string")
  pm.setAttr('%s.fileTextureName' %roughess_file_node, roughness_filename , type="string")
  pm.setAttr('%s.fileTextureName' %f0_file_node, f0_file_nanme , type="string")
  pm.setAttr('%s.fileTextureName' %normal_file_node, normal_filename , type="string")

  pm.rename(roughess_file_node, 'roughness')
  pm.rename(diffuse_file_node, 'diffuse')
  pm.rename(sepcular_file_node, 'specular')
  pm.rename(normal_file_node, 'normal')
  pm.rename(f0_file_node, 'f0')
  pm.rename(shader_node, mesh_textureSet_name)



def generateArnold(*args):
  shader_name = name.getText()
  ext_name = ext.getText()
  # create_substance_shader('dog_mt_wireframe')
  create_substance_shader(shader_name, ext_name)
  

def generateMentalRay(*args):
  print 'Mental ray'
#    Create a window with a some fields for entering text_name.
#
window = pm.window( title="substance to maya", iconName='Short Name', widthHeight=(350, 355) )
pm.rowColumnLayout( numberOfColumns=2, columnAttach=(1, 'right', 0), columnWidth=[(1, 100), (2, 250)] )
# Result: ui.RowColumnLayout('window1|rowColumnLayout8') #
pm.text( label='export_prefix' )
name = pm.textField()

pm.text( label='extension' )
ext = pm.textField(text='png')
# #    Attach commands to pass focus to the next field if the Enter
# #    key is pressed. Hitting just the Return key will keep focus
# #    in the current field.
# #
pm.textField( name, edit=True, enterCommand='defaultButtonPush' )

pm.rowColumnLayout( numberOfColumns=1, columnAttach=(1, 'left', 0), columnWidth=[(1, 350)] )
pm.button( label='Arnold shaders', command=generateArnold )
pm.button( label='mentalray shaders', command=generateMentalRay )

pm.showWindow(window)