#include "RandomOctreeBuilder.h"
#include "../Util/Stopwatch.h"
#include "OctreeLoader.h"
#include "../../inc/tbb/parallel_for.h"

RandomOctreeBuilder::RandomOctreeBuilder() :
		BaseMaterialOctreeBuilder(),
		mTree(NULL),
		mCurPassTree(NULL),
		mCurPassTreeCoord(glm::uvec3(0))
{}
RandomOctreeBuilder::~RandomOctreeBuilder() {}

std::string RandomOctreeBuilder::GetTreeType() { return "r"; }

void RandomOctreeBuilder::InitTree()
{
	mTree = new MaterialTree<BitsMaterial<1>, HashComparer<BitsMaterial<1>>>(GetTreeDepth());
}

void RandomOctreeBuilder::FinalizeTree()
{
	// Convert the octree to a DAG
	Stopwatch watch;
	if (verbose) printf("Converting final tree to DAG...\n");
	watch.Reset();
	mTree->ToDAG();
	if (verbose) printf("Done in %u ms, %llu nodes left.\n", (unsigned)(watch.GetTime() * 1000), (unsigned64)mTree->GetNodeCount());
}

void RandomOctreeBuilder::TerminateTree()
{
	OctreeLoader::WriteCache(mTree, GetTreeType(), GetOutputFile(), verbose);
	delete mTree;
}

void RandomOctreeBuilder::InitCurPassTree(glm::uvec3 coord)
{
	mCurPassTree = new MaterialTree<Color, ColorCompare>(GetSinglePassTreeDepth());
	mCurPassTreeCoord = coord;
}

void RandomOctreeBuilder::FinalizeCurPassTree(glm::uvec3 coord)
{
	Stopwatch watch;
	unsigned8 mainTreeLevel = mTree->GetMaxLevel() - mCurPassTree->GetMaxLevel();
	if (mCurPassTree->GetNodeCount() > 1) // Only append the tree (and compress) if it is not empty
	{
		//std::vector<BitsMaterial<1>> randomMaterials(mCurPassTree->GetNodeCount());
		//tbb::parallel_for(size_t(0), randomMaterials.size(), [&](size_t i) { randomMaterials[i] = BitsMaterial<1>((std::rand() < (RAND_MAX / 2)) ? 1 : 0); });
		//mCurPassTree->SetMaterials(randomMaterials);

		mCurPassTree->PropagateMaterials([](const std::vector<Color>& materials, const std::vector<float>& weights){ return Color::WeightedAverage(materials, weights); });
		std::vector<Color> colors = mCurPassTree->GetMaterials();
		std::vector<BitsMaterial<1>> bitMaterials(mCurPassTree->GetNodeCount());
		unsigned8 shift = 0;
		tbb::parallel_for(size_t(0), bitMaterials.size(), [&](size_t i) { bitMaterials[i] = BitsMaterial<1>((colors[i].GetColor().r & (1 << shift)) >> shift); });
		auto mCurPassBitTree = new MaterialTree<BitsMaterial<1>, HashComparer<BitsMaterial<1>>>(mCurPassTree->GetMaxLevel());
		mCurPassBitTree->MoveShallow(mCurPassTree);
		mCurPassBitTree->SetMaterials(bitMaterials);

		delete mCurPassTree;

		if (mainTreeLevel == 0) // Means we just constructed the root, so no need to append
		{
			delete mTree;
			mTree = mCurPassBitTree;
		}
		else
		{
			 //Convert the subtree to a DAG first, this saved time when appending and converting the total tree
			 //Benchmark (Total build time for subtrees of depth 10, final tree of depth 12, Release mode with pool):
			 //    209.922 seconds without early converting
			 //    163.645 seconds with early converting
			if (verbose) printf("Converting subtree to DAG...\n");
			watch.Reset();
			mCurPassBitTree->ToDAG();
			if (verbose) printf("Converting took %d ms.\n", (int)(watch.GetTime() * 1000));
			
			if (verbose) printf("Appending subtree... ");
			mTree->Append(mCurPassTreeCoord, mainTreeLevel, mCurPassBitTree);
			delete mCurPassBitTree;

			if (verbose) printf("Converting current tree to DAG...\n");
			mTree->ToDAG(mainTreeLevel);
		}

	}
	else
	{
		delete mCurPassTree;
	}
}

void RandomOctreeBuilder::AddNode(const glm::uvec3& coordinate, const Color& color)
{
	mCurPassTree->AddLeafNode(coordinate, color);
}

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

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