#pragma once
#include "Tree.h"
#include "MaterialNode.h"
#include "../../inc/tbb/parallel_for_each.h"
#include "../../core/BitHelper.h"
#include "../../core/CollectionHelper.h"
#include "../../core/Serializer.h"
#include <unordered_map>
//#include <map>
#include <stack>

template<typename T, typename Comparer = std::less<T>>
class MaterialTree : public Tree<MaterialNode<T, Comparer>>
{
private:
	std::unordered_map<T, Node*>	mLeafMap;
	bool							mUseLeafMap;
public:
	// Creates a Material Tree with "maxLevel" levels.
	MaterialTree(unsigned8 maxLevel) : Tree<MaterialNode<T, Comparer>>(maxLevel), mUseLeafMap(false)
	{
		mLeafsAreEqual = false;
	}

	~MaterialTree() override { }

	// When using the leaf map, leaf nodes with the same material will only be created once (e.g. the leaf level is already a DAG)
	void UseLeafMap(bool value)
	{
		mUseLeafMap = value;
		
		// Build the leaf map for the nodes that already exist
		if (mUseLeafMap)
		{
			for (unsigned32 i = 0; i < GetNodeCount(); i++)
			{
				MaterialNode<T, Comparer>* node = GetTypedNode(i);
				if (node->GetLevel() == GetMaxLevel())
					mLeafMap.insert(std::pair<T, Node*>(node->GetMaterial(), node));
			}
			// Convert the leaf nodes to a DAG
			ToDAG(GetMaxLevel() - 1);
		}
	}

	// Assuming all leaf nodes contain pointers to materials, propagates those materials up in the tree, leaving the average material everywhere.
	// It is advised to call this method after DAG conversion, as it is still valid at that point, and it will be cheaper.
	void PropagateMaterials(T(*Average)(const std::vector<T>& materials, const std::vector<float>& weights))
	{
		// Bottom up go through the nodes to propagate the materials
		auto levelIndices = SortOnLevel();
		// Calculate how many levels actually contain data
		unsigned8 filledLevels = 0;
		for (unsigned8 level = 0; level <= GetMaxLevel(); level++)
		{
			if (levelIndices[level] != levelIndices[level + 1]) filledLevels++;
			else break;
		}
		
		// If the levelIndices.size() <= 2, that means there's only the root level, thus no propagation is needed
		if (filledLevels < 2) return;

		// Set node weights for weighted average calculation
		std::vector<float> lastLevelNodeWeights;
		std::vector<float> levelNodeWeights;

		// Bottom-up calculate the weighted average material for each node in the tree, and store them in perfectMaterialPerNode
		for (unsigned8 level = filledLevels - 1; level-- > 0;)
		{
			auto levelStart = levelIndices[level];
			auto levelEnd = levelIndices[level + 1];
			levelNodeWeights = std::vector<float>(levelEnd - levelStart);
			tbb::parallel_for(levelStart, levelEnd, [&](const unsigned32 i)
			{
				// Get the current node
				MaterialNode<T, Comparer>* node = GetTypedNode(i);
				std::vector<T> childMaterials;
				std::vector<float> childWeights;
				float nodeWeight = 0;
				// Find all materials the children use
				unsigned32* children = node->GetChildren();
				unsigned8 childCount = node->GetChildCount();
				for (ChildIndex c = 0; c < childCount; c++)
				{
					const unsigned32& i = children[c];
					MaterialNode<T, Comparer>* child = GetTypedNode(i);
					float childWeight = (i - levelEnd < lastLevelNodeWeights.size()) ? lastLevelNodeWeights[i - levelEnd] : 1.f;
					childMaterials.push_back(child->GetMaterial());
					childWeights.push_back(childWeight);
					nodeWeight += childWeight;
				}
				// Calculate the average material and retrieve the closest material to that from the library
				node->SetMaterial(Average(childMaterials, childWeights));
				// Store the weighted average in perfectMaterialPerNode
				levelNodeWeights[i - levelStart] = nodeWeight;
			});
			// Update the last level node weights to the node weights of the current level
			lastLevelNodeWeights = std::move(levelNodeWeights);
		}
	}

	void ReplaceMaterials(const std::unordered_map<T, T>& replacementMap)
	{
		tbb::parallel_for((unsigned32)0, GetNodeCount(), [&](unsigned32 i)
		{
			auto node = GetTypedNode(i);
			node->SetMaterial(replacementMap.at(node->GetMaterial()));
		});
	}

	void SetMaterials(std::vector<T> materials)
	{
		assert(materials.size() == GetNodeCount());
		tbb::parallel_for((unsigned32)0, GetNodeCount(), [&](const unsigned32 i)
		{
			MaterialNode<T, Comparer>* node = GetTypedNode(i);
			node->SetMaterial(materials[i]);
		});
	}

	T GetMaterial(MaterialNode<T, Comparer>* node)
	{
		return node->GetMaterial();
	}

	std::vector<T> GetMaterials() const
	{
		std::vector<T> nodeMaterials(GetNodeCount());
		tbb::parallel_for((unsigned32)0, GetNodeCount(), [&](const unsigned32 i)
		{
			nodeMaterials[i] = GetTypedNode(i)->GetMaterial();
		});
		return nodeMaterials;
	}

	// Returns a list with all unique materials in the tree
	std::vector<T> GetUniqueMaterials() const
	{
		std::vector<T> nodeMaterials = GetMaterials();
		CollectionHelper::Unique(nodeMaterials, Comparer());
		return nodeMaterials;
	}

	// Sets the material of a node.
	void SetMaterial(glm::uvec3 coordinate, unsigned level, T material)
	{
		MaterialNode<T, Comparer>* node = AddNode(coordinate, level);
		node->SetMaterial(material);
	}

	// Creates a leaf node and sets its material
	void AddLeafNode(glm::uvec3 coordinate, T material)
	{
		if (mUseLeafMap)
		{
			// Check if there is already a leaf for this material (for auto reuse)
			auto existingLeaf = mLeafMap.find(material);
			Node* leaf = NULL;
			if (existingLeaf != mLeafMap.end())
			{
				// If there is a leaf, use it
				assert(material == ((MaterialNode<T, Comparer>*)(existingLeaf->second))->GetMaterial());
				leaf = existingLeaf->second;
			}
			else
			{
				// If no leaf with this material exists yet, create it
				leaf = Create(GetMaxLevel());
				((MaterialNode<T, Comparer>*)leaf)->SetMaterial(material);
				mLeafMap.insert(std::pair<T, Node*>(material, leaf));
			}
			// Create the parent node of the leaf
			MaterialNode<T, Comparer>* parentOfLeaf = Tree<MaterialNode<T, Comparer>>::AddNode(glm::uvec3(coordinate.x >> 1, coordinate.y >> 1, coordinate.z >> 1), GetMaxLevel() - 1);
			// The last bit of the coordinate can be used to find the childindex in this parent:
			ChildIndex index = (((coordinate.x & 1) == 1) ? 1 : 0)
				+ (((coordinate.y & 1) == 1) ? 2 : 0)
				+ (((coordinate.z & 1) == 1) ? 4 : 0);

			// Make sure the leaf points to this child
			parentOfLeaf->SetChild(index, leaf);
		}
		else
		{
			MaterialNode<T, Comparer>* leaf = Tree<MaterialNode<T, Comparer>>::AddLeafNode(coordinate);
			leaf->SetMaterial(material);
		}
	}

	bool HasAdditionalPool() const override { return false; }
protected:
	unsigned8 GetAdditionalBytesPerNode(unsigned8 level) const override
	{
		return sizeof(T);
	}
	std::vector<unsigned8> GetAdditionalNodeBytes(const Node* node) const override
	{
		return ((MaterialNode<T, Comparer>*)node)->GetMaterial().Serialize();
	}
};