#pragma once
#include "HierarchicalMaterialMultiRoot.h"
#include "../Material/BitsMaterial.h"
#include "IMaterialTexture.h"
#include "../../inc/tbb/parallel_for_each.h"
#include "../Material/BlockBasedMaterialLibrary.h"
#include "../../core/BitHelper.h"
#include "../../core/CollectionHelper.h"
#include "../../core/Util/BoolArray.h"
#include "NodeReplacementFinder.h"
#include "Tree.h"
#include <queue>
#include <unordered_set>

// Usage:
// This tree can only be built correctly if it is based on some material tree that has materials throughout 
// (Such as HierarchicalRoot<T> or MaterialRoot<T>). To build this tree, create this root object and then call
// the "BaseOn(tree)" method with the material tree.
template<typename T, typename Comparer = std::less<T>, unsigned8 channelsPerPixel = 3>
class MaterialLibraryMultiRoot : public HierarchicalMaterialMultiRoot<BitsMaterial<8>>, public IMaterialTexture
{
private:
	// After the tree is finalized, the material library will be used to contain the actual materials
	MaterialLibrary<T, Comparer, channelsPerPixel>* mMaterialLibrary;
	std::vector<unsigned8> mBitMap;

	std::vector<unsigned char> mMaterialTexture;
	unsigned short mMaterialTextureSize;
	MaterialLibraryPointer mMaxTextureIndex;

	inline void WriteMaterialTexture(std::ostream& file)
	{
		if (mMaterialTexture.empty())
			GetMaterialTexture();
		// Pack the texture size and the biggest texture index in one 32 bit unsigned int (for historic reasons...)
		unsigned materialTextureSizeSummary = (mMaxTextureIndex.x << 20) | (mMaxTextureIndex.y << 10) | (mMaterialTextureSize - 1);
		Serializer<unsigned>::Serialize(materialTextureSizeSummary, file);
		Serializer<unsigned8*>::Serialize(&mMaterialTexture[0], (size_t)mMaterialTextureSize * (size_t)mMaterialTextureSize * (size_t)channelsPerPixel, file);
	}

	inline void ReadMaterialTexture(std::istream& file)
	{
		unsigned materialTextureSizeSummary;
		Serializer<unsigned>::Deserialize(materialTextureSizeSummary, file);
		unsigned mask1 = BitHelper::GetLSMask<unsigned32>(20, 30);
		unsigned mask2 = BitHelper::GetLSMask<unsigned32>(10, 20);
		unsigned mask3 = BitHelper::GetLSMask<unsigned32>(0, 10);
		unsigned short maxTextureIndexX = (mask1 & materialTextureSizeSummary) >> 20;
		unsigned short maxTextureIndexY = (mask2 & materialTextureSizeSummary) >> 10;
		unsigned materialTextureSize = mask3 & materialTextureSizeSummary;
		mMaterialTextureSize = materialTextureSize + 1;
		mMaxTextureIndex = MaterialLibraryPointer(maxTextureIndexX, maxTextureIndexY);

		size_t textureArraySize = (size_t)mMaterialTextureSize * (size_t)mMaterialTextureSize * (size_t)channelsPerPixel;

		mMaterialTexture.resize(textureArraySize);
		Serializer<unsigned8*>::Deserialize(&mMaterialTexture[0], textureArraySize, file);
	}

	void AddMaterial(const T& material)
	{
		assert(!mMaterialLibrary->IsFinalized());
		mMaterialLibrary->AddMaterial(material);
	}

	void FinalizeMaterials()
	{
		assert(!mMaterialLibrary->IsFinalized());
		mMaterialLibrary->Finalize();
		unsigned requiredXBits = BitHelper::Log2Ceil(mMaterialLibrary->GetTextureSize());
		unsigned requiredYBits = BitHelper::Log2Ceil(mMaterialLibrary->GetMaxTextureIndex().y);
		unsigned requiredBits = requiredXBits + requiredYBits;
		unsigned32 mask = BitHelper::GetLSMask<unsigned32>(16, 16 + requiredXBits) | BitHelper::GetLSMask<unsigned32>(0, requiredYBits);
		mBitMap = BitHelper::GetBitMapHS(mask);
		AddSlaveRoots(requiredBits); // The main root can be used for the first bit, the rest of the bits require slave roots
	}

	bool CheckNodesToRemove(MultiRootMaterialNode<BitsMaterial<8>>* node, bool curValue, BoolArray& nodesCanBeShaved)
	{
		if (!node->HasChildren()) return true;
		auto mat = node->GetMaterial();
		bool childrenEqual = true;
		for (ChildIndex c = 0; c < 8; c++)
		{
			if (node->HasChild(c))
			{
				bool nodeMat = mat.GetLS(c);
				MultiRootMaterialNode<BitsMaterial<8>>* child = (MultiRootMaterialNode<BitsMaterial<8>>*)GetNode(node->GetChildIndex(c));
				bool nodeHasSameMat = CheckNodesToRemove(child, nodeMat, nodesCanBeShaved);
				if (nodeMat != curValue || !nodeHasSameMat) childrenEqual = false;
			}
		}
		if (!childrenEqual) nodesCanBeShaved.Set(node->GetIndex(), false);
		return childrenEqual;
	}

	static std::vector<unsigned32> HashChildren(const MultiRootMaterialNode<BitsMaterial<8>>* node)
	{
		assert(!node->GetIsGeometry());
		std::vector<unsigned32> hash(8, 0);
		auto mat = node->GetMaterial();
		for (ChildIndex c = 0; c < 8; c++)
		{
			if (node->HasChild(c))
				hash[c] = (node->GetChildIndex(c) << 1) | (mat.GetLS(c) ? 1 : 0);
		}
		return hash;
	}
public:
	MaterialLibraryMultiRoot(unsigned8 maxLevel) : HierarchicalMaterialMultiRoot<BitsMaterial<8>>(maxLevel, 0)
	{
		mMaterialLibrary = new BlockBasedMaterialLibrary<T, Comparer, channelsPerPixel>();
	}

	~MaterialLibraryMultiRoot() override {
		if (mMaterialLibrary != NULL)
			delete mMaterialLibrary;
	}

	void CreateMaterialLibrary(const std::vector<T>& materials)
	{
		for (const T& material : materials) AddMaterial(material);
		FinalizeMaterials();
	}

	// Copies the geometry and material information of the given tree to this tree
	template <typename MaterialTree>
	void BaseOn(MaterialTree* tree, bool autoDAG)
	{
		// Make sure the leaf nodes are the last nodes in the tree, so that we can skip them and only construct one.
		std::vector<unsigned32> levelIndices = tree->SortOnLevel();

		// Clear the existing tree
		Clear();
		auto root = Create(0);
		root->SetIsGeometry(true);

		// Create the single leaf node that can exist for the geometry tree (for memory efficiency)
		auto geometryLeaf = Create(GetMaxLevel());
		geometryLeaf->SetIsGeometry(true);

		// Copy the geometry information to a geometry tree
		unsigned32 destChildrenCache[8];
		unsigned32 offset = GetNodeCount();
		for (unsigned32 i = 0; i < levelIndices[GetMaxLevel()]; i++)
		{
			Node* source = tree->GetNode(i);
			unsigned32* sourceChildren = source->GetChildren();
			if (source->GetLevel() == GetMaxLevel() - 1)
			{
				for (ChildIndex child = 0; child < source->GetChildCount(); child++)
					destChildrenCache[child] = geometryLeaf->GetIndex();
			}
			else
			{
				for (ChildIndex child = 0; child < source->GetChildCount(); child++)
					destChildrenCache[child] = offset + sourceChildren[child] - 1; // The root is reused, so that index shouldn't be counted towards the offset
			}
			auto dest = source->GetLevel() == 0 ? root : Create(source->GetLevel());
			dest->SetIsGeometry(true);
			dest->SetChildren(source->GetChildmask(), destChildrenCache);
		}

		// Create and fill the material library (if this hasn't been done yet)
		if (!mMaterialLibrary->IsFinalized())
		{
			auto materials = tree->GetUniqueMaterials();
			CreateMaterialLibrary(materials);
		}
		else
		{ // Re-add the slaveroots
			AddSlaveRoots(mBitMap.size());
		}

		unsigned32 bitCount = (unsigned32)mBitMap.size();

		// Go through all nodes that aren't leaf nodes
		// And construct their (bit-based) material tree (e.g. the not-geometry part of the tree)
		auto leaf = Create(GetMaxLevel());
		leaf->SetIsGeometry(false);
		for (unsigned32 bit = 0; bit < bitCount; bit++)
		{
			offset = GetNodeCount();
			for (unsigned32 i = 0; i < levelIndices[GetMaxLevel()]; i++)
			{
				auto source = tree->GetTypedNode(i);

				// Create the material the new node has to have
				BitsMaterial<8> destMaterial;
				for (ChildIndex childIdx = 0; childIdx < 8; childIdx++)
				{
					if (source->HasChild(childIdx))
					{
						auto child = tree->GetTypedNode(source->GetChildIndex(childIdx));
						T sourceMaterial = tree->GetMaterial(child);
						unsigned32 sourceMaterialPointer = (unsigned32)mMaterialLibrary->GetTextureIndex(sourceMaterial);
						destMaterial.SetLS(childIdx, BitHelper::GetHS(sourceMaterialPointer, mBitMap[bit]));
					}
				}

				// Create the children pointers the new node has to have
				unsigned32* sourceChildren = source->GetChildren();
				if (source->GetLevel() == GetMaxLevel() - 1)
				{
					for (ChildIndex child = 0; child < source->GetChildCount(); child++)
						destChildrenCache[child] = leaf->GetIndex();
				}
				else
				{
					for (ChildIndex child = 0; child < source->GetChildCount(); child++)
						destChildrenCache[child] = offset + sourceChildren[child] - 1; // The root gets reused so it shouldn't be counted towards the offset
				}

				// Create the new node
				auto dest = source->GetLevel() == 0 ? GetSlaveRoot(bit) : Create(source->GetLevel());
				dest->SetMaterial(destMaterial);
				dest->SetIsGeometry(false);
				dest->SetChildren(source->GetChildmask(), destChildrenCache);
			}

			if (autoDAG)
			{
				ToDAG(1, false);
				printf(".");
			}
		}


	}

	// Bottom-up remove node that only have the same bit value (e.g. all 1 or all 0 for the non-geometry nodes).
	void ShaveEquals()
	{
		// For all the nodes on level 1, RemoveChildrenWithSameBit
		std::vector<unsigned32> levelIndices = SortOnLevel();
		BoolArray nodesToShave(GetNodeCount());
		// Assume that all nodes that arent geometry (or roots) can be shaved until the opposite has been proven
		for (unsigned32 i = 0; i < GetNodeCount(); i++)
		{
			Node* node = GetNode((unsigned32)i);
			if (node->GetLevel() > 0)
			{
				MultiRootMaterialNode<BitsMaterial<8>>* matNode = (MultiRootMaterialNode<BitsMaterial<8>>*)node;
				nodesToShave.Set(matNode->GetIndex(), !matNode->GetIsGeometry());
			}
		}

		// Now try to find proof not to shave certain nodes
		for (unsigned32 i = levelIndices[1]; i < levelIndices[2]; i++)
		{
			MultiRootMaterialNode<BitsMaterial<8>>* node =GetTypedNode(i);
			if (!node->GetIsGeometry()) CheckNodesToRemove(node, false, nodesToShave); // Since nodes on the second level don't have material properties, assume the bit value is false
		}

		// Shave all nodes of which no evidence is found not to shave them
		size_t nodesRemoved = 0;
		for (unsigned32 i = 0; i < GetNodeCount(); i++)
		{
			if (nodesToShave.Get(i))
			{
				nodesRemoved++;
				MultiRootMaterialNode<BitsMaterial<8>>* node = (MultiRootMaterialNode<BitsMaterial<8>>*)GetNode((unsigned32)i);
				node->SetMaterial(BitsMaterial<8>());
				node->SetChildren(0, NULL);
			}
		}

		printf("Connections removed from %llu nodes.\n", (unsigned64)nodesRemoved);
		ClearOrphans();
	}

	// Since the bit values only need to be correct for parts of the scene that are defined, we can randomly fill in the rest to be correct
	void FillEmptySpace(bool full = true)
	{
		// TODO: For the layer above the leaf node, just add all 8 children for each node (all pointing to the leaf node), and make
		// the only difference the bit node. Then connect all nodes in the layer above it correctly.
		std::vector<unsigned32> levelOffsets = SortOnLevel();

		// Find the leaf node that is not geometry
		MultiRootMaterialNode<BitsMaterial<8>>* leaf = NULL;
		for (unsigned32 i = levelOffsets[GetMaxLevel()]; i < levelOffsets[GetMaxLevel() + 1]; i++)
		{
			auto node = GetTypedNode(i);
			if (!node->GetIsGeometry()) { leaf = node; break; }
		}
		assert(leaf != NULL);

		// Create properties for pointers to this leaf node
		unsigned32 leafPointer = leaf->GetIndex();
		unsigned32 leafChildren[8];
		for (size_t i = 0; i < 8; i++) leafChildren[i] = leafPointer;
		ChildMask leafMask(255);

		// Make sure all nodes above the leaf level have full geometry, allowing more merging
		for (unsigned32 i = levelOffsets[GetMaxLevel() - 1]; i < levelOffsets[GetMaxLevel()]; i++)
		{
			auto node = GetTypedNode(i);
			if (!node->GetIsGeometry() && node->HasChildren())
				node->SetChildren(leafMask, leafChildren);
		}

		if (!full) return;

		// TODO: After every layer is made smaller, call "ToDAG" up to the next level to process :P
		// Now use a lookup table to quickly find all feasible nodes for a merge. 
		std::function<std::vector<unsigned32>(const MultiRootMaterialNode<BitsMaterial<8>>*)> childrenHasher = &MaterialLibraryMultiRoot::HashChildren;

		for (auto level = GetMaxLevel() - 1; level-- > 1;)
		{
			ToDAG(level - 1);
			levelOffsets = SortOnLevel();
			std::vector<size_t> parentsPerNode = GetParentCounts();
			NodeReplacementFinder<unsigned32, MultiRootMaterialNode<BitsMaterial<8>>*> finder(childrenHasher);

			auto levelStart = levelOffsets[level];
			auto levelEnd = levelOffsets[level + 1];
			std::vector<MultiRootMaterialNode<BitsMaterial<8>>*> nodesToTryVec;
			// Add all nodes to the finder:
			for (size_t i = levelStart; i < levelEnd; i++)
			{
				MultiRootMaterialNode<BitsMaterial<8>>* node = (MultiRootMaterialNode<BitsMaterial<8>>*)GetNode((unsigned32)i);
				if (!node->GetIsGeometry() && node->HasChildren())
				{
					finder.Add(node);
					nodesToTryVec.push_back(node);
				}
			}
			// Sort nodes on number of parents. Merging nodes with many parents is preferred as it has a high probability of leading to more merges higher up in the tree
			std::sort(nodesToTryVec.begin(), nodesToTryVec.end(), [&](MultiRootMaterialNode<BitsMaterial<8>>* a, MultiRootMaterialNode<BitsMaterial<8>>* b)
			{
				return parentsPerNode[a->GetIndex()] > parentsPerNode[b->GetIndex()];
			});

			std::queue<MultiRootMaterialNode<BitsMaterial<8>>*> nodesToTry;
			for (auto nodeToTry : nodesToTryVec) nodesToTry.push(nodeToTry);

			std::unordered_set<MultiRootMaterialNode<BitsMaterial<8>>*> allMergedNodes;
			// Now replace as much as possible
			while (!nodesToTry.empty())
			{
				MultiRootMaterialNode<BitsMaterial<8>>* node = nodesToTry.front();
				if (allMergedNodes.find(node) == allMergedNodes.end())
				{
					std::vector<MultiRootMaterialNode<BitsMaterial<8>>*> mergeOptions = finder.Find(node);
					// Prefer the merge options with most children :)
					std::sort(mergeOptions.begin(), mergeOptions.end(), [&](Node* a, Node* b) 
					{
						if (parentsPerNode[a->GetIndex()] != parentsPerNode[b->GetIndex()]) return parentsPerNode[a->GetIndex()] > parentsPerNode[b->GetIndex()];
						return a->GetChildCount() > b->GetChildCount();
					});
					// All merge options for this node will be explored, so remove it
					if (mergeOptions.size() > 1)
					{
						// Keep track of which nodes have been merged, so that we can set them to be equal to this node in the end
						ChildMask combinedMask = node->GetChildmask();
						unsigned32* combinedChildIndices = new unsigned32[8];
						for (ChildIndex c = 0; c < 8; c++) if (node->HasChild(c)) combinedChildIndices[c] = node->GetChildIndex(c);
						unsigned8 combinedMaterial = (unsigned8)node->GetMaterial().GetValue();

						std::vector<MultiRootMaterialNode<BitsMaterial<8>>*> mergedNodes(1, node);
						for (auto option : mergeOptions)
						{
							if (option != node)
							{
								// Check if the merge is still valid
								unsigned8 optionMaterial = (unsigned8)option->GetMaterial().GetValue();
								unsigned8 optionMask = (unsigned8)option->GetChildmask().mask;

								// The material mask should be the same for children that both nodes have:
								bool valid = (optionMaterial & (optionMask & combinedMask.mask)) == (combinedMaterial & (optionMask & combinedMask.mask));
								if (valid)
									for (ChildIndex c = 0; c < 8; c++)
										if (combinedMask.Get(c) && option->HasChild(c) && combinedChildIndices[c] != node->GetChildIndex(c))
										{
											valid = false;
											break;
										}

								// If the merge is still valid, updated the combined Mask, combined material and combinedChildIndices
								if (valid)
								{
									for (ChildIndex c = 0; c < 8; c++)
									{
										if (!combinedMask.Get(c) && option->HasChild(c))
										{
											combinedChildIndices[c] = option->GetChildIndex(c);
											combinedMask.Set(c, true);
										}
									}
									combinedMaterial |= optionMaterial;
									mergedNodes.push_back(option);
								}
							}
						}
						// If more nodes then the current node are merged,
						// Update all merged nodes to be equal. Also remove the (old) merged nodes from the finder and add the merged version
						if (mergedNodes.size() > 1)
						{
							unsigned8 i = 0;
							unsigned32* combinedChildren = new unsigned32[combinedMask.GetSet()];
							for (ChildIndex c = 0; c < 8; c++) if (combinedMask.Get(c)) combinedChildren[i++] = combinedChildIndices[c];
							BitsMaterial<8> combinedMaterialProp((size_t)combinedMaterial);

							for (auto mergedNode : mergedNodes)
							{
								finder.Remove(mergedNode);
								mergedNode->SetChildren(combinedMask, combinedChildren);
								mergedNode->SetMaterial(combinedMaterialProp);
								allMergedNodes.insert(mergedNode);
							}
							delete combinedChildren;
							finder.Add(node);
						}
						else
						{
							finder.Remove(node);
						}
					}
					else
					{
						finder.Remove(node);
					}
				}
				nodesToTry.pop();

			}

		}

	}

	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; }

	bool HasAdditionalPool() const override { return true; }
protected:
	void AppendPostProcess(glm::uvec3 coordinates, unsigned8 level, Tree<MultiRootMaterialNode<BitsMaterial<8>>>* tree) override
	{
		MaterialLibraryMultiRoot<T, Comparer, channelsPerPixel>* other = (MaterialLibraryMultiRoot<T, Comparer, channelsPerPixel>*)tree;
		if (other->mMaterialLibrary != NULL)
		{
			// Copy the material library of the appended tree
			if ((!this->mMaterialLibrary->IsFinalized()) && (*(other->mMaterialLibrary)) != (*(this->mMaterialLibrary)))
			{
				// Use copy constructor to copy the library of the other tree
				delete mMaterialLibrary;
				this->mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>(*(other->mMaterialLibrary));
			}

			assert((*(this->mMaterialLibrary)) == (*(other->mMaterialLibrary)));
		}
	}

	void WriteProperties(std::ostream& file) override {
		WriteMaterialTexture(file);
		HierarchicalMaterialMultiRoot<BitsMaterial<8>>::WriteProperties(file);
	}
	void ReadProperties(std::istream& file) override {
		// Reat the material texture
		ReadMaterialTexture(file);
		// Restore the material library from the texture
		if (mMaterialLibrary != NULL)
			delete mMaterialLibrary;
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>(mMaterialTexture, mMaterialTextureSize, mMaxTextureIndex);
		mMaterialLibrary->Finalize();

		HierarchicalMaterialMultiRoot<BitsMaterial<8>>::ReadProperties(file);
	}
	void WriteAdditionalPoolProperties(std::ostream& file) override { WriteMaterialTexture(file); HierarchicalMaterialMultiRoot<BitsMaterial<8>>::WriteAdditionalPoolProperties(file); }
	void ReadAdditionalPoolProperties(std::istream& file) override { ReadMaterialTexture(file); HierarchicalMaterialMultiRoot<BitsMaterial<8>>::ReadAdditionalPoolProperties(file); }
};