--ShaderAnalysisTool.ms
--2013 Andy Davis
--Rockstar London
--Script displays information about multi-layered shaders

FileIn (RsConfigGetWildWestDir() + "script/3dsMax/_config_files/Wildwest_header.ms")

--UI Xaml Markup	
UIMarkupFilePath = @"X:\gta5\tools\wildwest\script\3dsMax\UI\ShaderAnalysisTool.xaml"

--Custom dotNet Class
CSharp.CompileToMemory #(::RsConfigGetWildwestDir() + @"script\3dsMax\Maps\Materials\ShaderAnalysisTool.cs")

global ShaderAnalysisToolUI
global ToolFunctions

struct MaterialInfo
(
	mat,
	faceList,
	shaderClass,
	rageType, 
	isLayered = false,
	diffuseMapList = #()
)

--------------------------------------------------------------------------------------------------------------------------------
--functions that are candidates to be made generic
--------------------------------------------------------------------------------------------------------------------------------

--function returns the shader layer Influences when passed the green and blue components of a map value
--InMapValue: send the value for map channel 9 to the function
--Returns: 4 part array with the influences of each of the four shaders
fn RSMaterial_Get2LayerShaderInfluences InMapValue =
(
	local G = InMapValue[2]	--green component of map value
	local B = InMapValue[3]	--blue component of map value
	local Influences = #()
	
	if (G > B) then
	(
		Influences[1] = 1 - G
		Influences[4] = B
	)
	else
	(
		Influences[1] = 1 - B
		Influences[4] = G
	)
	
	if (1 - G > B) then
	(
		Influences[2] = B
		Influences[3] = G
	)
	else
	(
		Influences[2] = 1 - G
		Influences[3] = 1 - B
	)
	
	return Influences
)

--returns a material from a multi-material when passed the id or undefined if nout found
fn RSMaterials_GetMaterialFromMulti inMulti inMatID =
(
	local outMaterial = undefined
	local notFound = true
	
	for i = 1 to inMulti.MaterialList.count while notFound do
	(
		if inMulti.MaterialIDList[i] == inMatID then
		(
			outMaterial = inMulti.MaterialList[i]
			notFound = false
		)
	)	
	
	return outMaterial
)

--returns an array of MaterialInfo objects
fn RSMaterials_GetMaterials InObject =
(
	local MaterialInfoList = #()
	local materialList = #()
	local mat = InObject.material
	local modAdded = false
	
	if (classOf InObject != editable_mesh) then
	(
		addModifier InObject (Edit_Mesh())
		modAdded = true
	)
	
	if (classOf InObject.material == MultiMaterial) then
	(	
		--get a list of each unique material id used on the object, and corresponding faces
		local faceCount = InObject.faces.count
		local matIDList = #()			--list of unique face ids
		local matIDFaceList = #()	--corresponding faces
		
		for face = 1 to faceCount do
		(
			local matID = getFaceMatID  InObject face
			local item = findItem matIDList matID
			
			if (item == 0) then
			(
				append matIDList matID
				append matIDFaceList #{face}
			)
			else
			(
				append matIDFaceList[item] face
			)
		)
		
		--work out the material corresponding to id's in the material id list
		for i = 1 to matIDList.count do
		(
			local thisMat = RSMaterials_GetMaterialFromMulti InObject.material matIDList[i]
			append materialList (dataPair material:thisMat faces:matIDFaceList[i])
		)
	)
	else
	(
		append materialList (dataPair material:InObject.Material faces:#{1..InObject.faces.count})
	)	

-- 	clearListener()
	
	for item in materialList do
	(
		local shaderType = undefined
		local diffuseMaps = #()
		local isLayered = false
		
		if (classOf item.material == Rage_Shader) then
		(
			shaderType = RstGetShaderName item.material
			local variableCount = RstGetVariableCount item.material
			
			if (findString shaderType "terrain_cb_") != undefined then
			(
				isLayered = true
				
				local firstMap = 0
				
				for i = 1 to variableCount while firstMap == 0 do
				(
					if (RstGetVariableType item.material i) == "texmap" then
					(
						--check this is a diffuse map
						if ((findString (RstGetVariableName item.material i) "Diffuse" != undefined)) then
						(
							firstMap = i
						)
					)
				)
				
				local thisIsNormalMap = false
				--maps alternate between diffuse and normal, use this logic operation to get alternate maps
				
				for i = firstMap to (firstMap + 7) do
				(
					if (thisIsNormalMap = not thisIsNormalMap) then
					(
						append diffuseMaps (RstGetVariable item.material i)
					)
				)
			)
			else
			(
				local mapNotFound = true
				
				for i = 1 to variableCount while mapNotFound do
				(
					if (RstGetVariableType item.material i) == "texmap" then
					(
						mapNotFound = false
						append diffuseMaps (RstGetVariable item.material i)
					)
				)
			)	
		)
		
		local newMaterialInfo = MaterialInfo()
		newMaterialInfo.mat = item.material
		newMaterialInfo.faceList = item.faces
		newMaterialInfo.shaderClass = (classOf item.material)
		newMaterialInfo.rageType = shaderType
		newMaterialInfo.diffuseMapList = diffuseMaps
		newMaterialInfo.isLayered = isLayered
		append MaterialInfoList newMaterialInfo
	)
	
	if (modAdded) then deleteModifier InObject InObject.Edit_Mesh
	
	return MaterialInfoList
)

--------------------------------------------------------------------------------------------------------------------------------
--tool functionality
--------------------------------------------------------------------------------------------------------------------------------

struct ToolFunctionsStruct
(
	ShaderList = #(),
	AnalyseButton,
	ShaderAnalysisTool,
	ObjectList = #(),
	
	fn PrintMaterialInfoList InShaderList =
	(
		for item in InShaderList do
		(
			format "Material: %\n" item.mat.name
			format "Face list: %\n" item.faceList
			format "Shader class: %\n" item.shaderClass
			format "Rage type: %\n" item.rageType
			format "Layered shader? %\n" item.isLayered
			for map in item.diffuseMapList do format "diffuse map: %\n" map
			format "\n"
		)
	),	
	
	
	fn GetShaderLayerInfo InObject InMaterialInfo =
	(
		local ShaderLayerInfo = #(0, 0, 0, 0)
		local faceList = InMaterialInfo.FaceList
		
		--get the total influences of each of the four materials in the layered shader by getting the map channel data on each subvert
		
		for face in faceList do
		(
			local vertList = (meshop.GetVertsUsingFace InObject face) as array
			
			for vert in vertList do
			(
				local shaderLayerMapChannel = 9
				local mapValue = RSMesh_GetSubVert_MapValue InObject vert face shaderLayerMapChannel
				for i = 1 to 3 do this.ClampValueZeroOne &mapValue[i]		--adjust values not in the zero to one value range
				local influence =	RSMaterial_Get2LayerShaderInfluences mapValue				
				for i = 1 to 4 do ShaderLayerInfo[i] += influence[i]
			)
		)
		
		--convert the influences into percentage ratios by dividing by the total and multiplying by 100
		
		local divisor = ShaderLayerInfo[1] + ShaderLayerInfo[2] + ShaderLayerInfo[3] + ShaderLayerInfo[4]
		
		for i = 1 to 4 do ShaderLayerInfo[i] *= 100.0 / divisor
		
		return ShaderLayerInfo
	),
	
	fn ClampValueZeroOne &InValue =
	(
		if (InValue > 1) then
			InValue = InValue - (floor InValue)
		else if (InValue < 0) then
			InValue = (ceil InValue) - InValue
	),
	
	fn GetObjectData =
	(
		ToolFunctions.ObjectList.ObjectDataList.Clear()
		
		local ProgressValue = 0
		
		objList = for item in selection where SuperClassOf item == GeometryClass collect item
		
		--Create our top progress window
		ProgressWindow = RSProgressWindow Title:"ProgressWindow" StartStep:0 EndStep:objList.count
		ProgressWindow.Start()
		
		
		for obj in objList do
		(
			AddModifier obj (Edit_Mesh())
			local newObjDataItem = dotNetObject "ShaderAnalysis.ObjectData"
			newObjDataItem.Handle = obj.Handle
			newObjDataItem.Name = obj.Name
			newObjDataItem.FaceCount = obj.faces.count
			local ShaderList = RSMaterials_GetMaterials obj
				
			for thisShader in ShaderList do
			(
				local newDesc = dotNetObject "ShaderAnalysis.MaterialDescription"

				newDesc.ObjectHandle = obj.handle
				
				if thisShader.mat == undefined then
				(
					newDesc.Name = "no material"
					newDesc.Handle = -1
					newDesc.TextureMapList = undefined
					newDesc.ShaderType = "\t"
				)
				else
				(				
					newDesc.Name = thisShader.mat.name
					newDesc.Handle = GetHandleByAnim thisShader.mat
					newDesc.ShaderType = thisShader.RageType
					newDesc.Tint = (dotNetClass "System.Windows.Media.Color").FromRGB (random 80 255) (random 80 255) (random 80 255)
					
					local TextureMapList = #()
					
					for map in thisShader.diffuseMapList where (isKindOf map String) do
					(
						local newMap = dotNetObject "ShaderAnalysis.TextureMap"
						newMap.Path = map
						append TextureMapList newMap
					)
					
					newDesc.TextureMapList = TextureMapList
					
					if thisShader.isLayered then
					(
						local shaderLayerPercentages = this.GetShaderLayerInfo obj thisShader
						
						for i = 1 to 4 do
							newDesc.TextureMapList[i].Usage = shaderLayerPercentages[i]
					)
					else
					(
						newDesc.TextureMapList[1].Usage = 100
					)
				)
				
				newDesc.FaceList = (thisShader.faceList) as array
				newObjDataItem.CalculatePercentage newDesc
				newObjDataItem.ShaderList.Add newDesc					
			)
			
 			ToolFunctions.ObjectList.ObjectDataList.Add newObjDataItem
			DeleteModifier obj obj.Edit_Mesh
			ProgressWindow.PostProgressStep()
		)
		
		--End progress
		ProgressWindow.End()
	),

	fn AnalyseButtonPressed =
	(
		ToolFunctions.DeleteShaderProxy()
		ToolFunctions.GetObjectData()
	),
	
	fn InitStuff =
	(
		--analyse button
		this.AnalyseButton = ShaderAnalysisToolUI.Content.FindName("AnalyseSelectedButton")
		dotnet.AddEventHandler AnalyseButton "click" AnalyseButtonPressed
		
		--select and zoom button
		this.ShaderAnalysisTool = dotNetObject "ShaderAnalysis.ShaderAnalysisTool"
		ShaderAnalysisTool.BindElement ShaderAnalysisToolUI.content
		
		--ObjectList and bind
		ToolFunctions.ObjectList = dotNetObject "ShaderAnalysis.BindingList"
		ToolFunctions.GetObjectData()
		ShaderAnalysisToolUI.Content.DataContext = ObjectList.ObjectDataList
	),
	
	--returns an array of MaterialInfo objects
	fn GetMaterials InObject =
	(
		local MaterialInfoList = #()
		local materialList = #()
		local mat = InObject.material
		local modAdded = false
		
		if (classOf InObject != editable_mesh) then
		(
			addModifier InObject (Edit_Mesh())
			modAdded = true
		)
		
		if (classOf InObject.material == MultiMaterial) then
		(	
			--get a list of each unique material id used on the object, and corresponding faces
			local faceCount = InObject.faces.count
			local matIDList = #()			--list of unique face ids
			local matIDFaceList = #()	--corresponding faces
			
			for face = 1 to faceCount do
			(
				local matID = getFaceMatID  InObject face
				local item = findItem matIDList matID
				
				if (item == 0) then
				(
					append matIDList matID
					append matIDFaceList #{face}
				)
				else
				(
					append matIDFaceList[item] face
				)
			)
			
			--work out the material corresponding to id's in the material id list
			for i = 1 to matIDList.count do
			(
				local thisMat = RSMaterials_GetMaterialFromMulti InObject.material matIDList[i]
				append materialList (dataPair material:thisMat faces:matIDFaceList[i])
			)
		)
		else
		(
			append materialList (dataPair material:InObject.Material faces:#{1..InObject.faces.count})
		)	

-- 		clearListener()
		
		for item in materialList do
		(
			local shaderType = undefined
			local diffuseMaps = #()
			local isLayered = false
			local newMaterialInfo = MaterialInfo()
			
			if (classOf item.material == Rage_Shader) then
			(
				shaderType = RstGetShaderName item.material
				local variableCount = RstGetVariableCount item.material
				
				if (findString shaderType "terrain_cb_") != undefined then
				(
					isLayered = true
					
					local firstMap = 0
					
					for i = 1 to variableCount while firstMap == 0 do
					(
						if (RstGetVariableType item.material i) == "texmap" then
						(
							firstMap = i
						)
					)
					
					local thisIsNormalMap = false
					--maps alternate between diffuse and normal, use this logic operation to get alternate maps
					
					for i = firstMap to (firstMap + 7) do
					(
						if (thisIsNormalMap = not thisIsNormalMap) then
						(
							append diffuseMaps (RstGetVariable item.material i)
						)
					)
				)
				else
				(
					local mapNotFound = true
					
					for i = 1 to variableCount while mapNotFound do
					(
						if (RstGetVariableType item.material i) == "texmap" then
						(
							mapNotFound = false
							append diffuseMaps (RstGetVariable item.material i)
						)
					)
				)
			)
			
			newMaterialInfo.mat = item.material
			newMaterialInfo.faceList = item.faces
			newMaterialInfo.shaderClass = (classOf item.material)
			newMaterialInfo.rageType = shaderType
			newMaterialInfo.diffuseMapList = diffuseMaps
			newMaterialInfo.isLayered = isLayered
			append MaterialInfoList newMaterialInfo
		)
		
		if (modAdded) then deleteModifier InObject InObject.Edit_Mesh
		
		return MaterialInfoList
	),
	
	------------------------------------------------------------------------------------------------------
	--function creates a tinted proxy object representing the shaders in the target object
	------------------------------------------------------------------------------------------------------
	fn CreateShaderProxy InObject =
	(
		local nodeArray = undefined
		local NewMaterialList = #()		
		max modify mode
		local SubObjectMode = SubObjectLevel
		
		local proxyObjects = for item in objects where (item.name == "ShaderAnalysis_Proxy") collect item
		
		if not this.DeleteShaderProxy() then
		(
			if (classOf InObject.material == MultiMaterial) then
			(
				maxops.cloneNodes InObject cloneType:#reference newNodes:&nodeArray
				local ShaderProxy = nodeArray[1]
				ShaderProxy.name = "ShaderAnalysis_Proxy"
				ShaderProxy.showFrozenInGray = off
				ShaderProxy.showVertexColors = off
				AddModifier ShaderProxy (Push Push_Value:0.1)
				freeze ShaderProxy
				setAttr ShaderProxy (getAttrIndex "Gta Object" "Dont Export") true
				setAttr ShaderProxy (getAttrIndex "Gta Object" "Dont Add To IPL") true
				
				--find the object in the ObjectDataList
				--get the material description list
				local numObjs = ToolFunctions.ObjectList.ObjectDataList.Count
				local MaterialData = undefined
				
				for i = 0 to (numObjs - 1) do
				(
					if (InObject.Handle == ToolFunctions.ObjectList.ObjectDataList.Item[i].Handle) then
					(
						MaterialData = ToolFunctions.ObjectList.ObjectDataList.Item[i].ShaderList
					)
				)
				
				--for each material description, get the tint value and the material name
				if (MaterialData != undefined) then
				(
					local MaterialIDList = #()
					local numMaterials = MaterialData.count
					NewMulti = MultiMaterial name:"ShaderProxyMaterial"
					meditMaterials[activeMeditSlot] = NewMulti
					
					for i = 0 to (numMaterials - 1) do
					(
						local shaderNode = MaterialData.Item[i]
						local NewColor = color (shaderNode.Tint.R) (shaderNode.Tint.G) (shaderNode.Tint.B)
						local currentMaterial = GetAnimByHandle shaderNode.Handle
						local NewMaterial = StandardMaterial name:currentMaterial.name diffuse:NewColor opacity:50
						append MaterialIDList (RSMaterials_GetMaterialIDFromMulti InObject.material currentMaterial)
						append NewMaterialList NewMaterial
					)
					
					NewMulti.MaterialList = NewMaterialList
					NewMulti.MaterialIDList = MaterialIDList
					ShaderProxy.Material = NewMulti
					ShaderProxy
				)
				else
				(
					messagebox "Object not found" title:"Shader Analysis Tool"
				)
			)
			else
			(
				messagebox "Single material on object" title:"Shader Analysis Tool"
			)
		)
		
		SubObjectLevel = SubObjectMode
	),
	
	fn DeleteShaderProxy =
	(
		local proxyObjects = for item in objects where (item.name == "ShaderAnalysis_Proxy") collect item	
		if proxyObjects.count > 0 then
		(
			delete proxyObjects
			return true
		)
		else return false
	)
)
	
fn SelectAndZoomAction s e = 
(
	local obj = MaxOps.getNodeByHandle e.Data.ObjectHandle
	local faceList = e.Data.FaceList as bitarray
	
	if obj == undefined then
	(
		messagebox "Object no longer in scene" title:"Shader Analysis Tool"
	)
	else
	(
		select obj
		
		if (obj.modifiers.count > 0) then
		(
			if (classOf obj.modifiers[1] != Edit_Mesh) then
			(
				AddModifier obj (Edit_Mesh())
			)
		)
		else if classOf obj == Editable_Poly then
		(
			convertToMesh obj
		)
		
		
		if (classOf obj != Editable_Mesh) then
			AddModifier obj (Edit_Mesh())
		
		max modify mode
		SubObjectLevel = 4
		
		if (classOf obj == Editable_Mesh) then
		(
			SetFaceSelection obj faceList
		)
		else
		(
			PolyOp.SetFaceSelection obj faceList
		)
		
		max zoomext sel
	)
)

fn TintShadersAction s e =
(
	local obj = MaxOps.getNodeByHandle e.Data.Handle
	
	if obj == undefined then
	(
		messagebox "Object no longer in scene" title:"Shader Analysis Tool"
	)
	else
	(		
		if obj.material != undefined then
		(
			ToolFunctions.CreateShaderProxy obj
		)
		else
		(
			messagebox "No shaders found to tint" title:"Shader Analysis Tool"
		)
	)
)

fn OnCloseAction s e =
(
	ToolFunctions.DeleteShaderProxy()
)

ToolFunctions = ToolFunctionsStruct()

--Setup our tool window
if ShaderAnalysisToolUI != undefined then ShaderAnalysisToolUI.Form.Close()

ShaderAnalysisToolUI = RSToolWindow Name:"Shader Analysis Tool" UISource:UIMarkupFilePath
ToolFunctions.InitStuff()
dotnet.AddEventHandler ToolFunctions.ShaderAnalysisTool "SelectAndZoom" SelectAndZoomAction
dotnet.AddEventHandler ToolFunctions.ShaderAnalysisTool "TintShaders" TintShadersAction
dotnet.AddEventHandler ShaderAnalysisToolUI.Form "Closing" OnCloseAction
ShaderAnalysisToolUI.Open()