#pragma once
#include "BaseMaterialOctreeBuilder.h"
#include "../../scene/Octree/LeafMaterialMultiRootTree.h"
#include "../Util/Stopwatch.h"
#include "OctreeLoader.h"

template <typename T>
class MaterialMultiRootOctreeBuilder : public BaseMaterialOctreeBuilder<Color>
{
private:
	LeafMaterialMultiRootTree<T>* mTree;

	LeafMaterialMultiRootTree<T>* mCurPassTree;

	unsigned8 mBitsPerChannel;

	std::string mTreeType;

public:
	MaterialMultiRootOctreeBuilder(std::string treeType, unsigned8 bitsPerChannel) :
			BaseMaterialOctreeBuilder(),
			mTree(NULL),
			mCurPassTree(NULL),
			mBitsPerChannel(bitsPerChannel),
			mTreeType(treeType)
	{}

	MaterialMultiRootOctreeBuilder(std::string treeType) : MaterialMultiRootOctreeBuilder(treeType, 8) {}
	~MaterialMultiRootOctreeBuilder() override {}

	std::string GetTreeType() override { bool explicitSpecializationGiven = false; assert(explicitSpecializationGiven); return ""; }
protected:
	// Initialize the main tree
	void InitTree() override
	{
		mTree = (LeafMaterialMultiRootTree<T>*)OctreeLoader::CreateTree(mTreeType, GetTreeDepth());
	}

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

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

	// Initialize the tree for the current pass.
	void InitCurPassTree(glm::uvec3 coord) override
	{
		mCurPassTree = (LeafMaterialMultiRootTree<T>*)OctreeLoader::CreateTree(mTreeType, GetSinglePassTreeDepth());
	}

	// 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 (mCurPassTree->GetNodeCount() > 1) // Only append the tree (and compress) if it is not empty
		{
			if (IsSinglePass()) // Means we just constructed the root, so no need to append
			{
				delete mTree;
				mTree = mCurPassTree;
			}
			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();
				mCurPassTree->ToDAG();
				if (verbose) printf("Converting took %d ms.\n", (int)(watch.GetTime() * 1000));

				if (verbose) printf("Appending subtree... ");
				mTree->Append(coord, GetAppendedTreeLevel(), mCurPassTree);
				delete mCurPassTree;

				if (verbose) printf("Converting current tree to DAG...\n");
				mTree->ToDAG(GetAppendedTreeLevel());
			}
		}
		else
		{
			delete mCurPassTree;
		}
	}

	void AddMissingNode(const glm::uvec3& coordinate, const Color& color) override
	{
		if (!mCurPassTree->HasLeaf(coordinate))
			AddNode(coordinate, color);
	}

	// Should add a node to the current pass tree at the given coordinate and color
	void AddNode(const glm::uvec3& coordinate, const Color& color) override
	{
		((Root*)mCurPassTree)->AddLeafNode(coordinate);
	}

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