-- Weapon automatic render tool 
-- Rick Stirling
-- Rockstar North
-- July 2012

-- Uses an XML definition file to load weapons and accessories and render them out 
filein (RsConfigGetWildWestDir() + "script/3dsMax/_config_files/Wildwest_header.ms")


-- Data section 
weaponRenderList = #()
attachmentRenderList = #()
	



-- Structs

-- Weaponset holds the details of each weapon and the attachaments
struct weaponsset
(
	weaponName = undefined,
	weaponMesh = undefined,
	weaponFile = undefined,
	attachment_Scop = undefined,
	attachment_Supp = undefined,
	attachment_Clip = undefined,
	attachment_Stck = undefined,
	attachment_Lasr = undefined,
	attachment_Grip = undefined,
	attachment_Flsh = undefined,
	attachment_SeWp = undefined
)	


struct attachmentdetails
(
	attachmentName = undefined,
	attachmentMesh = undefined,
	attachmentFile = undefined,
	attachmentParent = undefined
)	





-- **************************************************************
--Functions


-- Load the XML config file and get the weapon and attachment data from it
fn load_weapon_data weaponConfigFile = (
	
	-- Define an XML file and make the path safe:
	local theXMLfile = weaponConfigFile
	theXMLfile = RSMakeSafeSlashes theXMLfile 

	-- Sync that file from Perforce 
	try 
	(
		depotPath = gRsPerforce.local2depot theXMLfile 
		gRsPerforce.p4.run "sync" #(depotPath) silent:true
	)
	catch(print "Couldnt sync the XML file")

	-- Initialise the XML file for processing using Xpath
	ExpClass = dotNetClass "System.Xml.XPath.XPathExpression"
	stream = dotNetObject "System.IO.StreamReader" theXMLfile
	xmldoc = dotNetObject "System.Xml.XmlDocument"
	xmldoc.load stream
	xmlroot = xmldoc.documentElement

	
	-- Clear the data structues 
	weaponRenderList = #()
	attachmentRenderList = #()
	
	-- Grab the weapon infomation and create a struct for each
	searchstring = ".//weaponassets/Item"
	nodelist = xmlroot.Selectnodes (searchstring)
	
	for nlindex = 0 to (nodelist.count-1) do
	(
		-- get the basic weapon information
		thenode = nodelist.item[nlindex]
		xmlentry = thenode.innerxml

		innernodelist = thenode.Selectnodes "weaponname"
		innernode = innernodelist.itemof[0]
		
		str_weapon = weaponsset weaponName: innernode.innerxml
		
		innernodelist = thenode.Selectnodes "weaponmesh"
		innernode = innernodelist.itemof[0]
		str_weapon.weaponMesh = innernode.innerxml
		
		innernodelist = thenode.Selectnodes "weaponfile"
		innernode = innernodelist.itemof[0]
		str_weapon.weaponFile = innernode.innerxml
		
		
		-- Now search for the attachment lists for that weapon
		weaponNodelist = xmlroot.Selectnodes  "//weaponattachments/item" 
		
		-- Loop through all the weapon attachment sections looking for this weapon 
		for weaponIndex = 0 to (weaponnodelist.count-1) do
		(
			weaponNode = weaponNodelist.item[weaponindex]
			theWeaponParameter = weaponNode.Selectnodes "@weapon"
			attlistWeaponName = theWeaponParameter.itemof[0].value
			
			-- If we find the weapon, build up the attachment list for it 
			if attlistWeaponName == str_weapon.weaponMesh then
			(
				-- Set up some arrays for storing the attachment data
				ar_AAPScop = #()
				ar_AAPSupp = #()
				ar_AAPClip = #()
				ar_AAPStck = #()
				ar_AAPLasr = #()
				ar_AAPGrip = #()
				ar_AAPFlsh = #()
				ar_AAPSeWp = #()
				
				-- How many children are there of this weaponNode?
				-- Each child is an attachment entity
				attachmentNodelist = weaponNode.Selectnodes "attachmentoverride"
				for attachmentIndex = 0 to (attachmentNodelist.count-1) do 
				(
					attachmentNode = attachmentNodelist.item[attachmentIndex]
					theAttachmentTypeParamater = attachmentNode.Selectnodes "@attachmenttype"
					attachmentType = theAttachmentTypeParamater.itemof[0].value
					
					theAttachmentNameParamater = attachmentNode.Selectnodes "@attachmentname"
					attachmentName = theAttachmentNameParamater.itemof[0].value
					
					case attachmentType of
					(
						"AAPScop": append ar_AAPScop attachmentName
						"AAPSupp": append ar_AAPSupp attachmentName
						"AAPClip": append ar_AAPClip attachmentName
						"AAPStck": append ar_AAPStck attachmentName
						"AAPLasr": append ar_AAPLasr attachmentName
						"AAPGrip": append ar_AAPGrip attachmentName
						"AAPFlsh": append ar_AAPFlsh attachmentName
						"AAPSeWp": append ar_AAPSeWp attachmentName
					)	
					
				) -- End attachment gathering loop 
				
				-- Add this attachment data to the weapon struct 
				-- Each element in the struct is an array (though usually only with 1 entry)
				str_weapon.attachment_Scop = ar_AAPScop
				str_weapon.attachment_Supp = ar_AAPSupp
				str_weapon.attachment_Clip = ar_AAPClip
				str_weapon.attachment_Stck = ar_AAPStck
				str_weapon.attachment_Lasr = ar_AAPLasr
				str_weapon.attachment_Grip = ar_AAPGrip
				str_weapon.attachment_Flsh = ar_AAPFlsh
				str_weapon.attachment_SeWp = ar_AAPSeWp
				
			)-- end weapon attachment  search loop				
		)
		
		
		-- Store this struct in the attachments array 
		append weaponRenderList str_weapon
	)  
		
	
	-- Grab the attachment infomation and create a struct for each
	searchstring = ".//attachmentassets/Item"
	nodelist = xmlroot.Selectnodes (searchstring)
	
	for nlindex = 0 to (nodelist.count-1) do
	(
		
		thenode = nodelist.item[nlindex]
		xmlentry = thenode.innerxml

		innernodelist = thenode.Selectnodes "attachmentname"
		innernode = innernodelist.itemof[0]
		
		str_attachment = attachmentdetails attachmentName: innernode.innerxml
		
		innernodelist = thenode.Selectnodes "attachmentmesh"
		innernode = innernodelist.itemof[0]
		str_attachment.attachmentMesh = innernode.innerxml
		
		innernodelist = thenode.Selectnodes "attachmentfile"
		innernode = innernodelist.itemof[0]
		str_attachment.attachmentFile = innernode.innerxml
		
		innernodelist = thenode.Selectnodes "attachmentnode"
		innernode = innernodelist.itemof[0]
		str_attachment.attachmentParent = innernode.innerxml
		
		
		-- Store this struct in the attachments array 
		append attachmentRenderList  str_attachment
	)  
	
	
	
) -- end of load_weapon_data 




-- Given the mesh name of an attachment, find the details in the attachment struct 
fn find_attachment attachmentMeshName =
(

	condensedAttachment =#()
	
	for attachmentIndex = 1 to attachmentRenderList.count do 
	(
		theAttachment = attachmentRenderList[attachmentIndex]
		structMeshName = theAttachment.attachmentMesh
		
		-- Populate the return array
		if structMeshName == attachmentMeshName then
		(
			append condensedAttachment theAttachment.attachmentName
			append condensedAttachment theAttachment.attachmentMesh
			append condensedAttachment theAttachment.attachmentFile
			append condensedAttachment theAttachment.attachmentParent
		)
	)
	
	condensedAttachment -- return the attachment data
)




-- Taking the name of an attachment point, try to get the corresponding weapon attachment point
fn get_wap_node aapname =
(
	wapnode = undefined
	
	-- Usually the node is the same name as the aap, with the first letter replaced
	wapname = "W" + (substring aapname 2 255)
	try	(wapnode = getNodeByName wapname) catch (print "no joy")
	
	-- If that failed, try some of the overrides
	if wapnode == undefined then
	(		
		if aapname == "AAPLasr" then wapname = "WAPFlshLasr"
		if aapname == "AAPFlsh" then wapname = "WAPFlshLasr"	
	)	
	
	try	(wapnode = getNodeByName wapname) catch (print "no joy")
	
	wapnode -- return this node
	
)	



-- Merge in the attachments for the weapon and position them correctly in the scene 
fn merge_attachment attachmentData = 
(
	clearSelection()
	-- Check that there is valid data
	if attachmentData.count > 1 then 
	(	
		-- Merge the file - max object and the attachment helper 
		-- There might be many attachment helpers, but that doesn't matter 
		mergemaxfile attachmentData[3] #(attachmentData[2],attachmentData[4]) #mergeDups #useSceneMtlDups #select
		
		-- Store these incoming objects, we'll need to access them for positioning later
		attachmentSelectionSet = getCurrentSelection()
		
		-- Figure out the name of the weapon attachment point 
		aapname = attachmentData[4]
		wapnode = undefined
		wapnode = get_wap_node aapname 
		
		

		-- Position the helper object (skinning will take care of the attachment)	
		if wapnode != undefined then
		(	
			for obj in attachmentSelectionSet do 
			(
				if obj.name == aapname then obj.pos = wapnode.pos
			)
		)

		else 
		(
			messagebox ("Could not find an attachment point for this object: " + aapname) 
			max delete
		)
	)
	else
	(
		print "Incoming attachment data was empty"
	)
)




-- Load or merge the maxfiles and build the datasets
-- Takes a weapon and does a lookup to know what to load
fn load_weapon_meshes chosenWeapon =
(
	theWeapon = undefined
	weaponComponents = #()
	
	
	-- Load the weapon first
	-- Find the weapon in the weapon struct array
	for weaponindex = 1 to weaponRenderList.count do 
	(
		theWeaponEntry = weaponRenderList[weaponindex]
		theWeaponMesh = theWeaponEntry.weaponMesh 
		if theWeaponMesh == chosenWeapon then theWeapon = theWeaponEntry
	)	
	
	-- Load that max file
	theWeaponMesh = theWeapon.weaponMesh 
	theWeaponFile = theWeapon.weaponFile 
	loadMaxFile theWeaponFile quiet:true
	append weaponComponents theWeaponMesh
	
	
	-- Sort out the attachments 
	-- First find the first instance of each attachement and match that against the attachment list 
	AAPScop = find_attachment theWeapon.attachment_Scop[1] 
	AAPSupp = find_attachment theWeapon.attachment_Supp[1]
	AAPClip = find_attachment theWeapon.attachment_Clip[1]
	AAPStck = find_attachment theWeapon.attachment_Stck[1]
	AAPLasr = find_attachment theWeapon.attachment_Lasr[1]
	AAPGrip = find_attachment theWeapon.attachment_Grip[1]
	AAPFlsh = find_attachment theWeapon.attachment_Flsh[1]
	AAPSeWp = find_attachment theWeapon.attachment_SeWp[1]
	
	
	-- Merge each item from their parent max file 
	merge_attachment AAPScop
	merge_attachment AAPSupp
	merge_attachment AAPClip
	merge_attachment AAPStck
	merge_attachment AAPLasr
	merge_attachment AAPGrip
	merge_attachment AAPFlsh
	merge_attachment AAPSeWp
	
	
	-- Store the names of the data elements that we need for the rendering 
	append weaponComponents AAPScop[2]
	append weaponComponents AAPSupp[2]
	append weaponComponents AAPClip[2]
	append weaponComponents AAPStck[2]
	append weaponComponents AAPLasr[2]
	append weaponComponents AAPGrip[2]
	append weaponComponents AAPFlsh[2]
	append weaponComponents AAPSeWp[2]
	
	
	
	-- Return the data set for the rendering part
	weaponComponents
)



-- This fucntion writes the image to disk
-- Made as a function so that I can render it or screengrab it, or whatever
fn render_weapon_image weaponRenderFolder weaponName objectName =
(
	weaponImageFile = weaponName+"_"+objectName
	
	grabbedImg = gw.getViewportDib()
	grabbedImg.filename = weaponRenderFolder + weaponImageFile + ".bmp"
	save grabbedImg
	close grabbedImg
)




-- Do the renders loop
fn render_scene_elements objectlist =
(	
	-- turn bits on and off
	for weaponIndex = 1 to objectlist.count do
	(
		weaponElement = objectlist[weaponIndex]
		
		if weaponElement != undefined then
		(	
			clearSelection()
			max hide inv
			weaponObj = getNodeByName weaponElement
			weaponObj.showVertexColors = on
			weaponObj.vertexColorType = 0	
			unhide weaponObj
			redrawViews() 
			print weaponObj.name
			
			-- Save the image 
			render_weapon_image objectlist[1] weaponObj.name
			
			
		)
	)

	
	clearSelection()
	max hide inv
	
	-- turn bits on and off
	for weaponIndex = 1 to objectlist.count do
	(
		weaponElement = objectlist[weaponIndex]
		
		if weaponElement != undefined then
		(	
			weaponObj = getNodeByName weaponElement
			unhide weaponObj
			weaponObj.showVertexColors = on
			weaponObj.vertexColorType = 0	
		)
	)
	
	redrawViews() 
	render_weapon_image objectlist[1] "complete_set"
	
)



-- Set up the scene for renders 
fn set_up_grabber chosenWeapon = 
(
	-- Check that the weapon exists here 
	
	-- Use the weapons pivot point for the camera target 
	cameraPos = [0.1,-1.5,0] 
	cameraTargetPos = [0.1,0,0] 
	
	-- set up cameras 
	weaponcam = Targetcamera fov:45 nearclip:0.01 farclip:100 nearrange:0 farrange:1000 transform:(matrix3 [1,0,0] [0,0,1] [0,-1,0] cameraPos) isSelected:on target:(Targetobject transform:(matrix3 [1,0,0] [0,1,0] [0,0,1] cameraTargetPos))
	weaponcam.orthoProjection = true 
	weaponcam.name ="Weapon_Cam"
	viewport.setCamera weaponcam	

	-- Turn off the grid	
	viewport.setGridVisibility 1 false	

	-- Enable polyview
	if viewport.isWire() then actionMan.executeAction 0 "272"
	-- Turn off wireframe overlay
	if (viewport.GetShowEdgeFaces() == true) then viewport.SetShowEdgeFaces false
		
	displayColor.shaded = #material
	

	-- Create a chroma key Plane
	chromaPlane = Plane length:4 width:4 transform:(matrix3 [1,0,0] [0,0,1] [0,-1,0] [0,2,0]) isSelected:off lengthsegs:1 widthsegs:1
	chromaPlane.wirecolor = color 0 255 0
	chromaPlane.name = "chromaPlane"

)


fn cleanup_scene =
(
	try (delete $weapon_cam) catch()
	try (delete $chromaPlane) catch()
)





-- UI Section 


try( destroyDialog weaponRenderTool )catch()
rollout weaponRenderTool "Weapon Render Tool"
(
	dotNetControl rsBannerPanel "Panel" pos:[0,0] height:32 width:weaponRenderTool.width
	local banner = makeRsBanner dn_Panel:rsBannerPanel wiki:"weaponRenderTool" filename:(getThisScriptFilename())

	on weaponRenderTool open do banner.setup()

	button btnLoadConfigData "Load Config Data" width:160 height:30 tooltip:"Load Config Data and populate UI"
	dropDownList ddlWeaponList "Weapon List" width:160 height:40 items:#() 

	button btnLoadSelectedWeapon "Load Selected Weapon" width:160 height:30 tooltip:"Loads the selected weapon and attachments"

	button btnSetUpRenderScene "Setup Render Scene" width:160 height:30 tooltip:"Setup Render Scene - camera, etc."
	
	button btnRenderScene "Setup Render Scene" width:160 height:30 tooltip:"Render weapon and attachments"

	
	-- THIS SHOULD NOT BE HARDCODED!
	local weaponConfigFile = (RsConfigGetWildwestDir() + "script/3dsMax/_config_files/weapons/weaponrenderconfig.xml")


	local weaponListUINames = #()
	local weaponListMapping = #()
	local chosenWeaponMesh = undefined
	
	
	-- Populate the data structures
	on btnLoadConfigData pressed do 	
	(
		
		load_weapon_data weaponConfigFile
		
		
		for weaponEntry in weaponRenderList do 
		(
			append weaponListUINames weaponEntry.weaponName
			append weaponListMapping #(weaponEntry.weaponName, weaponEntry.weaponMesh)
		)	
		
		--populate the ddlWeaponList
		weaponRenderTool.ddlWeaponList.items = weaponListUINames
	)

	
	-- Load the weapon and attachments
	on btnLoadSelectedWeapon pressed do 
	(
		local chosenUIweapon = weaponRenderTool.ddlWeaponList.selected
		
		-- convert the UI name into the weapon name
		chosenWeaponMesh = undefined
		for weaponIndex = 1 to weaponListUINames.count do 
		(
			if weaponListUINames[weaponIndex] == chosenUIweapon then chosenWeaponMesh =  weaponListMapping[weaponIndex][2]
		)
		
		if chosenWeaponMesh != undefined then 
		(	
			objectlist = load_weapon_meshes chosenWeaponMesh
		)	
		
		else 
		(
			messagebox "No weapon selected to load"
		)
	)	
	
	
	
	on btnSetUpRenderScene pressed do 
	(
		
		if chosenWeaponMesh != undefined then 
		(	
			set_up_grabber chosenWeapon
		)	
		
		else 
		(
			messagebox "No weapon selected/loaded"
		)
	)	
	
	
	
	on btnRenderScene pressed do 
	(
		-- hardcoded stuff to kill 
		weaponRenderFolder =  "c:/temp/weapon_renders/"
		makeDir (weaponRenderFolder)
		render_scene_elements objectlist
	)
	
	
)

createDialog weaponRenderTool 180 300