#include "MaterialLibraryMultiRootOctreeBuilder.h"
#include "../Util/Stopwatch.h"
#include "OctreeLoader.h"
#include "../CollectionHelper.h"
#include "../../inc/tbb/parallel_for.h"
#include "../../inc/tbb/concurrent_queue.h"

MaterialLibraryMultiRootOctreeBuilder::MaterialLibraryMultiRootOctreeBuilder(BaseQuantizer<Color, ColorCompare>* quantizer) :
		BaseMaterialOctreeBuilder(),
		mTree(NULL),
		mIntermediateTree(NULL),
		mQuantizeColors(quantizer != NULL),
		mQuantizer(quantizer),
		mSceneColors(std::vector<Color>()),
		mColorReplacers(std::unordered_map<Color, Color>())
{
	// TODO: when building in steps, use preprocessing to find the materials in the upper levels
}
MaterialLibraryMultiRootOctreeBuilder::~MaterialLibraryMultiRootOctreeBuilder() {}

bool MaterialLibraryMultiRootOctreeBuilder::IsLimitedColors() const { return mQuantizeColors || !IsSinglePass(); }

bool MaterialLibraryMultiRootOctreeBuilder::UsePreprocessing() const { return IsLimitedColors(); }

std::string MaterialLibraryMultiRootOctreeBuilder::GetTreeType() { return "blc" + (mQuantizer == NULL ? "" : mQuantizer->GetQuantizerDescriptor()); }

void MaterialLibraryMultiRootOctreeBuilder::InitTree()
{
	mTree = new MaterialLibraryMultiRoot<Color, ColorCompare>(GetTreeDepth());
}

void MaterialLibraryMultiRootOctreeBuilder::FinalizeTree()
{
	// Store the amount of nodes the octree had for the first few levels
	auto nodesPerLevel = mTree->GetOctreeNodesPerLevel();
}

void MaterialLibraryMultiRootOctreeBuilder::TerminateTree()
{
	OctreeLoader::WriteCache(mTree, GetTreeType(), GetOutputFile(), verbose);
	delete mTree;
	mSceneColors.clear();
	mColorReplacers.clear();
}

void MaterialLibraryMultiRootOctreeBuilder::InitCurPassTree(glm::uvec3 coord)
{
	mIntermediateTree = new MaterialTree<Color, ColorCompare>(GetSinglePassTreeDepth());
	mIntermediateTree->UseLeafMap(true);

	//for (auto color : mSceneColors)
	//	mIntermediateTree->AddMaterial(Color(color));
	//mIntermediateTree->FinalizeMaterials();
}

void MaterialLibraryMultiRootOctreeBuilder::FinalizeCurPassTree(glm::uvec3 coord)
{
	Stopwatch watch;

	if (mIntermediateTree->GetNodeCount() > 1) // Only append the tree (and compress) if it is not empty
	{
		// Convert the intermediate tree to a DAG before building the material library multiroot tree:
		if (verbose) printf("Converting intermediate tree to a DAG...\n");
		watch.Reset();
		mIntermediateTree->ToDAG();
		if (verbose) printf("Converting took %d ms.\n", (int)(watch.GetTime() * 1000));

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

		// Quantize the colors in the material tree
		if (IsLimitedColors())
		{
			if (verbose) printf("Replacing materials by their quantized counterparts.");
			std::vector<Color> curPassMaterials = mIntermediateTree->GetUniqueMaterials();
			CollectionHelper::Unique(curPassMaterials, ColorCompare());
			std::vector<Color> curPassQuantizedMaterials(curPassMaterials.size());

			if (verbose) printf(".");
			// For each material in the curPassMaterials, find the closest material
			tbb::concurrent_queue<size_t> newReplacers;
			tbb::parallel_for(size_t(0), curPassMaterials.size(), [&](const size_t& i)
			{
				auto cachedReplacer = mColorReplacers.find(curPassMaterials[i]);
				if (cachedReplacer != mColorReplacers.end())
					curPassQuantizedMaterials[i] = cachedReplacer->second;
				else
				{
					curPassQuantizedMaterials[i] = NearestFinder<Color>()(curPassMaterials[i], mSceneColors);
					newReplacers.push(i);
				}
			});

			// Add the newly found replacers to the cache
			for (auto i = newReplacers.unsafe_begin(); i != newReplacers.unsafe_end(); ++i)
				mColorReplacers.insert(std::pair<Color, Color>(curPassMaterials[*i], curPassQuantizedMaterials[*i]));

			// Update the current scene color map
			std::unordered_map<Color, Color> quantizedColors;
			for (size_t i = 0; i < curPassMaterials.size(); i++)
				quantizedColors.insert(std::pair<Color, Color>(curPassMaterials[i], curPassQuantizedMaterials[i]));

			if (verbose) printf(".");
			mIntermediateTree->ReplaceMaterials(quantizedColors);
			if (verbose) printf("Replaced in %d ms.\n", (int)(watch.GetTime() * 1000));
			if (verbose) printf("%llu new materials quantized, cache now contains %llu materials.\n", (unsigned64)newReplacers.unsafe_size(), (unsigned64)mColorReplacers.size());
		}

		// Build the actual MaterialLibaryMultiRootTree based on the intermediate tree
		if (verbose) printf("Building multiroot tree based on intermediate tree...");
		watch.Reset();
		auto curPassTree = new MaterialLibraryMultiRoot<Color, ColorCompare>(mIntermediateTree->GetMaxLevel());
		if (IsLimitedColors()) curPassTree->CreateMaterialLibrary(mSceneColors);
		curPassTree->BaseOn(mIntermediateTree, false);
		delete mIntermediateTree;

		curPassTree->ShaveEquals();
		curPassTree->ToDAG();
		curPassTree->FillEmptySpace(false);
		if (verbose) printf("Multiroot tree build in %d ms.\n", (int)(watch.GetTime() * 1000));

		// Convert the current pass tree to a DAG
		if (verbose) printf("Converting the current pass multiroot tree to a DAG...\n");
		watch.Reset();
		curPassTree->ToDAG();
		if (verbose) printf("Converted in %d ms.\n", (int)(watch.GetTime() * 1000));

		auto octreeNodesPerLevel = curPassTree->GetOctreeNodesPerLevel();

		if (IsSinglePass()) // Means we just constructed the root, so no need to append
		{
			delete mTree;
			mTree = curPassTree;
		}
		else
		{
			if (verbose) printf("Appending subtree... ");
			watch.Reset();
			mTree->Append(coord, GetAppendedTreeLevel(), curPassTree);
			delete curPassTree;
			if (verbose) printf("Appending took %d ms.\n", (int)(watch.GetTime() * 1000));

			if (verbose) printf("Converting current tree to DAG...\n");
			watch.Reset();
			mTree->ToDAG(GetAppendedTreeLevel());
			if (verbose) printf("Converted in %d ms.\n", (int)(watch.GetTime() * 1000));
		}
	}
	else
	{
		delete mIntermediateTree;
	}
}

void MaterialLibraryMultiRootOctreeBuilder::AddMissingNode(const glm::uvec3& coordinate, const Color& color)
{
	if (!mIntermediateTree->HasLeaf(coordinate))
		AddNode(coordinate, color);
}
void MaterialLibraryMultiRootOctreeBuilder::AddNode(const glm::uvec3& coordinate, const Color& color)
{
	if (IsLimitedColors())
	{
		auto replacer = mColorReplacers.find(color);
		if (replacer == mColorReplacers.end())
		{
			Color replacement = NearestFinder<Color>()(color, mSceneColors);
			mColorReplacers.insert(std::pair<Color, Color>(color, replacement));
			replacer = mColorReplacers.find(color);
		}
		assert(replacer != mColorReplacers.end());
		mIntermediateTree->AddLeafNode(coordinate, replacer->second);
	}
	else
	{
		mIntermediateTree->AddLeafNode(coordinate, Color(color));
	}
}

//-----------------------------------------------------------------------------------
// ---------------------------------PREPROCESSING------------------------------------
//-----------------------------------------------------------------------------------
void MaterialLibraryMultiRootOctreeBuilder::PreProcessNode(const glm::uvec3& coordinate, const Color& color)
{
	mSceneColors.push_back(color);
}

void MaterialLibraryMultiRootOctreeBuilder::FinalizePreprocessing()
{
	CalculateUniqueSceneColors();
	if (mQuantizer != NULL)
	{
		if (verbose) printf("Quantizing %llu materials...", (unsigned64)mSceneColors.size());
		Stopwatch watch; watch.Reset();
		auto quantizedMaterials = mQuantizer->QuantizeMaterials(mSceneColors);
		mSceneColors.clear();
		for (auto quantizedMaterial : *quantizedMaterials)
		{
			mColorReplacers.insert(std::pair<Color, Color>(quantizedMaterial.first, quantizedMaterial.second));
			mSceneColors.push_back(quantizedMaterial.second);
		}
		CalculateUniqueSceneColors();
		if (verbose) printf("Quantized in %u ms, %llu materials left\n", (unsigned)(watch.GetTime() * 1000), (unsigned64)mSceneColors.size());
	}
	else
	{
		for (Color color : mSceneColors)
			mColorReplacers.insert(std::pair<Color, Color>(color, color));
	}
}

void MaterialLibraryMultiRootOctreeBuilder::CalculateUniqueSceneColors()
{
	// Remove duplicates from scene colors
	CollectionHelper::Unique(mSceneColors, ColorCompare());
}

//-----------------------------------------------------------------------------------
// -----------------------------------STATISTICS-------------------------------------
//-----------------------------------------------------------------------------------
std::vector<size_t> MaterialLibraryMultiRootOctreeBuilder::GetOctreeNodesPerLevel()
{
	return mTree->GetOctreeNodesPerLevel();
}

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