#include "BaseTreePoolBuilder.h"
#include <algorithm>
#include "../../inc/tbb/parallel_sort.h"

//************************************
// Calculates the minimum size of the node pool texture, so that all nodes + pointers fit
//************************************
size_t BaseTreePoolBuilder::GetPoolSize(const BaseTree* tree)
{
	return GetMinimumNodePoolTexelCount(tree);
}

size_t BaseTreePoolBuilder::GetMinimumNodePoolTexelCount(const BaseTree* tree)
{
	assert(tree != NULL);
	if (!mIsBuilding)
		InitBuild(tree);
	size_t texelCount = GetTreeInfoBytesSize(tree);
	for (unsigned32 i = 0; i < (unsigned32)tree->GetNodeCount(); i++)
		texelCount += GetNodeSize(tree, i);
	if (!mIsBuilding)
		FinishBuild(tree);
	return texelCount;
}

unsigned32 BaseTreePoolBuilder::GetTreeInfoBytesSize(const BaseTree* tree) const
{
	assert(tree != NULL);
	bool hasAdditionalPointerBytes = HasAdditionalBytesPerPointer(tree);
	bool hasAdditionalBytesPerNode = HasAdditionalBytesPerNode(tree);

	unsigned8 depth = tree->GetMaxLevel();
	size_t treeInfoSize = (hasAdditionalPointerBytes ? (depth + 1) : 0) // additional pointer sizes(if not 0, 8 bits = 1 byte per level)
		+ (hasAdditionalBytesPerNode ? (depth + 1) : 0) //and additional node sizes(if not 0, 8 bits = 1 byte per level)
		+ tree->GetAdditionalTreeInfoSize()
		+ (unsigned32)GetPoolInfoSize(tree)
		+ (unsigned32)GetAdditionalPoolInfoSize(tree);
	if (WordAligned()) RoundToWords(treeInfoSize);
	return (unsigned32)treeInfoSize;
}

//************************************
// Insert all nodes into final node pool and updates pointers
//************************************
bool BaseTreePoolBuilder::BuildPool(const BaseTree* tree, std::vector<unsigned8>& pool) {
	if (tree == NULL) return false;

	// Notify subclasses that the build is initiating (so they can do pre-calculations)
	InitBuild(tree);
	mIsBuilding = true;

	// Initialize the final node pool
	size_t poolSize = GetMinimumNodePoolTexelCount(tree);
	pool.resize(poolSize, 0);

	// Get information about the node sizes per level

	std::vector<unsigned8> additionalBytesPerPointer = tree->GetAdditionalBytesPerPointer();
	std::vector<unsigned8> additionalBytesPerNode = tree->GetAdditionalBytesPerNode();
	bool hasAdditionalPointerBytes = HasAdditionalBytesPerPointer(tree);
	bool hasAdditionalBytesPerNode = HasAdditionalBytesPerNode(tree);
	bool lastChildHasAdditionalBytesPerPointer = LastChildHasAdditionalBytesPerPointer(tree);

	// Find an ordering of the nodes such that all children of a node appear after that node in memory.
	unsigned32 nodeCount = (unsigned32)tree->GetNodeCount();
	unsigned8 depth = tree->GetMaxLevel();

	std::vector<unsigned32> nodeMap(nodeCount);
	for (unsigned32 i = 0; i < nodeCount; i++) nodeMap[i] = i;
	OrderNodes(tree, nodeMap);
	std::vector<unsigned32> reverseNodeMap(nodeCount);
	for (unsigned32 i = 0; i < nodeCount; i++) reverseNodeMap[nodeMap[i]] = i;

	// Calculate all node indices beforehand (also makes calculating the level offsets easy)
	std::vector<size_t> nodePointers(nodeCount);
	size_t treeInfoSize = GetTreeInfoBytesSize(tree);
	size_t curIndex = treeInfoSize;
	for (unsigned32 i = 0; i < nodeCount; i++)
	{
		unsigned32 nodeId = nodeMap[i];
		nodePointers[nodeId] = curIndex;
		curIndex += GetNodeSize(tree, nodeId);
	}

	// Start building the pool.
	curIndex = 0;
	// First the (subclass specific) tree/pool information
	size_t poolInfoSize = GetPoolInfoSize(tree);
	if (poolInfoSize != 0)
	{
		std::vector<unsigned8> poolInfo = GetPoolInfo(tree, nodePointers, nodeMap);
		assert(poolInfo.size() == poolInfoSize);
		std::move(poolInfo.begin(), poolInfo.end(), pool.begin());
		curIndex += poolInfoSize;
	}

	// Write the additional node bytes (size per level)
	if (hasAdditionalBytesPerNode)
		for (unsigned8 level = 0; level <= depth; level++)
			pool[curIndex++] = additionalBytesPerNode[level];

	// Write the additional pointer bytes (size per level)
	if (hasAdditionalPointerBytes)
		for (unsigned8 level = 0; level <= depth; level++)
			pool[curIndex++] = additionalBytesPerPointer[level];

	// Write the additional tree info from the tree itself (MultiRoot uses this for root pointers)
	// TODO: Find a better way to do this, don't let the tree itself determine part of how to pool looks
	assert(GetAdditionalTreeInfoStart(tree) == curIndex);
	std::vector<unsigned8> additionalTreeInfo = tree->GetAdditionalTreeInfo(nodePointers);
	assert(additionalTreeInfo.size() == tree->GetAdditionalTreeInfoSize());
	std::move(additionalTreeInfo.begin(), additionalTreeInfo.end(), pool.begin() + curIndex);
	curIndex += additionalTreeInfo.size();

	// Write additional pool info (e.g. lookup tables)
	assert(GetAdditionalPoolInfoStart(tree) == curIndex);
	size_t additionalPoolInfoSize = GetAdditionalPoolInfoSize(tree);
	if (additionalPoolInfoSize != 0)
	{
		std::vector<unsigned8> additionalPoolInfo = GetAdditionalPoolInfo(tree, nodePointers, nodeMap);
		assert(additionalPoolInfo.size() == additionalPoolInfoSize);
		std::move(additionalPoolInfo.begin(), additionalPoolInfo.end(), pool.begin() + curIndex);
		curIndex += additionalPoolInfoSize;
	}

	// Assert that the nodes start at the expected position
	if (WordAligned()) RoundToWords(curIndex);
	assert(curIndex == nodePointers[nodeMap[0]]);

	// Write the nodes
#ifdef _DEBUG
	for (unsigned32 i = 0; i < nodeCount; i++)
#else
	tbb::parallel_for((unsigned32)0, nodeCount, [&](const unsigned32& i)
#endif

	{
		unsigned32 nodeId = nodeMap[i];
		std::vector<unsigned8> bytesForNode = GetBytesForNode(tree, nodeId, nodePointers, reverseNodeMap, additionalBytesPerNode, additionalBytesPerPointer);
		assert(bytesForNode.size() == GetNodeSize(tree, nodeId));
		size_t nodeIndex = nodePointers[nodeId];
		std::move(bytesForNode.begin(), bytesForNode.end(), pool.begin() + nodeIndex);
#ifdef _DEBUG
	}
#else 
	});
#endif
	FinishBuild(tree);
	mIsBuilding = false;
	return true;
}

size_t BaseTreePoolBuilder::GetAdditionalTreeInfoStart(const BaseTree* tree) const
{
	std::vector<unsigned8> additionalBytesPerPointer = tree->GetAdditionalBytesPerPointer();
	std::vector<unsigned8> additionalBytesPerNode = tree->GetAdditionalBytesPerNode();
	bool hasAdditionalPointerBytes = false; for (auto abpp : additionalBytesPerPointer) if (abpp != 0) hasAdditionalPointerBytes = true;
	bool hasAdditionalBytesPerNode = false; for (auto abpn : additionalBytesPerNode) if (abpn != 0) hasAdditionalBytesPerNode = true;
	unsigned8 depth = tree->GetMaxLevel();
	size_t poolInfoSize = GetPoolInfoSize(tree);
	size_t additionalInfoSize = ((hasAdditionalBytesPerNode ? 1 : 0) + (hasAdditionalPointerBytes ? 1 : 0)) * (depth + 1);
	return poolInfoSize + additionalInfoSize;
}

size_t BaseTreePoolBuilder::GetAdditionalPoolInfoStart(const BaseTree* tree) const
{
	return GetAdditionalTreeInfoStart(tree) + tree->GetAdditionalTreeInfoSize();
}

void BaseTreePoolBuilder::OrderNodes(const BaseTree* tree, std::vector<unsigned32>& nodeOrder) const
{
	tbb::parallel_sort(nodeOrder.begin(), nodeOrder.end(), [&](const unsigned32& i1, const unsigned32& i2)
	{
		const Node* a = tree->GetNode(i1);
		const Node* b = tree->GetNode(i2);
		return a->GetLevel() < b->GetLevel();
	});
}

std::vector<unsigned8> BaseTreePoolBuilder::GetBytesForNode(const BaseTree* tree, const unsigned32& nodeId, const std::vector<size_t>& nodePointers, const std::vector<unsigned32>& reverseNodeMap,
	const std::vector<unsigned8>& fullAdditionalBytesPerNode, const std::vector<unsigned8>& fullAdditionalBytesPerPointer) const
{
	const Node* node = tree->GetNode(nodeId);
	unsigned8 childCount = node->GetChildCount();
	unsigned32* children = node->GetChildren();
	unsigned8 level = node->GetLevel();
	unsigned8 additionalBytesPerNode = fullAdditionalBytesPerNode[level];
	unsigned8 additionalBytesPerPointer = fullAdditionalBytesPerPointer[level];

	size_t nodeSize = GetNodeSize(tree, nodeId);
	std::vector<unsigned8> bytes(nodeSize);
	size_t curIndex = 0;

	// Add the childmask
	bytes[0] = node->GetChildmask().mask;
	curIndex++;

	if (GetAdditionalPoolInfoForNodeSize(tree, nodeId) != 0)
	{
		auto additionalPoolInfoForNode = GetAdditionalPoolInfoForNode(tree, nodeId, reverseNodeMap[nodeId]);
		std::move(additionalPoolInfoForNode.begin(), additionalPoolInfoForNode.end(), bytes.begin() + curIndex);
		assert(additionalPoolInfoForNode.size() == GetAdditionalPoolInfoForNodeSize(tree, nodeId));
		curIndex += additionalPoolInfoForNode.size();
	}

	// Add the additional node bytes
	if (additionalBytesPerNode != 0)
	{
		std::vector<unsigned8> additionalNodeBytes = tree->GetAdditionalNodeBytes(node);
		assert(additionalNodeBytes.size() == additionalBytesPerNode);
		std::move(additionalNodeBytes.begin(), additionalNodeBytes.end(), bytes.begin() + curIndex);
		curIndex += additionalNodeBytes.size();
	}
	if (WordAligned()) RoundToWords(curIndex);

	// Write the node pointers
	for (ChildIndex c = 0; c < childCount; c++)
	{
		unsigned32 childNodeIndex = children[c];
		size_t pointer = nodePointers[childNodeIndex];
		std::vector<unsigned8> pointerBytes = WrapPointer(tree, childNodeIndex, reverseNodeMap[childNodeIndex], (unsigned32)pointer);
		std::move(pointerBytes.begin(), pointerBytes.end(), bytes.begin() + curIndex);
		curIndex += pointerBytes.size();
		if (WordAligned()) RoundToWords(curIndex);
	}

	// Followed by the additional pointer info
	if (additionalBytesPerPointer != 0)
	{
		ChildIndex i = 0;
		for (ChildIndex c = 0; c < 8; c++)
		{
			if (node->HasChild(c) && (i < childCount - 1 || LastChildHasAdditionalBytesPerPointer(tree)))
			{
				i++;
				std::vector<unsigned8> additionalBytesForPointer = tree->GetAdditionalPointerBytes(node, c);
				if (WordAligned())
				{
					curIndex += additionalBytesForPointer.size();
					RoundToWords(curIndex);
					curIndex -= additionalBytesForPointer.size();
				}
				std::move(additionalBytesForPointer.begin(), additionalBytesForPointer.end(), bytes.begin() + curIndex);
				curIndex += additionalBytesForPointer.size();
				if (WordAligned()) RoundToWords(curIndex);
			}
		}
	}
	return bytes;
}

bool BaseTreePoolBuilder::VerifyPool(std::vector<unsigned8>& pool, const unsigned8& treeDepth) const
{
	// Verify that the last level offset is not bigger than the size of the pool
	for (unsigned8 level = 0; level <= treeDepth; level++)
	{
		unsigned32 levelOffset = 0;
		BitHelper::JoinBytes(pool, levelOffset, level * 4);
		if (levelOffset > pool.size()) return false;
	}
	return true;
}