#include "BVHBinnedBuilder.h"
#include "flexilog.h"
#include "common.h"
#include <wx/stopwatch.h>
#include <stack>

using namespace VRUT;

BVHBottom * BVHBinnedBuilder::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();

	// empty geometry
	if (numberOfTriangles == 0)
	{
		// 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].offset = 0;
		bvh->nodes[0].t = 0;
		bvh->numberOfNodes = 1;
	}

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

		// 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->cMin = centroidMin;
		bn->cMax = centroidMax;
		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;

			// find domain axis
			unsigned char domain = 1;
			VRUT_ALIGNED(16) Point d;
			d._p = _mm_sub_ps(bn->cMax._p, bn->cMin._p);
			if (d.v[3] < d.v[2])
			{
				if (d.v[2] < d.v[1])
					domain = 1;
				else
					domain = 2;
			}
			else
			{
				if (d.v[3] < d.v[1])
					domain = 1;
				else
					domain = 3;
			}

			// node surface area
			d._p = _mm_sub_ps(bn->bMin._p, bn->bMax._p);
			const float sa = 2.0f * (d.v[1] * d.v[2] + d.v[1] * d.v[3] + d.v[2] * d.v[3]);

			// initialize count of triangles in bin and bins AABB as negative box
			for (unsigned int i = 0; i < NUMBER_OF_BINS; ++i)
			{
				binMin[i]._p = _mm_set1_ps(FLT_MAX);
				binMax[i]._p = _mm_set1_ps(-FLT_MAX);
				binCounter[i] = 0;
			}
			// precompute constant to binning
			static const float EPSILON = 0.0001f;
			const float k0 = bn->cMin.v[domain];
			const float k1 = (NUMBER_OF_BINS * (1.0f- EPSILON)) / (bn->cMax.v[domain] - bn->cMin.v[domain]);

			// find bin projection for all triangles
			for (unsigned int i = node->offset; i < node->offset + node->t; ++i)
			{
				aabbs[i].bin = k1 * (aabbs[i].centroid.v[domain] - k0);
				++(binCounter[aabbs[i].bin]);
				binMin[aabbs[i].bin]._p = _mm_min_ps(binMin[aabbs[i].bin]._p, aabbs[i].min._p);
				binMax[aabbs[i].bin]._p = _mm_max_ps(binMax[aabbs[i].bin]._p, aabbs[i].max._p);
			}

			unsigned int bestIndex = 0;
			unsigned int nL[NUMBER_OF_BINS];
			unsigned int nR[NUMBER_OF_BINS];
			float aL[NUMBER_OF_BINS];
			float bestCost = FLT_MAX;
			nL[0] = 0;
			nR[NUMBER_OF_BINS - 1] = 0;
			VRUT_ALIGNED(16) Point bMin;
			VRUT_ALIGNED(16) Point bMax;
			
			// prefix scan from left
			bMin._p = _mm_set1_ps(FLT_MAX);
			bMax._p = _mm_set1_ps(-FLT_MAX);
			for (unsigned int i = 0; i < NUMBER_OF_BINS - 1; ++i)
			{
				nL[i + 1] = nL[i] + binCounter[i];
				bMin._p = _mm_min_ps(bMin._p, binMin[i]._p);
				bMax._p = _mm_max_ps(bMax._p, binMax[i]._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];
				aL[i + 1] = 2.0f * (x*y + x*z + y*z);
			}

			// prefix scan from right
			bMin._p = _mm_set1_ps(FLT_MAX);
			bMax._p = _mm_set1_ps(-FLT_MAX);
			for (unsigned int i = NUMBER_OF_BINS - 1; i > 0; --i)
			{
				nR[i - 1] = nR[i] + binCounter[i];
				bMin._p = _mm_min_ps(bMin._p, binMin[i]._p);
				bMax._p = _mm_max_ps(bMax._p, binMax[i]._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];
				const float aR = 2.0f * (x*y + x*z + y*z);
				const float cost = (aL[i] * nL[i] + aR * nR[i-1]) / sa;
				if (cost < bestCost)
				{
					bestCost = cost;
					bestIndex = i;
				}
			}

			// split
			if (node->t > maxTrianglesPerLeaf)
			{

				// SAH
				const float SAH_C = 1.0f;
				if ((maxTrianglesPerLeaf > 0 && nL[bestIndex] > 0 && nR[bestIndex - 1] > 0) 
					|| (maxTrianglesPerLeaf == 0 && SAH_C + bestCost < node->t))
				{
					BVHBottom::BVHNode * leftNode = bvh->nodes + bvh->numberOfNodes;
					leftNode->offset = node->offset;
					leftNode->t = nL[bestIndex];
					BVHBuildNode * lbn = (BVHBuildNode*) VrutAlignedMalloc(sizeof(BVHBuildNode), 16);
					lbn->cMin._p = _mm_set1_ps(FLT_MAX);
					lbn->cMax._p = _mm_set1_ps(-FLT_MAX);
					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 + nL[bestIndex];
					rightNode->t = nR[bestIndex - 1];
					BVHBuildNode * rbn = (BVHBuildNode*) VrutAlignedMalloc(sizeof(BVHBuildNode), 16);
					rbn->cMin._p = _mm_set1_ps(FLT_MAX);
					rbn->cMax._p = _mm_set1_ps(-FLT_MAX);
					rbn->bMin._p = _mm_set1_ps(FLT_MAX);
					rbn->bMax._p = _mm_set1_ps(-FLT_MAX);
					rbn->node = rightNode;

					node->offset = -1;
					node->t = bvh->numberOfNodes;
					bvh->numberOfNodes += 2;

					TriangleAABB * left = aabbs + leftNode->offset;
					TriangleAABB * right = aabbs + rightNode->offset + rightNode->t - 1;
					while (left <= right)
					{
						if (left->bin < bestIndex)
						{
							lbn->cMin._p = _mm_min_ps(lbn->cMin._p, left->centroid._p);
							lbn->cMax._p = _mm_max_ps(lbn->cMax._p, left->centroid._p);
							lbn->bMin._p = _mm_min_ps(lbn->bMin._p, left->min._p);
							lbn->bMax._p = _mm_max_ps(lbn->bMax._p, left->max._p);
							++left;
							continue;
						}
						if (right->bin >= bestIndex)
						{
							rbn->cMin._p = _mm_min_ps(rbn->cMin._p, right->centroid._p);
							rbn->cMax._p = _mm_max_ps(rbn->cMax._p, right->centroid._p);
							rbn->bMin._p = _mm_min_ps(rbn->bMin._p, right->min._p);
							rbn->bMax._p = _mm_max_ps(rbn->bMax._p, right->max._p);
							--right;
							continue;
						}
						// both bad side
						// swap, in next step use it.
						TriangleAABB tmp(*left);
						*left = *right;
						*right = tmp;
					}

					// 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]);

					// push new nodes
					stack.push(lbn);
					stack.push(rbn);

				}

				// object median
				else if (maxTrianglesPerLeaf > 0)
				{
					BVHBottom::BVHNode * leftNode = bvh->nodes + bvh->numberOfNodes;
					leftNode->offset = node->offset;
					leftNode->t = node->t / 2;
					BVHBuildNode * lbn = (BVHBuildNode*) VrutAlignedMalloc(sizeof(BVHBuildNode), 16);
					lbn->cMin._p = _mm_set1_ps(FLT_MAX);
					lbn->cMax._p = _mm_set1_ps(-FLT_MAX);
					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 = leftNode->offset + leftNode->t;
					rightNode->t = node->t - leftNode->t;
					BVHBuildNode * rbn = (BVHBuildNode*) VrutAlignedMalloc(sizeof(BVHBuildNode), 16);
					rbn->cMin._p = _mm_set1_ps(FLT_MAX);
					rbn->cMax._p = _mm_set1_ps(-FLT_MAX);
					rbn->bMin._p = _mm_set1_ps(FLT_MAX);
					rbn->bMax._p = _mm_set1_ps(-FLT_MAX);
					rbn->node = rightNode;

					node->offset = -1;
					node->t = bvh->numberOfNodes;
					bvh->numberOfNodes += 2;

					int l = leftNode->offset;
					int r = rightNode->offset + rightNode->t - 1;
					int m = (l + r) / 2;
					while (l < r)
					{
						int i = l;
						int j = r;
						float pivot = aabbs[(i + j) / 2].centroid.v[domain];
						while (i < j)
						{
							while (aabbs[i].centroid.v[domain] < pivot) ++i;
							while (aabbs[j].centroid.v[domain] > pivot) --j;
							if (i <= j)
							{
								std::swap(aabbs[i], aabbs[j]);
								++i;
								--j;
							}
						}
						if (m <= j)
							r = j;
						else
							l = i;
					}

					// recompute aabbs
					for (unsigned int i = leftNode->offset; i < leftNode->offset + leftNode->t; ++i)
					{
						lbn->cMin._p = _mm_min_ps(lbn->cMin._p, aabbs[i].centroid._p);
						lbn->cMax._p = _mm_max_ps(lbn->cMax._p, aabbs[i].centroid._p);
						lbn->bMin._p = _mm_min_ps(lbn->bMin._p, aabbs[i].min._p);
						lbn->bMax._p = _mm_max_ps(lbn->bMax._p, aabbs[i].max._p);
					}
					for (unsigned int i = rightNode->offset; i < rightNode->offset + rightNode->t; ++i)
					{
						rbn->cMin._p = _mm_min_ps(rbn->cMin._p, aabbs[i].centroid._p);
						rbn->cMax._p = _mm_max_ps(rbn->cMax._p, aabbs[i].centroid._p);
						rbn->bMin._p = _mm_min_ps(rbn->bMin._p, aabbs[i].min._p);
						rbn->bMax._p = _mm_max_ps(rbn->bMax._p, aabbs[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]);

					// push new nodes
					stack.push(lbn);
					stack.push(rbn);

				}

			}

			// delete node
			VrutAlignedFree(bn);

		}

	}

	// post setup triangles
	PostSetupTriangles(geometry);

	sw.Pause();
	LOGDEBUG(wxString::Format("<BVHBinnedBuilder> 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;

}
