#include "BvhConstructor.h"
#include "Triangle3.h"
#include "SceneEntity.h"
#include "Geometry.h"
#include "gzstream.h"

#include <queue>
#include <stack>
#include <fstream>

#define MAX_FLOAT 1e20f


#ifdef _CRT_SET
	#define _CRTDBG_MAP_ALLOC
	#include <stdlib.h>
	#include <crtdbg.h>

	// redefine new operator
	#define DEBUG_NEW new(_NORMAL_BLOCK, __FILE__, __LINE__)
	#define new DEBUG_NEW
#endif


using namespace std;


namespace CHCDemoEngine
{

BvhConstructor::BvhConstructor(SceneEntity **entities, 
							   int first, 
							   int last):
mEntities(entities),
mFirst(first),
mLast(last),
mMaxDepth(20),
mMaxObjects(1),
mMaxTriangles(1),
mNumNodes(0),
mSplitType(SAH)
{
}


float BvhConstructor::EvaluateSahCost(BvhLeaf *leaf, int axis, float position)
{
	// count triangles on the left / right
	int left = 0;
	int right = 0;
	AxisAlignedBox3 leftBox, rightBox;

	leftBox.Initialize();
	rightBox.Initialize();

	const int candidates = std::max(50, (int)(leaf->CountPrimitives() / 20));
	const float finc = leaf->CountPrimitives() / (float)candidates;

	int i = leaf->mFirst;
	float fi = leaf->mFirst + 0.5f;

	AxisAlignedBox3 box;

	for (; i <= leaf->mLast;) 
	{
		box = mEntities[i]->GetWorldBoundingBox();

		if (box.Center(axis) < position) 
		{
			++ left;
			// update the box
			leftBox.Include(box);
		} 
		else 
		{
			++ right;
			rightBox.Include(box);
		}
		 
		fi += finc;
		i = (int)fi;
	}

	float bW = 1.0f;
	float leftRatio = left / (float)leaf->CountPrimitives();
	float rightRatio = right / (float)leaf->CountPrimitives();

	float saLeft = 0.0f;
	float saRight = 0.0f;

	// not a valid split
	if (!left || !right)
		return 1e25f;

	saLeft = leftBox.SurfaceArea();
	saRight = rightBox.SurfaceArea();


	return
		saLeft  * ((1.0f - bW) + bW * leftRatio) + 
		saRight * ((1.0f - bW) + bW * rightRatio);
}


float BvhConstructor::SelectPlaneSah(BvhLeaf *leaf, int &axis, float &minCost)
{
	minCost = MAX_FLOAT;
	float bestPos = minCost;
	int bestAxis = 0;

	//  cout<<"Evaluating axis..."<<endl<<flush;

	const int initialPlanes = 3;

	// initiate the costs
	for (axis = 0; axis < 3; ++ axis) 
	{
		const float size = leaf->mBox.Size(axis);

		for (int i = 0; i < initialPlanes; ++ i) 
		{
			const float pos = leaf->mBox.Min(axis) + (i + 1) * size / (initialPlanes + 1);
			const float cost = EvaluateSahCost(leaf, axis, pos);

			if (cost < minCost) 
			{
				minCost = cost;
				bestPos = pos;
				bestAxis = axis;
			}
		}
	}

	axis = bestAxis;

	//  cout<<axis<<endl<<flush;
	const float shrink = 0.5f;

	// now hierarchically refine the initial interval
	float size = shrink * 
		(leaf->mBox.Max(axis) - leaf->mBox.Min(axis)) / (initialPlanes + 1);

	const int steps = 4;
	float cost;

	for (int i = 0; i < steps; ++ i) 
	{
		const float left = bestPos - size;
		const float right = bestPos + size;

		cost = EvaluateSahCost(leaf, axis, left);

		if (cost < minCost) 
		{
			minCost = cost;
			bestPos = left;
		}

		cost = EvaluateSahCost(leaf, axis, right);

		if (cost < minCost) 
		{
			minCost = cost;
			bestPos = right;
		}

		size = shrink * size;
	}

	//OUT1("best axis: " << axis << " " << bestPos << " " << minCost);

	return bestPos;
}


int BvhConstructor::SortTriangles(BvhLeaf *leaf, 
								  int axis, 
								  float position
								  )
{
	int i = leaf->mFirst;
	int j = leaf->mLast;

    while (1)
	{
		//while (mEntities[i]->GetWorldCenter()[axis] < position) ++ i;
		//while (position < mEntities[j]->GetWorldCenter()[axis]) -- j;

		while ((i < leaf->mLast) && (mEntities[i]->GetWorldCenter()[axis] < position)) ++ i;
		while ((j > leaf->mFirst) && (position < mEntities[j]->GetWorldCenter()[axis])) -- j;

		// sorting finished
		if (i >= j) break;

		// swap entities
		swap(mEntities[i], mEntities[j]);
			
		++ i;
		-- j;
	}

	return j;
}


void BvhConstructor::UpdateBoundingBoxes(BvhNode *node)
{
	if (node->IsLeaf())
	{
		int numEntities = node->CountPrimitives();
		SceneEntity **entities = mEntities + node->mFirst;

		node->mBox = SceneEntity::ComputeBoundingBox(entities, numEntities);
		//cout << "box: " << node->mBox << endl;
	}
	else
	{
		BvhNode *f = static_cast<BvhInterior *>(node)->GetFront();
		BvhNode *b = static_cast<BvhInterior *>(node)->GetBack();

		UpdateBoundingBoxes(f);
		UpdateBoundingBoxes(b);

		node->mBox = f->mBox;
		node->mBox.Include(b->mBox);
	}
}


int BvhConstructor::SortTrianglesSpatialMedian(BvhLeaf *leaf, 
											   int axis)
{
	// spatial median
	float m = leaf->GetBox().Center()[axis];
	return SortTriangles(leaf, axis, m);
}


int BvhConstructor::SortTrianglesObjectMedian(BvhLeaf *leaf, int axis, float &pos)
{
	// Now distribute the objects into the subnodes
	// Use QuickMedian sort
	int l = leaf->mFirst;
	int r = leaf->mLast;
	int k = (l + r) / 2;

	while (l < r) 
	{
		int i = l;
		int j = r;

		// get some estimation of the pivot
		pos = mEntities[(l + r) / 2]->GetBoundingBox().Center(axis);

		while (1) 
		{
			while ((i <= leaf->mLast) && (mEntities[i]->GetWorldBoundingBox().Center(axis) < pos))
				++ i;

			while((j >= leaf->mFirst) && (pos < mEntities[j]->GetWorldBoundingBox().Center(axis)))
				-- j;

			if (i <= j) 
			{
				std::swap(mEntities[i], mEntities[j]);
				++ i;
				-- j;
			}
			else
			{
				break;
			}
		}

		// now check the extents
		if (i >= k)
			r = j;
		else
			l = i;
	}

	return k;
}


BvhNode *BvhConstructor::SubdivideLeaf(BvhLeaf *leaf, 
									   int parentAxis
									   )
{
	if (TerminationCriteriaMet(leaf))
	{
		//if (leaf->CountPrimitives() <= 0)
		//cout << "leaf constructed:" << leaf->mBox << " " << leaf->mFirst << " " << leaf->mLast << endl;
		return leaf;
	}

	//const int axis = (parentAxis + 1) % 3;
	int axis = leaf->mBox.MajorAxis();
	const int scale = 20;

	// position of the split in the partailly sorted array of triangles
	// corresponding to this leaf
	int split = -1;
	float pos = -1.0f;

	// Spatial median subdivision
	switch (mSplitType) 
	{
	case SPATIAL_MEDIAN:
		split = SortTrianglesSpatialMedian(leaf, axis);
		pos = leaf->mBox.Center()[axis];
		break;
	case OBJECT_MEDIAN:
		// Object median subdivision: approximately the same number
		// of objects on the left of the the splitting point.
		split = SortTrianglesObjectMedian(leaf, axis, pos);
		break;
	case SAH:
		{
			float cost;
			pos = SelectPlaneSah(leaf, axis, cost);

			if (pos != MAX_FLOAT)
			{
				split = SortTriangles(leaf, axis, pos);
			}

			if ((pos == MAX_FLOAT) || (split == leaf->mLast))
			{
				split = -1;
				split = SortTrianglesObjectMedian(leaf, axis, pos);
			}
		}
		break;
	default:
		cerr << "should not come here" << endl;
		break;
	}
	
	if (split == leaf->mLast)
	{
		// no split could be achieved => just halve number of objects
		split = (leaf->mLast + leaf->mFirst) / 2;
		//cerr << "no reduction " << leaf->CountPrimitives() << " " << leaf->mFirst << " " << leaf->mLast << endl;
	}

	// create two more nodes
	mNumNodes += 2;
	BvhInterior *parent = new BvhInterior(leaf->GetParent());
	parent->mFirst = leaf->mFirst;
	parent->mLast = leaf->mLast;
	
	parent->mAxis = axis;
	parent->mBox = leaf->mBox;
	parent->mDepth = leaf->mDepth;

	BvhLeaf *front = new BvhLeaf(parent);

	parent->mBack = leaf;
	parent->mFront = front;
		
	// now assign the triangles to the subnodes
	front->mFirst = split + 1;
	front->mLast = leaf->mLast;
	front->mDepth = leaf->mDepth + 1;

	leaf->mLast = split;
	leaf->mDepth = front->mDepth;
	leaf->mParent = parent;
	
	front->mBox = SceneEntity::ComputeBoundingBox(mEntities + front->mFirst, front->CountPrimitives());
	leaf->mBox = SceneEntity::ComputeBoundingBox(mEntities + leaf->mFirst, leaf->CountPrimitives());

	// recursively continue subdivision
	parent->mBack = SubdivideLeaf(static_cast<BvhLeaf *>(parent->mBack), axis);
	parent->mFront = SubdivideLeaf(static_cast<BvhLeaf *>(parent->mFront), axis);
	
	return parent;
}


int BvhConstructor::CountTriangles(BvhNode *node) const
{
	int numTriangles = 0;

	for (int i = node->mFirst; i <= node->mLast; ++ i)
	{
		numTriangles += mEntities[i]->CountNumTriangles(0);
	}

	return numTriangles;
}


bool BvhConstructor::TerminationCriteriaMet(BvhLeaf *leaf) const
{
	const bool criteriaMet = 
		(leaf->mDepth > mMaxDepth) ||
		(leaf->CountPrimitives() <= mMaxObjects) ||
		(CountTriangles(leaf) <= mMaxTriangles);

	return criteriaMet;
}


BvhNode *BvhConstructor::Construct(int &numNodes)
{
	BvhNode *root;
	mNumNodes = 1;

	BvhLeaf *l = new BvhLeaf(NULL);
	
	l->mFirst = mFirst;
	l->mLast = mLast;

	cout << "\n================================" << endl;
	cout << "constructing bvh for " << l->mLast - l->mFirst + 1 << " entities" << endl;

	l->mBox = SceneEntity::ComputeBoundingBox(mEntities + mFirst, mLast - mFirst + 1);
	l->mArea = l->mBox.SurfaceArea();
	
	root = SubdivideLeaf(l, 0);

	UpdateBoundingBoxes(root);

	numNodes = mNumNodes;

	return root;
}


}