#pragma once
#include "UniqueIndexTree.h"
#include "MaterialTree.h"
#include "IMaterialTexture.h"
#include "../../inc/tbb/parallel_for_each.h"
#include "../Material/MaterialLibrary.h"
#include "../../core/BitHelper.h"
#include "../../core/Serializer.h"
#include "../../core/Util/BoolArray.h"
#include "../../inc/lodepng/lodepng.h"
#include <unordered_map>
#include <map>
#include <stack>

//#define PRINT_ERROR_DATA

template<typename T, typename Comparer = std::less<T>, unsigned8 channelsPerPixel = 3>
class MaterialLibraryUniqueIndexTree : public UniqueIndexTree<MaterialLibraryPointer>, 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>		mMaterialTexture;
	unsigned16					mMaterialTextureSize;
	MaterialLibraryPointer		mMaxTextureIndex;

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

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

	// Creates a UniqueIndexRoot with "maxLevel" levels.  
	// Note that the nodeMaterialsTexture will be deleted if the tree is deleted.
	MaterialLibraryUniqueIndexTree(unsigned8 maxLevel, CompressedTexture<MaterialLibraryPointer>* nodeMaterialsTexture)
		: MaterialLibraryUniqueIndexTree(maxLevel, nodeMaterialsTexture, 0)
	{}

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

	void AppendPostProcess(glm::uvec3 coordinates, unsigned8 level, Tree* tree) override
	{
		MaterialLibraryUniqueIndexTree<T, Comparer, channelsPerPixel>* uniqueIndexTree = (MaterialLibraryUniqueIndexTree<T, Comparer, channelsPerPixel>*)tree;

		// Copy the color information from the blocks of the other tree. Make sure that the blocks are inserted in the correct position.
		auto oldLibrary = mMaterialLibrary;
		auto existingMaterials = GetUniqueMaterials();
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		for (auto material : existingMaterials)
			mMaterialLibrary->AddMaterial(material);
		auto appendedLibrary = uniqueIndexTree->GetMaterialLibrary();
		auto appendedMaterials = uniqueIndexTree->GetUniqueMaterials();
		for (auto material : appendedMaterials)
			mMaterialLibrary->AddMaterial(material);
		mMaterialLibrary->Finalize();
		mMaterialTexture = std::vector<unsigned char>();
		mMaterialTextureSize = 0;
		mMaxTextureIndex =  mMaterialLibrary->GetMaxTextureIndex();

		// Replace the existing materials
		std::unordered_map<MaterialLibraryPointer, MaterialLibraryPointer> ownPointerReplacers;
		for (auto material : existingMaterials)
			ownPointerReplacers[oldLibrary->GetTextureIndex(material)] = mMaterialLibrary->GetTextureIndex(material);
		UniqueIndexTree::ReplaceMaterials(ownPointerReplacers);
		delete oldLibrary;

		// Create a map from the old tree to the new tree replacers:
		std::unordered_map<MaterialLibraryPointer, MaterialLibraryPointer> appendedPointerReplacers;
		for (auto material : appendedMaterials)
			appendedPointerReplacers[appendedLibrary->GetTextureIndex(material)] = mMaterialLibrary->GetTextureIndex(material);

		UniqueIndexTree::AppendPostProcess(coordinates, level, uniqueIndexTree, appendedPointerReplacers);
	}

	void ReplaceMaterials(const std::unordered_map<T, T>& replacementMap)
	{
		assert(mMaterialLibrary != NULL);
		auto oldLibrary = mMaterialLibrary;
		std::vector<T> oldMaterials = oldLibrary->GetMaterials();
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		std::unordered_map<T, T> actualReplacers;
		for (auto material : oldMaterials)
		{
			auto replacerIt = replacementMap.find(material);
			T replacer = material;
			if (replacerIt == replacementMap.end())
				printf("Material not found in replacementMap");
			else
				replacer = replacerIt->second;
			mMaterialLibrary->AddMaterial(replacer);
			actualReplacers.insert(std::pair<T, T>(material, replacer));
		}
		mMaterialLibrary->Finalize();

		std::unordered_map<MaterialLibraryPointer, MaterialLibraryPointer> pointerReplacers;
		for (auto material : oldMaterials)
		{
			auto oldPointer = oldLibrary->GetTextureIndex(material);
			auto newPointer = mMaterialLibrary->GetTextureIndex(actualReplacers[material]);
			pointerReplacers.insert(std::pair<MaterialLibraryPointer, MaterialLibraryPointer>(oldPointer, newPointer));
		}

		UniqueIndexTree::ReplaceMaterials(pointerReplacers);

		delete oldLibrary;
	}

	// Will build a unique index tree with the same material information as the given material tree. 
	// This will consume (and delete) the given tree!
	template<typename MaterialTree>
	void BaseOn(MaterialTree* tree)
	{
		// Create a material library for all original materials
		assert(mMaterialLibrary == NULL);
		mMaterialLibrary = new MaterialLibrary<T, Comparer, channelsPerPixel>();
		std::vector<T> uniqueMaterials = tree->GetUniqueMaterials();
		for (auto material : uniqueMaterials)
			mMaterialLibrary->AddMaterial(material);
		mMaterialLibrary->Finalize();

		// Create a new material tree containing the material pointers:
		size_t nodeCount = tree->GetNodeCount();
		std::vector<MaterialLibraryPointer> pointers(nodeCount);
		tbb::parallel_for(size_t(0), nodeCount, [&](size_t i)
		{
			pointers[i] = mMaterialLibrary->GetTextureIndex(tree->GetMaterial(tree->GetTypedNode((unsigned32)i)));
		});
		
		UniqueIndexTree::BaseOn(tree, pointers);
	}

	// Returns a list with all unique materials in the tree
	std::vector<T> GetUniqueMaterials() const
	{
		return mMaterialLibrary->GetMaterials();
	}

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

	// Returns the material of the node at index i
	T GetMaterial(const size_t& i) const
	{
		return mMaterialLibrary->GetMaterial(GetNodeValue(i));
	}

	std::vector<T> GetMaterials(size_t fromIndex = 0) const
	{
		std::vector<MaterialLibraryPointer> pointersTexture = GetNodeValues(fromIndex);
		std::vector<T> res(pointersTexture.size());
		tbb::parallel_for(size_t(0), res.size(), [&](size_t i)
		{
			res[i] = mMaterialLibrary->GetMaterial(pointersTexture[i]);
		});
		return res;
	}

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

	void PrintDebugInfo() const override
	{
		UniqueIndexTree::PrintDebugInfo();
	}

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

#ifdef PRINT_ERROR_DATA
#include "../../core/OctreeBuilder/ColorQuantizerFactory.h"
#include "../Material/MaterialQuantizer/ColorQuantizer/BaseColorQuantizer.h"

struct ColorErrorInfo
{
	glm::u64vec3 sumError = glm::u64vec3(0);
	unsigned maxError = 0;
	glm::uvec3 maxErrorValue = glm::uvec3(0);
	long double sumDeltaE = 0;
	float maxDeltaE = 0;
};

template<>
void MaterialLibraryUniqueIndexTree<Color, ColorCompare>::PrintDebugInfo() const
{
	// If the colors are not quantized, print some information
	std::vector<Color> uniqueMaterials = mMaterialLibrary->GetMaterials();
	printf("Calculating quantization errors over the full scene (%u unique colors, %llu values)...\n", (unsigned32)uniqueMaterials.size(), (unsigned64)this->GetMaterialCount());
	std::string quantizationTypes[4] = { "lab256", "lab1024", "lab4096", "lab16384" };
	for (std::string type : quantizationTypes)
	{
		printf("%s: ", type.c_str());
		BaseColorQuantizer* quantizer = ColorQuantizerFactory::Create(type);
		auto quantizationMap = quantizer->QuantizeMaterials(uniqueMaterials);

		//glm::dvec3 sumError(0);
		//unsigned maxError = 0;
		//glm::uvec3 maxErrorValue(0);
		//long double sumDeltaE = 0;
		//float maxDeltaE = 0;
		ColorErrorInfo initial;
		//tbb::parallel_for((size_t)0, this->GetMaterialCount(), [&](const size_t& i)
		//for (size_t i = 0; i < this->GetMaterialCount(); i++)
		ColorErrorInfo error = tbb::parallel_reduce(tbb::blocked_range<size_t>(0, this->GetMaterialCount()), initial, [&](const tbb::blocked_range<size_t>& r, ColorErrorInfo errorInfo)
		{
			for (size_t i = r.begin(); i != r.end(); i++)
			{
				const Color& actual = GetMaterial(i);
				const Color& quantized = quantizationMap->at(actual);
				glm::ivec3 error = glm::abs(glm::ivec3(actual.GetColor()) - glm::ivec3(quantized.GetColor()));
				errorInfo.sumError += error;
				unsigned errorU = error.r + error.g + error.b;
				if (errorU > errorInfo.maxError)
				{
					errorInfo.maxError = errorU;
					errorInfo.maxErrorValue = error;
				}
				float deltaE = ColorHelper::GetDeltaEFromRGB(actual.GetColor(), quantized.GetColor());
				if (deltaE == deltaE) // Only sum if it is not NaN...
					errorInfo.sumDeltaE += deltaE;
				if (deltaE > errorInfo.maxDeltaE) errorInfo.maxDeltaE = deltaE;
			}
			return errorInfo;
		}, [&](ColorErrorInfo a, ColorErrorInfo b)
		{
			ColorErrorInfo res;
			res.sumError = a.sumError + b.sumError;
			res.maxError = a.maxError > b.maxError ? a.maxError : b.maxError;
			res.maxErrorValue = a.maxError > b.maxError ? a.maxErrorValue : b.maxErrorValue;
			res.sumDeltaE = a.sumDeltaE + b.sumDeltaE;
			res.maxDeltaE = a.maxDeltaE > b.maxDeltaE ? a.maxDeltaE : b.maxDeltaE;
			return res;
		});
		//);
		glm::dvec3 sumError = glm::dvec3(error.sumError);
		glm::dvec3 meanError = sumError / double(GetMaterialCount());
		double meanDeltaE = error.sumDeltaE / double(GetMaterialCount());
		printf("Mean errors: (%f, %f, %f), Max errors: (%u, %u, %u), Mean delta-E: %f, Max delta-E: %f\n", meanError.x, meanError.y, meanError.z, error.maxErrorValue.x, error.maxErrorValue.y, error.maxErrorValue.z, meanDeltaE, error.maxDeltaE);

		delete quantizationMap;
	}
}

#include "../Material/MaterialQuantizer/NormalQuantizer.h"
#include "../../inc/glm/gtx/vector_angle.hpp"

struct NormalErrorInfo
{
	float maxAngle = 0;
	long double sumAngle = 0;
};

template<>
void MaterialLibraryUniqueIndexTree<SmallNormal, NormalCompare, 4>::PrintDebugInfo() const
{
	// If the colors are not quantized, print some information 
	std::vector<SmallNormal> uniqueMaterials = mMaterialLibrary->GetMaterials();
	printf("Calculating normal quantization errors over the full scene (%u unique normals, %llu values)...\n", (unsigned32)uniqueMaterials.size(), (unsigned64)this->GetMaterialCount());
	unsigned8 bits[] = { 12 };
	for (unsigned8 quantizeToBits : bits)
	{
		printf("%u bits: ", quantizeToBits);
		NormalQuantizer quantizer(quantizeToBits);

		NormalErrorInfo initial;
		NormalErrorInfo error = tbb::parallel_reduce(tbb::blocked_range<size_t>(0, this->GetMaterialCount()), initial, [&](const tbb::blocked_range<size_t>& r, NormalErrorInfo errorInfo)
		{
			for (size_t i = r.begin(); i != r.end(); i++)
			{
				const SmallNormal& actual = GetMaterial(i);
				SmallNormal quantized = quantizer.Quantize(actual);
				float angle = glm::angle(actual.Get(), quantized.Get());
				if (angle > errorInfo.maxAngle) errorInfo.maxAngle = angle;
				errorInfo.sumAngle += angle;
			}
			return errorInfo;
		}, [&](NormalErrorInfo a, NormalErrorInfo b)
		{
			NormalErrorInfo res;
			res.maxAngle = a.maxAngle > b.maxAngle ? a.maxAngle : b.maxAngle;
			res.sumAngle = a.sumAngle + b.sumAngle;
			return res;
		});
		//);
		long double meanError = error.sumAngle / long double(GetMaterialCount());
		printf("Mean angle error: %f, Max angle error: %f\n", meanError * (360.0 / (2 * mPi)), error.maxAngle * (360.0 / (2 * mPi)));
	}
}
#endif