#pragma once
#include "Tree.h"
#include "MaterialNode.h"
#include "IMaterialTexture.h"
#include "../../inc/tbb/parallel_for_each.h"
#include "../../inc/tbb/concurrent_queue.h"
#include "../Material/MaterialLibrary.h"
#include "../../core/BitHelper.h"
#include <unordered_map>
#include <map>
#include <set>

typedef MaterialNode<MaterialLibraryPointer> MaterialLibraryNode;

template<typename T, typename Comparer = std::less<T>, unsigned8 channelsPerPixel = 3>
class MaterialLibraryTree : public Tree<MaterialLibraryNode>, public IMaterialTexture
{
protected:
	MaterialLibrary<T, Comparer, channelsPerPixel>* mMaterialLibrary;
	bool mOwnsLibrary;
	std::vector<unsigned char> mMaterialTexture;
	MaterialLibraryPointer mMaxTextureIndex;
	unsigned short mMaterialTextureSize;

	std::unordered_map<T, MaterialLibraryNode*> leafMap;

	inline void WriteMaterialTexture(std::ostream& file)
	{
		assert(mMaterialLibrary != NULL);
		mMaterialLibrary->Serialize(file);
	}

	inline void ReadMaterialTexture(std::istream& file)
	{
		if (mMaterialLibrary == NULL)
			mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		mMaterialLibrary->Deserialize(file);
		GetMaterialTexture();
	}

	inline void CreateAllMaterialLeafs()
	{
		std::vector<std::pair<T, MaterialLibraryPointer>> materialsAndIndices = mMaterialLibrary->GetMaterialTextureIndices();
		for (auto materialAndIndex : materialsAndIndices)
		{
			T material = materialAndIndex.first;
			MaterialLibraryPointer textureIndex = materialAndIndex.second;
			MaterialLibraryNode* leaf = (MaterialLibraryNode*)Create(GetMaxLevel());
			leaf->SetMaterial(textureIndex);
			leafMap.insert(std::pair<T, MaterialLibraryNode*>(material, leaf));
		}
	}
public:
	// Creates a material library tree. 
	MaterialLibraryTree(unsigned8 maxLevel) :
		Tree<MaterialLibraryNode>(maxLevel),
		mOwnsLibrary(true),
		mMaterialLibrary(new MaterialLibrary<T, Comparer, channelsPerPixel>())
	{
		mLeafsAreEqual = false;
	}

	// Creates a material library tree with the given (finalized!) material library. The library is not owned by this tree and therefore will not
	// be deleted when this tree is deleted.
	MaterialLibraryTree(unsigned8 maxLevel, MaterialLibrary<T, Comparer, channelsPerPixel>* materialLibrary) : 
		Tree<MaterialLibraryNode>(maxLevel),
		mOwnsLibrary(false),
		mMaterialLibrary(materialLibrary)
	{
		mLeafsAreEqual = false;
		CreateAllMaterialLeafs();
	}

	~MaterialLibraryTree() override {
		if(mOwnsLibrary)
			delete mMaterialLibrary;
	}

	MaterialLibrary<T, Comparer, channelsPerPixel>* GetMaterialLibrary()
	{
		return mMaterialLibrary;
	}

	static void PassLibraryOwnership(MaterialLibraryTree* tree1, MaterialLibraryTree* tree2)
	{
		tree1->mOwnsLibrary = false;
		tree2->mOwnsLibrary = true;
	}

	void AddMaterial(T material)
	{
		if (mMaterialLibrary->IsFinalized())
			return;
		mMaterialLibrary->AddMaterial(material);
	}

	void FinalizeMaterials()
	{
		mMaterialLibrary->Finalize();
		CreateAllMaterialLeafs();
	}

	// 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()
	{
		// Bottom up go through the nodes to propagate the materials
		auto levelIndices = SortOnLevel();

		// Set node weights for weighted average calculation
		std::vector<float> lastLevelNodeWeights;
		std::vector<T> perfectMaterialPerNode(GetNodeCount());
		std::vector<float> levelNodeWeights;

		auto leafStart = levelIndices[GetMaxLevel()];
		auto leafEnd = levelIndices[GetMaxLevel() + 1];
		lastLevelNodeWeights.resize(leafEnd - leafStart, 1.f);
		// Initialize the vectors for leaf nodes, assuming leaf nodes have weight 1 and perfect materials
		for (auto i = leafStart; i != leafEnd; i++)
			perfectMaterialPerNode[i] = mMaterialLibrary->GetMaterial(GetTypedNode(i)->GetMaterial());

		// Bottom-up calculate the weighted average material for each node in the tree, and store them in perfectMaterialPerNode
		for (unsigned8 level = GetMaxLevel(); level-- > 0;)
		{
			unsigned32 levelStart = levelIndices[level];
			unsigned32 levelEnd = levelIndices[level + 1];
			levelNodeWeights = std::vector<float>(levelEnd - levelStart);
			tbb::parallel_for(levelStart, levelEnd, [&](unsigned32 i)
			{
				// Get the current node
				MaterialLibraryNode* 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++)
				{
					unsigned32 i = children[c];
					MaterialLibraryNode* child = GetTypedNode(i);
					MaterialLibraryPointer matPtr = child->GetMaterial();
					float childWeight = lastLevelNodeWeights[i - levelEnd];
					childMaterials.push_back(perfectMaterialPerNode[i]);
					childWeights.push_back(childWeight);
					nodeWeight += childWeight;
				}
				// Calculate the average material and retrieve the closest material to that from the library
				T nodeMaterial = T::WeightedAverage(childMaterials, childWeights);
				// Store the weighted average in perfectMaterialPerNode
				perfectMaterialPerNode[i] = nodeMaterial;
				levelNodeWeights[i - levelStart] = nodeWeight;
			});
			// Update the last level node weights to the node weights of the current level
			lastLevelNodeWeights = std::move(levelNodeWeights);
		}
		
		// Find all materials required for this scene
		std::vector<T> perfectMaterials(GetNodeCount());
		tbb::parallel_for((size_t)0, perfectMaterialPerNode.size(), [&](size_t i)
		{
			perfectMaterials[i] = perfectMaterialPerNode[i];
		});

		// Reduce the materials by only using unique ones
		tbb::parallel_sort(perfectMaterials, Comparer());
		perfectMaterials.erase(std::unique(perfectMaterials.begin(), perfectMaterials.end()), perfectMaterials.end());
		perfectMaterials.shrink_to_fit();

		// Create a new material library based on these materials
		if (mOwnsLibrary) delete mMaterialLibrary;
		mMaterialTexture.clear();
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		for (auto material : perfectMaterials)
			mMaterialLibrary->AddMaterial(material);
		mMaterialLibrary->Finalize();

		// Update all nodes in the tree to point to the new material library
		tbb::parallel_for((unsigned32)0, GetNodeCount(), [&](unsigned32 i)
		{
			MaterialLibraryPointer materialPointer = mMaterialLibrary->GetTextureIndex(perfectMaterialPerNode[i]);
			auto node = GetTypedNode(i);
			node->SetMaterial(materialPointer);
		});

		// Update the material texture
		GetMaterialTexture();
	}

	std::vector<unsigned8> GetMaterialTexture() override
	{
		if (!mMaterialTexture.empty())
			return mMaterialTexture;
		assert(mMaterialLibrary->IsFinalized());
		mMaterialTextureSize = mMaterialLibrary->GetTextureSize();
		mMaterialTexture = mMaterialLibrary->GetTexture();
		mMaxTextureIndex = mMaterialLibrary->GetMaxTextureIndex();
		return mMaterialTexture;
	}

	unsigned GetMaterialTextureSize() override
	{
		GetMaterialTexture();
		return mMaterialTextureSize;
	}

	unsigned8 GetMaterialTextureChannelsPerPixel() override { return channelsPerPixel; }

	std::vector<T> GetLeafMaterials()
	{
		auto levelIndices = SortOnLevel();
		unsigned32 leafsStart = levelIndices[GetMaxLevel()];
		unsigned32 leafsEnd =  levelIndices[GetMaxLevel() + 1];
		std::vector<T> leafMaterials(levelIndices[GetMaxLevel() + 1] - levelIndices[GetMaxLevel()]);
		tbb::parallel_for(leafsStart, leafsEnd, [&](const unsigned32 i)
		{
			leafMaterials[i - leafsStart] = GetMaterial(GetTypedNode(i));
		});
		return leafMaterials;
	}

	// Replaces the current leaf materials (and thereby also the material library) by the given material.
	// In this, it assumes the leafMaterials list is ordered in the same way the leafs are ordered in the current tree.
	// It is advised to follow this command by "ToDAG()" and "PropagateMaterials()"
	void SetLeafMaterials(const std::vector<T> leafMaterials)
	{
		// Create the new material library
		delete mMaterialLibrary;
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		for (auto material : leafMaterials)
			mMaterialLibrary->AddMaterial(material);
		mMaterialLibrary->Finalize();

		// Set the leaf node materials
		size_t i = 0;
		for (unsigned32 i = 0; i < GetNodeCount(); i++)
			MaterialLibraryNode* node = GetTypedNode(i);
			if (node->GetLevel() == GetMaxLevel())
			{
				node->SetMaterial(mMaterialLibrary->GetTextureIndex(leafMaterials[i]));
				i++;
			}
	}

	T GetMaterial(const MaterialLibraryNode* node) const
	{
		assert(mMaterialLibrary->IsFinalized());
		return mMaterialLibrary->GetMaterial(node->GetMaterial());
	}

	std::vector<T> GetMaterials() const
	{
		// Read all materials from all nodes
		std::vector<T> materials(GetNodeCount());
		//for (unsigned32 i = 0; i < GetNodeCount(); i++)
		tbb::parallel_for((unsigned32)0, GetNodeCount(), [&](const unsigned32 i)
		{
			materials[i] = GetMaterial(GetTypedNode(i));
		});
		return materials;
	}

	void SetMaterials(const std::vector<T>& materials)
	{
		// Not a valid material vector
		if (materials.size() != GetNodeCount())
			return;

		// Create a new material library containing the new materials
		if (mOwnsLibrary) delete mMaterialLibrary;
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		for (auto material : materials)
		{
			mMaterialLibrary->AddMaterial(material);
		}
		mMaterialLibrary->Finalize();
		mMaterialTexture.clear();

		// Update all nodes to point to this new library
		tbb::parallel_for((unsigned32)0, GetNodeCount(), [&](unsigned32 i)
		{
			MaterialLibraryNode* node = GetTypedNode(i);
			node->SetMaterial(mMaterialLibrary->GetTextureIndex(materials[i]));
		});
	}

	std::vector<T> GetUniqueMaterials() const
	{
		return mMaterialLibrary->GetMaterials();
	}

	// Clears all material pointers from non-leaf nodes
	void ClearPropagation()
	{
		auto levelIndices = SortOnLevel();
		tbb::parallel_for((unsigned32)0, levelIndices[GetMaxLevel()], [&](const unsigned32 i)
		{
			MaterialLibraryNode* node = GetTypedNode(i);
			node->SetMaterial(MaterialLibraryPointer(0));
		});
	}

	// Replaces the leaf materials by the given leaf materials. It is advised to follow this command by "ToDAG()" and "PropagateMaterials()" 
	void ReplaceLeafMaterials(const std::map<T, T, Comparer>& leafMaterialReplacers)
	{
		auto newMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		for (auto material : leafMaterialReplacers)
			newMaterialLibrary->AddMaterial(material.second);
		newMaterialLibrary->Finalize();
		std::unordered_map<MaterialLibraryPointer, MaterialLibraryPointer> materialLibraryPointerReplacers;
		for (auto material : leafMaterialReplacers)
			materialLibraryPointerReplacers.insert(std::make_pair(mMaterialLibrary->GetTextureIndex(material.first), newMaterialLibrary->GetTextureIndex(material.second)));

		auto levelIndices = SortOnLevel();
		tbb::parallel_for(levelIndices[GetMaxLevel()], levelIndices[GetMaxLevel() + 1], [&](const unsigned32 i)
		{
			MaterialLibraryNode* node = GetTypedNode(i);
			node->SetMaterial(materialLibraryPointerReplacers[node->GetMaterial()]);
		});
		delete mMaterialLibrary;
		mMaterialLibrary = newMaterialLibrary;
	}

	unsigned8 GetAdditionalBytesPerNode(unsigned8 level) const override
	{
		return 3; // For now, assume that material pointers are always 12+12 bits
	}
	std::vector<unsigned8> GetAdditionalNodeBytes(const Node* node) const override
	{
		MaterialLibraryPointer materialPointer = ((MaterialLibraryNode*)node)->GetMaterial();
		std::vector<unsigned8> res(3);
		res[0] = (unsigned8)(materialPointer.x >> 4);
		res[1] = (unsigned8)((materialPointer.x << 4) | ((materialPointer.y & (0x0FFF)) >> 8));
		res[2] = (unsigned8)materialPointer.y;
		return res;
	}

	void AddLeafNode(glm::uvec3 coordinate, T material)
	{
		// Get the material pointer
		assert(mMaterialLibrary->IsFinalized());
		// Check if there is already a leaf for this material (for auto reuse)
		auto existingLeaf = leafMap.find(material);
		if (existingLeaf != leafMap.end())
		{
			// If the leaf node already exists, reuse it:
			// Create the parent node of the leaf
			MaterialLibraryNode* parentOfLeaf = Tree::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:
			ChildIndex index = (((coordinate.x & 1) == 1) ? 1 : 0)
				+ (((coordinate.y & 1) == 1) ? 2 : 0)
				+ (((coordinate.z & 1) == 1) ? 4 : 0);

			parentOfLeaf->SetChild(index, existingLeaf->second);
		}
		else
		{
			MaterialLibraryPointer textureIndex = mMaterialLibrary->GetTextureIndex(material);
			MaterialLibraryNode* leaf = Tree::AddLeafNode(coordinate);
			leaf->SetMaterial(textureIndex);
			leafMap.insert(std::pair<T, MaterialLibraryNode*>(material, leaf));
		}
	}

	bool HasAdditionalPool() const override { return true; }
protected:
	void WriteProperties(std::ostream& file) override {
		// Write the material texture
		WriteMaterialTexture(file);
	}
	void ReadProperties(std::istream& file) override {
		// Reat the material texture
		ReadMaterialTexture(file);
		// Restore the material library from the texture
		delete mMaterialLibrary;
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>(mMaterialTexture, mMaterialTextureSize, mMaxTextureIndex);
		mMaterialLibrary->Finalize();

		//// Refresh the material texture for debug purposes
		//mMaterialTexture.clear();
		//GetMaterialTexture();
	}
	void WriteAdditionalPoolProperties(std::ostream& file) override { WriteMaterialTexture(file); }
	void ReadAdditionalPoolProperties(std::istream& file) override { ReadMaterialTexture(file); }
};