#pragma once
#include "Tree.h"
#include "EdgeMaterialNode.h"
#include "IBlockTexture.h"
#include "IMaterialTexture.h"
#include "IAdditionalProperties.h"
#include "../../inc/tbb/parallel_for_each.h"
#include "../../core/BitHelper.h"
#include "../../core/Serializer.h"
#include "../../core/Util/BoolArray.h"
#include "../TextureCompressor/CompressedTexture.h"
#include "../../inc/lodepng/lodepng.h"
#include <unordered_map>
#include <map>
#include <stack>

typedef EdgeMaterialNode<unsigned64> UniqueIndexNode;

// Comment or uncomment these defines to enable or disable offset/shift compression
#define offset_sizes_per_level
#define implicit_store_first_offset

template<typename T>
class UniqueIndexTree : public Tree<UniqueIndexNode>, public IBlockTexture, public IAdditionalProperties
{
private:
	std::vector<unsigned8>* mShiftBytesPerLevel;

	// After the tree is finalized, this vector contains the materialPointers for all nodes in the original octree
	CompressedTexture<T>* mNodeMaterials;
	unsigned8 mUniqueShiftLevel;

	// Textures and pools
	std::vector<unsigned8>		mBlockPointerPool;
	bool						mBlockPointerPoolLoaded;
	std::vector<unsigned8>		mBlockPool;
	bool						mBlockPoolLoaded;
	std::map<std::string, std::string> mAdditionalProperties;
	bool						mAdditionalPropertiesLoaded;

	inline void WriteBlockPointerPool(std::ostream& file)
	{
		// Make sure the block pointer pool exists
		GetBlockPointerPool();

		// Write it to the filestream
		Serializer<std::vector<unsigned8>, unsigned64>::Serialize(mBlockPointerPool, file);
	}

	inline void ReadBlockPointerPool(std::istream& file)
	{
		// Read from filestream
		Serializer<std::vector<unsigned8>, unsigned64>::Deserialize(mBlockPointerPool, file);
		mBlockPointerPoolLoaded = true;
	}

	inline void WriteBlockPool(std::ostream& file)
	{
		// Make sure the block pool exists
		GetBlockPool();

		Serializer<std::vector<unsigned8>, unsigned64>::Serialize(mBlockPool, file);
	}

	inline void ReadBlockPool(std::istream& file)
	{
		Serializer<std::vector<unsigned8>, unsigned64>::Deserialize(mBlockPool, file);
		mBlockPoolLoaded = true;
	}

	inline void WriteAdditionalProperties(std::ostream& file)
	{
		// Make sure the additional properties are set.
		GetAdditionalProperties();

		// Write them to a file
		Serializer<std::map<std::string, std::string>>::Serialize(mAdditionalProperties, file);
	}

	inline void ReadAdditionalProperties(std::istream& file)
	{
		Serializer<std::map<std::string, std::string>>::Deserialize(mAdditionalProperties, file);
		mAdditionalPropertiesLoaded = true;
	}

	void ClearPooledNodeMaterials()
	{
		if (mBlockPoolLoaded && mBlockPointerPoolLoaded && mAdditionalPropertiesLoaded) ClearNodeMaterials();
	}

	void ClearNodeMaterials()
	{
		if (mNodeMaterials != NULL)
		{
			delete mNodeMaterials;
			mNodeMaterials = NULL;
		}
	}

	unsigned64 GetShift(const UniqueIndexNode* node, const ChildIndex& child) const
	{
		return node->GetEdgeMaterial(child);
	}

	// Calculates for each node whether the edge pointing to it has a shift of zero on it
	// (incorrect if the tree is a DAG, as the results might be ambiguous if nodes have multiple parents).
	BoolArray GetNodesWithZeroShiftParents()
	{
		BoolArray hasZeroShiftParent(GetNodeCount());
		for (unsigned32 i = 1; i < GetNodeCount(); i++)
		{
			UniqueIndexNode* cur = (UniqueIndexNode*)GetNode(i);
			unsigned64* childShifts = cur->GetEdgeMaterials();
			unsigned32* children = cur->GetChildren();
			ChildIndex childCount = cur->GetChildCount();
			for (ChildIndex c = 0; c < childCount; c++)
				if (childShifts[c] == 0) hasZeroShiftParent.Set(children[c], true);
		}
		for (ChildIndex c = 0; c < GetChildCount(); c++)
			if (mChildShifts[c] == 0) hasZeroShiftParent.Set(mChildren[c], true);
		return hasZeroShiftParent;
	}

	size_t GetHighestShiftSum(const UniqueIndexNode* node) const
	{
		// The first child is always the child with the highest shift
		if (node->HasChildren()) return node->GetEdgeMaterials()[0] + GetHighestShiftSum(GetTypedNode(node->GetChildren()[0]));
		else return 0;
	}

	// Create a vector containing all materials from the original tree
	std::vector<T> GetMappedMaterialTexture(const BaseTree* tree, std::vector<T>& materialPerNode)
	{
		std::vector<T> materialTexture;
		// Go depth-first through all nodes in the tree and copy the materials as they are reached.
		std::stack<unsigned32> curStack;
		curStack.push(0); // Push the root index to the stack
		bool skipNext = false;
		while (!curStack.empty())
		{
			// Go to the next child
			unsigned node = curStack.top();
			const Node* cur = tree->GetNode(node);
			curStack.pop();
			if (!skipNext)
				materialTexture.push_back(materialPerNode[node]);
			else
				skipNext = false;

			if (cur->GetLevel() < mUniqueShiftLevel)
			{
				// if a node has one child, don't update the index of the next node to process (as it will have the same color)
				//if (cur->GetChildCount() == 1)
				//	skipNext = true;

				// Push the children to the stack
				unsigned32* children = cur->GetChildren();
				unsigned8 childCount = cur->GetChildCount();
				for (ChildIndex c = 0; c < childCount; c++)
					curStack.push(children[c]);
			}
		}
		return materialTexture;
	}

	// This method will move the layout of the given (material) tree to this tree.
	// Since UniqueIndexTrees only need one leaf, all other leafs will be discarded. 
	// The tree will be deleted by this method (as it will be left in an undefined state otherwise)
	void MoveShallowWithOneLeaf(BaseTree* tree)
	{
		Clear();

		std::vector<unsigned32> levelOffsets = tree->SortOnLevel();
		// Find the first leaf, the one that will replace all other leafs
		unsigned32 leafsStart = levelOffsets[GetMaxLevel()];
		if (leafsStart < tree->GetNodeCount())
		{
			// Delete all leafs except the first
			for (unsigned32 i = leafsStart + 1; i < tree->GetNodeCount(); i++) tree->Destroy(i);

			// Now replace all pointers to leaf nodes with pointers to the first leaf.
			tbb::parallel_for(levelOffsets[GetMaxLevel() - 1], levelOffsets[GetMaxLevel()], [&](unsigned32 i)
			{
				Node* node = tree->GetNode(i);
				unsigned32* children = node->GetChildren();
				unsigned8 childCount = node->GetChildCount();
				for (ChildIndex c = 0; c < childCount; c++)
					children[c] = (unsigned32)leafsStart;
			});
		}

		// Create new nodes to replace all nodes in the tree. Since we're doing this in order, the pointers will always be correct
		for (unsigned32 i = 0; i < std::min(leafsStart + 1, tree->GetNodeCount()); i++)
		{
			Node* node = tree->GetNode(i);
			UniqueIndexNode* replacement = Create(node->GetLevel());
			replacement->SetChildren(node->GetChildmask(), node->GetChildren());
			tree->Destroy(node->GetIndex());
		}

		// Delete the given tree as it is now obsolete
		delete tree;
	}

	void UpdateShifts(const unsigned8& untilLevel)
	{
		// Go bottom up through the tree, calculate the shifts in each node
		std::vector<unsigned64> highestShiftSums(GetNodeCount(), 0);
		if (untilLevel != GetMaxLevel()) // If the until level is not the max level, calculate the highest shift sums in the until level before calculating the shifts
			for (unsigned32 i = 0; i < GetNodeCount(); i++)
			{
				UniqueIndexNode* cur = GetTypedNode(i);
				if (cur->GetLevel() == untilLevel)
					highestShiftSums[i] = GetHighestShiftSum(cur);
			}

		for (unsigned8 level = untilLevel; level-- > 0;)
		{
			//tbb::parallel_for((unsigned32)0, GetNodeCount(), [&](const unsigned32 i)
			for (unsigned32 i = 0; i < GetNodeCount(); i++)
			{
				UniqueIndexNode* cur = GetTypedNode(i);
				if (cur->GetLevel() == level)
				{
					unsigned8 childCount = cur->GetChildCount();
					unsigned64 curShift = 1;
					unsigned64* childShifts = new unsigned64[childCount];
					unsigned* curChildren = cur->GetChildren();
					// Only calculate non-zero shifts for the nodes that need this
					if (/*childCount == 1 ||*/ level >= mUniqueShiftLevel) {
						for (ChildIndex c = childCount; c-- > 0;)
						{
							childShifts[c] = 0;
							curShift += highestShiftSums[curChildren[c]];
						}
					}
					else {
						for (ChildIndex c = childCount; c-- > 0;)
						{
							childShifts[c] = curShift;
							curShift += highestShiftSums[curChildren[c]] + 1;
						}
					}
					cur->SetEdgeMaterials(childShifts);
					delete childShifts;
					highestShiftSums[i] = curShift - 1;
				}
			}//);
		}
	}

	std::vector<unsigned64> GetUniqueNodeIndices(const unsigned8& untilLevel) const
	{
		// Now go top down through the tree and calculate the indices of each node
		std::vector<unsigned64> uniqueNodeIndices(GetNodeCount(), 0);
		for (unsigned8 level = 0; level <= untilLevel; level++)
		{
			//tbb::parallel_for(unsigned32(1), GetNodeCount(), [&](const unsigned32& i)
			for (unsigned32 i = 0; i < GetNodeCount(); i++)
			{
				const UniqueIndexNode* cur = GetTypedNode(i);
				if (cur->GetLevel() == level)
				{
					unsigned8 childCount = cur->GetChildCount();
					unsigned64* childShifts = cur->GetEdgeMaterials();
					unsigned32* curChildren = cur->GetChildren();
					for (ChildIndex c = 0; c < childCount; c++)
					{
						unsigned64 childIndex = uniqueNodeIndices[i] + childShifts[c];
						uniqueNodeIndices[curChildren[c]] = childIndex;
					}
				}
			}//);
		}

		return uniqueNodeIndices;
	}

public:
	// Creates a UniqueIndexRoot with "maxLevel" levels.
	// collapsedMaterialLevels are used to decide how many levels of the tree will have the same materials as their parents.
	// Note that the nodeMaterialsTexture will be deleted if the tree is deleted.
	UniqueIndexTree(unsigned8 maxLevel, CompressedTexture<T>* nodeMaterialsTexture, unsigned32 collapsedMaterialLevels) :
		Tree<UniqueIndexNode>(maxLevel),
		mShiftBytesPerLevel(NULL),
		mNodeMaterials(nodeMaterialsTexture),
		mUniqueShiftLevel(maxLevel - collapsedMaterialLevels),
		mBlockPointerPool(std::vector<unsigned8>()),
		mBlockPointerPoolLoaded(false),
		mBlockPool(std::vector<unsigned8>()),
		mBlockPoolLoaded(false),
		mAdditionalProperties(std::map<std::string, std::string>()),
		mAdditionalPropertiesLoaded(false)
	{
		mLeafsAreEqual = true;
	}

	// Creates a UniqueIndexRoot with "maxLevel" levels.  
	// Note that the nodeMaterialsTexture will be deleted if the tree is deleted.
	UniqueIndexTree(unsigned8 maxLevel, CompressedTexture<T>* nodeMaterialsTextureL)
		: UniqueIndexRoot(maxLevel, nodeMaterialsTexture, 0)
	{}

	~UniqueIndexTree() override {
		if (mShiftBytesPerLevel != NULL)
			delete mShiftBytesPerLevel;
		ClearNodeMaterials();
	}

	void SetNodeValues(const std::vector<T>& materialPointers, const size_t& fromIndex = 0)
	{
		mNodeMaterials->SetTexture(materialPointers, fromIndex);
	}

	std::vector<T> GetNodeValues(size_t fromIndex = 0) const
	{
		return mNodeMaterials->GetTexture(fromIndex);
	}

	T GetNodeValue(size_t i) const
	{
		return mNodeMaterials->operator[](i);
	}

	size_t GetMaterialCount() const
	{
		return GetHighestShiftSum(GetRoot());
	}

	// Call this method after appending. 
	// Replacers should be a map from materials in the appended tree to materials in this tree. 
	// If a certain item is not found in the replacer map, it will be ignored
	void AppendPostProcess(glm::uvec3 coordinates, unsigned8 level, UniqueIndexTree* tree, std::unordered_map<T, T> replacers)
	{
		// Unpack all the blocks after where the new blocks should be inserted
		std::vector<T> newNodeMaterials = tree->GetNodeValues(1);
		// Replace the material pointers from the other tree to pointers to the same materials in this tree
		tbb::parallel_for(size_t(0), newNodeMaterials.size(), [&](size_t i)
		{
			auto replacer = replacers.find(newNodeMaterials[i]);
			if (replacer != replacers.end())
				newNodeMaterials[i] = replacer->second;
		});


		// Copy the shifts from the root of the appended tree to the node in the new tree
		UniqueIndexNode* appendedRootCopy = (UniqueIndexNode*)this->AddNode(coordinates, level);

		// Update the shifts in the current tree on the levels above the appended tree
		UpdateShifts(level + 1);
		auto uniqueIndices = GetUniqueNodeIndices(level + 1);

		unsigned64 pointerIndex = uniqueIndices[appendedRootCopy->GetIndex()] + 1;
		std::vector<T> existingPointers = this->GetNodeValues(pointerIndex);
		assert(existingPointers.size() <= (size_t)((1 << level) * (1 << level) * (1 << level)));
		newNodeMaterials.insert(newNodeMaterials.end(), existingPointers.begin(), existingPointers.end());
		mNodeMaterials->SetTexture(newNodeMaterials, pointerIndex);

		// Update the shift sizes per level to be the maximum of the current and the new tree
		for (unsigned8 l = 0; l < tree->GetMaxLevel(); l++)
			this->mShiftBytesPerLevel->at(l + level) = std::max(tree->mShiftBytesPerLevel->at(l), this->mShiftBytesPerLevel->at(l + level));
	}

	void ReplaceMaterials(const std::unordered_map<T, T>& replacers)
	{
		mNodeMaterials->ReplaceMaterials(replacers);
	}

	// Takes the current material texture and compresses it again
	void Recompress()
	{
		mNodeMaterials->Recompress();
	}

	// Will build a unique index tree with the structure taken from root. The materials
	// should be given as an additional array that contains at the index of the node 
	// it's material.
	// This will consume (and delete) both the given tree and the given material array!
	void BaseOn(BaseTree* tree, std::vector<T>& materialPerNode)
	{
		std::vector<T> nodeValues = GetMappedMaterialTexture(tree, materialPerNode);
		// Free some memory
		materialPerNode.clear();
		materialPerNode.shrink_to_fit();

		MoveShallowWithOneLeaf(tree);
		UpdateShifts(GetMaxLevel());

		// Assert that the maximum index that can be reached using the shifts is equal to the maximum index in the material texture
		assert(GetHighestShiftSum(GetRoot()) == nodeValues.size() - 1);

		// Add all materials to the material table (BIG texture)
		SetNodeValues(nodeValues);

		
		// Let's check if the compression works correctly
		// Method 1: Unpack the whole material at once and check that:
		//auto check = GetNodeValues();
		//assert(check.size() == nodeMaterialPointers.size());
		//for (size_t i = 0; i < check.size(); i++)
		//	assert(nodeMaterialPointers[i] == check[i]);

		// Method 2: Unpack per value
		assert(mNodeMaterials->size() == nodeValues.size());
		for (size_t i = 0; i < mNodeMaterials->size(); i++)
		{
			T value = mNodeMaterials->operator[](i);
			assert(nodeValues[i] == value);
		}

		CalculateShiftBytesPerLevel();
	}

	void ReplaceCompressedTexture(CompressedTexture<T>* replacementCompressor)
	{
		std::vector<T> materials = GetNodeValues();
		delete mNodeMaterials;
		mNodeMaterials = replacementCompressor;
		SetNodeValues(materials);
	}

	std::vector<unsigned8> GetBlockPointerPool() override
	{
		if (mBlockPointerPoolLoaded) return mBlockPointerPool;

		GetBlockPointerPoolSize();
		mBlockPointerPool = mNodeMaterials->GetAdditionalTexturePool();
		mBlockPointerPoolLoaded = true;
		ClearPooledNodeMaterials();
		return mBlockPointerPool;
	}

	// Assuming 3D texture with one pointer per node (for example GL_R32UI)
	size_t GetBlockPointerPoolSize() override
	{
		if (mBlockPointerPoolLoaded) return mBlockPointerPool.size();
		return mNodeMaterials->GetAdditionalTexturePoolSize();
	}

	std::vector<unsigned8> GetBlockPool() override
	{
		if (mBlockPoolLoaded) return mBlockPool;

		GetBlockPoolSize();
		mBlockPool = mNodeMaterials->GetTexturePool();
		mBlockPoolLoaded = true;
		ClearPooledNodeMaterials();
		return mBlockPool;
	}

	size_t GetBlockPoolSize() override
	{
		if (mBlockPoolLoaded) return mBlockPool.size();
		return mNodeMaterials->GetTexturePoolSize();
	}

	std::map<std::string, std::string> GetAdditionalProperties() override
	{
		if (mAdditionalPropertiesLoaded) return mAdditionalProperties;
		mAdditionalProperties = mNodeMaterials->GetAdditionalProperties();
		mAdditionalPropertiesLoaded = true;
		ClearPooledNodeMaterials();
		return mAdditionalProperties;
	}

	bool HasAdditionalPool() const override{ return true; }
protected:
	void CalculateShiftBytesPerLevel()
	{
		if (mShiftBytesPerLevel != NULL) delete mShiftBytesPerLevel;
#ifdef offset_sizes_per_level
		mShiftBytesPerLevel = new std::vector<unsigned8>(GetMaxLevel() + 1);
		std::vector<unsigned64> maxShiftPerLevel(GetMaxLevel() + 1);
		for (unsigned32 i = 0; i < GetNodeCount(); i++)
		{
			UniqueIndexNode* node = GetTypedNode(i);
			for (ChildIndex child = 0; child < 8; child++)
			{
				if (node->HasChild(child))
				{
					unsigned64 shift = GetShift(node, child);
					unsigned8 level = node->GetLevel();
					if (shift > maxShiftPerLevel[level]) maxShiftPerLevel[level] = shift;
				}
			}
		}

		for (unsigned8 level = 0; level <= GetMaxLevel(); level++)
			(*mShiftBytesPerLevel)[level] = BitHelper::RoundToBytes(BitHelper::Log2Ceil(maxShiftPerLevel[level] + 1)) / 8;
#else
		mShiftBytesPerLevel = new std::vector<unsigned8>(GetMaxLevel() + 1, 4);
#endif
	}

	unsigned8 GetAdditionalBytesPerPointer(unsigned8 level) const override
	{
		assert(mShiftBytesPerLevel != NULL);
		return mShiftBytesPerLevel->operator[](level);
	}

	bool LastChildHasAdditionalBytes() const override { 
#ifdef implicit_store_first_offset
		return false;
#else
		return true;
#endif
	}

	std::vector<unsigned8> GetAdditionalPointerBytes(const Node* node, ChildIndex child) const override
	{
		return BitHelper::SplitInBytes(GetShift((UniqueIndexNode*)node, child), GetAdditionalBytesPerPointer(node->GetLevel()));
	}

	virtual void WriteAdditionalUniqueIndexTreeProperties(std::ostream& file) {}
	virtual void ReadAdditionalUniqueIndexTreeProperties(std::istream& file) {}

	void WriteProperties(std::ostream& file) override {
		WriteAdditionalUniqueIndexTreeProperties(file);
		mNodeMaterials->WriteToFile(file);

	}
	void ReadProperties(std::istream& file) override {
		ReadAdditionalUniqueIndexTreeProperties(file);
		mNodeMaterials->ReadFromFile(file);
	}

	void FinalizeReadTree() override
	{
		CalculateShiftBytesPerLevel();
	}

	void WriteAdditionalPoolProperties(std::ostream& file) override { WriteBlockPointerPool(file); WriteBlockPool(file); WriteAdditionalProperties(file); }
	void ReadAdditionalPoolProperties(std::istream& file) override { ReadBlockPointerPool(file); ReadBlockPool(file); ReadAdditionalProperties(file); }

	void PrintDebugInfo() const override
	{
		printf("Leaf voxel count: %llu\n", GetLeafVoxelCount());
		//printf("Total voxel count: %llu\n", GetOctreeNodeCount());
		printf("Max node index: %llu\n", GetHighestShiftSum(GetRoot()));
		auto nodesPerLevel = GetNodesPerLevel();
		for (unsigned8 level = 0; level <= GetMaxLevel(); level++)
			printf("Shifts in level %u with %llu nodes: %u bytes\n", level, (unsigned64)nodesPerLevel[level], GetAdditionalBytesPerPointer(level));

		// Print the childshifts of the root for debug purposes
		unsigned32* rootChildren = GetRoot()->GetChildren();
		unsigned64* rootEdgeMaterials = GetRoot()->GetEdgeMaterials();
		for (ChildIndex c = 0; c < GetRoot()->GetChildCount(); c++)
		{
			printf("mChildShifts[%u] = %9llu; ->", c, rootEdgeMaterials[c]);
			const UniqueIndexNode* child = GetTypedNode(rootChildren[c]);
			for (unsigned j = 0; j < child->GetChildCount(); j++)
				printf("  %9llu", child->GetEdgeMaterials()[j]);
			printf("\n");
		}

		//printf("Node count = %u\nOctreeNodeCount = %u\n", GetNodeCount(), GetOctreeNodeCount());
		//printf("Leaf node count = %u\n", GetLeafVoxelCount());
	}
};