#include "BVHOptimizer.h"
#include "BinTreesUtils.h"
#include "flexilog.h"
#include <stack>
#include <wx/stopwatch.h>

using namespace VRUT;

typedef std::pair<BVHOptimizer::BVHOptNode*, int> OptPair;

void BVHOptimizer::BVHOptNode::UpdateBVNodeStatic(BVHOptNode * node)
{
	node->UpdateBVNode();
}

BVHOptimizer::BVHOptNode::BVHOptNode(BVHOptNode * parent, BVHOptNode * left, BVHOptNode * right)
		: bTree::BinTreeNodeWithAABB<BVHOptNode>(parent, left, right), offset(-1), t(0)
{
}

BVHOptimizer::BVHOptNode::~BVHOptNode()
{
	SafeDeleteChilds();
}

void BVHOptimizer::BVHOptNode::UpdateBVNode()
{
	if (!IsLeaf())
	{
		aabb=lChild->aabb;
		aabb.Merge(rChild->aabb);
	}
	SetAABBUndirty();
}

BVHOptimizer::BVHOptNode * BVHOptimizer::BVHToOpt(BVHBottom * bvh)
{
	std::stack<OptPair> stack;
	BVHOptNode * root = new BVHOptNode(nullptr, nullptr, nullptr);
	stack.push(OptPair(root, 0));
	while (!stack.empty())
	{
		BVHBottom::BVHNode * node = bvh->nodes + stack.top().second;
		BVHOptNode * optNode = stack.top().first;
		stack.pop();
		optNode->aabb = node->box;
		if (node->offset >= 0)
		{
			optNode->offset = node->offset;
			optNode->t = node->t;
		}
		else
		{
			BVHOptNode * left = new BVHOptNode(optNode, nullptr, nullptr);
			optNode->lChild = left;
			stack.push(OptPair(left, node->t));
			BVHOptNode * right = new BVHOptNode(optNode, nullptr, nullptr);
			optNode->rChild = right;
			stack.push(OptPair(right, node->t + 1));
		}
	}
	root->UpdateBVSubtree();
	return root;
}

void BVHOptimizer::OptToBVH(BVHOptNode * optRoot, BVHBottom * bvh)
{
	bvh->numberOfNodes = ComputeNodesCount(optRoot);
	VrutAlignedFree(bvh->nodes);
	bvh->nodes = (BVHBottom::BVHNode*) VrutAlignedMalloc(bvh->numberOfNodes * sizeof(BVHBottom::BVHNode), 16);
	std::queue<OptPair> queue;
	queue.push(OptPair(optRoot, 0));
	int i = 1;
	while (!queue.empty())
	{
		BVHBottom::BVHNode * node = bvh->nodes + queue.front().second;
		BVHOptNode * optNode = queue.front().first;
		queue.pop();
		node->box = optNode->aabb;
		if (optNode->IsLeaf())
		{
			node->offset = optNode->offset;
			node->t = optNode->t;
		}
		else
		{
			node->offset = optNode->offset;
			node->t = i;
			queue.push(OptPair(optNode->GetLChild(), node->t));
			queue.push(OptPair(optNode->GetRChild(), node->t + 1));
			i += 2;
		}
	}
}

float BVHOptimizer::GetCost(BVHOptNode * node)
{
	if (node)
	{
		float cost = 0.0f;
		if (!node->IsLeaf())
		{ 
			BVHOptNode * left = node->lChild;
			BVHOptNode * right = node->rChild;
			float sa = node->aabb.SurfaceArea();
			float leftCost = (sa != 0.0f) ? left->aabb.SurfaceArea() / sa * GetCost(left) : 0.0f;
			float rightCost = (sa != 0.0f) ? right->aabb.SurfaceArea() / sa * GetCost(right) : 0.0f;
			cost = CT + leftCost + rightCost; 
		} 
		else 
		{ 
			cost =  CI * node->t; 
		}
		return cost;
	}
	return 0.0f;
}

unsigned int BVHOptimizer::ComputeTrisCount(BVHOptNode * node)
{
	if (node)
	{
		if (!node->IsLeaf())
		{ 
			node->t = ComputeTrisCount(node->lChild);
			node->t += ComputeTrisCount(node->rChild);
			int leftOffset = node->lChild->offset < 0 ? ~node->lChild->offset : node->lChild->offset;
			int rightOffset = node->rChild->offset < 0 ? ~node->rChild->offset : node->rChild->offset;
			node->offset = ~std::min(leftOffset, rightOffset);
		}
		return node->t;
	}
	return 0;
}

unsigned int BVHOptimizer::ComputeNodesCount(BVHOptNode * node)
{
	if (node)
	{
		unsigned int nodeCount = 1;
		if (!node->IsLeaf())
		{ 
			nodeCount += ComputeNodesCount(node->lChild);
			nodeCount += ComputeNodesCount(node->rChild);
		} 
		return nodeCount;
	}
	return 0;
}

void BVHOptimizer::ReorderTriangles(BVHOptNode * root, BVHBottom * bvh)
{
	std::vector<BVHBottom::TriangleCol> trianglesCol(bvh->triangleCols.size());
	std::vector<BVHBottom::TriangleNorm> trianglesNorm(bvh->triangleNorms.size());
	std::vector<BVHBottom::TriangleTex> trianglesTex(bvh->triangleTexs.size());
	int j = 0;

	std::stack<BVHOptNode*> stack;
	stack.push(root);
	while (!stack.empty())
	{
		BVHOptNode * node = stack.top();
		stack.pop();
		if (node->IsLeaf())
		{
			int begin = node->offset;
			int end = node->offset + node->t;
			node->offset = j;
			for (int i = begin; i < end; ++i)
			{
				trianglesCol[j] = bvh->triangleCols[i];
				trianglesNorm[j] = bvh->triangleNorms[i];
				trianglesTex[j] = bvh->triangleTexs[i];
				++j;
			}
		}
		else
		{
			stack.push(node->GetLChild());
			stack.push(node->GetRChild());
		}
	}
	bvh->triangleCols = trianglesCol;
	bvh->triangleNorms = trianglesNorm;
	bvh->triangleTexs = trianglesTex;
}

void BVHOptimizer::CollapseLeaf(BVHOptNode * node)
{
	if (!node->IsLeaf())
	{
		BVHOptNode * left = node->GetLChild();
		if (!left->IsLeaf())
		{
			CollapseLeaf(left);
			float leafCost = CI * left->t;
			float subtreeCost = GetCost(left);
			if (leafCost <= subtreeCost)
			{
				delete left->GetLChild();
				delete left->GetRChild();
				left->lChild = nullptr;
				left->rChild = nullptr;
				left->offset = ~left->offset;
			}
			
		}
		BVHOptNode * right = node->GetRChild();
		if (!right->IsLeaf())
		{
			CollapseLeaf(right);
			float leafCost = CI * right->t;
			float subtreeCost = GetCost(right);
			if (leafCost <= subtreeCost)
			{
				delete right->GetLChild();
				delete right->GetRChild();
				right->lChild = nullptr;
				right->rChild = nullptr;
				right->offset = ~right->offset;
			}
			
		}
	}
}

void BVHOptimizer::CollapseLeaves(BVHOptNode * root, BVHBottom * bvh)
{
	ReorderTriangles(root, bvh);
	ComputeTrisCount(root);
	CollapseLeaf(root);
}

void BVHOptimizer::Optimize(BVHBottom * bvh, int ratio, int maxIterations, int threshold, int termThreshold)
{
	if ((bvh != nullptr) && (bvh->numberOfNodes > 1) && (maxIterations!=0) && ((threshold!=0) || (termThreshold!=0)))
	{
		wxStopWatch sw;
		sw.Start();
		BVHOptNode * optRoot = BVHToOpt(bvh);
		//int depth0 = bTree::BinTreeGetDepth(optRoot);
		float cost0 = GetCost(optRoot);
		bTree::BinTreeOptimize<BVHOptNode, bTree::MsumGetter<BVHOptNode > >(optRoot, ratio, maxIterations, threshold, termThreshold);
		float cost1 = GetCost(optRoot);
		//int depth1 = bTree::BinTreeGetDepth(optRoot);
		//int nodes1 = bTree::BinTreeCountNodes(optRoot);
		//int leaves1 = bTree::BinTreeCountLeaves(optRoot);
		//LOG(wxString::Format(wxT("<BVHOptimizer> depth (%i, %i), cost (%f, %f), nodes %i, leaves %i"), depth1, depth0, cost1, cost0, nodes1, leaves1));
		if (cost1 < cost0) 
		{
			CollapseLeaves(optRoot, bvh);
			//float cost2 = GetCost(optRoot);
			//int depth2 = bTree::BinTreeGetDepth(optRoot);
			OptToBVH(optRoot, bvh);
			sw.Pause();
			//LOG(wxString::Format(wxT("<BVHOptimizer> depth %i (%i, %i), cost %f (%f, %f) %f percent in %ld ms"), depth2, depth1, depth0, cost2, cost1, cost0, cost2 /cost0 * 100.0f, sw.Time()));
		}
		delete optRoot;
	}
}
