#include "dxstdafx.h"
#include ".\meshexploder.h"
#include "TriangleMesh.h"
#include "Vector.hpp"
#include "L.h"
#include <vector>
#include <fstream>
#include "xmlwriter.h"
#include "FlexVertex.h"

using namespace xmlw;

MeshExploder::MeshExploder(LPDIRECT3DDEVICE9 device, LPD3DXMESH originalMesh, D3DXMATERIAL* originalMaterials, wchar_t** materialNames, D3DXEFFECTINSTANCE* originalShaders, unsigned int nPieces)
{
	this->device = device;
	this->originalMesh = originalMesh;
	this->originalMaterials = originalMaterials;
	this->originalShaders = originalShaders;
	this->nPieces = nPieces;

	this->materialNames = materialNames;

	fragmentMeshes = NULL;
	fragmentAttribIndices = NULL;
	composedMaterials = NULL;
	composedShaders = NULL;
	composedMaterialNames = NULL;
}

MeshExploder::~MeshExploder(void)
{
	if(fragmentMeshes)
	{
		for(int f=0; f<nPieces; f++)
		{
			fragmentMeshes[f]->Release();
		}
		delete [] fragmentMeshes;
	}
	if(fragmentAttribIndices)
		delete [] fragmentAttribIndices;

	//we need these as output
/*	if(composedMaterials)
		delete [] composedMaterials;
	if(composedShaders)
		delete [] composedShaders;*/

	if(composedMaterialNames)
		delete [] composedMaterialNames;
}

struct Tri
{
	unsigned short origMeshIndex;
	Vector pos;
	Vector normal;
};

void MeshExploder::explode()
{
	fragmentMeshes = new LPD3DXMESH[nPieces];
	fragmentAttribIndices = new DWORD[nPieces];

	originalMesh->OptimizeInplace(D3DXMESHOPT_ATTRSORT, NULL, NULL, NULL, NULL);

	LPDIRECT3DVERTEXBUFFER9 vertexBuffer;
	originalMesh->GetVertexBuffer(&vertexBuffer);

	FlexVertexArray vertexData(originalMesh->GetNumBytesPerVertex());
	D3DVERTEXELEMENT9 elem[64];	originalMesh->GetDeclaration(elem);	 vertexData.setFormat(elem);
	vertexBuffer->Lock(0,originalMesh->GetNumVertices()*originalMesh->GetNumBytesPerVertex(),vertexData.getDataPointerReference(),0);

	LPDIRECT3DINDEXBUFFER9 indexBuffer;
	originalMesh->GetIndexBuffer(&indexBuffer);
	
	DWORD nSubsets;
	originalMesh->GetAttributeTable(NULL, &nSubsets);
	D3DXATTRIBUTERANGE* subsetRanges = new D3DXATTRIBUTERANGE[nSubsets];
	originalMesh->GetAttributeTable(subsetRanges, &nSubsets);

	unsigned int piecesToGo =  nPieces;
	unsigned int facesToGo =  originalMesh->GetNumFaces();
	int iPiece=0;

	for(int iSubset=0; iSubset<nSubsets; iSubset++)
	{

		unsigned int nTris = subsetRanges[iSubset].FaceCount;

		unsigned int pcsPerSubset = piecesToGo * nTris / facesToGo;
		while(piecesToGo - pcsPerSubset < nSubsets - 1 - iSubset )
			pcsPerSubset--;
		if(pcsPerSubset < 1)
			pcsPerSubset = 1;
		piecesToGo -= pcsPerSubset;
		facesToGo -= nTris;

		unsigned short* indexData;
		indexBuffer->Lock(subsetRanges[iSubset].FaceStart*3*sizeof(unsigned short),
						nTris*3*sizeof(unsigned short),
						(void**)&indexData,0);

		Tri* tris = new Tri[nTris];
		//partition to small meshes
		for(int g=0; g<nTris; g++)
		{
			tris[g].origMeshIndex = g;
			tris[g].pos =	(vertexData[indexData[g * 3]].pos() + 
							vertexData[indexData[g * 3 + 1]].pos() + 
							vertexData[indexData[g * 3 + 2]].pos()) * 0.33333333333333;
			tris[g].normal =	(vertexData[indexData[g * 3]].normal() + 
							vertexData[indexData[g * 3 + 1]].normal() + 
							vertexData[indexData[g * 3 + 2]].normal());
			tris[g].normal.normalize();
		}

		int* slabNTris = new int[pcsPerSubset];
		clusterTrisKMeans(tris, nTris, pcsPerSubset, slabNTris);
		int triClusterStart = 0;

		for(int ict=0; ict<pcsPerSubset; ict++)
		{
			// count vertices
			unsigned short* ill = new unsigned short[originalMesh->GetNumVertices()];
			for(int zzz=0; zzz<originalMesh->GetNumVertices(); zzz++)
				ill[zzz] = 0;
			for(int rmi = triClusterStart; rmi<triClusterStart + slabNTris[ict]; rmi++)
			{
				ill[indexData[tris[rmi].origMeshIndex*3]] = 1;
				ill[indexData[tris[rmi].origMeshIndex*3+1]] = 1;
				ill[indexData[tris[rmi].origMeshIndex*3+2]] = 1;
			}

			int nVertices = 0;
			for(int ccc=0; ccc<originalMesh->GetNumVertices(); ccc++)
				if(ill[ccc])
					nVertices++;

			LPD3DXMESH fragMesh;
			// create new mesh
			HRESULT dhre = D3DXCreateMeshFVF(
				slabNTris[ict],
				nVertices,
				0,
				originalMesh->GetFVF(),
				device,
				&fragMesh);
			// fill with tris
			LPDIRECT3DVERTEXBUFFER9 svertexBuffer;
			fragMesh->GetVertexBuffer(&svertexBuffer);
			FlexVertexArray svertexData(fragMesh->GetNumBytesPerVertex());
			svertexBuffer->Lock(0, nVertices * fragMesh->GetNumBytesPerVertex(),svertexData.getDataPointerReference(),0);
			LPDIRECT3DINDEXBUFFER9 sindexBuffer;
			fragMesh->GetIndexBuffer(&sindexBuffer);
			unsigned short* sindexData;
			sindexBuffer->Lock(0,slabNTris[ict]*3*sizeof(unsigned short),(void**)&sindexData,0);

			int ufi = 0;
			for(int ifi=0; ifi<originalMesh->GetNumVertices(); ifi++)
				if(ill[ifi])
				{
					//svertexData[ufi] = vertexData[ifi];
					svertexData.assign(ufi, vertexData[ifi]);
					ill[ifi] = ufi;
					ufi++;
				}
				else
					ill[ifi] = 0xffff;

			for(int trifi=0; trifi<slabNTris[ict]; trifi++)
			{
				sindexData[trifi*3] = ill[indexData[tris[trifi+triClusterStart].origMeshIndex*3]];
				sindexData[trifi*3+1] = ill[indexData[tris[trifi+triClusterStart].origMeshIndex*3+1]];
				sindexData[trifi*3+2] = ill[indexData[tris[trifi+triClusterStart].origMeshIndex*3+2]];
			}

			svertexBuffer->Unlock();
			sindexBuffer->Unlock();
			svertexBuffer->Release();
			sindexBuffer->Release();

			triClusterStart += slabNTris[ict];

			fragmentMeshes[iPiece] = fragMesh;
			fragmentAttribIndices[iPiece] = iSubset;
			iPiece++;
		}

		indexBuffer->Unlock();
	}

	vertexBuffer->Unlock();
	vertexBuffer->Release();
	indexBuffer->Release();

	delete subsetRanges;
}


void MeshExploder::clusterTrisKMeans(Tri* triArray, int nTris, int nClusters, int* clusterLenghts)
{
	Tri* centroids = new Tri[nClusters];
	std::vector<Tri>* result = new std::vector<Tri>[nClusters];

	//initial clusters: uniform length
	for(unsigned int i=0; i<nClusters; i++)
		clusterLenghts[i] = nTris / nClusters;
	clusterLenghts[nClusters - 1] = nTris - (nClusters - 1) * (nTris / nClusters);
	for(unsigned int iKMeans = 0; iKMeans < 128; iKMeans++)
	{
		//find centroids
		unsigned int iTri = 0;
		for(unsigned int i=0; i<nClusters; i++)
		{
			centroids[i].pos = Vector::RGBBLACK;
			centroids[i].normal = Vector::RGBBLACK;
			for(unsigned int iTriInCluster=0; iTriInCluster < clusterLenghts[i]; iTriInCluster++, iTri++)
			{
				centroids[i].pos += triArray[iTri].pos;
				centroids[i].normal += triArray[iTri].normal;
			}
			centroids[i].pos *= 1.0f / (float)clusterLenghts[i];
			centroids[i].normal.normalize();
		}
		
		for(unsigned int i=0; i<nClusters; i++)
			result[i].clear();

		//sort radions to centroids
		iTri = 0;
		for(unsigned int i=0; i<nClusters; i++)
		{
			for(unsigned int iTriInCluster=0; iTriInCluster < clusterLenghts[i]; iTriInCluster++, iTri++)
			{
				unsigned int minimumDistanceClusterIndex = 0;
				float minimumDistance = FLT_MAX;
				for(unsigned int u=0; u<nClusters; u++)
				{
					float dirDist = (triArray[iTri].normal * centroids[u].normal);
					float dist = (triArray[iTri].pos - centroids[u].pos).norm() * pow(2.0, 1.0 - (double)dirDist);
					if(dist < minimumDistance)
					{
						minimumDistanceClusterIndex = u;
						minimumDistance = dist;
					}
				}
				result[minimumDistanceClusterIndex].push_back( triArray[iTri] );
			}
		}
		//refill bushStarters, set cluster lengths
		iTri = 0;
		for(unsigned int i=0; i<nClusters; i++)
		{
			clusterLenghts[i] = result[i].size();
			for(unsigned int iTriInCluster=0; iTriInCluster < clusterLenghts[i]; iTriInCluster++, iTri++)
			{
				triArray[iTri] = result[i].at(iTriInCluster);
			}
		}
		//eliminate empty clusters
		for(unsigned int i=1; i<nClusters; i++)
		{
			if(clusterLenghts[i] == 0)
			{
				clusterLenghts[i] = clusterLenghts[i-1] >> 1;
				clusterLenghts[i-1] -= clusterLenghts[i-1] >> 1;
			}
		}
		for(int i=nClusters - 1; i >= 0; i--)
		{
			if(clusterLenghts[i] == 0)
			{
				clusterLenghts[i] = clusterLenghts[i+1] >> 1;
				clusterLenghts[i+1] -= clusterLenghts[i+1] >> 1;
			}
		}
	}
	delete [] centroids;
	delete [] result;
}

void MeshExploder::saveFragMeshes()
{
	for(int i=0; i<nPieces; i++)
	{
		char sname[256];
		sprintf(sname, "media/explode_fragment%d.x", i);
		HRESULT hr = D3DXSaveMeshToX(L::l+sname,	//LPCTSTR pFilename,
						fragmentMeshes[i],	//LPD3DXMESH pMesh,
						NULL,	//CONST DWORD * pAdjacency,
						NULL,	//CONST D3DXMATERIAL * pMaterials,
						NULL,	//CONST D3DXEFFECTINSTANCE * pEffectInstances,
						0,	//DWORD NumMaterials,
  						D3DXF_FILEFORMAT_TEXT	//DWORD Format
						);
	}
}

void MeshExploder::generateAtlas(int atlasSize)
{
	for(int i=0; i<nPieces; i++)
	{

		LPD3DXMESH tempMesh;

		float maxStretchApplied;
		unsigned int nCharts;

		unsigned int atlasIndex = (originalMesh->GetFVF() & D3DFVF_TEXCOUNT_MASK) / D3DFVF_TEX1 - 1;

		HRESULT hr =
			D3DXUVAtlasCreate(
				fragmentMeshes[i],		//	LPD3DXMESH pMesh,
				0,			//	UINT dwMaxChartNumber,
				1.0f,		//	FLOAT fMaxStretch,
				atlasSize - 4,		//	UINT dwWidth,
				atlasSize - 4,		//	UINT dwHeight,
				2.0f,		//	FLOAT fGutter,
				atlasIndex,	//	DWORD dwTextureIndex,
				NULL,		//	CONST DWORD * pdwAdjacency,
				NULL,		//	CONST CHAR * pcFalseEdges,
				NULL,		//	FLOAT * pfIMTArray,
				NULL,		//	LPD3DXUVATLASCB pCallback,
				0.001f,		//	FLOAT fCallbackFrequency,
				NULL,		//	LPVOID pUserContent,
				0,	// options
				&tempMesh,	//	LPD3DXMESH * ppMeshOut,
				NULL,		//	LPD3DXBUFFER * ppFacePartitioning,
				NULL,		//	LPD3DXBUFFER * ppVertexRemapArray,
				&maxStretchApplied,	//	FLOAT * pfMaxStretchOut,
				&nCharts			//	UINT * pdwNumChartsOut
			);

		if(hr == S_OK)
		{
			//forget the old mesh
			fragmentMeshes[i]->Release();
			fragmentMeshes[i] = tempMesh;
		}
		else
		{
			MessageBoxA( NULL, "Sorry!", "Impossible to create UVatlas, please increase PRM resolution!", MB_OK);
		}

		LPDIRECT3DVERTEXBUFFER9 vertexBuffer;
		fragmentMeshes[i]->GetVertexBuffer(&vertexBuffer);
		FlexVertexArray vertexData(fragmentMeshes[i]->GetNumBytesPerVertex());
		D3DVERTEXELEMENT9 elem[64];	fragmentMeshes[i]->GetDeclaration(elem);	 vertexData.setFormat(elem);
		vertexBuffer->Lock(0,fragmentMeshes[i]->GetNumVertices()*fragmentMeshes[i]->GetNumBytesPerVertex(),vertexData.getDataPointerReference(),0);

		for(unsigned int u=0; u < fragmentMeshes[i]->GetNumVertices(); u++ )
			vertexData[u].tex(atlasIndex) = (vertexData[u].tex(atlasIndex) * (atlasSize - 4) + D3DXVECTOR2(2.0,2.0)) / (float)atlasSize;

		vertexBuffer->Unlock();
		vertexBuffer->Release();
	}
}

void MeshExploder::merge()
{
	int sumVertices=0;
	int sumFaces=0;
	for(int i=0; i<nPieces; i++)
	{
		sumVertices += fragmentMeshes[i]->GetNumVertices();
		sumFaces += fragmentMeshes[i]->GetNumFaces();
	}

	composedMaterials = new D3DXMATERIAL[nPieces];
	composedShaders = new D3DXEFFECTINSTANCE[nPieces];
	composedMaterialNames = new wchar_t*[nPieces];

	for(int t=0; t<nPieces; t++)
	{
		composedMaterials[t] = originalMaterials[fragmentAttribIndices[t]];
		composedShaders[t] = originalShaders[fragmentAttribIndices[t]];
		composedMaterialNames[t] = materialNames[fragmentAttribIndices[t]];
	}

	HRESULT dhre = D3DXCreateMeshFVF(
		sumFaces,
		sumVertices,
		0,
		originalMesh->GetFVF(),
		device,
		&composedMesh);

	LPDIRECT3DVERTEXBUFFER9 composedVertexBuffer;
	composedMesh->GetVertexBuffer(&composedVertexBuffer);
	FlexVertexArray composedVertexData(composedMesh->GetNumBytesPerVertex());
	D3DVERTEXELEMENT9 elem[64];	composedMesh->GetDeclaration(elem);	 composedVertexData.setFormat(elem);
	composedVertexBuffer->Lock(0,composedMesh->GetNumVertices()*composedMesh->GetNumBytesPerVertex(),composedVertexData.getDataPointerReference(),0);

	LPDIRECT3DINDEXBUFFER9 composedIndexBuffer;
	composedMesh->GetIndexBuffer(&composedIndexBuffer);
	unsigned short* composedIndexData;
	composedIndexBuffer->Lock(0,
						composedMesh->GetNumFaces()*3*sizeof(unsigned short),
						(void**)&composedIndexData,0);

	DWORD* composedAttributeData;
	composedMesh->LockAttributeBuffer(0, &composedAttributeData);

	int vertexOffset = 0;
	int faceOffset = 0;

	for(int f=0; f<nPieces; f++)
	{
		LPDIRECT3DVERTEXBUFFER9 fragmentVertexBuffer;
		fragmentMeshes[f]->GetVertexBuffer(&fragmentVertexBuffer);
		FlexVertexArray fragmentVertexData(fragmentMeshes[f]->GetNumBytesPerVertex());
		fragmentVertexBuffer->Lock(0,fragmentMeshes[f]->GetNumVertices()*fragmentMeshes[f]->GetNumBytesPerVertex(),fragmentVertexData.getDataPointerReference(),0);

		LPDIRECT3DINDEXBUFFER9 fragmentIndexBuffer;
		fragmentMeshes[f]->GetIndexBuffer(&fragmentIndexBuffer);
		unsigned short* fragmentIndexData;
		fragmentIndexBuffer->Lock(0,
							fragmentMeshes[f]->GetNumFaces()*3*sizeof(unsigned short),
							(void**)&fragmentIndexData,0);

		for(int iv=0; iv<fragmentMeshes[f]->GetNumVertices(); iv++)
			composedVertexData.assign(iv + vertexOffset, fragmentVertexData[iv]);
			//composedVertexData[iv + vertexOffset] = fragmentVertexData[iv];
		for(int ic=0; ic<fragmentMeshes[f]->GetNumFaces(); ic++)
		{
			composedAttributeData[ic + faceOffset] = f;
			composedIndexData[(ic + faceOffset) * 3] = fragmentIndexData[ic * 3] + vertexOffset;
			composedIndexData[(ic + faceOffset) * 3 + 1] = fragmentIndexData[ic * 3 + 1] + vertexOffset;
			composedIndexData[(ic + faceOffset) * 3 + 2] = fragmentIndexData[ic * 3 + 2] + vertexOffset;
		}
		fragmentVertexBuffer->Unlock();
		fragmentVertexBuffer->Release();
		fragmentIndexBuffer->Unlock();
		fragmentIndexBuffer->Release();

		vertexOffset += fragmentMeshes[f]->GetNumVertices();
		faceOffset += fragmentMeshes[f]->GetNumFaces();
	}

	composedMesh->UnlockAttributeBuffer();

	composedVertexBuffer->Unlock();
	composedVertexBuffer->Release();
	composedIndexBuffer->Unlock();
	composedIndexBuffer->Release();
}

void MeshExploder::saveComposedMesh(LPCSTR name)
{
	char sname[256];
	sprintf(sname, "processedMeshes/%s.x", name);
	HRESULT hr = D3DXSaveMeshToX(L::l+sname,	//LPCTSTR pFilename,
					composedMesh,	//LPD3DXMESH pMesh,
					NULL,	//CONST DWORD * pAdjacency,
					composedMaterials,	//CONST D3DXMATERIAL * pMaterials,
					composedShaders,	//CONST D3DXEFFECTINSTANCE * pEffectInstances,
					nPieces,	//DWORD NumMaterials,
  					D3DXF_FILEFORMAT_TEXT	//DWORD Format
					);
}

void MeshExploder::saveComposedMeshToOgreXML(LPCSTR name)
{
	char sname[256];
	sprintf(sname, "processedMeshes/%s.mesh.xml", name);

	std::ofstream f(sname);
	XmlStream xml(f);

	DWORD mfvf = composedMesh->GetFVF();
	unsigned int nTexCoords = ( mfvf & D3DFVF_TEXCOUNT_MASK) / D3DFVF_TEX1;

	xml << tag("mesh");
	xml << tag("sharedgeometry")
		<< attr("vertexcount") << composedMesh->GetNumVertices();
		
		xml << tag("vertexbuffer");
		if(mfvf & D3DFVF_XYZ)
			xml << attr("positions") << "true";
		else
			xml << attr("positions") << "false";
		if(mfvf & D3DFVF_NORMAL)
			xml << attr("normals") << "true";
		else
			xml << attr("normals") << "false";
		if(mfvf & D3DFVF_DIFFUSE)
			xml << attr("colours_diffuse") << "true";
		else
			xml << attr("colours_diffuse") << "false";
		if(mfvf & D3DFVF_SPECULAR)
			xml << attr("colours_specular") << "true";
		else
			xml << attr("colours_specular") << "false";
		xml << attr("texture_coords") << nTexCoords;

	LPDIRECT3DVERTEXBUFFER9 composedVertexBuffer;
	composedMesh->GetVertexBuffer(&composedVertexBuffer);
	FlexVertexArray composedVertexData(composedMesh->GetNumBytesPerVertex());
	D3DVERTEXELEMENT9 elem[64];	composedMesh->GetDeclaration(elem);	 composedVertexData.setFormat(elem);
	composedVertexBuffer->Lock(0,composedMesh->GetNumVertices()*composedMesh->GetNumBytesPerVertex(),composedVertexData.getDataPointerReference(),0);

	LPDIRECT3DINDEXBUFFER9 composedIndexBuffer;
	composedMesh->GetIndexBuffer(&composedIndexBuffer);
	unsigned short* composedIndexData;
	composedIndexBuffer->Lock(0,
						composedMesh->GetNumFaces()*3*sizeof(unsigned short),
						(void**)&composedIndexData,0);

	DWORD* composedAttributeData;
	composedMesh->LockAttributeBuffer(0, &composedAttributeData);

	for(int xvi = 0; xvi < composedMesh->GetNumVertices(); xvi++)
	{
		xml << tag("vertex");
			if(mfvf & D3DFVF_XYZ)
			{
				xml << tag("position");
					xml << attr("x") << composedVertexData[xvi].pos().x;
					xml << attr("y") << composedVertexData[xvi].pos().y;
					xml << attr("z") << composedVertexData[xvi].pos().z;
				xml << endtag();
			}
			if(mfvf & D3DFVF_NORMAL)
			{
				xml << tag("normal");
					xml << attr("x") << composedVertexData[xvi].normal().x;
					xml << attr("y") << composedVertexData[xvi].normal().y;
					xml << attr("z") << composedVertexData[xvi].normal().z;
				xml << endtag();
			}
			if(mfvf & D3DFVF_DIFFUSE)
			{
				char colval[256];
				D3DCOLOR dc = composedVertexData[xvi].diffuse();
				unsigned char* pdc = (unsigned char*)&dc;
				sprintf(colval, "%f %f %f %f",
					pdc[3],
					pdc[0],
					pdc[1],
					pdc[2]
					);
				xml << tag("colours_diffuse");
					xml << attr("value") << colval;
				xml << endtag();
			}
			if(mfvf & D3DFVF_SPECULAR)
			{
				char colval[256];
				D3DCOLOR dc = composedVertexData[xvi].specular();
				unsigned char* pdc = (unsigned char*)&dc;
				sprintf(colval, "%f %f %f %f",
					pdc[3],
					pdc[0],
					pdc[1],
					pdc[2]
					);
				xml << tag("colours_specular");
					xml << attr("value") << colval;
				xml << endtag();
			}
			for(int taxi=0; taxi<nTexCoords; taxi++)
			{
				xml << tag("texcoord");
					xml << attr("u") << composedVertexData[xvi].tex(taxi).x;
					xml << attr("v") << composedVertexData[xvi].tex(taxi).y;
				xml << endtag();
			}

		xml << endtag();
	}

	xml << endtag(); // end vertexbuffer
	xml << endtag(); // end sharedgeometry

	xml << tag("submeshes");
	int faceOffset = 0;
	for(int xai = 0; xai < nPieces; xai++)
	{
		xml << tag("submesh")
			<< attr("material") << LC::c-composedMaterialNames[xai]
			<< attr("use32bitindexes") << "true"
			<< attr("usesharedvertices") << "true"
			<< attr("operationtype") << "triangle_list";
			
			xml << tag("faces")
				<< attr("count") << fragmentMeshes[xai]->GetNumFaces();
			for(int ic = 0; ic < fragmentMeshes[xai]->GetNumFaces(); ic++)
			{
				xml << tag("face")
					<< attr("v1") << composedIndexData[(ic + faceOffset) * 3]
					<< attr("v2") << composedIndexData[(ic + faceOffset) * 3 + 1]
					<< attr("v3") << composedIndexData[(ic + faceOffset) * 3 + 2];
				xml << endtag();
			}
			faceOffset += fragmentMeshes[xai]->GetNumFaces();
			xml << endtag();
		xml << endtag();			
	}
	xml << endtag(); // submeshes
	xml << tag("submeshnames");
	xml << endtag();

	composedMesh->UnlockAttributeBuffer();

	composedVertexBuffer->Unlock();
	composedVertexBuffer->Release();
	composedIndexBuffer->Unlock();
	composedIndexBuffer->Release();
}