path: root/NW4RTools/ColladaWriter.cs

using System;
using System.IO;
using System.Collections.Generic;
using System.Text;
using NW4RTools.Models;
using Collada141;

namespace NW4RTools {
	public class ColladaWriter {
		public static void WriteModel(Stream outputStream, ResFile file, string modelName) {
			new ColladaWriter(file).SaveModel(outputStream, modelName);
		}



		ResFile CurrentFile;
		Models.Model CurrentModel;
		COLLADA Collada;

		library_geometries LibGeometries;
		library_visual_scenes LibVisualScenes;

		private ColladaWriter(ResFile file) {
			CurrentFile = file;
		}

		public void SaveModel(Stream outputStream, string modelName) {
			CurrentModel = CurrentFile.GetGroup<Model>("3DModels(NW4R)")[modelName];
			Collada = new COLLADA();
			
			Collada.asset = new asset();
			
			Collada.asset.contributor = new assetContributor[1];
			Collada.asset.contributor[0] = new assetContributor();
			Collada.asset.contributor[0].authoring_tool = "NW4RTools Collada exporter by Treeki";
			Collada.asset.contributor[0].source_data = "NW4R model: " + modelName;
			Collada.asset.created = DateTime.Now;
			Collada.asset.modified = DateTime.Now;
			Collada.asset.unit = new assetUnit();
			Collada.asset.unit.meter = 1.0;
			Collada.asset.up_axis = UpAxisType.Y_UP;
			
			List<object> ColladaItems = new List<object>();
			
			LibGeometries = new library_geometries();
			ColladaItems.Add(LibGeometries);
			
			var GeometryList = new List<geometry>();
			foreach (var kv in CurrentModel.Shapes) {
				GeometryList.Add(CreateGeometryFromShape(kv.Key, kv.Value));
			}
			
			LibGeometries.geometry = GeometryList.ToArray();

			// SHAPES ARE DONE.
			// Next up: Visual Scenes (I will just create one atm)
			
			LibVisualScenes = new library_visual_scenes();
			ColladaItems.Add(LibVisualScenes);
			
			LibVisualScenes.visual_scene = new visual_scene[1];
			var mainScene = LibVisualScenes.visual_scene[0] = new visual_scene();
			
			// TODO: Change this so it doesn't have the possibility of name collisions with shapes
			mainScene.id = "RootNode";
			mainScene.name = "RootNode";
			var mainSceneNodeList = new List<node>();

			// OK, so here's what's up: first off, we must create a definition for every node
			var NodeDefs = new Dictionary<Node, node>();

			foreach (var kv in CurrentModel.Nodes) {
				string nodeName = kv.Key;
				Node origNode = kv.Value;
				var cNode = new node();

				cNode.id = nodeName;
				cNode.name = nodeName;
				//cNode.type = NodeType.JOINT;
				cNode.node1 = new node[0];

				/*var origMatrix = origNode.NodeMatrix;
				var nodeMatrix = new matrix();
				nodeMatrix.Values = new double[] {
					origMatrix.v00, origMatrix.v01, origMatrix.v02, origMatrix.v03,
					origMatrix.v10, origMatrix.v11, origMatrix.v12, origMatrix.v13,
					origMatrix.v20, origMatrix.v21, origMatrix.v22, origMatrix.v23,
					0.0, 0.0, 0.0, 1.0
				};

				cNode.Items = new object[] { nodeMatrix };
				cNode.ItemsElementName = new ItemsChoiceType2[] { ItemsChoiceType2.matrix };*/

				var tScale = new TargetableFloat3();
				var tRotate = new rotate();
				var tTranslate = new TargetableFloat3();

				tScale.Values = new double[] {
					origNode.Scale.x, origNode.Scale.y, origNode.Scale.z
				};

				tRotate.Values = new double[] {
					0, 0, 0, 0 // TODO
				};

				tTranslate.Values = new double[] {
					origNode.Translation.x, origNode.Translation.y, origNode.Translation.z
				};

				cNode.Items = new object[] { tScale, tRotate, tTranslate };
				cNode.ItemsElementName = new ItemsChoiceType2[] {
					ItemsChoiceType2.scale, ItemsChoiceType2.rotate,
					ItemsChoiceType2.translate
				};

				NodeDefs[origNode] = cNode;
			}

			// Now add them to the hierarchy
			foreach (var kv in NodeDefs) {
				Node origNode = kv.Key;
				node cNode = kv.Value;

				if (origNode.Parent == null) {
					mainSceneNodeList.Add(cNode);
				} else {
					var parentNode = NodeDefs[origNode.Parent];

					// this is stupid, thanks C#
					node[] nodeArrayCopy = parentNode.node1;
					Array.Resize<node>(ref nodeArrayCopy, nodeArrayCopy.Length + 1);
					nodeArrayCopy[nodeArrayCopy.Length - 1] = cNode;
					parentNode.node1 = nodeArrayCopy;
				}
			}

			// Apply shapes to nodes
			foreach (var kv in CurrentModel.Shapes) {
				Shape shape = kv.Value;
				Node origNode = CurrentModel.Nodes[CurrentModel.MatrixIDtoNodeID[shape.MatrixID]];
				node cNode = NodeDefs[origNode];

				var newGeoEntry = new instance_geometry();
				newGeoEntry.name = kv.Key;
				newGeoEntry.url = String.Format("#{0}-lib", kv.Key);

				instance_geometry[] geoArrayCopy = cNode.instance_geometry;
				if (geoArrayCopy == null)
					geoArrayCopy = new instance_geometry[1];
				else
					Array.Resize<instance_geometry>(ref geoArrayCopy, geoArrayCopy.Length + 1);
				geoArrayCopy[geoArrayCopy.Length - 1] = newGeoEntry;
				cNode.instance_geometry = geoArrayCopy;
			}

			/*foreach (var kv in CurrentModel.Shapes) {
				var thisNode = new node();

				thisNode.id = kv.Key;
				thisNode.name = kv.Key;
				thisNode.instance_geometry = new instance_geometry[1];
				thisNode.instance_geometry[0] = new instance_geometry();
				thisNode.instance_geometry[0].url = String.Format("#{0}-lib", kv.Key);

				mainSceneNodeList.Add(thisNode);
			}*/
			
			mainScene.node = mainSceneNodeList.ToArray();

			
			// Finally, create a scene
			Collada.scene = new COLLADAScene();
			Collada.scene.instance_visual_scene = new InstanceWithExtra();
			Collada.scene.instance_visual_scene.url = "#RootNode";
			
			Collada.Items = ColladaItems.ToArray();
			Collada.Save(outputStream);
		}





		private geometry CreateGeometryFromShape(string name, Shape shape) {
			var geo = new geometry();
			
			geo.id = name + "-lib";
			geo.name = name + "Mesh";
			
			var m = new mesh();
			geo.Item = m;
			
			// Vertex settings
			var firstDL = new InputStream(shape.DisplayList1);
			firstDL.Seek(0x0C);
			UInt32 vtxDesc1 = firstDL.ReadUInt32();
			firstDL.Seek(0x12);
			UInt32 vtxDesc2 = firstDL.ReadUInt32();
			firstDL.Seek(0x22);
			UInt32 vtxAttr1 = firstDL.ReadUInt32();
			firstDL.Seek(0x28);
			UInt32 vtxAttr2 = firstDL.ReadUInt32();
			firstDL.Seek(0x2E);
			UInt32 vtxAttr3 = firstDL.ReadUInt32();
			
			var vs = new VertexSettings();
			vs.SetDesc(vtxDesc1, vtxDesc2);
			vs.SetAttrFmt(vtxAttr1, vtxAttr2, vtxAttr3);
			
			// Figure out how many elements we need in the Source array
			// Position data ALWAYS exists
			int sourceCount = 1;
			sourceCount += (shape.NrmData != null) ? 1 : 0;
			for (int i = 0; i < 8; i++)
				sourceCount += (shape.TexCoordData[i] != null) ? 1 : 0;
			
			m.source = new source[sourceCount];
			int currentSource = 0;
			
			// TODO: Refactor this messy code!
			
			int dest;
			
			// Write position data
			var posData = shape.PosData;
			var posSource = new source();
			posSource.id = name + "-lib-Position";
			m.source[currentSource++] = posSource;
			
			var posArray = new float_array();
			posArray.id = name + "-lib-Position-array";
			posArray.count = (ulong)(posData.GetRealCount() * posData.EntryCount);
			posArray.Values = new double[posArray.count];
			
			dest = 0;
			for (int i = 0; i < posData.EntryCount; i++) {
				float[] data = posData.GetEntry(i);
				for (int j = 0; j < data.Length; j++) {
					posArray.Values[dest++] = data[j];
				}
			}
			
			posSource.Item = posArray;
			
			// Write position technique
			posSource.technique_common = new sourceTechnique_common();
			var posAcc = posSource.technique_common.accessor = new accessor();
			posAcc.source = String.Format("#{0}-lib-Position-array", name);
			posAcc.count = posData.EntryCount;
			posAcc.stride = (ulong)posData.GetRealCount();
			
			posAcc.param = new param[posData.GetRealCount()];
			string[] posParamNames = new string[] { "X", "Y", "Z" };
			
			for (int i = 0; i < posAcc.param.Length; i++) {
				posAcc.param[i] = new param();
				posAcc.param[i].name = posParamNames[i];
				posAcc.param[i].type = "float";
			}
			
			
			// Write normal data
			if (shape.NrmData != null) {
				var nrmData = shape.NrmData;
				var nrmSource = new source();
				nrmSource.id = name + "-lib-Normal";
				m.source[currentSource++] = nrmSource;
				
				var nrmArray = new float_array();
				nrmArray.id = name + "-lib-Normal-array";
				nrmArray.count = (ulong)(nrmData.GetRealCount() * nrmData.EntryCount);
				nrmArray.Values = new double[nrmArray.count];
				
				dest = 0;
				for (int i = 0; i < nrmData.EntryCount; i++) {
					float[] data = nrmData.GetEntry(i);
					for (int j = 0; j < data.Length; j++) {
						nrmArray.Values[dest++] = data[j];
					}
				}

				nrmSource.Item = nrmArray;
				
				// Write normal technique
				nrmSource.technique_common = new sourceTechnique_common();
				var nrmAcc = nrmSource.technique_common.accessor = new accessor();
				nrmAcc.source = String.Format("#{0}-lib-Normal-array", name);
				nrmAcc.count = nrmData.EntryCount;
				nrmAcc.stride = (ulong)nrmData.GetRealCount();
				
				nrmAcc.param = new param[nrmData.GetRealCount()];
				string[] nrmParamNames = new string[] { "X", "Y", "Z" };
				
				for (int i = 0; i < nrmAcc.param.Length; i++) {
					nrmAcc.param[i] = new param();
					nrmAcc.param[i].name = nrmParamNames[i];
					nrmAcc.param[i].type = "float";
				}
			}
			
			
			// Write TexCoord data
			for (int tcIndex = 0; tcIndex < 8; tcIndex++) {
				if (shape.TexCoordData[tcIndex] != null) {
					var tcData = shape.TexCoordData[tcIndex];
					var tcSource = new source();
					tcSource.id = String.Format("{0}-lib-TexCoord{1}", name, tcIndex);
					m.source[currentSource++] = tcSource;
					
					var tcArray = new float_array();
					tcArray.id = String.Format("{0}-lib-TexCoord{1}-array", name, tcIndex);
					tcArray.count = (ulong)(tcData.GetRealCount() * tcData.EntryCount);
					tcArray.Values = new double[tcArray.count];
					
					dest = 0;
					for (int i = 0; i < tcData.EntryCount; i++) {
						float[] data = tcData.GetEntry(i);
						for (int j = 0; j < data.Length; j++) {
							tcArray.Values[dest++] = data[j];
						}
					}
					
					tcSource.Item = tcArray;
					
					// Write texcoord technique
					tcSource.technique_common = new sourceTechnique_common();
					var tcAcc = tcSource.technique_common.accessor = new accessor();
					tcAcc.source = String.Format("#{0}-lib-TexCoord{1}-array", name, tcIndex);
					tcAcc.count = tcData.EntryCount;
					tcAcc.stride = (ulong)tcData.GetRealCount();
					
					tcAcc.param = new param[tcData.GetRealCount()];
					string[] tcParamNames = new string[] { "S", "T" };
					
					for (int i = 0; i < tcAcc.param.Length; i++) {
						tcAcc.param[i] = new param();
						tcAcc.param[i].name = tcParamNames[i];
						tcAcc.param[i].type = "float";
					}
				}
			}
			
			
			
			// Ok, we've written all the raw float data, now set up vertices
			// TODO: Vertex colours
			m.vertices = new vertices();
			m.vertices.id = String.Format("{0}-lib-Vertex", name);
			m.vertices.input = new InputLocal[1];
			m.vertices.input[0] = new InputLocal();
			m.vertices.input[0].semantic = "POSITION";
			m.vertices.input[0].source = String.Format("#{0}-lib-Position", name);
			
			// And before we finish, write the polygon data of course
			var dl = new InputStream(shape.DisplayList2);
			
			List<object> meshItems = new List<object>();
			
			// create the Input array -- we can reuse sourceCount!
			var inputArray = new InputLocalOffset[sourceCount];
			currentSource = 0;
			
			var posInput = inputArray[currentSource] = new InputLocalOffset();
			posInput.semantic = "VERTEX";
			posInput.offset = (ulong)currentSource;
			posInput.source = String.Format("#{0}-lib-Vertex", name);
			currentSource++;
			
			if (shape.NrmData != null) {
				var nrmInput = inputArray[currentSource] = new InputLocalOffset();
				nrmInput.semantic = "NORMAL";
				nrmInput.offset = (ulong)currentSource;
				nrmInput.source = String.Format("#{0}-lib-Normal", name);
				currentSource++;
			}
			
			for (int i = 0; i < 8; i++) {
				if (shape.TexCoordData[i] != null) {
					var tcInput = inputArray[currentSource] = new InputLocalOffset();
					tcInput.semantic = "TEXCOORD";
					tcInput.offset = (ulong)currentSource;
					tcInput.@set = (ulong)i;
					tcInput.source = String.Format("#{0}-lib-TexCoord{1}", name, i);
					currentSource++;
				}
			}
			
			
			// Create a list for tristrips beforehand, because they're THE most common
			List<string> triStrips = new List<string>();
			
			
			// Now go through the display list
			while (true) {
				if (dl.AtEnd)
					break;
				
				byte cmd = dl.ReadByte();
				if (cmd == 0)
					break;
				
				PrimitiveType prim = (PrimitiveType)((cmd >> 3) & 7);
				int vtxCount = dl.ReadUInt16();
				
				// first, parse it into a list of vertices
				GXIndexedVertex[] vtxs = new GXIndexedVertex[vtxCount];
				string[] pVtxs = new string[vtxCount];
				
				for (int i = 0; i < vtxCount; i++) {
					vtxs[i].LoadFrom(dl, vs);
					
					pVtxs[i] = vtxs[i].Position.ToString();
					
					if (vs.NormalDesc != VertexSettings.DescType.None)
						pVtxs[i] += " " + vtxs[i].Normal.ToString();
					
					for (int j = 0; j < 8; j++) {
						if (vs.TexCoordDesc[j] != VertexSettings.DescType.None) {
							pVtxs[i] += " " + vtxs[i].TexCoords[j].ToString();
						}
					}
				}
				
				switch (prim) {
				case PrimitiveType.Triangles:
					var pTri = new triangles();
					pTri.count = (ulong)(vtxCount / 3);
					// should be 1? dunno
					pTri.input = inputArray;
					
					StringBuilder pTriData = new StringBuilder();
					
					for (int i = 0; i < vtxCount; i++) {
						pTriData.AppendFormat("{0} ", pVtxs[i]);
					}

					
					pTri.p = pTriData.ToString();
					
					meshItems.Add(pTri);
					break;
				
				case PrimitiveType.TriangleStrip:
					StringBuilder pTriStripData = new StringBuilder();
					
					for (int i = 0; i < vtxCount; i++) {
						pTriStripData.AppendFormat("{0} ", pVtxs[i]);
					}

					
					triStrips.Add(pTriStripData.ToString());
					break;
				default:
					
					
					Console.WriteLine("UNIMPLEMENTED PRIMITIVE TYPE");
					return geo;
				}
			}
			
			
			// If any tristrips were found, add them!
			if (triStrips.Count > 0) {
				var pTriStrips = new tristrips();
				pTriStrips.input = inputArray;
				pTriStrips.count = (ulong)triStrips.Count;
				pTriStrips.p = triStrips.ToArray();
				meshItems.Add(pTriStrips);
			}
			
			
			m.Items = meshItems.ToArray();
			
			// FINALLY DONE!
			return geo;
		}
	}
}