# -*- coding: utf-8 -*-
"""
Spyder Editor

Author: Jared Athias

PYTHON
for use in Autodesk Maya

Auto texture tool

This tool will automatically assign PBR textures to the selected model.

Base Color
Metalness
Roughness
Normal

-------------------------------------------------------------------------------

-------------------------------------------------------------------------------
"""

import maya.cmds as cmds
import mtoa.core as core
from functools import partial
import os
import os.path

SELECTED = cmds.ls(sl=True)[0]
FILE_TYPE = ['baseColor', 'roughness', 'metallic', 'normal', 'emissive']

#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/auto_texture_tool_script/"
FILES = []


# Returns list of files in directory
def RETURN_FILE(fieldname,*arg):
    FILE = cmds.textField(fieldname, query=True, text=True)
    textures = os.listdir(FILE)
    
    for item in textures:
        file_name = os.path.basename(item)
        FILES.append(file_name)
    print (FILES)


# Create File Node network Function
def createFileTexture(fileTextureName, p2dName):
    tex = cmds.shadingNode('file', name=fileTextureName, asTexture=True, isColorManaged=True)
    if not cmds.objExists(p2dName):
        p2d = cmds.shadingNode('place2dTexture', name=p2dName, asUtility=True)
        
        cmds.connectAttr(p2d + ".outUV", tex + ".uvCoord")
        cmds.connectAttr(p2d + ".outUvFilterSize", tex + ".uvFilterSize")
        cmds.connectAttr(p2d + ".vertexCameraOne", tex + ".vertexCameraOne")
        cmds.connectAttr(p2d + ".vertexUvOne", tex + ".vertexUvOne")
        cmds.connectAttr(p2d + ".vertexUvThree", tex + ".vertexUvThree")
        cmds.connectAttr(p2d + ".vertexUvTwo", tex + ".vertexUvTwo")
        cmds.connectAttr(p2d + ".coverage", tex + ".coverage")
        cmds.connectAttr(p2d + ".mirrorU", tex + ".mirrorU")
        cmds.connectAttr(p2d + ".mirrorV", tex + ".mirrorV")
        cmds.connectAttr(p2d + ".noiseUV", tex + ".noiseUV")
        cmds.connectAttr(p2d + ".offset", tex + ".offset")
        cmds.connectAttr(p2d + ".repeatUV", tex + ".repeatUV")
        cmds.connectAttr(p2d + ".rotateFrame", tex + ".rotateFrame")
        cmds.connectAttr(p2d + ".rotateUV", tex + ".rotateUV")
        cmds.connectAttr(p2d + ".stagger", tex + ".stagger")
        cmds.connectAttr(p2d + ".translateFrame", tex + ".translateFrame")
        cmds.connectAttr(p2d + ".wrapU", tex + ".wrapU")
        cmds.connectAttr(p2d + ".wrapV", tex + ".wrapV")
        return tex

# Create Node tree
def create_node_tree(fieldname,*arg):
    FILE = cmds.textField(fieldname, query=True, text=True)
    textures = os.listdir(FILE)
    key = '.png'
    for item in textures:
        file_name = os.path.basename(item)
        new_name = file_name.replace(key,'')
        FILES.append(new_name)
    
    SHADER = "aiSS_" + str(SELECTED)
    print (SHADER)
    surface_name = core.createArnoldNode("aiStandardSurface", name=SHADER)
    texture_file = str(FILE) + '/' + str(file_name)
    cmds.sets(SELECTED, edit=True, forceElement=str(SHADER) + "SG")
    for i in FILES:
        NAME = i
        FPN = str(FILE) + "/" + str(NAME) + str(key) 
        
    
        if 'BaseColor' in NAME:
            TFILE = createFileTexture(str(NAME), 'p2d' + str(NAME))
            cmds.connectAttr(str(NAME) + ".outColor", str(SHADER) + ".baseColor")
            cmds.setAttr(str(NAME) + ".fileTextureName", str(FPN), type="string")
            
        elif 'Roughness' in NAME:
            TFILE = createFileTexture(str(NAME), 'p2d' + str(NAME))
            cmds.connectAttr(str(NAME) + ".outAlpha", str(SHADER) + ".specularRoughness")
            cmds.setAttr(str(NAME) + ".fileTextureName", str(FPN), type="string")
            cmds.setAttr(str(NAME) + ".colorSpace", "Raw", type="string")
        
        elif 'Metallic' in NAME:
            TFILE = createFileTexture(str(NAME), 'p2d' + str(NAME))
            cmds.connectAttr(str(NAME) + ".outAlpha", str(SHADER) + ".metalness")
            cmds.setAttr(str(NAME) + ".fileTextureName", str(FPN), type="string")
            cmds.setAttr(str(NAME) + ".colorSpace", "Raw", type="string")
     
        elif 'Normal' in NAME:
            TFILE = createFileTexture(str(NAME), 'p2d' + str(NAME))
            cmds.shadingNode("bump2d", name="b2d" + str(NAME), asUtility=True)
            cmds.setAttr("b2d" + str(NAME) + ".bumpInterp", 1)
            cmds.connectAttr(str(NAME) + ".outAlpha", "b2d" + str(NAME) + ".bumpValue")
            cmds.connectAttr("b2d" + str(NAME) + ".outNormal", str(SHADER) + ".normalCamera")
            cmds.setAttr(str(NAME) + ".fileTextureName", str(FPN), type="string")
            cmds.setAttr(str(NAME) + ".colorSpace", "Raw", type="string")
        
        elif 'Emissive' in NAME:
            TFILE = createFileTexture(str(NAME), 'p2d' + str(NAME))
            cmds.connectAttr(str(NAME) + ".outColor", str(SHADER) + ".emissionColor")
            cmds.setAttr(str(NAME) + ".fileTextureName", str(FPN), type="string")
        
        else:
            note = str(NAME) + " is not connected"
            print(note)
            pass

    
            
def assign_shader(*arg):
    SELECTED = cmds.ls(sl=True)[0]
    SHADER = "aiSS_" + str(SELECTED)
    mel.sets(str(SELECTED), edit=True, forceElement=str(SHADER))      
        
# Auto Texture Tool Function
def auto_texeture_tool_UI():
    ui_title = 'auto texture tool'
    
    if cmds.window(ui_title, exists=True):
        print('CLOSE DUPLICATE WINDOW')
        cmds.deleteUI(ui_title)
        
    window = cmds.window(title='Auto Texture Tool', widthHeight=(500,300))
    cmds.columnLayout(adjustableColumn=True)
    
#    cmds.image(image=IMG_PATH + 'Title_Image.jpg')
    cmds.separator(height=50)
    
    "File Location"
    cmds.rowColumnLayout( numberOfColumns=2, columnAttach=(1, 'right', 0), columnWidth=[(1, 100), (2, 250)] )
    cmds.text(label= "File Location")
    FILE = cmds.textField()
    
    cmds.button(label="Launch", en=True, c=partial(create_node_tree, FILE))
    
#    cmds.button(label="Assign Shader", en=True, c=assign_shader)
    
    cmds.showWindow(window)


auto_texeture_tool_UI()