////////////////////////////////////////////////////////////////////////////////
// submesh.cpp
// Author       : Francesco Giordana
// Sponsored by : Anygma N.V. (http://www.nazooka.com)
// Start Date   : January 13, 2005
// Copyright    : (C) 2006 by Francesco Giordana
// Email        : fra.giordana@tiscali.it
////////////////////////////////////////////////////////////////////////////////

/*********************************************************************************
*                                                                                *
*   This program is free software; you can redistribute it and/or modify         *
*   it under the terms of the GNU Lesser General Public License as published by  *
*   the Free Software Foundation; either version 2 of the License, or            *
*   (at your option) any later version.                                          *
*                                                                                *
**********************************************************************************/

#include "submesh.h"

namespace OgreMayaExporter
{
	/***** Class Submesh *****/
	// constructor
	Submesh::Submesh(const MString& name)
	{
		m_pBlendShape = NULL;
		clear();
		m_name = name;
	}

	// destructor
	Submesh::~Submesh()
	{
		clear();
	}

	// clear data
	void Submesh::clear()
	{
		m_name = "";
		m_numTriangles = 0;
		m_pMaterial = NULL;
		m_indices.clear();
		m_vertices.clear();
		m_faces.clear();
		m_uvsets.clear();
		m_use32bitIndexes = false;
		m_pBlendShape = NULL;
	}

	// return number of triangles composing the mesh
	long Submesh::numTriangles()
	{
		return m_numTriangles;
	}

	// return number of vertices composing the mesh
	long Submesh::numVertices()
	{
		return m_numVertices;
	}

	// return submesh name
	MString& Submesh::name()
	{
		return m_name;
	}

/***** load data *****/
	MStatus Submesh::loadMaterial(MObject& shader,MStringArray& uvsets,ParamList& params)
	{
		MPlug plug;
		MPlugArray srcplugarray;
		bool foundShader = false;
		MStatus stat;
		MFnDependencyNode* pShader;
		//get shader from shading group
		MFnDependencyNode shadingGroup(shader);
		plug = shadingGroup.findPlug("surfaceShader");
		plug.connectedTo(srcplugarray,true,false,&stat);
		for (int i=0; i<srcplugarray.length() && !foundShader; i++)
		{
			if (srcplugarray[i].node().hasFn(MFn::kLambert) || srcplugarray[i].node().hasFn(MFn::kSurfaceShader) || 
				srcplugarray[i].node().hasFn(MFn::kPluginHwShaderNode) )
			{
				pShader = new MFnDependencyNode(srcplugarray[i].node());
				foundShader = true;
			}
		}
		if (foundShader)
		{
			std::cout << "Found material: " << pShader->name().asChar() << "\n";
			std::cout.flush();
			//check if this material has already been created
			//fix material name, adding the requested prefix
			MString tmpStr = params.matPrefix;
			if (tmpStr != "")
				tmpStr += "/";
			tmpStr += pShader->name();
			MStringArray tmpStrArray;
			tmpStr.split(':',tmpStrArray);
			MString name = "";
			for (int i=0; i<tmpStrArray.length(); i++)
			{
				name += tmpStrArray[i];
				if (i < tmpStrArray.length()-1)
					name += "_";
			}
			Material* pMaterial = MaterialSet::getSingleton().getMaterial(name);
			//if the material has already been created, update the pointer
			if (pMaterial)
				m_pMaterial = pMaterial;
			//else create it and add it to the material set
			else
			{
				pMaterial = new Material();
				pMaterial->load(pShader,uvsets,params);
				m_pMaterial = pMaterial;
				MaterialSet::getSingleton().addMaterial(pMaterial);
			}
			//delete temporary shader
			delete pShader;
		}
		else
		{
			std::cout << "Unsupported material, replacing with default lambert\n";
			std::cout.flush();
			m_pMaterial = MaterialSet::getSingleton().getDefaultMaterial();
		}
		
		//loading complete
		return MS::kSuccess;
	}

	MStatus Submesh::load(const MDagPath& dag,std::vector<face>& faces, std::vector<vertexInfo>& vertInfo, MPointArray& points, 
		MFloatVectorArray& normals, MStringArray& texcoordsets,ParamList& params,bool opposite)
	{
		//save the dag path of the maya node from which this submesh will be created
		m_dagPath = dag;
		//create the mesh Fn
		MFnMesh mesh(dag);
		std::cout << "Loading submesh : " << m_name.asChar() << "...";
		std::cout.flush();
		//save uvsets info
		for (int i=m_uvsets.size(); i<texcoordsets.length(); i++)
		{
			uvset uv;
			uv.size = 2;
			m_uvsets.push_back(uv);
		}
		//iterate over faces array, to retrieve vertices info
		std::vector<int> vtx_mapping;
		vtx_mapping.resize(vertInfo.size());
		for (int i=0; i<vtx_mapping.size(); i++)
			vtx_mapping[i] = -1;
		for (int i=0; i<faces.size(); i++)
		{
			face newFace;
			// If we are using shared geometry, indexes refer to the vertex buffer of the whole mesh
			if (params.useSharedGeom)
			{
				if(opposite)
				{	// reverse order of face vertices for correct culling
					newFace.v[0] = faces[i].v[2];
					newFace.v[1] = faces[i].v[1];
					newFace.v[2] = faces[i].v[0];
				}
				else
				{
					newFace.v[0] = faces[i].v[0];
					newFace.v[1] = faces[i].v[1];
					newFace.v[2] = faces[i].v[2];
				}
			}
			// Otherwise we create a vertex buffer for this submesh.
			// First we create a list of the vertices that we need for this submesh
			// then we get the vertex info for those vertices.
			// The indices are stored in a set so we don't get duplicates
			else
			{	// faces are triangles, so retrieve index of the three vertices
				for (int j=0; j<3; j++)
				{
					long new_idx;
					// Store the vertex index, if it hasn't already been stored in the set
					long idx = faces[i].v[j];
					if (vtx_mapping[idx] == -1)
					{
						m_indices.push_back(idx);
						new_idx = m_indices.size() - 1;
						vtx_mapping[idx] = new_idx;
					}
					else
						new_idx = vtx_mapping[idx];
					// Save the index in the set to the face info
					if (opposite)	// reverse order of face vertices to get correct culling
						newFace.v[2-j] = new_idx;
					else
						newFace.v[j] = new_idx;
				}
			}
			m_faces.push_back(newFace);
		}
		// Get the vertices info (if there are vertices to store in this submesh)
		for (int i=0; i<m_indices.size(); i++)
		{
			long vtx_idx = m_indices[i];
			vertex v;
			// get the vertex info
			vertexInfo vInfo = vertInfo[vtx_idx];
			// save vertex coordinates (rescale to desired length unit)
			MPoint point = points[vInfo.pointIdx] * params.lum;
			if (fabs(point.x) < PRECISION)
				point.x = 0;
			if (fabs(point.y) < PRECISION)
				point.y = 0;
			if (fabs(point.z) < PRECISION)
				point.z = 0;
			v.x = point.x;
			v.y = point.y;
			v.z = point.z;
			// save vertex normal
			MFloatVector normal = normals[vInfo.normalIdx];
			if (fabs(normal.x) < PRECISION)
				normal.x = 0;
			if (fabs(normal.y) < PRECISION)
				normal.y = 0;
			if (fabs(normal.z) < PRECISION)
				normal.z = 0;
			if (opposite)
			{
				v.n.x = -normal.x;
				v.n.y = -normal.y;
				v.n.z = -normal.z;
			}
			else
			{
				v.n.x = normal.x;
				v.n.y = normal.y;
				v.n.z = normal.z;
			}
			v.n.normalize();
			// save vertex color
			v.r = vInfo.r;
			v.g = vInfo.g;
			v.b = vInfo.b;
			v.a = vInfo.a;
			// save vertex bone assignements
			for (int k=0; k<vInfo.vba.size(); k++)
			{
				vba newVba;
				newVba.jointIdx = vInfo.jointIds[k];
				newVba.weight = vInfo.vba[k];
				if (newVba.weight > PRECISION)
					v.vbas.push_back(newVba);
			}
			// save texture coordinates
			for (int k=0; k<vInfo.u.size(); k++)
			{
				texcoord newTexCoords;
				newTexCoords.u = vInfo.u[k];
				newTexCoords.v = vInfo.v[k];
				newTexCoords.w = 0;
				v.texcoords.push_back(newTexCoords);
			}
			// save vertex index in maya mesh, to retrieve future positions of the same vertex
			v.index = vInfo.pointIdx;
			// add newly created vertex to vertex list
			m_vertices.push_back(v);
		}
		// set use32bitIndexes flag
		if (params.useSharedGeom || (m_vertices.size() > 65535) || (m_faces.size() > 65535))
			m_use32bitIndexes = true;
		else
			m_use32bitIndexes = false;
		// get submesh bounding box
		MPoint min = mesh.boundingBox().min();
		MPoint max = mesh.boundingBox().max();
		MBoundingBox bbox(min,max);
		if (params.exportWorldCoords)
			bbox.transformUsing(dag.inclusiveMatrix());
		min = bbox.min() * params.lum;
		max = bbox.max() * params.lum;
		MBoundingBox newbbox(min,max);
		m_boundingBox = newbbox;
		// add submesh pointer to parameters list
		params.loadedSubmeshes.push_back(this);

		std::cout << "DONE\n";
		std::cout.flush();
		return MS::kSuccess;
	}


	// Load a keyframe for this submesh
	MStatus Submesh::loadKeyframe(Track& t,float time,ParamList& params)
	{
		// create a new keyframe
		vertexKeyframe k;
		// set keyframe time
		k.time = time;
		// get the mesh Fn
		MFnMesh mesh(m_dagPath);
		// get vertex positions
		MFloatPointArray points;
		if (params.exportWorldCoords)
			mesh.getPoints(points,MSpace::kWorld);
		else
			mesh.getPoints(points,MSpace::kObject);
		// calculate vertex offsets
		for (int i=0; i<m_vertices.size(); i++)
		{
			vertexPosition pos;
			vertex v = m_vertices[i];
			pos.x = points[v.index].x * params.lum;
			pos.y = points[v.index].y * params.lum;
			pos.z = points[v.index].z * params.lum;
			if (fabs(pos.x) < PRECISION)
					pos.x = 0;
			if (fabs(pos.y) < PRECISION)
				pos.y = 0;
			if (fabs(pos.z) < PRECISION)
				pos.z = 0;
			k.positions.push_back(pos);
		}
		// add keyframe to given track
		t.addVertexKeyframe(k);
		if (params.vertBB)
		{
			// update bounding box
			MPoint min = mesh.boundingBox().min();
			MPoint max = mesh.boundingBox().max();
			MBoundingBox bbox(min,max);
			if (params.exportWorldCoords)
				bbox.transformUsing(m_dagPath.inclusiveMatrix());
			min = bbox.min() * params.lum;
			max = bbox.max() * params.lum;
			MBoundingBox newbbox(min,max);
			m_boundingBox.expand(newbbox);
		}
		// keyframe successfully loaded
		return MS::kSuccess;
	}



/***** Export data *****/
	// Write submesh data to an Ogre compatible mesh
	MStatus Submesh::createOgreSubmesh(Ogre::MeshPtr pMesh,const ParamList& params)
	{
		MStatus stat;
		// Create a new submesh
		Ogre::SubMesh* pSubmesh;
		if (m_name != "")
			pSubmesh = pMesh->createSubMesh(m_name.asChar());
		else
			pSubmesh = pMesh->createSubMesh();
		// Set material
        pSubmesh->setMaterialName(m_pMaterial->name().asChar());
        // Set use shared geometry flag
		pSubmesh->useSharedVertices = params.useSharedGeom;
		// Create vertex data for current submesh
		pSubmesh->vertexData = new Ogre::VertexData();
        // Set number of indexes
        pSubmesh->indexData->indexCount = 3*m_faces.size();
		pSubmesh->vertexData->vertexCount = m_vertices.size();
		// Check if we need to use 32 bit indexes
		bool use32BitIndexes = false;
		if (m_vertices.size() > 65536 || params.useSharedGeom)
		{
			use32BitIndexes = true;
		}
		// Create a new index buffer
		pSubmesh->indexData->indexBuffer = 
			Ogre::HardwareBufferManager::getSingleton().createIndexBuffer(
				use32BitIndexes ? Ogre::HardwareIndexBuffer::IT_32BIT : Ogre::HardwareIndexBuffer::IT_16BIT,
				pSubmesh->indexData->indexCount,
				Ogre::HardwareBuffer::HBU_STATIC_WRITE_ONLY);
		// Fill the index buffer with faces data
		if (use32BitIndexes)
        {
			Ogre::uint32* pIdx = static_cast<Ogre::uint32*>(
				pSubmesh->indexData->indexBuffer->lock(Ogre::HardwareBuffer::HBL_DISCARD));
			for (int i=0; i<m_faces.size(); i++)
			{
				*pIdx++ = static_cast<Ogre::uint32>(m_faces[i].v[0]);
				*pIdx++ = static_cast<Ogre::uint32>(m_faces[i].v[1]);
				*pIdx++ = static_cast<Ogre::uint32>(m_faces[i].v[2]);
			}
			pSubmesh->indexData->indexBuffer->unlock();
        }
        else
        {
            Ogre::uint16* pIdx = static_cast<Ogre::uint16*>(
				pSubmesh->indexData->indexBuffer->lock(Ogre::HardwareBuffer::HBL_DISCARD));
            for (int i=0; i<m_faces.size(); i++)
			{
				*pIdx++ = static_cast<Ogre::uint16>(m_faces[i].v[0]);
				*pIdx++ = static_cast<Ogre::uint16>(m_faces[i].v[1]);
				*pIdx++ = static_cast<Ogre::uint16>(m_faces[i].v[2]);
			}
			pSubmesh->indexData->indexBuffer->unlock();
		}
		// Define vertex declaration (only if we're not using shared geometry)
		if(!params.useSharedGeom)
		{
			Ogre::VertexDeclaration* pDecl = pSubmesh->vertexData->vertexDeclaration;
			unsigned buf = 0;
			size_t offset = 0;
			// Add vertex position
			pDecl->addElement(buf, offset, Ogre::VET_FLOAT3, Ogre::VES_POSITION);
			offset += Ogre::VertexElement::getTypeSize(Ogre::VET_FLOAT3);
			// Add vertex normal
			if (params.exportVertNorm)
			{
				pDecl->addElement(buf, offset, Ogre::VET_FLOAT3, Ogre::VES_NORMAL);
				offset += Ogre::VertexElement::getTypeSize(Ogre::VET_FLOAT3);
			}
			// Add vertex colour
			if(params.exportVertCol)
			{
				pDecl->addElement(buf, offset, Ogre::VET_COLOUR, Ogre::VES_DIFFUSE);
				offset += Ogre::VertexElement::getTypeSize(Ogre::VET_COLOUR);
			}
			// Add texture coordinates
			for (int i=0; i<m_vertices[0].texcoords.size(); i++)
			{
				Ogre::VertexElementType uvType = Ogre::VertexElement::multiplyTypeCount(Ogre::VET_FLOAT1, 2);
				pDecl->addElement(buf, offset, uvType, Ogre::VES_TEXTURE_COORDINATES, i);
				offset += Ogre::VertexElement::getTypeSize(uvType);
			}
			Ogre::VertexDeclaration* pOptimalDecl = pDecl->getAutoOrganisedDeclaration(
				params.exportVBA,params.exportBlendShapes || params.exportVertAnims);
			// Fill the vertex buffer using the newly created vertex declaration
			stat = createOgreVertexBuffer(pSubmesh,pDecl,m_vertices);
			// Write vertex bone assignements list
			if (params.exportVBA)
			{
				// Create a new vertex bone assignements list
				Ogre::SubMesh::VertexBoneAssignmentList vbas;
				// Scan list of shared geometry vertices
				for (int i=0; i<m_vertices.size(); i++)
				{
					vertex v = m_vertices[i];
					// Add all bone assignemnts for every vertex to the bone assignements list
					for (int j=0; j<v.vbas.size(); j++)
					{
						Ogre::VertexBoneAssignment vba;
						vba.vertexIndex = i;
						vba.boneIndex = v.vbas[j].jointIdx;
						vba.weight = v.vbas[j].weight;
						vbas.insert(Ogre::SubMesh::VertexBoneAssignmentList::value_type(i, vba));
					}
				}
				// Rationalise the bone assignements list
				pSubmesh->parent->_rationaliseBoneAssignments(pSubmesh->vertexData->vertexCount,vbas);
				// Add bone assignements to the submesh
				for (Ogre::SubMesh::VertexBoneAssignmentList::iterator bi = vbas.begin(); bi != vbas.end(); bi++)
				{
					pSubmesh->addBoneAssignment(bi->second);
				}
				pSubmesh->_compileBoneAssignments();
			}
			pSubmesh->vertexData->reorganiseBuffers(pOptimalDecl);
		}
		return MS::kSuccess;
	}


	// Create an Ogre compatible vertex buffer
	MStatus Submesh::createOgreVertexBuffer(Ogre::SubMesh* pSubmesh,Ogre::VertexDeclaration* pDecl,const std::vector<vertex>& vertices)
	{
		Ogre::HardwareVertexBufferSharedPtr vbuf = 
			Ogre::HardwareBufferManager::getSingleton().createVertexBuffer(pDecl->getVertexSize(0),
			pSubmesh->vertexData->vertexCount, 
			Ogre::HardwareBuffer::HBU_STATIC_WRITE_ONLY);
		pSubmesh->vertexData->vertexBufferBinding->setBinding(0, vbuf);
		size_t vertexSize = pDecl->getVertexSize(0);
		char* pBase = static_cast<char*>(vbuf->lock(Ogre::HardwareBuffer::HBL_DISCARD));
		Ogre::VertexDeclaration::VertexElementList elems = pDecl->findElementsBySource(0);
		Ogre::VertexDeclaration::VertexElementList::iterator ei, eiend;
		eiend = elems.end();
		float* pFloat;
		Ogre::RGBA* pRGBA;
		// Fill the vertex buffer with shared geometry data
		Ogre::ColourValue col;
		float ucoord, vcoord;
		for (long vi=0; vi<vertices.size(); vi++)
		{
			int iTexCoord = 0;
			vertex v = vertices[vi];
			for (ei = elems.begin(); ei != eiend; ++ei)
			{
				Ogre::VertexElement& elem = *ei;
				switch(elem.getSemantic())
				{
				case Ogre::VES_POSITION:
					elem.baseVertexPointerToElement(pBase, &pFloat);
					*pFloat++ = v.x;
					*pFloat++ = v.y;
					*pFloat++ = v.z;
					break;
				case Ogre::VES_NORMAL:
					elem.baseVertexPointerToElement(pBase, &pFloat);
					*pFloat++ = v.n.x;
					*pFloat++ = v.n.y;
					*pFloat++ = v.n.z;
					break;
				case Ogre::VES_DIFFUSE:
					{
						elem.baseVertexPointerToElement(pBase, &pRGBA);
						Ogre::ColourValue col(v.r, v.g, v.b, v.a);
						*pRGBA = Ogre::VertexElement::convertColourValue(col, 
							Ogre::VertexElement::getBestColourVertexElementType());
					}
					break;
				case Ogre::VES_TEXTURE_COORDINATES:
					elem.baseVertexPointerToElement(pBase, &pFloat);
					ucoord = v.texcoords[iTexCoord].u;
					vcoord = v.texcoords[iTexCoord].v;
					*pFloat++ = ucoord;
					*pFloat++ = vcoord;
					iTexCoord++;
					break;
				}
			}
			pBase += vertexSize;
		}
		vbuf->unlock();
		return MS::kSuccess;
	}

}; //end of namespace