#pragma once
#include "BaseMaterialOctreeBuilder.h"
#include <unordered_set>
#include "../../inc/glm/common.hpp"
#include "../../core/Hashers.h"
#include "../../scene/Material/MaterialLibraryPointer.h"
#include "../../scene/Octree/MaterialLibraryUniqueIndexTree.h"
#include "../../scene/Octree/MaterialTree.h"

template<typename T, typename Comparer>
class UniqueIndexMaterialOctreeBuilder : public BaseMaterialOctreeBuilder<T>
{
private:
	static std::string GetSubtreeTextureCompressionType() { return "b"; } // Basic texture compression (e.g. no texture compression)
	typedef MaterialLibraryUniqueIndexTree<T, Comparer, T::CHANNELSPERPIXEL> FinalTreeType;
	typedef MaterialTree<T, Comparer> IntermediateTreeType;
public:
	UniqueIndexMaterialOctreeBuilder(std::string textureCompressionType, BaseQuantizer<T, Comparer>* quantizer = NULL, unsigned32 levelsWithoutMaterials = 0) :
		BaseMaterialOctreeBuilder(),
		mTree(NULL),
		mReduceMaterials(quantizer != NULL),
		mTextureCompressionType(textureCompressionType),
		mLevelsWithoutMaterials(levelsWithoutMaterials),
		mSceneMaterials(std::vector<T>()),
		mMaterialReplacers(std::unordered_map<T, T>()),
		mCurPreprocessPassMaterials(std::vector<T>()),
		mMainTreeMaterials(std::vector<std::pair<glm::uvec3, T>>()),
		mIntermediateTree(NULL),
		mQuantizer(quantizer)
	{}

	UniqueIndexMaterialOctreeBuilder(std::string textureCompressionType, unsigned32 levelsWithoutMaterials = 0) :
		UniqueIndexMaterialOctreeBuilder(textureCompressionType, NULL, levelsWithoutMaterials)
	{}

	~UniqueIndexMaterialOctreeBuilder() override
	{
		if (mIntermediateTree != NULL) delete mIntermediateTree;
		if (mTree != NULL) delete mTree;
	}

	std::string GetTreeType() override
	{
		return "u" +
			(mLevelsWithoutMaterials == 0 ? "" : (std::to_string(mLevelsWithoutMaterials) + "lod"))
			+ mTextureCompressionType + MaterialAbbreviation<T>()() +
			(mReduceMaterials ? mQuantizer->GetQuantizerDescriptor() : "");
	}

	std::string GetSubtreeType()
	{
		return "u" +
			(mLevelsWithoutMaterials == 0 ? "" : (std::to_string(mLevelsWithoutMaterials) + "lod"))
			+ GetSubtreeTextureCompressionType() + MaterialAbbreviation<T>()() +
			(mReduceMaterials ? mQuantizer->GetQuantizerDescriptor() : "");
	}
protected:
	bool		SubTreeCompare(const glm::uvec3& coord1, const glm::uvec3& coord2) const override
	{
		// Sort them so that the trees with the lowest index will be processed first. 
		// Since indexes are given in a depth-first order, and the lowest indexes are given to the trees with the highest ChildIndex in each level,
		for (unsigned8 bit = GetAppendedTreeLevel(); bit > 0; bit--)
		{
			unsigned mask = 1 << (bit - 1);
			if ((coord1.z & mask) != (coord2.z & mask))
				return (coord1.z & mask) > (coord2.z & mask);
			if ((coord1.y & mask) != (coord2.y & mask))
				return (coord1.y & mask) > (coord2.y & mask);
			if ((coord1.x & mask) != (coord2.x & mask))
				return (coord1.x & mask) > (coord2.x & mask);
		}
		return true;
	}
	
	// Initialize the main tree
	void		InitTree() override
	{
		auto texture = CompressedTextureFactory<MaterialLibraryPointer>::GetCompressedTexture(mTextureCompressionType);
		mTree = new FinalTreeType(GetTreeDepth(), texture, mLevelsWithoutMaterials);
		IntermediateTreeType* tempTree = new IntermediateTreeType(GetTreeDepth());

		if (!IsSinglePass())
		{
			for (auto coordMaterial : mMainTreeMaterials)
			{
				tempTree->SetMaterial(coordMaterial.first, GetAppendedTreeLevel(), coordMaterial.second);
			}
			tempTree->PropagateMaterials(T::WeightedAverage);
			mTree->BaseOn(tempTree);
		}

		mFirstPass = true;
	}

	bool		UsePreprocessing() const override { return !IsSinglePass(); }
	void		InitPreprocessing() override
	{
		mSceneMaterials.clear();
		mMaterialReplacers.clear();
		mMainTreeMaterials.clear();
	}
	void		InitCurPreprocessPass(glm::uvec3 coordinate) override
	{
		mCurPreprocessPassMaterials.clear();
	}
	void		PreProcessNode(const glm::uvec3& coordinate, const T& color) override { mCurPreprocessPassMaterials.push_back(color); }
	void		FinalizeCurPreprocessPass(glm::uvec3 coord) override
	{
		if (mCurPreprocessPassMaterials.empty())
			return;
		tbb::parallel_sort(mCurPreprocessPassMaterials, Comparer());
		std::vector<T> uniqueMaterials;
		std::vector<float> uniqueMaterialWeights;
		unsigned curMaterialCount = 0;
		T lastSeenMaterial = mCurPreprocessPassMaterials[0];
		for (auto color : mCurPreprocessPassMaterials)
		{
			if (!(lastSeenMaterial == color))
			{
				uniqueMaterials.push_back(lastSeenMaterial);
				uniqueMaterialWeights.push_back((float)curMaterialCount);
				lastSeenMaterial = color;
				curMaterialCount = 0;
			}
			curMaterialCount++;
		}
		T avgMaterial = T::WeightedAverage(uniqueMaterials, uniqueMaterialWeights);
		mMainTreeMaterials.push_back(std::pair<glm::uvec3, T>(coord, avgMaterial));

		// Append the unique colors to the scene colors and compress them to keep the memory usage acceptable
		mSceneMaterials.insert(mSceneMaterials.end(), uniqueMaterials.begin(), uniqueMaterials.end());
		tbb::parallel_sort(mSceneMaterials, Comparer());
		mSceneMaterials.erase(std::unique(mSceneMaterials.begin(), mSceneMaterials.end()), mSceneMaterials.end());
	}

	void QuantizeSceneMaterials()
	{
		if (mQuantizer == NULL || !mReduceMaterials) return;
		if (verbose) printf("Quantizing/merging %llu %s...", (unsigned64)mSceneMaterials.size(), MaterialName<T>()());
		Stopwatch watch; watch.Reset();
		auto quantizedSceneMaterials = mQuantizer->QuantizeMaterials(mSceneMaterials);

		// Replace the list of scene colors with the quantized scene colors
		mSceneMaterials.clear();
		for (auto color : *quantizedSceneMaterials)
			mSceneMaterials.push_back(color.second);
		tbb::parallel_sort(mSceneMaterials, Comparer());
		mSceneMaterials.erase(std::unique(mSceneMaterials.begin(), mSceneMaterials.end()), mSceneMaterials.end());
		mSceneMaterials.shrink_to_fit();

		// Build a dictionary for quick lookup of original scene colors and their quantized counterparts
		mMaterialReplacers = std::unordered_map<T, T>();
		for (auto color : *quantizedSceneMaterials)
			mMaterialReplacers.insert(std::make_pair(color.first, color.second));

		// Clear the cur preprocessMaterials (to free up memory during tree construction)
		mCurPreprocessPassMaterials = std::vector<T>();
		delete quantizedSceneMaterials;

		// Replace the old colors by the new ones
		if (verbose) printf("Quantized %s in %d ms\n", MaterialName<T>()(), (int)(watch.GetTime() * 1000));
	}

	void		FinalizeTree() override
	{
		Stopwatch watch;
		if (!IsSinglePass())
		{
			unsigned32 i = 1;
			for (const glm::uvec3& coord : GetValidCoords())
			{
				{ // Scope subTree variable
					if (verbose) printf("Reading subtree %u / %u at (%u, %u, %u) from cache...\n", i, (unsigned32)(GetValidCoords().size()), coord.x, coord.y, coord.z);
					FinalTreeType* subTree = (FinalTreeType*)OctreeLoader::ReadCache(GetSubtreeType(), GetSinglePassTreeDepth(), GetSinglePassTreeFilename(coord), verbose);
					if (subTree != NULL)
					{
						// Append the subtree to the main tree
						if (verbose) printf("Appending subtree... ");;
						watch.Reset();
						mTree->Append(coord, GetAppendedTreeLevel(), subTree);
						delete subTree;
						if (verbose) printf("Appending took %d ms.\n", (int)(watch.GetTime() * 1000));
					}
				}

				// Convert the new part of the main tree to a DAG
				if (verbose) printf("Converting current tree to DAG...\n");
				watch.Reset();
				mTree->ToDAG(GetAppendedTreeLevel());
				if (verbose) printf("Converting took %d ms.\n", (int)(watch.GetTime() * 1000));

				i++;
			}
		}
		
		// Generate the material texture
		watch.Reset();
		if (verbose) printf("Generating material texture... ");
		mTree->GetMaterialTexture();
		if (verbose) printf("Material texture generated in %d ms\n", (int)(watch.GetTime() * 1000));

		// Delete the cache files
		if (!IsSinglePass())
		{
			if (verbose) printf("Deleting cache...");
			watch.Reset();
			for (const glm::uvec3& coord : GetValidCoords())
				OctreeLoader::DeleteCache(GetSubtreeType(), GetSinglePassTreeDepth(), GetSinglePassTreeFilename(coord));
			if (verbose) printf("Cache deleted in %d ms\n", (int)(watch.GetTime() * 1000));
		}
	}

	// Step to finalize the main tree (for example storing it to a file)
	void TerminateTree() override
	{
		OctreeLoader::WriteCache(mTree, GetTreeType(), GetOutputFile(), verbose);
		delete mTree;
		mTree = NULL;
		mSceneMaterials.clear();
		mMaterialReplacers.clear();
		mMainTreeMaterials.clear();
	}

	// Don't build this subtree again if a cache file exists for it
	bool CancelCurPassTree(const glm::uvec3& coord) override
	{
		return OctreeLoader::VerifyCache(GetSubtreeType(), GetSinglePassTreeDepth(), GetSinglePassTreeFilename(coord));
	}

	// Initialize the tree for the current pass.
	void		InitCurPassTree(glm::uvec3 coord) override
	{
		mIntermediateTree = new IntermediateTreeType(GetSinglePassTreeDepth());
		mIntermediateTree->UseLeafMap(false);
	}
	// Terminate the tree in the current pass. This means it should also be appended to the main tree and deleted
	void		FinalizeCurPassTree(glm::uvec3 coord) override
	{
		Stopwatch watch;
		if (mIntermediateTree->GetNodeCount() > 1) // Only append the tree (and compress) if it is not empty
		{
			mIntermediateTree->ToDAG();

			// Propagate the materials
			watch.Reset();
			if (verbose) printf("Propagating materials in subtree... ");
			mIntermediateTree->PropagateMaterials(T::WeightedAverage);
			if (verbose) printf("Materials propagated in %d ms.\n", (int)(watch.GetTime() * 1000));

			if (mFirstPass)
			{
				std::vector<T> uniqueMaterials = mIntermediateTree->GetUniqueMaterials();
				mSceneMaterials.insert(mSceneMaterials.end(), uniqueMaterials.begin(), uniqueMaterials.end());
				QuantizeSceneMaterials();
				mFirstPass = false;
			}

			// Replace the materials by their quantized counterparts
			if (mQuantizer != NULL)
			{
				if (verbose) printf("Replacing materials by their quantized counterparts.");
				watch.Reset();
				std::vector<T> curPassMaterials = mIntermediateTree->GetMaterials();
				std::vector<T> curPassQuantizedMaterials(curPassMaterials.size());
				auto quickQuantizer = dynamic_cast<QuickQuantizer<T>*>(mQuantizer);

				if (verbose) printf(".");
				// For each material in the curPassMaterials, find the closest material
				tbb::parallel_for(size_t(0), curPassMaterials.size(), [&](size_t i)
				{
					auto cachedReplacer = mMaterialReplacers.find(curPassMaterials[i]);
					if (cachedReplacer != mMaterialReplacers.end())
						curPassQuantizedMaterials[i] = cachedReplacer->second;
					else
					{
						// If we can use the quick quantizer, try it
						if (quickQuantizer != NULL)
						{
							auto quickQuantizedValue = quickQuantizer->Quantize(curPassMaterials[i]);
							auto quickQuantizedCachedReplacer = mMaterialReplacers.find(quickQuantizedValue);
							if (quickQuantizedCachedReplacer != mMaterialReplacers.end())
								curPassQuantizedMaterials[i] = quickQuantizedCachedReplacer->second;
							return;
						}
						curPassQuantizedMaterials[i] = NearestFinder<T>()(curPassMaterials[i], mSceneMaterials);
					}
				});
				// Update the current scene color map
				std::unordered_map<T, T> quantizedMaterials;
				for (size_t i = 0; i < curPassMaterials.size(); i++)
					quantizedMaterials.insert(std::make_pair(curPassMaterials[i], curPassQuantizedMaterials[i]));

				if (verbose) printf(".");
				mIntermediateTree->ReplaceMaterials(quantizedMaterials);
				if (verbose) printf("Replaced in %d ms.\n", (int)(watch.GetTime() * 1000));
			}

			// Create the UniqueIndexTree for this current pass
			std::string subtreeCompressionType = mTextureCompressionType;
			if (!IsSinglePass())
				subtreeCompressionType = GetSubtreeTextureCompressionType();
			auto texture = CompressedTextureFactory<MaterialLibraryPointer>::GetCompressedTexture(subtreeCompressionType);
			auto curPassTree = new FinalTreeType(GetSinglePassTreeDepth(), texture, mLevelsWithoutMaterials);

			// Finalize the current pass tree
			if (verbose) printf("Finalizing subtree...");
			watch.Reset();
			curPassTree->BaseOn(mIntermediateTree); // Note that BaseOn will delete the intermediate tree, no need to do that manually
			mIntermediateTree = NULL;
			if (verbose) printf("Finalized in %d ms.\n", (int)(watch.GetTime() * 1000));

			// Convert the subtree to a DAG first, this saved time when appending and converting the total tree
			if (verbose) printf("Converting subtree to DAG...\n");
			watch.Reset();
			curPassTree->ToDAG();
			if (verbose) printf("Converting took %d ms.\n", (int)(watch.GetTime() * 1000));

			if (IsSinglePass()) // Means we just constructed the root, so no need to append
			{
				delete mTree;
				mTree = curPassTree;
			}
			else
			{
				OctreeLoader::WriteCache(curPassTree, GetSubtreeType(), GetSinglePassTreeFilename(coord), verbose);
				delete curPassTree;
			}
		}
		else
		{
			delete mIntermediateTree;
		}
	}

	// Should add a node to the current pass tree at the given coordinate and color
	void AddNode(const glm::uvec3& coordinate, const T& color) override { mIntermediateTree->AddLeafNode(coordinate, color); }
	void AddMissingNode(const glm::uvec3& coordinate, const T& color) override { if (!mIntermediateTree->HasLeaf(coordinate)) AddNode(coordinate, color); }

	std::vector<size_t> GetOctreeNodesPerLevel() override { return mTree->GetOctreeNodesPerLevel(); }
	std::vector<size_t> GetNodesPerLevel() override { return mTree->GetNodesPerLevel(); }
private:
	std::string GetSinglePassTreeFilename(glm::uvec3 coord)
	{
		char buffer[255];
		sprintf(buffer, "%s_%u_(%u_%u_%u)", GetOutputFile().c_str(), GetTreeDepth(), coord.x, coord.y, coord.z);
		return std::string(buffer);
	}

	FinalTreeType*			mTree;
	bool					mReduceMaterials;
	bool					mFirstPass;
	std::string				mTextureCompressionType;
	unsigned32				mLevelsWithoutMaterials;

	std::vector<T>		    mSceneMaterials;
	std::unordered_map<T, T> mMaterialReplacers;
	std::vector<T>		    mCurPreprocessPassMaterials;
	std::vector<std::pair<glm::uvec3, T>> mMainTreeMaterials;

	IntermediateTreeType* mIntermediateTree;

	BaseQuantizer<T, Comparer>*		mQuantizer;
};