#include "OctreeConverter.h"
#include "OctreeLoader.h"
#include "SettingsParser.h"
#include "TreeTypeParser.h"
#include "ColorQuantizerFactory.h"
#include "CompressedTextureFactory.h"
#include "../Util/Stopwatch.h"
#include "../../scene/Material/MaterialQuantizer/ColorQuantizer/BaseColorQuantizer.h"
#include "../../PropertyLoader.h"
#include <regex>

#include "../../scene/Octree/BaseTree.h"
#include "../../scene/Octree/Tree.h"
#include "../../scene/Octree/MaterialLibraryTree.h"
#include "../../scene/Octree/HierarchicalShiftingColoredTree.h"
#include "../../scene/Octree/HierarchicalColorsOnlyTree.h"
#include "../../scene/Octree/MaterialLibraryUniqueIndexTree.h"

#ifdef _WIN32
#include <Windows.h>
#endif

bool OctreeConverter::Convert()
{
	PropertyLoader::Create();
	auto propertyLoader = PropertyLoader::Instance();
	if (!propertyLoader->GetBoolProperty("octreebuilder_tryconvert") || !propertyLoader->GetBoolProperty("octreebuilder_usecache"))
		return false;

	bool verbose = propertyLoader->GetProperty("octreebuilder_verbose") != "0";
	if (verbose) printf("Trying to convert an existing tree to the wanted treetype.\n");

	std::string octreeType = propertyLoader->GetProperty("octree_type");
	unsigned8 maxLevel = propertyLoader->GetIntProperty("shader_max_level");
	std::string filename = propertyLoader->GetProperty("dag_file");

	bool res = ConvertTo(octreeType, maxLevel, filename, verbose);
	if (verbose)
	{
		if (res) printf("Conversion succeeded\n");
		if (!res) printf("No valid conversion found\n");
	}
	return res;
}

bool OctreeConverter::ConvertTo(std::string destinationType, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
#ifdef _WIN32
	// Go through all .oct files that can support the given dag file. 
	// For each of the files found, try conversion.
	std::string dagStrippedFilename = dagFileName;
	const size_t last_slash_idx = dagStrippedFilename.find_last_of("\\/");
	if (std::string::npos != last_slash_idx)
	{
		dagStrippedFilename.erase(0, last_slash_idx + 1);
	}

	std::string validFileWildcard = std::string(dagFileName + "*.oct");
	std::regex fileNameMatcher(dagStrippedFilename + "_([0-9]+)_(.+).oct");
	std::smatch sm;

	WIN32_FIND_DATA FindFileData;
	HANDLE hFind = FindFirstFile(validFileWildcard.c_str(), &FindFileData);
	LARGE_INTEGER filesize;
	// Find all convertable files
	std::vector<FileInfo> convertableFiles;
	while (hFind != INVALID_HANDLE_VALUE)
	{
		// Try conversion
		std::string foundFilename(FindFileData.cFileName);
		filesize.LowPart = FindFileData.nFileSizeLow;
		filesize.HighPart = FindFileData.nFileSizeHigh;
		// Extract the level and type from the filename
		if (std::regex_match(foundFilename, sm, fileNameMatcher))
		{
			unsigned8 sourceDepth = std::stoi(sm.str(1));
			std::string sourceType = sm.str(2);
			if (CanConvert(sourceType, destinationType, sourceDepth, destinationDepth))
				convertableFiles.push_back(FileInfo(foundFilename, sourceType, sourceDepth, filesize.QuadPart));
		}

		if (!FindNextFile(hFind, &FindFileData))
		{
			FindClose(hFind);
			hFind = INVALID_HANDLE_VALUE;
		}
	}
	// Sort them on filesize, assuming that smaller files are always quicker to convert
	std::sort(convertableFiles.begin(), convertableFiles.end(), [&](const FileInfo& a, const FileInfo& b)
	{
		return a.filesize < b.filesize;
	});
	for (const FileInfo& info : convertableFiles)
	{
		if (Convert(info.treeType, destinationType, info.depth, destinationDepth, dagFileName, verbose))
			return true;
	}
#endif
	return false;
}

// Check if a certain conversion is valid (theore
bool OctreeConverter::CanConvert(std::string sourceType, std::string destinationType, unsigned8 sourceDepth, unsigned8 destinationDepth)
{
	// Converting to the same type always works
	if (sourceType == destinationType && sourceDepth == destinationDepth)
		return true;

	if (sourceDepth != destinationDepth)
		return false; // Can't change the depth...

	TreeTypeDescriptor destDescr = TreeTypeParser::GetTreeTypeDescriptor(destinationType);
	TreeTypeDescriptor sourceDescr = TreeTypeParser::GetTreeTypeDescriptor(sourceType);
	std::string normalizedSourceQuantizerType = ColorQuantizerFactory::GetNormalizedType(sourceDescr.quantizer);
	std::string normalizedDestinationQuantizerType = ColorQuantizerFactory::GetNormalizedType(destDescr.quantizer);

	// Any tree type can be converted to standard if the source depth is at least as deep as the destination depth
	if (destDescr.globalType == STANDARD) return true;
	else if (destDescr.globalType == HIERARCHICAL)
	{
		return sourceDepth == destinationDepth
			&& sourceDescr.globalType == HIERARCHICAL 
			&& sourceDescr.additionalTypeInfo == "" && (destDescr.additionalTypeInfo == "" || destDescr.additionalTypeInfo == "s")
			&& destDescr.material == sourceDescr.material
			&& (normalizedSourceQuantizerType == "" || normalizedSourceQuantizerType == normalizedDestinationQuantizerType);
	}
	else if (destDescr.globalType == UNIQUEINDEX)
	{
		return sourceDepth == destinationDepth
			&& sourceDescr.globalType == UNIQUEINDEX && sourceDescr.material == "c"
			&&  destDescr.material == "c"
			&& (normalizedSourceQuantizerType == "" || normalizedSourceQuantizerType == normalizedDestinationQuantizerType)
			&& sourceDescr.additionalTypeInfo == "" && destDescr.additionalTypeInfo == "";
	}
	else if (destDescr.globalType == ONLYMATERIAL)
	{
		return sourceDepth == destinationDepth
			&& destDescr.material == "c" && sourceDescr.material == "c"
			&& sourceDescr.globalType == HIERARCHICAL
			&& (normalizedSourceQuantizerType == "" || normalizedSourceQuantizerType == normalizedDestinationQuantizerType);
	}
	return false;
}

bool OctreeConverter::Convert(std::string sourceType, std::string destinationType, unsigned8 sourceDepth, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
	if (!CanConvert(sourceType, destinationType, sourceDepth, destinationDepth)) return false;
	TreeTypeDescriptor destDescr = TreeTypeParser::GetTreeTypeDescriptor(destinationType);
	switch (destDescr.globalType)
	{
	case HIERARCHICAL:
		if (destDescr.material == "c")
		{
			if (destDescr.additionalTypeInfo == "s")
				return ConvertToHierarchicalColorShift(sourceType, destinationType, sourceDepth, destinationDepth, dagFileName, verbose);
			else
				return ConvertToHierarchicalColored(sourceType, destinationType, sourceDepth, destinationDepth, dagFileName, verbose);
		}
		break;
	case UNIQUEINDEX: return ConvertToUniqueIndex(sourceType, destinationType, sourceDepth, destinationDepth, dagFileName, verbose);
	case STANDARD: return ConvertToStandard(sourceType, sourceDepth, destinationDepth, dagFileName, verbose);
	case ONLYMATERIAL: return ConvertToOnlyMaterials(sourceType, destinationType, sourceDepth, destinationDepth, dagFileName, verbose);
	}
	return false;
}

bool OctreeConverter::ConvertToStandard(std::string sourceType, unsigned8 sourceDepth, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
	if (verbose) printf("Found valid octree file: %s, converting...\n", OctreeLoader::GetFullFilename(sourceType, sourceDepth, dagFileName).c_str());
	BaseTree* originalTree = OctreeLoader::ReadCache(sourceType, sourceDepth, dagFileName);
	if (originalTree == NULL) return false;
	Tree<>* standardTree = (Tree<>*)OctreeLoader::CreateTree("s", destinationDepth);
	originalTree->Shave(destinationDepth);
	standardTree->MoveShallow(originalTree);
	delete originalTree;
	standardTree->ToDAG();
	OctreeLoader::WriteCache(standardTree, "s", dagFileName, verbose);
	delete standardTree;
	return true;
}

static void CompressHierarchicalTreeColors(MaterialLibraryTree<Color, ColorCompare>* tree, std::string compressor)
{
	std::vector<Color> leafMaterials = tree->GetLeafMaterials();
	tbb::parallel_sort(leafMaterials.begin(), leafMaterials.end(), ColorCompare());
	leafMaterials.erase(std::unique(leafMaterials.begin(), leafMaterials.end()), leafMaterials.end());
	std::vector<glm::u8vec3> leafColors(leafMaterials.size());
	tbb::parallel_for(size_t(0), leafMaterials.size(), [&](size_t i)
	{
		leafColors[i] = leafMaterials[i].GetColor();
	});
	tbb::parallel_sort(leafColors.begin(), leafColors.end(), u8vec3comparer());
	leafColors.erase(std::unique(leafColors.begin(), leafColors.end()), leafColors.end());

	auto colorQuantizer = ColorQuantizerFactory::Create(compressor);
	auto colorReplacers = colorQuantizer->QuantizeColors(leafColors);
	std::map<Color, Color, ColorCompare> materialReplacers;
	for (auto colorReplacer = colorReplacers->begin(); colorReplacer != colorReplacers->end(); std::advance(colorReplacer, 1))
		materialReplacers.insert(std::make_pair(Color(colorReplacer->first), Color(colorReplacer->second)));
	delete colorReplacers;
	delete colorQuantizer;
	tree->ReplaceLeafMaterials(materialReplacers);
}

bool OctreeConverter::ConvertToHierarchicalColored(std::string sourceType, std::string destinationType, unsigned8 sourceDepth, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
	// Only works if trees are of the same depth. Higher depth will mean the colors are averaged which might lead to wrong results
	if (sourceDepth != destinationDepth)
		return false;

	TreeTypeDescriptor sourceDescr = TreeTypeParser::GetTreeTypeDescriptor(sourceType);
	TreeTypeDescriptor destDescr = TreeTypeParser::GetTreeTypeDescriptor(destinationType);

	// Make sure that the destination type is actually what is expected
	if (destDescr.globalType != HIERARCHICAL || destDescr.material != "c" || destDescr.additionalTypeInfo != "")
		return false;
	// Only colored and hierarchical trees can be converted, so they need to match this regex
	if (sourceDescr.globalType != HIERARCHICAL || sourceDescr.additionalTypeInfo != "" || sourceDescr.material != "c")
		return false;

	std::string normalizedSourceQuantizerType = ColorQuantizerFactory::GetNormalizedType(sourceDescr.quantizer);
	std::string normalizedDestinationQuantizerType = ColorQuantizerFactory::GetNormalizedType(destDescr.quantizer);

	// Only trees with full colors or the same compression scheme as the destination type can be compressed
	if (normalizedSourceQuantizerType != "" && normalizedSourceQuantizerType != normalizedDestinationQuantizerType)
		return false;

	// Now there are 4 options:
	// Source is hierarchical		 fully colored
	// Source is not-hierarchical	 fully colored
	// Source is hierarchical        compressed colored
	// Source is not-hierarchical    compressed colored
	if (sourceDescr.globalType == HIERARCHICAL)
	{
		// Apparently we already have the desired type, so just return true
		if (normalizedDestinationQuantizerType == normalizedSourceQuantizerType)
			return true;

		// Convertert fully colored hierarchical root to compressed colored hierarchical root
		if (verbose) printf("Converting fully colored tree in %s to compressed colored tree...\n", OctreeLoader::GetFullFilename(sourceType, sourceDepth, dagFileName).c_str());
		MaterialLibraryTree<Color, ColorCompare>* source = (MaterialLibraryTree<Color, ColorCompare>*)OctreeLoader::ReadCache(sourceType, sourceDepth, dagFileName);
		if (source == NULL)
		{
			if (verbose) printf("Reading the source tree failed...\n");
			return false;
		}
		if (verbose) printf("Quantizing original materials...\n");
		CompressHierarchicalTreeColors(source, normalizedDestinationQuantizerType);
		if (verbose) printf("Clearing propagation...\n");
		source->ClearPropagation();

		if (verbose) printf("Recalculating DAG (since the materials are now compressed).\n");
		source->ToDAG();
		if (verbose) printf("Propagating materials...\n");
		source->PropagateMaterials();

		OctreeLoader::WriteCache(source, destinationType, dagFileName, verbose);
		delete source;
		return true;
	}
	return false;
}

bool OctreeConverter::ConvertToUniqueIndex(std::string sourceType, std::string destinationType, unsigned8 sourceDepth, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
	// Conversion currently only allowed from other UniqueIndexRoots
	if (sourceDepth != destinationDepth) return false;

	TreeTypeDescriptor sourceDescr = TreeTypeParser::GetTreeTypeDescriptor(sourceType);
	TreeTypeDescriptor destDescr = TreeTypeParser::GetTreeTypeDescriptor(destinationType);
	std::string normalizedSourceQuantizerType = ColorQuantizerFactory::GetNormalizedType(sourceDescr.quantizer);
	std::string normalizedDestinationQuantizerType = ColorQuantizerFactory::GetNormalizedType(destDescr.quantizer);

	// Check the source and destination type to be valid
	if (destDescr.globalType != UNIQUEINDEX || destDescr.material != "c")
		return false;
	if ((/*sourceDescr.globalType != HIERARCHICAL && sourceDescr.globalType != COLORED && */sourceDescr.globalType != UNIQUEINDEX) || sourceDescr.material != "c")
		return false;
	// Only trees with full colors or the same compression scheme as the destination type can be compressed
	if (normalizedSourceQuantizerType != "" && normalizedSourceQuantizerType != normalizedDestinationQuantizerType)
		return false;
	if (sourceDescr.additionalTypeInfo != "" || destDescr.additionalTypeInfo != "")
		return false;

	if (sourceDescr.globalType == UNIQUEINDEX)
	{
		// Check if the current type is not already the desired type
		if (sourceDescr.textureCompressor == destDescr.quantizer && normalizedDestinationQuantizerType == normalizedSourceQuantizerType) return true;

		// If not, load the source tree
		if (verbose) printf("Reading tree in %s...\n", OctreeLoader::GetFullFilename(sourceType, sourceDepth, dagFileName).c_str());
		MaterialLibraryUniqueIndexTree<Color, ColorCompare>* source = (MaterialLibraryUniqueIndexTree<Color, ColorCompare>*)OctreeLoader::ReadCache(sourceType, sourceDepth, dagFileName);
		if (source == NULL)
		{
			if (verbose) printf("Reading the source tree failed...\n");
			return false;
		}

		if (normalizedSourceQuantizerType != normalizedDestinationQuantizerType)
		{
			// Quantize all materials
			if (verbose) printf("Quantizing materials...");
			Stopwatch watch; watch.Reset();
			std::vector<Color> materials = source->GetUniqueMaterials();
			BaseColorQuantizer* colorQuantizer = ColorQuantizerFactory::Create(normalizedDestinationQuantizerType);
			std::map<Color, Color, ColorCompare>* colorReplacers = colorQuantizer->QuantizeMaterials(materials);
			delete colorQuantizer;
			std::unordered_map<Color, Color> colorReplacersUnorderedMap;
			for (auto colorReplacer : *colorReplacers) colorReplacersUnorderedMap.insert(colorReplacer);
			delete colorReplacers;
			if (verbose) printf("Materials quantized in %u ms.\n Replacing materials...", (unsigned)(watch.GetTime() * 1000.0));
			watch.Reset();
			source->ReplaceMaterials(colorReplacersUnorderedMap);
			if (verbose) printf("Materials replaced in %u ms.\n", (unsigned)(watch.GetTime() * 1000.0));
		}

		if (sourceDescr.textureCompressor != destDescr.textureCompressor)
		{
			// Replace the compressed texture type
			if (verbose) printf("Replacing texture compressor...");
			Stopwatch watch; watch.Reset();
			source->ReplaceCompressedTexture(CompressedTextureFactory<MaterialLibraryPointer>::GetCompressedTexture(destDescr.textureCompressor));
			if (verbose) printf("Texture compressor replaced in %u ms.\n", (unsigned)(watch.GetTime() * 1000.0));
		}
		else
		{
			if (verbose) printf("Recompressing texture with quantized materials...");
			Stopwatch watch; watch.Reset();
			source->Recompress();
			if (verbose) printf("Done in %u ms.\n", (unsigned)(watch.GetTime() * 1000.0));
		}

		OctreeLoader::WriteCache(source, destinationType, dagFileName, verbose);
		delete source;
		return true;
	}

	return false;
}

bool OctreeConverter::ConvertToHierarchicalColorShift(std::string sourceType, std::string destinationType, unsigned8 sourceDepth, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
	if (sourceDepth != destinationDepth)
		return false;
	TreeTypeDescriptor sourceDescr = TreeTypeParser::GetTreeTypeDescriptor(sourceType);
	TreeTypeDescriptor destDescr = TreeTypeParser::GetTreeTypeDescriptor(destinationType);

	// Make sure the source and destination type is valid
	assert(destDescr.globalType == HIERARCHICAL && destDescr.material == "c" && destDescr.additionalTypeInfo == "s");
	assert(sourceDescr.globalType == HIERARCHICAL && sourceDescr.material == "c");

	std::string normalizedSourceQuantizerType = ColorQuantizerFactory::GetNormalizedType(sourceDescr.quantizer);
	std::string normalizedDestinationQuantizerType = ColorQuantizerFactory::GetNormalizedType(destDescr.quantizer);

	// Only trees with full colors or the same compression scheme as the destination type can be compressed
	if (normalizedSourceQuantizerType != "" && normalizedSourceQuantizerType != normalizedDestinationQuantizerType)
		return false;

	// Remove the second characted, as that is the one that indicates that we want a shifting tree
	std::string hierarchicalColoredDestinationType = std::string(destinationType).erase(1, 1);

	// If we're dealing with a normal "colored octree" that isn't hierarchical, or has a different compression, convert it to hierarchical first:
	if (TreeTypeParser::GetGlobalType(sourceType) != HIERARCHICAL || normalizedDestinationQuantizerType != normalizedSourceQuantizerType)
	{
		if (verbose) printf("Checking cache for hierarchical tree with correct compression...");
		if (OctreeLoader::VerifyCache(hierarchicalColoredDestinationType, sourceDepth, dagFileName))
		{
			if (verbose) printf("Cache file found.\n");
		}
		else
		{
			if (verbose) printf("Nothing found.\nConverting non-hierarchical tree %s to hierarchical tree with correct compression.\n", OctreeLoader::GetFullFilename(sourceType, sourceDepth, dagFileName).c_str());
			if (OctreeConverter::ConvertToHierarchicalColored(sourceType, hierarchicalColoredDestinationType, sourceDepth, destinationDepth, dagFileName, verbose))
				return false;
		}
	}

	// By now we're sure that a fully colored or compressed colored hierarchical octree with the same compression exists
	// Let's convert it to a shifter:
	if (verbose) printf("Reading source tree to new shifting tree...");
	auto source = (MaterialLibraryTree<Color, ColorCompare>*)OctreeLoader::ReadCache(hierarchicalColoredDestinationType, destinationDepth, dagFileName);
	auto dest = (HierarchicalShiftingColoredTree*)OctreeLoader::CreateTree(destinationType, destinationDepth);
	std::vector<Color> sourceMaterials = source->GetMaterials();
	dest->MoveShallow(source);
	delete source;
	dest->SetMaterials(sourceMaterials);
	if (verbose) printf("\nReplacing original colors by color shifts...\n");
	dest->ReplaceColorsByShifts();
	if (verbose) printf("Converting the tree containing color shifts to a DAG...\n");
	dest->ToDAG();

	OctreeLoader::WriteCache(dest, destinationType, dagFileName, verbose);
	delete dest;
	if (verbose) printf("Conversion succeeded");
	return true;
}



bool OctreeConverter::ConvertToOnlyMaterials(std::string sourceType, std::string destinationType, unsigned8 sourceDepth, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
	if (sourceDepth != destinationDepth)
		return false;

	TreeTypeDescriptor destDescr = TreeTypeParser::GetTreeTypeDescriptor(destinationType);
	if (destDescr.globalType != ONLYMATERIAL)
		return false;

	if (destDescr.material == "c")
		return ConvertToOnlyColors(sourceType, destinationType, destinationDepth, dagFileName, verbose);

	// No valid type found
	return false;

}

bool OctreeConverter::ConvertToOnlyColors(std::string sourceType, std::string destinationType, unsigned8 destinationDepth, std::string dagFileName, bool verbose)
{
	TreeTypeDescriptor destDescr = TreeTypeParser::GetTreeTypeDescriptor(destinationType);
	TreeTypeDescriptor sourceDescr = TreeTypeParser::GetTreeTypeDescriptor(sourceType);

	// Make sure that the destination type is actually what is expected
	if (destDescr.globalType != ONLYMATERIAL || destDescr.material != "c") // Make sure the destination is actual the type we want
		return false;


	// Only colored and hierarchical trees can be converted to only color trees
	if (sourceDescr.globalType != HIERARCHICAL || sourceDescr.material != "c")
		return false;
	std::string normalizedSourceQuantizerType = ColorQuantizerFactory::GetNormalizedType(sourceDescr.quantizer);
	std::string normalizedDestinationQuantizerType = ColorQuantizerFactory::GetNormalizedType(destDescr.quantizer);
	
	// Only trees with full colors or the same compression scheme as the destination type can be compressed
	if (normalizedSourceQuantizerType != "" && normalizedSourceQuantizerType != normalizedDestinationQuantizerType)
		return false;

	// Remove the second characted, as that is the one that indicates that we want a shifting tree
	std::string hierarchicalColoredDestinationType = "h" + std::string(destinationType).erase(0, 1);

	// If we're dealing with a normal "colored octree" that isn't hierarchical, or has a different compression, convert it to hierarchical first:
	if (TreeTypeParser::GetGlobalType(sourceType) != HIERARCHICAL || normalizedDestinationQuantizerType != normalizedSourceQuantizerType)
	{
		if (verbose) printf("Checking cache for hierarchical tree with correct compression...");
		if (OctreeLoader::VerifyCache(hierarchicalColoredDestinationType, destinationDepth, dagFileName))
		{
			if (verbose) printf("Cache file found.\n");
		}
		else
		{
			if (verbose) printf("Nothing found.\nConverting non-hierarchical tree %s to hierarchical tree with correct compression.\n", OctreeLoader::GetFullFilename(sourceType, destinationDepth, dagFileName).c_str());
			if (OctreeConverter::ConvertToHierarchicalColored(sourceType, hierarchicalColoredDestinationType, destinationDepth, destinationDepth, dagFileName, verbose))
				return false;
		}
	}


	// By now we're sure that a fully colored or compressed colored hierarchical octree with the same compression exists
	// Let's convert it to a tree with only colors:
	if (verbose) printf("Reading source tree and moving to new tree with colors only...");
	auto source = (MaterialLibraryTree<Color, ColorCompare>*)OctreeLoader::ReadCache(hierarchicalColoredDestinationType, destinationDepth, dagFileName);
	auto dest = (HierarchicalColorsOnlyTree*)OctreeLoader::CreateTree(destinationType, destinationDepth);
	std::vector<Color> sourceMaterials = source->GetMaterials();

	dest->MoveShallow(source);
	delete source;
	dest->SetMaterials(sourceMaterials);
	if (verbose) printf("\nConverting to DAG while merging empty nodes...\n");
	auto quantizer = ColorQuantizerFactory::Create(normalizedDestinationQuantizerType);
	dest->ToDAG(quantizer);

	OctreeLoader::WriteCache(dest, destinationType, dagFileName, verbose);
	delete dest;

	if (verbose) printf("Conversion succeeded");
	return true;
}