#pragma once
#include "UniqueIndexTree.h"
#include "../Material/SignedIntMaterial.h"
#include "MaterialTree.h"
#include "../Material/MaterialLibrary.h"
#include "../../inc/tbb/parallel_for.h"
#include "IMaterialTexture.h"
#include<vector>

template<typename T, typename Comparer = std::less<T>, unsigned8 channelsPerPixel = 3>
class UniqueIndexShiftTree : public UniqueIndexTree<SignedIntMaterial>, public IMaterialTexture
{
private:
	MaterialLibrary<T, Comparer, channelsPerPixel>* mMaterialLibrary;
	std::vector<unsigned8>		mMaterialTexture;
	unsigned16					mMaterialTextureSize;
	MaterialLibraryPointer		mMaxTextureIndex;

	void ClearMaterialLibrary()
	{
		if (mMaterialLibrary != NULL)
		{
			delete mMaterialLibrary;
			mMaterialLibrary = NULL;
		}
	}

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

	void BuildMaterialLibrary(const std::vector<T>& materials)
	{
		ClearMaterialLibrary();
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		for (T m : materials) mMaterialLibrary->AddMaterial(m);
		mMaterialLibrary->Finalize();
	}

	bool CheckMaterialLibrary(const std::vector<T>& materials)
	{
		for (T m : materials) if (!mMaterialLibrary->Contains(m)) return false;
		return true;
	}

public:
	UniqueIndexShiftTree(unsigned8 maxLevel, CompressedTexture<SignedIntMaterial>* nodeMaterialsTexture, unsigned32 collapsedMaterialLevels) :
		UniqueIndexTree(maxLevel, nodeMaterialsTexture, collapsedMaterialLevels),
		mMaterialLibrary(NULL),
		mMaterialTexture(std::vector<unsigned8>()),
		mMaterialTextureSize(0),
		mMaxTextureIndex(MaterialLibraryPointer(0))
	{}

	UniqueIndexShiftTree(unsigned8 maxLevel, CompressedTexture<SignedIntMaterial>* nodeMaterialsTexture)
		: UniqueIndexShiftTree(maxLevel, nodeMaterialsTexture, 0)
	{}

	~UniqueIndexShiftTree() override {
		ClearMaterialLibrary();
	}

	void AppendPostProcess(glm::uvec3 coordinates, unsigned8 level, Tree<UniqueIndexNode>* tree) override
	{
		UniqueIndexShiftTree<T, Comparer, channelsPerPixel>* typedTree = (UniqueIndexShiftTree<T, Comparer, channelsPerPixel>*)tree;
		// If the material libraries are equal, appending should work correctly
		assert(*typedTree->GetMaterialLibrary() == *mMaterialLibrary);
		UniqueIndexTree::AppendPostProcess(coordinates, level, typedTree, std::unordered_map<SignedIntMaterial, SignedIntMaterial>());
	}

	// Prepare the tree for a certain set of materials. This should be done before calling "BaseOn" or otherwise adding nodes to the tree.
	void PrepareForMaterials(const std::vector<T>& materials)
	{
		assert(GetNodeCount() == 1);
		BuildMaterialLibrary(materials);
	}
	
	template<typename MaterialTree>
	void BaseOn(MaterialTree* tree)
	{
		// First we build a new material tree that contains the differences in the materials compared to the parents for each node (kind of watch hsc is)
		// To do this, we need an octree (as the same color doesn't mean it will have the same difference to it's parent)
		tree->ToOctree();
		std::vector<T> uniqueMaterials = tree->GetUniqueMaterials();

		// Create the actual material library
		if (mMaterialLibrary == NULL)
			BuildMaterialLibrary(uniqueMaterials);
		else
			assert(CheckMaterialLibrary(uniqueMaterials));

		// Now we calculate the shift for each node in the tree
		std::vector<int> intMaterialPointer(tree->GetNodeCount());
		tbb::parallel_for((unsigned32)0, tree->GetNodeCount(), [&](unsigned32 i)
		{
			auto node = tree->GetTypedNode(i);
			MaterialLibraryPointer matPointer = mMaterialLibrary->GetTextureIndex(tree->GetMaterial(node));
			intMaterialPointer[i] = (matPointer.y * mMaterialLibrary->GetTextureSize()) + matPointer.x;
		});
		std::vector<SignedIntMaterial> shiftMaterials(tree->GetNodeCount());
		shiftMaterials[0] = intMaterialPointer[0];
		tbb::parallel_for((unsigned32)0, tree->GetNodeCount(), [&](unsigned32 i)
			//for (size_t i = 0; i < tree->GetNodeCount(); i++)
		{
			Node* node = tree->GetNode(i);
			int curIntIndex = intMaterialPointer[i];
			unsigned32* children = node->GetChildren();
			unsigned8 childCount = node->GetChildCount();
			for (ChildIndex c = 0; c < childCount; c++)
			{
				int childIntIndex = intMaterialPointer[children[c]];
				shiftMaterials[children[c]] = SignedIntMaterial(childIntIndex - curIntIndex);
			}
		});

		std::vector<int> shiftsCopy(shiftMaterials.size());
		for (size_t i = 0; i < shiftMaterials.size(); i++) shiftsCopy[i] = shiftMaterials[i];

		CollectionHelper::PrintStats(shiftsCopy);
		for (size_t i = 0; i < shiftsCopy.size(); i++) shiftsCopy[i] = abs(shiftsCopy[i]);
		CollectionHelper::PrintStats(shiftsCopy);
		
		UniqueIndexTree::BaseOn(tree, shiftMaterials);
	}

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

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

	// Returns the texture containing all materials once
	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; }

protected:
	void WriteAdditionalUniqueIndexTreeProperties(std::ostream& file) override {
		WriteMaterialTexture(file);
	}

	void ReadAdditionalUniqueIndexTreeProperties(std::istream& file) override {
		ReadMaterialTexture(file);
	}

	void WriteAdditionalPoolProperties(std::ostream& file) override { WriteMaterialTexture(file);  UniqueIndexTree::WriteAdditionalPoolProperties(file); }
	void ReadAdditionalPoolProperties(std::istream& file) override { ReadMaterialTexture(file); UniqueIndexTree::ReadAdditionalPoolProperties(file); }

	void PrintDebugInfo() const override
	{
		std::vector<SignedIntMaterial> shifts = GetNodeValues();
		std::vector<int> shiftValues(shifts.size());
		for (size_t i = 0; i < shifts.size(); i++) shiftValues[i] = (int)shifts[i];
		CollectionHelper::PrintStats(shiftValues);
	}
};