#pragma once
#include "Tree.h"
#include "IAdditionalProperties.h"

template<typename NodeType = Node>
class MultiRootTree : public Tree<NodeType>, public IAdditionalProperties
{
public:
	MultiRootTree(unsigned8 maxLevel, unsigned32 slaveRootCount) : Tree(maxLevel)
	{
		// Initialize the slaves
		AddSlaveRoots(slaveRootCount);
	}
	~MultiRootTree() override {}

	void Clear() override
	{
		mSlaveRoots.clear();
		Tree<NodeType>::Clear();
	}

	unsigned32 GetSlaveRootCount() const { return (unsigned32)mSlaveRoots.size(); }
	const NodeType* GetSlaveRoot(unsigned32 i) const { return GetTypedNode(mSlaveRoots[i]); }
	NodeType* GetSlaveRoot(unsigned32 i) { return GetTypedNode(mSlaveRoots[i]); }
	void AddSlaveRoots(size_t count)
	{
		size_t originalSlaveRootCount = GetSlaveRootCount();
		mSlaveRoots.resize(originalSlaveRootCount + count);
		for (size_t i = 0; i < count; i++)
			mSlaveRoots[originalSlaveRootCount + i] = Create(0)->GetIndex();
	}

	void Append(glm::uvec3 coordinates, unsigned8 level, MultiRootTree* tree)
	{
		// Let some possible inhereting class do preprocessing on the current tree before append
		AppendPreProcess(coordinates, level, tree);

		// If the tree that is to be appended has more roots, add roots to this tree first:
		if (tree->GetSlaveRootCount() > this->GetSlaveRootCount())
			AddSlaveRoots(tree->GetSlaveRootCount() - this->GetSlaveRootCount());

		std::vector<unsigned32> equivalents(tree->GetNodeCount());

		// First create/get the node that acts as the root of the tree to append
		NodeType* treeRoot = this->AddNode(coordinates, level);
		treeRoot->CopyProperties(tree->GetRoot());
		equivalents[0] = treeRoot->GetIndex();

		// Make copies of all nodes in tree, and add them to our own nodepool (with the correct level and root)
		NodeType* copy = NULL;
		for (unsigned32 nodeId = 1; nodeId < tree->GetNodeCount(); nodeId++)
		{
			NodeType* node = tree->GetTypedNode(nodeId);
			if (node->GetLevel() == 0)
			{
				// Find which slave this root belongs to
				for (unsigned i = 0; i < tree->GetSlaveRootCount(); i++)
					if (tree->GetSlaveRoot(i) == node)
						// And add it
						copy = (NodeType*)this->GetSlaveRoot(i)->AddNode(coordinates, level);
			}
			else
				copy = this->Create(node->GetLevel() + level);
			assert(copy != NULL);
			copy->CopyProperties(node);
			equivalents[nodeId] = copy->GetIndex();
		}
		// Restore all child pointers
		unsigned32* newChildren = new unsigned32[8];
		for (unsigned32 i = 0; i < tree->GetNodeCount(); i++)
		{
			NodeType* copy = GetTypedNode(equivalents[i]);
			NodeType* node = tree->GetTypedNode(i);

			unsigned32* children = node->GetChildren();
			unsigned8 childCount = node->GetChildCount();

			for (ChildIndex c = 0; c < childCount; c++) newChildren[c] = equivalents[children[c]];

			copy->SetChildren(node->GetChildmask(), newChildren);
		}
		delete[] newChildren;

		// Let some possible inhereting class do postprocessing on the added nodes
		AppendPostProcess(coordinates, level, tree);
	}

	// Adds a leaf node to the given slave root ID. Use AddLeafNode without additional arguments to add leafs to the main root.
	NodeType* AddLeafNode(glm::uvec3 coordinates)
	{
		return Tree<NodeType>::AddLeafNode(coordinates);
	}
	NodeType* AddLeafNode(glm::uvec3 coordinates, unsigned32 slaveRootID)
	{
		return (NodeType*)GetSlaveRoot(slaveRootID)->AddNode(coordinates, GetMaxLevel());
	}
	bool SlaveHasLeaf(glm::uvec3 coordinates, unsigned32 slaveRootID) const
	{
		return GetSlaveRoot(slaveRootID)->HasNode(coordinates, GetMaxLevel());
	}

	std::vector<size_t> GetOctreeNodesPerLevel() const override
	{
		std::vector<size_t> octreeNodesPerLevel(GetMaxLevel() + 1);
		std::function<void(const Node*)> nodeCounter = [&octreeNodesPerLevel](const Node* node) -> void
		{
			octreeNodesPerLevel[node->GetLevel()]++;
		};
		this->Traverse(nodeCounter);
		for (unsigned32 i = 0; i < GetSlaveRootCount(); i++) GetSlaveRoot(i)->Traverse(nodeCounter);
		return octreeNodesPerLevel;
	}

	//size_t				GetMinimumNodePoolSize() const override;
	//std::vector<unsigned8>& GetNodePool() override;

	void WriteProperties(std::ostream& file) override
	{
		// Write the number of slave roots, and their indexes
		Serializer<std::vector<unsigned32>, unsigned32>::Serialize(mSlaveRoots, file);
	}
	void ReadProperties(std::istream& file) override
	{
		// By this time all original slaves have been deleted, so we need to rebuild them
		Serializer<std::vector<unsigned32>, unsigned32>::Deserialize(mSlaveRoots, file);
	}

	std::map<std::string, std::string> GetAdditionalProperties() override
	{
		if (!mAdditionalProperties.empty()) return mAdditionalProperties;
		else
		{
			mAdditionalProperties.insert(std::pair<std::string, std::string>("RootCount", std::to_string(1 + mSlaveRoots.size())));
			return mAdditionalProperties;
		}
	}

	bool HasAdditionalPool() const override { return true; }
protected:
	unsigned8 GetAdditionalBytesPerNode(unsigned8 level) const override { return level == GetMaxLevel() ? 1 : 0; }
	virtual std::vector<unsigned8> GetAdditionalNodeBytes(const Node* node) const override
	{
		if (node->GetLevel() != GetMaxLevel()) return std::vector<unsigned8>();
		return std::vector<unsigned8>(1, 1);
	}

	unsigned8			GetAdditionalTreeInfoSize() const override { return (unsigned8)(mSlaveRoots.size() + 1) * 2; }
	std::vector<unsigned8> GetAdditionalTreeInfo(const std::vector<size_t>& nodePointers) const override
	{
		std::vector<unsigned8> res(GetAdditionalTreeInfoSize());
		// Start with the main root:
		std::vector<unsigned8> rootPointer = BitHelper::SplitInBytes(nodePointers[0], 2);
		std::move(rootPointer.begin(), rootPointer.end(), res.begin());
		for (unsigned32 i = 0; i < GetSlaveRootCount(); i++)
		{
			size_t slavePointer = nodePointers[GetSlaveRoot(i)->GetIndex()];
			std::vector<unsigned8> slavePointerBytes = BitHelper::SplitInBytes(slavePointer, 2);
			std::move(slavePointerBytes.begin(), slavePointerBytes.end(), res.begin() + (1 + i) * 2);
		}
		return res;
	}

	void UpdateLocalReferences(const std::vector<unsigned32>& indexMap)
	{
		for (size_t i = 0; i < mSlaveRoots.size(); i++)
			mSlaveRoots[i] = indexMap[mSlaveRoots[i]];
		Tree<NodeType>::UpdateLocalReferences(indexMap);
	}

	void  WriteAdditionalPoolProperties(std::ostream& file) override
	{
		GetAdditionalProperties();
		Serializer<std::map<std::string, std::string>>::Serialize(mAdditionalProperties, file);
	}
	void ReadAdditionalPoolProperties(std::istream& file) override
	{
		Serializer<std::map<std::string, std::string>>::Deserialize(mAdditionalProperties, file);
	}

	std::vector<unsigned32>				mSlaveRoots;
	std::map<std::string, std::string>	mAdditionalProperties;
};