#include "BVHSAHBuilder.h"
#include "flexilog.h"
#include <wx/stopwatch.h>
#include <stack>
#include <vector>

using namespace VRUT;

BVHSAHBuilder::Comparator::Comparator(unsigned char axis, TriangleAABB * aabbs) : axis(axis) , aabbs(aabbs)
{
}

bool BVHSAHBuilder::Comparator::operator()(int left, int right)
{
	return aabbs[left].centroid.v[axis] < aabbs[right].centroid.v[axis];
}

BVHBottom * BVHSAHBuilder::Build(const smart_pointer<Geometry> & geometry, unsigned int maxTrianglesPerLeaf)
{

	wxStopWatch sw;
	sw.Start();

	bvh = new BVHBottom();

	// setup triangles
	PreSetupTriangles(geometry);

	// compute triangle bounding boxes
	ComputeAABBs();

	// too few triangles => create leaf
	if (numberOfTriangles <= maxTrianglesPerLeaf)
	{
		// we will do one leaf node with empty AABB
		// alocate triangle AABB array and nodes
		bvh->nodes = (BVHBottom::BVHNode*) VrutAlignedMalloc(sizeof(BVHBottom::BVHNode), 16);
		bvh->nodes[0].box.MinBound = Vector3(sceneMin.v[3], sceneMin.v[2], sceneMin.v[1]);
		bvh->nodes[0].box.MaxBound = Vector3(sceneMax.v[3], sceneMax.v[2], sceneMax.v[1]);
		bvh->nodes[0].offset = 0;
		bvh->nodes[0].t = numberOfTriangles;
		bvh->numberOfNodes = 1;
	}

	// non-empty geometry
	else
	{

		// nodes
		bvh->nodes = (BVHBottom::BVHNode*)VrutAlignedMalloc((2 * numberOfTriangles - 1) * sizeof(BVHBottom::BVHNode), 16);

		// aux. data
		float * aR = new float[numberOfTriangles];
		int * tmpIndices = new int[numberOfTriangles];
		bool * leftIndices = new bool[numberOfTriangles];
		VRUT_ALIGNED(16) Point bMin;
		VRUT_ALIGNED(16) Point bMax;

		// sort triangles
		std::vector<int> indices[3];
		for (unsigned char a = 0; a < 3; ++a)
		{
			indices[a].resize(numberOfTriangles);
			for (unsigned int i = 0; i < numberOfTriangles; ++i) 
				indices[a][i] = i;
			std::sort(indices[a].begin(), indices[a].end(), Comparator(a + 1, aabbs));
		}
		
		// root
		bvh->nodes[0].box.MinBound = Vector3(sceneMin.v[3], sceneMin.v[2], sceneMin.v[1]);
		bvh->nodes[0].box.MaxBound = Vector3(sceneMax.v[3], sceneMax.v[2], sceneMax.v[1]);
		bvh->nodes[0].offset = 0;
		bvh->nodes[0].t = numberOfTriangles;
		bvh->numberOfNodes = 1;

		// build node
		BVHBuildNode * bn = (BVHBuildNode*) VrutAlignedMalloc(sizeof(BVHBuildNode), 16);
		bn->node = &bvh->nodes[0];
		bn->bMin = sceneMin;
		bn->bMax = sceneMax;

		// stack
		std::stack<BVHBuildNode*> stack;
		stack.push(bn);
		while (!stack.empty())
		{

			// pop
			bn = stack.top();
			stack.pop();
			BVHBottom::BVHNode * node = bn->node;

			// try all splitting planes
			float bestCost = FLT_MAX;
			unsigned char bestAxis = 0;
			int bestIndex = -1;
			int bestLeftCount = 0;

			// node surface area
			VRUT_ALIGNED(16) Point d;
			d._p = _mm_sub_ps(bn->bMax._p, bn->bMin._p);
			const float sa = 2.0f * (d.v[1] * d.v[2] + d.v[1] * d.v[3] + d.v[2] * d.v[3]);
			
			for (unsigned char a = 0; a < 3; ++a) 
			{

				// right prefix scan.
				bMin._p = _mm_set1_ps(FLT_MAX);
				bMax._p = _mm_set1_ps(-FLT_MAX);
				for (int i = node->offset + node->t - 1; i >= node->offset; --i)
				{
					bMin._p =  _mm_min_ps(bMin._p, aabbs[indices[a][i]].min._p);
					bMax._p =  _mm_max_ps(bMax._p, aabbs[indices[a][i]].max._p);
					const float x = bMax.v[3] - bMin.v[3];
					const float y = bMax.v[2] - bMin.v[2];
					const float z = bMax.v[1] - bMin.v[1];
					aR[i] = 2.0f * (x*y + x*z + y*z);
				}

				// find the best splitting plane
				bMin._p = _mm_set1_ps(FLT_MAX);
				bMax._p = _mm_set1_ps(-FLT_MAX);
				for (unsigned int i = node->offset; i < node->offset + node->t - 1; ++i) 
				{
					int leftCount = (i + 1) - node->offset;
					int rightCount = node->offset + node->t - (i + 1);
					bMin._p =  _mm_min_ps(bMin._p, aabbs[indices[a][i]].min._p);
					bMax._p =  _mm_max_ps(bMax._p, aabbs[indices[a][i]].max._p);
					const float x = bMax.v[3] - bMin.v[3];
					const float y = bMax.v[2] - bMin.v[2];
					const float z = bMax.v[1] - bMin.v[1];
					float aL = 2.0f * (x*y + x*z + y*z);
					float cost = (aL * leftCount + aR[i + 1] * rightCount) / sa;
					if (cost < bestCost)
					{
						bestCost = cost;
						bestIndex = i + 1;
						bestAxis = a;
						bestLeftCount = leftCount;
					}
				}

			}

			//assert(bestIndex > node->offset);
			//assert(node->offset + node->t > unsigned int(bestIndex));

			if (node->t > maxTrianglesPerLeaf)
			{

				const float SAH_C = 1.0f;
				if (maxTrianglesPerLeaf > 0 || (maxTrianglesPerLeaf == 0 && SAH_C + bestCost < node->t))
				{

					// compute left index flags
					memset(leftIndices, 0, numberOfTriangles * sizeof(bool));
					for (int i = node->offset; i < bestIndex; ++i)
					{
						int index = indices[bestAxis][i];
						leftIndices[index] = true;
					}

					// reorder indices
					for (unsigned char a = 0; a < 3; ++a)
					{
						if (a != bestAxis) {
							int k = node->offset;
							int l = bestIndex;
							for (unsigned int i = node->offset; i < node->offset + node->t; ++i)
							{
								int index = indices[a][i];
								if (leftIndices[indices[a][i]]) 
									tmpIndices[k++] = index;
								else 
									tmpIndices[l++] = index;
							}
							assert(k == bestIndex);
							assert(l == node->offset + node->t);
							memcpy(&indices[a][node->offset], &tmpIndices[node->offset], node->t * sizeof(int));
						}
					}

					// create child nodes
					BVHBottom::BVHNode * leftNode = bvh->nodes + bvh->numberOfNodes;
					leftNode->offset = node->offset;
					leftNode->t = bestLeftCount;
					BVHBuildNode * lbn = (BVHBuildNode*) VrutAlignedMalloc(sizeof(BVHBuildNode), 16);
					lbn->bMin._p = _mm_set1_ps(FLT_MAX);
					lbn->bMax._p = _mm_set1_ps(-FLT_MAX);
					lbn->node = leftNode;

					BVHBottom::BVHNode * rightNode = leftNode + 1;
					rightNode->offset = node->offset + bestLeftCount;
					rightNode->t = node->t - bestLeftCount;
					BVHBuildNode * rbn = (BVHBuildNode*) VrutAlignedMalloc(sizeof(BVHBuildNode), 16);
					rbn->bMin._p = _mm_set1_ps(FLT_MAX);
					rbn->bMax._p = _mm_set1_ps(-FLT_MAX);
					rbn->node = rightNode;

					// compute bounding boxes.
					for (int i = node->offset; i < bestIndex; ++i) 
					{
						lbn->bMin._p = _mm_min_ps(lbn->bMin._p, aabbs[indices[bestAxis][i]].min._p);
						lbn->bMax._p = _mm_max_ps(lbn->bMax._p, aabbs[indices[bestAxis][i]].max._p);
					}

					for (unsigned int i = bestIndex; i < node->offset + node->t; ++i) 
					{
						rbn->bMin._p = _mm_min_ps(rbn->bMin._p, aabbs[indices[bestAxis][i]].min._p);
						rbn->bMax._p = _mm_max_ps(rbn->bMax._p, aabbs[indices[bestAxis][i]].max._p);
					}

					// node bounding boxes
					leftNode->box.MinBound = Vector3(lbn->bMin.v[3], lbn->bMin.v[2], lbn->bMin.v[1]);
					leftNode->box.MaxBound = Vector3(lbn->bMax.v[3], lbn->bMax.v[2], lbn->bMax.v[1]);
					rightNode->box.MinBound = Vector3(rbn->bMin.v[3], rbn->bMin.v[2], rbn->bMin.v[1]);
					rightNode->box.MaxBound = Vector3(rbn->bMax.v[3], rbn->bMax.v[2], rbn->bMax.v[1]);
			
					// interior
					node->offset = -1;
					node->t = bvh->numberOfNodes;
					bvh->numberOfNodes += 2;

					// push left node
					if (leftNode->t > 1) stack.push(lbn);
					else VrutAlignedFree(lbn);

					// psu hright child
					if (rightNode->t > 1) stack.push(rbn);
					else VrutAlignedFree(rbn);

				}

			}

			// delete node
			VrutAlignedFree(bn);

		}

		// Delete aux. data.
		delete [] aR;
		delete [] tmpIndices;
		delete [] leftIndices;

		// set bounding boxes indices
		for (unsigned int i = 0; i < numberOfTriangles; ++i)
			aabbs[i].index = indices[0][i];

	}

	// post setup triangles
	PostSetupTriangles(geometry);

	sw.Pause();
	LOGDEBUG(wxString::Format("<RayTracer> BVH (gid=%lu) has %lu triangles, built in %ld ms", (unsigned long)geometry->GetID(), (unsigned long)bvh->triangleCols.size(), sw.Time()));
	SAFE_DELETE(bvh->triangles);
	
	return bvh;

}
