#pragma once
#include "BaseMaterialOctreeBuilder.h"
#include "../../inc/glm/common.hpp"
#include <unordered_set>
#include "../../core/Hashers.h"
#include "../../scene/Octree/MaterialLibraryTree.h"
#include "../../scene/Material/MaterialQuantizer/BaseQuantizer.h"

template<typename T, typename Comparer>
class HierarchicalMaterialOctreeBuilder : public BaseMaterialOctreeBuilder<T>
{
public:
	HierarchicalMaterialOctreeBuilder(BaseQuantizer<T, Comparer>* quantizer = NULL) :
		mTree(NULL),
		mMaterialLibrary(NULL),
		mReduceMaterials(quantizer != NULL),
		mSceneMaterials(std::vector<T>()),
		mReplacers(NULL),
		mQuantizer(quantizer)
	{}

	~HierarchicalMaterialOctreeBuilder() override
	{
		if (mReplacers != NULL)
			delete mReplacers;
		if (mTree != NULL)
			delete mTree;
		if (mMaterialLibrary != NULL)
			delete mMaterialLibrary;
	}

	std::string GetTreeType() {
		return "h" + MaterialAbbreviation<T>()() +
			(mReduceMaterials ? mQuantizer->GetQuantizerDescriptor() : "");
	}
protected:
	bool MaterialHasDefaults() const { return !DefaultMaterials<T>()().empty(); }
	MaterialLibraryTree<T, Comparer, T::CHANNELSPERPIXEL>* CreateTree(unsigned8 depth) const
	{
		return new MaterialLibraryTree<T, Comparer, T::CHANNELSPERPIXEL>(depth);
	}

	MaterialLibraryTree<T, Comparer, T::CHANNELSPERPIXEL>* CreateTree(unsigned8 depth, MaterialLibrary<T, Comparer, T::CHANNELSPERPIXEL>* lib) const
	{
		return new MaterialLibraryTree<T, Comparer, T::CHANNELSPERPIXEL>(depth, lib);
	}

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

	bool		UsePreprocessing() const override { return !MaterialHasDefaults(); }
	void		PreProcessNode(const glm::uvec3& coordinate, const T& mat) override { mSceneMaterials.push_back(mat); }
	void		PreProcessMissingNode(const glm::uvec3& coordinate, const T& mat) override { mSceneMaterials.push_back(mat); }
	void		FinalizeCurPreprocessPass(glm::uvec3 coord) override { CalculateUniqueSceneMaterials(); }

	void InitTree() override
	{
		// Get scene materials if this is a material type that doesn't need preprocessing
		if (mSceneMaterials.empty())
			mSceneMaterials = DefaultMaterials<T>()();

		if (mReduceMaterials)
		{
			// Quantize the scene materials
			CalculateQuantizedMaterials();
			// Store the quantized values in the mSceneMaterials
			mSceneMaterials = std::vector<T>(); 
			for (auto it = mReplacers->begin(); it != mReplacers->end(); it++)
				mSceneMaterials.push_back(it->second);
			CalculateUniqueSceneMaterials();
		}

		// Initialize the material library
		Stopwatch watch; watch.Reset();
		if (verbose) printf("Initializing the material library...");
		mMaterialLibrary = new MaterialLibrary<T, Comparer, T::CHANNELSPERPIXEL>();
		for (auto mat : mSceneMaterials)
			mMaterialLibrary->AddMaterial(mat);
		mMaterialLibrary->Finalize();
		if (verbose) printf("Material library initialized in %u ms.\n", (int)(watch.GetTime() * 1000));
	}

	void FinalizeTree() override
	{
		Stopwatch watch;
		if (!IsSinglePass())
		{
			// If this tree was constructred in steps, we should have cache files by now to build the main tree from
			unsigned8 mainTreeDepth = GetTreeDepth();
			unsigned8 subTreeDepth = GetSinglePassTreeDepth();
			unsigned8 mainTreeLevel = mainTreeDepth - subTreeDepth;
			mTree = CreateTree(GetTreeDepth(), mMaterialLibrary);

			unsigned32 i = 1;
			for (const glm::uvec3& coord : GetValidCoords())
			{
				if (verbose) printf("Reading subtree at %u / %u (%u, %u, %u) from cache...\n", i, (unsigned32)GetValidCoords().size(), coord.x, coord.y, coord.z);
				MaterialLibraryTree<T, Comparer, T::CHANNELSPERPIXEL>* subTree = (MaterialLibraryTree<T, Comparer, T::CHANNELSPERPIXEL>*)OctreeLoader::ReadCache(GetTreeType(), GetSinglePassTreeDepth(), GetSinglePassTreeFilename(coord), verbose);
				if (verbose) printf("Appending subtree...");
				watch.Reset();
				mTree->AppendAndMerge(coord, mainTreeLevel, subTree);
				delete subTree;
				if (verbose) printf("Subtree appended in %d ms.\n", (int)(watch.GetTime() * 1000));
				i++;
			}

			mTree->ClearOrphans();
		}

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

		// 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())
		{
			for (const glm::uvec3& coord : GetValidCoords())
				OctreeLoader::DeleteCache(GetTreeType(), GetSinglePassTreeDepth(), GetSinglePassTreeFilename(coord));
		}
	}

	// 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();
		if (mReplacers != NULL) delete mReplacers;
		mReplacers = NULL;
	}

	bool CancelCurPassTree(const glm::uvec3& coord) override
	{
		return OctreeLoader::VerifyCache(GetTreeType(), GetSinglePassTreeDepth(), GetSinglePassTreeFilename(coord));
	}

	// Initialize the tree for the current pass.
	void InitCurPassTree(glm::uvec3 coord) override
	{
		mTree = CreateTree(GetSinglePassTreeDepth(), mMaterialLibrary);
	}
	// 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;

		// Remove unused leafs and stuff
		mTree->ClearOrphans();

		// Convert the subtree to a DAG
		if (verbose) printf("Converting subtree to DAG...\n");
		watch.Reset();
		mTree->ToDAG();
		if (verbose) printf("Converting took %u ms.\n", (unsigned)(watch.GetTime() * 1000.0));

		if (!IsSinglePass())
		{
			// Write the subtree to a cache file
			OctreeLoader::WriteCache(mTree, GetTreeType(), GetSinglePassTreeFilename(coord), verbose);

			delete mTree;
			mTree = NULL;
		}
	}

	// Should add a node to the current pass tree at the given coordinate and material
	void AddNode(const glm::uvec3& coordinate, const T& mat) override
	{
		auto replacer = mat;
		if (mReduceMaterials)
		{
			auto replacerPair = mReplacers->find(mat);
			if (replacerPair != mReplacers->end())
				replacer = replacerPair->second;
			else
			{
				replacer = ParallelNearestFinder<T>()(mat, mSceneMaterials);
				mReplacers->insert(std::make_pair(mat, replacer));
			}
		}
		mTree->AddLeafNode(coordinate, replacer);
	}
	void AddMissingNode(const glm::uvec3& coordinate, const T& mat) override
	{
		if (!mTree->HasLeaf(coordinate))
			AddNode(coordinate, mat);
	}

	std::vector<size_t> GetOctreeNodesPerLevel() override { return mTree->GetOctreeNodesPerLevel(); }
	std::vector<size_t> GetNodesPerLevel() override { return mTree->GetNodesPerLevel(); }

	void CalculateUniqueSceneMaterials() { CollectionHelper::Unique<T, Comparer>(mSceneMaterials); }
	void CalculateQuantizedMaterials()
	{
		if (verbose) printf("Quantizing/merging %llu materials...", (unsigned64)(mSceneMaterials.size()));
		Stopwatch watch; watch.Reset();

		DebugQuantize();
		mReplacers = mQuantizer->QuantizeMaterials(mSceneMaterials);

		if (verbose) printf("Quantized materials in %d ms.\n", (int)(watch.GetTime() * 1000));
		watch.Reset();
	}

	void DebugQuantize() {}

	MaterialLibraryTree<T, Comparer, T::CHANNELSPERPIXEL>*	mTree;

	bool									mReduceMaterials;
	std::vector<T>							mSceneMaterials;
	MaterialLibrary<T, Comparer, T::CHANNELSPERPIXEL>* mMaterialLibrary;
	std::map<T, T, Comparer>*				mReplacers;
	BaseQuantizer<T, Comparer>*				mQuantizer;
};

//template<> void HierarchicalMaterialOctreeBuilder<Color, ColorCompare>::DebugQuantize()
//{
//	std::vector<std::string> types = { "lab256", "lab1024", "lab4096", "lab16384" };
//	for (std::string type : types)
//	{
//		printf("%s: ", type.c_str());
//		BaseColorQuantizer* quantizer = ColorQuantizerFactory::Create(type);
//		auto temp = quantizer->QuantizeMaterials(mSceneMaterials);
//		delete temp;
//	}
//}

template<> bool HierarchicalMaterialOctreeBuilder<SmallNormal, NormalCompare>::MaterialHasDefaults() const { return false; }
template<> void HierarchicalMaterialOctreeBuilder<ColorAndNormal, ColorAndNormalCompare>::PreProcessMissingNode(const glm::uvec3& coordinate, const ColorAndNormal& mat)
{
	mSceneMaterials.push_back(mat);
}