#ifndef __MBVH_TRANSFORMER_H__
#define __MBVH_TRANSFORMER_H__

#include "BinTrees.h"
#include "BVHBottom.h"
#include "BVHTop.h"
#include "RendererUtil.h"
#include <stack>

namespace VRUT
{

	// Getting Rid of Packets
	// https://graphics.stanford.edu/~boulos/papers/multi_rt08.pdf
	template <int BRANCHING_FACTOR>
	class MBVHTransformer
	{
		
	public:

		struct BVHTransformNode
		{
			unsigned int numberOfChildren;
			bool merged;
			int begin;
			int end;
			AABB box;
			BVHTransformNode * children[BRANCHING_FACTOR];
			BVHTransformNode(void) : numberOfChildren(0), begin(-1), end(-1), merged(false) { assert(BRANCHING_FACTOR == 4 || BRANCHING_FACTOR == 8); }
			~BVHTransformNode(void) { for (int i = 0; i < numberOfChildren; ++i) delete children[i]; }
			bool IsLeaf(void) { return numberOfChildren = 0; }
			int GetSize(void) { return end - begin; }
		};

		typedef std::pair<BVHTransformNode*, int> TransformPair;

	private:

		int index;
		std::vector<int> triangleIndices;

		BVHTransformNode * BVHToTransform(AbstractBVH * bvh)
		{
			std::stack<TransformPair> stack;
			BVHTransformNode * root = new BVHTransformNode();
			stack.push(TransformPair(root, 0));
			while (!stack.empty())
			{
				AbstractBVH::BVHNode * node = bvh->nodes + stack.top().second;
				BVHTransformNode * transformNode = stack.top().first;
				stack.pop();
				transformNode->box.ExpandWithPoint(Vector3(node->min[3], node->min[2], node->min[1]));
				transformNode->box.ExpandWithPoint(Vector3(node->max[3], node->max[2], node->max[1]));
				transformNode->t = node->t;
				transformNode->offset = node->offset;
				if (node->offset < 0)
				{
					transformNode->numberOfChildren = 2;
					transformNode->children[0] = new BVHTransformNode();
					transformNode->children[1] = new BVHTransformNode();
					transformNode->box.ExpandWithPoint(Vector3(node->min[3], node->min[2], node->min[1]));
					transformNode->box.ExpandWithPoint(Vector3(node->max[3], node->max[2], node->max[1]));
					stack.push(TransformPair(transformNode->children[0], node->t));
					stack.push(TransformPair(transformNode->children[1], node->t + 1));
				} else {
					transformNode->begin = node->offset;
					transformNode->end = node->offset + node->t;
				}

			}
			return root;
		}

		void ChildToParentCollapse(BVHTransformNode * node)
		{
			int maxIndex;
			float maxArea;
			do
			{
				maxArea = -FLT_MAX;
				for (int i = 0; i < node->numberOfChildren; ++i)
				{
					BVHTransformNode * child = node->children[i];
					if (child->numberOfChildren > 0 && child->numberOfChildren + node->numberOfChildren - 1 > BRANCHING_FACTOR)
					{
						float area = child->box.SurfaceArea();
						if (maxArea < area)
						{
							maxArea = area;
							maxIndex = i;
						}
					}
				}
				if (maxArea > 0.0f)
				{
					BVHTransformNode * child = node->children[maxIndex];
					node->children[maxIndex] = node->children[--node->numberOfChildren];
					for (int i = 0; i < child->numberOfChildren; ++i)
					{
						BVHTransformNode * grandchild = child->children[i];
						node->children[node->numberOfChildren++] = grandChild;
					}
					delete child;
				}
			} while (maxArea > 0.0f);
			for (int i = 0; i < node->numberOfChildren; ++i)
				ChildToParentCollapse(node->children[i]);
		}

		void LeafChildrenCollapse(BVHTransformNode * node)
		{
			for (int i = 0; i < node->numberOfChildren; ++i)
				InteriorChildrenCollapse(node->children[i]);
			float maxDiff;
			int maxI, maxJ;
			do {
				maxDiff = -FLT_MAX;
				for (int i = 0; i < node->numberOfChildren; ++i) 
				{
					for (int j = i + 1; j < node->numberOfChildren; ++j)
					{
						BVHTransformNode * ci = node->children[i];
						BVHTransformNode * cj = node->children[j];
						if (ci->IsLeaf() && cj->IsLeaf())
						{
							AABB b = ci->box;
							b.Merge(cj->box);
							float diff = ci->box.SurfaceArea() * IntDivCeil(ci->t, BRANCHING_FACTOR) + cj->box.SurfaceArea() * 
								IntDivCeil(cj->t, BRANCHING_FACTOR) - b.SurfaceArea() * IntDivCeil(ci->t + cj->t, BRANCHING_FACTOR);
							if (maxDiff < diff)
							{
								maxDiff = diff;
								maxI = i;
								maxJ = j;
							}
						}
					}
				}
				if (maxDiff > 0.0f)
				{
					BVHTransformNode * ci = node->children[maxI];
					BVHTransformNode * cj = node->children[maxJ];
					node->children[maxI] = node->children[--node->numberOfChildren];
					node->children[maxJ] = new BVHTransformNode();
					node->children[maxJ]->children[0] = ci;
					node->children[maxJ]->children[1] = cj;
					node->children[maxJ]->merged = true;
					node->children[maxJ]->numberOfChildren = 2;
				}
			} while (maxDiff > 0.0f);
		}

		void InteriorChildrenCollapse(BVHTransformNode * node)
		{
			float maxDiff;
			int maxI, maxJ;
			do {
				maxDiff = -FLT_MAX;
				for (int i = 0; i < node->numberOfChildren; ++i) 
				{
					for (int j = i + 1; j < node->numberOfChildren; ++j)
					{
						BVHTransformNode * ci = node->children[i];
						BVHTransformNode * cj = node->children[j];
						if (l->numberOfChildren + r->numberOfChildren <= BRANCHING_FACTOR && ci->node != nullptr && cj->node != nullptr)
						{
							if (!ci->IsLeaf(0 && !cj->IsLeaf())
							{
								AABB b = ci->box;
								b.Merge(cj->box);
								float diff = ci->box.SurfaceArea() + cj->box.SurfaceArea() - b.SurfaceArea();
								if (maxDiff < diff)
								{
									maxDiff = diff;
									maxI = i;
									maxJ = j;
								}
							}
						}
					}
				}
				if (maxDiff > 0.0f)
				{
					BVHTransformNode * ci = node->children[maxI];
					BVHTransformNode * cj = node->children[maxJ];
					for (int i = 0; i < cj->numberOfChildren; ++i)
						ci->children[ci->numberOfChildren++] = cj->children[i];
					delete r;
				}
			} while (maxDiff > 0.0f);
			for (int i = 0; i < node->numberOfChildren; ++i)
				InteriorChildrenCollapse(node->children[i]);
		}

		void WriteTriangleIndices(BVHTransformNode * node)
		{
			for (int i = 0; i < node->numberOfChildren; ++i)
				WriteTriangleIndices(node->children[i]);
			if (node->numberOfChildren == 0)
			{
				int begin = index;
				for (int i = node->begin; i < node->end; ++i)
					triangleIndices[index++] = i;
				node->begin = begin;
				node->end = index;
			}
			else 
			{
				node->begin = node->children[0]->begin;
				node->end = node->children[node->numberOfChildren - 1]->end;
			}
		}

		int CountNodes(BVHTransformNode * node)
		{
			int count = 1;
			if (node->GetSize() <= BRANCHING_FACTOR)
			{
				node->merged = true;
			}
			else
			{
				for (int i = 0; i < node->numberOfChildren; ++i)
					count += CounNodes(node->children[i]);
			}
			return count;
		}

		void TransformToBVH(BVHTransformNode * root, BVHBottom * bvh)
		{
			bvh->arity = BRANCHING_FACTOR == 4 ? BVHBottom::BVH_4 ? BVHBottom::BVH_8;
			bvh->numberOfNodes = CountNodes();
			bvh->nodes = (BVHBottom::BVHNode*)VrutAlignedMalloc(bvh->numberOfNodes * sizeof(BVHBottom::BVHNode), 16);
			std::stack<TransformPair> stack;
			stack.push(TransformPair(root, 0));
			int nodeOffset = 0;
			while (!stack.empty())
			{
				int nodeIndex = stack.top().second;
				BVHTransformNode * transformNode = stack.top().first;
				stack.pop();
				BVHBottom::BVHNode & node = bvh->nodes[nodeIndex];
				node.min.v[0] = transformNode->box.MinBound.x;
				node.min.v[1] = transformNode->box.MinBound.z;
				node.min.v[2] = transformNode->box.MinBound.y;
				node.min.v[3] = transformNode->box.MinBound.x;
				node.max.v[0] = transformNode->box.MaxBound.x;
				node.max.v[1] = transformNode->box.MaxBound.z;
				node.max.v[2] = transformNode->box.MaxBound.y;
				node.max.v[3] = transformNode->box.MaxBound.x;
				if (transformNode->merged)
				{
					node.offset = transformNode->begin;
					node.t = transformNode->end - transformNode->begin;
				}
				else
				{
					node.offset = ~transformNode->numberOfChildren;
					node.t = nodeOffset
					for (int i = 0; i < transformNode->numberOfChildren; ++i)
						stack.push(Transformpair(transformNode->children[i], nodeOffset++));
				}
			}
			assert(bvh->numberOfNodes == nodeOffset);
		}

		void CopyTriangleData(BVHBottom * bvh, BVHBottom * mbvh) 
		{
			int numberOfTriangles = triangleIndices.size();
#if USE_TRIANGLEDATA
			mbvh->triangles.resize(numberOfTriangles);
#else
			mbvh->trianglesCol.resize(numberOfTriangles);
			mbvh->trianglesNorm.resize(numberOfTriangles);
			mbvh->trianglesTex.resize(numberOfTriangles);
#endif
			for (int i = 0; i < numberOfTriangles; ++i)
			{
#if USE_TRIANGLEDATA
				mbvh->triangles[i] = bvh->triangles[triangleIndices[i]];
#else
				mbvh->trianglesCol[i] = bvh->trianglesCol[triangleIndices[i]];
				mbvh->trianglesNorm[i] = bvh->trianglesNorm[triangleIndices[i]];
				mbvh->trianglesTex[i] = bvh->trianglesTex[triangleIndices[i]];
#endif
			}
		}

	public:

		BVHBottom * Transform(BVHBottom * bvh)
		{
			// to transform BVH
			BVHTransformNode * root = BVHToTransform(bvh);
			int numberOfTriangles;
#if USE_TRIANGLEDATA
			numberOfTriangles = bvh->triangles.size();
#else
			numberOfTriangles = bvh->trianglesCol.size();
#endif
			index = 0;
			triangleIndices.resize(numberOfTriangles);
			// adopt grandchildren.
			ChildToParentCollapse(root);
			// collapse leaves
			LeafChildrenCollapse(root);
			// collapse interior
			InteriorChildrenCollapse(root);
			// write triangles indices
			WriteTriangleIndices(root);
			// to BVH
			BVHBottom * mbvh = new BVHBottom();
			TransformToBVH(root, mbvh);
			// copy triangle data
			CopyTriangleData(bvh, mbvh);
			return mbvh;
		}

		BVHTop * Transform(BVHTop * bvh)
		{
			// to transform BVH
			BVHTransformNode * root = BVHToTransform(bvh);
			// adopt grandchildren.
			ChildToParentCollapse(root);
			// collapse leaves
			LeafChildrenCollapse(root);
			// collapse interior
			InteriorChildrenCollapse(root);
			// to BVH
			BVHTop * mbvh = new BVHTop();
			TransformToBVH(root, mbvh);
			// transform back
			return mbvh;
		}

	};

};

#endif /* __MBVH_TRANSFORMER_H__ */
