#include "BVHTop.h"
#include "BVHBinnedBuilder.h"
#include "BVHSAHBuilder.h"
#include "BVHOptimizer.h"

//Includes because of fonts...
#include "geometrytext.h"
#include "FontGlyphsTexture.h"
#include "imagecommon.h"

using namespace VRUT;

typedef std::pair<RayTracerBHVNode*, BVHTop::BVHCollapseNode*> CollapsePair;
typedef std::pair<BVHTop::BVHCollapseNode*, int> TopPair;

float BVHTop::GetSubtreeCost(int index) const 
{
	if (nodes[index].offset < 0)
	{
		// interior node
		BVHNode & left = nodes[nodes[index].t];
		BVHNode & right = nodes[nodes[index].t + 1];
		float sa = nodes[index].box.SurfaceArea();
		float leftCost = sa != 0.0f ? left.box.SurfaceArea() / sa * GetSubtreeCost(nodes[index].t) : 0.0f;
		float rightCost = sa != 0.0f ? right.box.SurfaceArea() / sa * GetSubtreeCost(nodes[index].t+1) : 0.0f;
		return CT + leftCost + rightCost; 
	}
	else
	{
		// leaf node
		NODE_ID nodeID = nodes[index].offset;
		GEOMETRY_ID gid = bvhNodeInfos[nodeID].geometryID;
		if (geometryAccelerators[gid])
			return geometryAccelerators[gid]->GetCost();
		return 0.0f;
	}
}

BVHTop::BVHCollapseNode * BVHTop::CollapseNodes(BVHCollapseNode * node) {
	if (node->active) 
	{
		if (node->leaf) 
		{
			return node;
		}
		else 
		{
			node->left = CollapseNodes(node->left);
			node->right = CollapseNodes(node->right);
			if (node->left != nullptr && node->right != nullptr)
				return node;
			else if (node->left != nullptr && node->right == nullptr)
				return node->left;
			else if (node->left == nullptr && node->right != nullptr)
				return node->right;
			else
				return nullptr;
		}
	}
	return nullptr;
}

void BVHTop::UpdateTopology(RayTracerBHV * bvh, SceneShadow * sceneShadow)
{

	// the ray casting has stack of depth 256, so we need to ensure that BVH has not bigger depth
	int depthBefore = bTree::BinTreeGetDepth(bvh->GetRoot());
	while (depthBefore > 200)
	{
		bvh->OptimizeIteration<bTree::MsumGetter<VRUT::RayTracerBHVNode> >(1);
		int depthAfter = bTree::BinTreeGetDepth(bvh->GetRoot());
		LOGWARNING(wxString::Format("<BVHTop> Too high BVH depth detected, optimize was used to decrease depth, depthBefore=%i, depthAfter=%i", depthBefore, depthAfter));
		depthBefore = depthAfter;
	}

	// aux. node buffer
	std::vector<BVHCollapseNode> auxNodes(bvh->GetNumberOfNodes());

	// number of actives nodes
	int numberOfActiveLeaves = 0;

	// copy BVH nodes
	std::stack<CollapsePair> collapseStack;
	collapseStack.push(CollapsePair(bvh->GetRoot(), &auxNodes[0]));
	size_t firstEmpty = 1;
	while (!collapseStack.empty())
	{

		// pop
		VRUT::RayTracerBHVNode * node = collapseStack.top().first;
		BVHCollapseNode * collapseNode = collapseStack.top().second;
		collapseStack.pop();

		// top node
		collapseNode->leaf = node->IsLeaf();
		collapseNode->box = node->GetAABB();
		collapseNode->offset = node->GetSceneNodeID();

		// leaf node
		if (node->IsLeaf())
		{
			collapseNode->active = sceneShadow->IsActive(node->GetSceneNodeID());
			if (collapseNode->active)
				++numberOfActiveLeaves;
		}

		// interior
		else
		{
			collapseNode->active = true;
			collapseNode->left = &auxNodes[firstEmpty++];
			collapseNode->right = &auxNodes[firstEmpty++];
			collapseStack.push(CollapsePair(node->GetLChild(), collapseNode->left));
			collapseStack.push(CollapsePair(node->GetRChild(), collapseNode->right));
		}

	}

	// empty bvh => return
	if (numberOfActiveLeaves == 0)
	{
		Clear();
		return;
	}

	// collapse inactive branches
	BVHCollapseNode * collapseRoot = CollapseNodes(&auxNodes[0]);

	// number of active nodes
	int numberOfActiveNodes = 2 * numberOfActiveLeaves - 1;

	// resize node buffer
	if (numberOfNodes != numberOfActiveNodes)
	{
		Clear();
		numberOfNodes = numberOfActiveNodes;
		nodes = (BVHNode*) VrutAlignedMalloc(numberOfNodes * sizeof(BVHNode), 16);
	}

	// copy collpase BVH
	std::stack<TopPair> topStack;
	topStack.push(TopPair(collapseRoot, 0));
	firstEmpty = 1;
	while (!topStack.empty())
	{

		// pop
		BVHCollapseNode * collapseNode = topStack.top().first;
		BVHNode * topNode = &nodes[topStack.top().second];
		topStack.pop();

		// bounding box
		topNode->box = collapseNode->box;

		// leaf node
		if (collapseNode->leaf)
		{
			topNode->offset = collapseNode->offset;
			topNode->t = 1;
		}

		// interior
		else
		{
			topNode->offset = ~2;
			topNode->t = firstEmpty;
			topStack.push(TopPair(collapseNode->left, firstEmpty++));
			topStack.push(TopPair(collapseNode->right, firstEmpty++));
		}

	}

}	

size_t BVHTop::GetNumberOfTriangles(smart_pointer<Geometry> g)
{
	size_t tris = 0;
	GeometryTriangles * geoTri = (GeometryTriangles*)g.getPtr();
	if (geoTri->type == Geometry::GEOMETRY_TRIANGLE)
	{
		for (size_t j = 0; j < geoTri->triDescList.size(); j++)
		{
			long count = long(((j + 1 < geoTri->triDescList.size()) ? geoTri->triDescList.at(j + 1).first:geoTri->indices.size()) - geoTri->triDescList.at(j).first);
			switch (geoTri->triDescList.at(j).second)
			{
			case GeometryTriangles::TRI_LIST:
				tris += count / 3;
				break;
			case GeometryTriangles::TRI_FAN:
			case GeometryTriangles::TRI_STRIP:
			case GeometryTriangles::POLYGON:
				tris += count - 2;
				break;
			case GeometryTriangles::QUADS:
				tris += count / 2;
				break;
			default:
				break;
			}
		}
	}
	return tris;
}

BVHTop::BVHTop() : workersBuildgeoAcc(false)
{
	numberOfNodes = 1;
	nodes = (BVHNode*) VrutAlignedMalloc(sizeof(BVHNode), 16);
#ifdef CLEANUP_ACCELERATORS
	resetGeometry = true;
#endif
}

BVHTop::~BVHTop()
{
	for (auto &geometryAccelerator : geometryAccelerators)
		SAFE_DELETE(geometryAccelerator);
}

void BVHTop::Update(RayTracerBHV * bvh, SceneShadow * sceneShadow, MaterialCache & materialCache)
{
	smart_pointer_v2<Scene> &scene = sceneShadow->GetScene();

	// resize leaf buffer
	bvhNodeInfos.resize(scene->GetNodeMaxID());

	// if there some new geometry - needed to build some new geometry accelerators
	size_t newSize = scene->GetGeometryMaxID();
	if (newSize != geometryAccelerators.size())
	{
		// resize the array used to store pointers to created geometry accelerators
		wxCriticalSectionLocker locker(geometricAccBuildCriticalSection);
		geometryAccelerators.resize(newSize, nullptr);
		resetGeometries.resize(newSize, true);
	}

	// push root
	std::stack<VRUT::RayTracerBHVNode *> stack;
	stack.push(nullptr);
	VRUT::RayTracerBHVNode * node = bvh->GetRoot();

	// nullptr root
	if (!node)
	{
		Clear();
		nodes = (BVHNode*)VrutAlignedMalloc(1 * sizeof(BVHNode), 16);
		nodes[0].box.SetInvalid();
		return;
	}

	materialCache.ResizeRTMaterials(scene->GetMaterialMaxID());

	// traversing scene BVH
	while (!stack.empty())
	{
		// if node is leaf copy data to bvhNodeInfo
		if (node->IsLeaf())
		{
			size_t nodeId = node->GetSceneNodeID();
			if (nodeId < bvhNodeInfos.size())
			{
				SceneNodeShadow *sNode=sceneShadow->GetNode(nodeId);
				if (sNode && (sNode->GetType() == SceneNode::GEOMETRY))
				{
					//bvhNodeInfos[nodeId].worldMat = *sceneShadow->GetWorldMatrix(nodeId);
					bvhNodeInfos[nodeId].worldMat = node->GetWorldMatrix();
					bvhNodeInfos[nodeId].invWorldMat = bvhNodeInfos[nodeId].worldMat;
					bvhNodeInfos[nodeId].isMatUnit = bvhNodeInfos[nodeId].worldMat.IsUnit();
					if (!bvhNodeInfos[nodeId].isMatUnit)
						bvhNodeInfos[nodeId].invWorldMat.InvertThisPrecise();
					bvhNodeInfos[nodeId].geometryID = sNode->GetGeometryID();
					if ((bvhNodeInfos[nodeId].geometryID != GEOMETRY_ID_NONE) && resetGeometries[bvhNodeInfos[nodeId].geometryID])
					{
						resetGeometries[bvhNodeInfos[nodeId].geometryID] = false;
						SAFE_DELETE(geometryAccelerators[bvhNodeInfos[nodeId].geometryID]);
						geometryAccRebuildQueue.Post(new GeometryAccRebuildQueueNode(bvhNodeInfos[nodeId].geometryID));
					}

					//Handling text geometry
					//Material must be merged from both assigned material and font material (which is not between materials)
					auto geo = scene->GetGeometry(bvhNodeInfos[nodeId].geometryID);
					//TODO: commented out because of problems with GLEW initialization
					//CopyDDSTexDataToRayTracerTexData(bvhNodeInfo[nodeId].material.getPtr());

					bvhNodeInfos[nodeId].material = materialCache.GetOrCreateRTMaterial(sNode->GetMaterialID(), scene, geo);
				}
			}
			node = stack.top();
			stack.pop();
		}
		// interior node
		else
		{
			if (node->GetRChild())
				stack.push(node->GetRChild());
			node = node->GetLChild();
		}
	}

#ifdef CLEANUP_ACCELERATORS
	if (resetGeometry)
	{
		for (size_t i = 0; i<resetGeometries.size(); ++i)
		{
			if (resetGeometries[i])
			{
				resetGeometries[i] = false;
				SAFE_DELETE(geometryAccelerators[i]);
			}
		}
		resetGeometry = false;
	}
#endif

	// update rayTracer simple BVH by scene BVH
	UpdateTopology(bvh, sceneShadow);

}

void BVHTop::DeleteGeometryCache()
{
	resetGeometries.clear();
#ifdef CLEANUP_ACCELERATORS
	resetGeometry = false;
#endif
}

void BVHTop::InvalidateGeometryCacheEntry(unsigned int index)
{
	if (resetGeometries.size() <= index)
		resetGeometries.resize(index + 1, true);
	resetGeometries[index] = true;
#ifdef CLEANUP_ACCELERATORS
	resetGeometry = true;
#endif
}

void BVHTop::WorkersBuildGeoAccStart()
{
	workersBuildgeoAcc = true;
}
		
void BVHTop::WorkersBuildGeoAccStop()
{
	workersBuildgeoAcc = false;
	// Wake up waiting threads
	geometryAccRebuildQueue.CancelTimeout();
}

bool BVHTop::BuildGeomAcc(smart_pointer_v2<Scene> & scene, int optRatio, int optMaxIterations, int optThreshold, int optRandomThreshold)
{
	unsigned long timeout = workersBuildgeoAcc ? 1 : 0; // Do not wait, when queue is filled
	GeometryAccRebuildQueueNode *queueNode = geometryAccRebuildQueue.ReceiveTimeout(timeout);
	if (queueNode)
	{
		smart_pointer<Geometry> g;
		scene->GetGeometry(g, queueNode->gid);
		if (g)
		{
			if (geometryAccelerators[queueNode->gid] == nullptr)
			{
				//LOG(wxString::Format(wxT("BVHTop::BuildGeomAcc start %lu"), (unsigned long)gid));
				size_t tris = GetNumberOfTriangles(g);
				BVHBottom * ge = nullptr;
				if (tris <= 150000)
					ge = BVHSAHBuilder().Build(g);
				else
					ge = BVHBinnedBuilder().Build(g);
				//LOG(wxString::Format(wxT("BVHTop::BuildGeomAcc build %lu"), (unsigned long)queueNode->gid));
				BVHOptimizer().Optimize(ge, optRatio, optMaxIterations, optThreshold, optRandomThreshold);
				//LOG(wxString::Format(wxT("BVHTop::BuildGeomAcc end %lu"), (unsigned long)gid));
				geometryAccelerators[queueNode->gid] = ge;
			}
		}
		SAFE_DELETE(queueNode);
		return true;
	}
	return workersBuildgeoAcc;
}

bool BVHTop::IsEmpty() const
{
	return numberOfNodes == 0;
}

float BVHTop::Radius(void) const
{
	if (IsEmpty())
		return 0.0f;
	else
		return nodes[0].box.Radius();
}

const std::vector<BVHTop::BVHNodeInfo> & BVHTop::GetNodeInfos() const 
{
	return bvhNodeInfos;
}

const std::vector<BVHBottom*> & BVHTop::GetGeomAccelerators() const
{
	return geometryAccelerators;
}
