#pragma once
#include <algorithm>
#include <unordered_map>
#include <unordered_set>
#include <math.h>

#include "CompressedTexture.h"
#include "BasicTexture.h"
#include "BlockHashers.h"

#include "../Material/Block.h"
#include "../../core/CollectionHelper.h"
#include "../../core/Util/BoolArray.h"
#include "../../core/Serializer.h"


#include "../../inc/tbb/parallel_for.h"
#include "../../inc/tbb/concurrent_queue.h"
#include "../../inc/tbb/parallel_for_each.h"

template<typename T>
class PaletteBlockTexture : public CompressedTexture < T >
{
private:
	const size_t								MAX_STEP_SIZE = 20000000;

	unsigned64									mOriginalSize;

	std::unordered_map<unsigned32, unsigned32>	mFirstBlockPointers;
	unsigned8									mFirstBlockPointerSize;
	// The pointers to the location in mBlockPalettePointers where the current block starts (one per block)
	std::vector<unsigned64>						mBlockPointers;
	// The pointers to the palettes (one per block)
	std::vector<unsigned32>						mPalettePointers;
	// All palettes
	std::vector<Block<T>>						mBlockPalettes;
	// Dictionary for quick finding of palette indices
	std::unordered_map<Block<T>, unsigned32>	mPaletteIndices;
	// All palette pointers (as big as the original texture). Use block pointers to find out the index of the block, and therefore the accompanying palette
	std::vector<std::vector<unsigned8>>			mPackedBlockPalettePointers;

	// Mask used for all original materials
	unsigned32									mPaletteMask;
	// Bitmap corresponding to the current paletteMask
	std::vector<unsigned8>						mPaletteBitMap;

	void RebuildPaletteIndices()
	{
		mPaletteIndices.clear();
		for (unsigned32 i = 0; i < mBlockPalettes.size(); i++)
			mPaletteIndices.insert(std::pair<Block<T>, unsigned32>(mBlockPalettes[i], i));
	}

	void SetPaletteMask(unsigned32 mask)
	{
		mPaletteMask = mask;
		mPaletteBitMap = BitHelper::GetBitMapHS(mask);
	}

	// Adds the palette to the block palettes (and dictionary)
	// If the palette already exists, the pointer to the existing palette is returned
	unsigned32 AddPalette(const Block<T>& palette)
	{
		// Check if the palette already exists and return a pointer to it if it does
		auto existingPalette = mPaletteIndices.find(palette);
		if (existingPalette != mPaletteIndices.end())
			return existingPalette->second;
		// Otherwise, add the palette
		unsigned32 index = (unsigned32)mBlockPalettes.size();
		mBlockPalettes.push_back(palette);
		mPaletteIndices.insert(std::pair<Block<T>, unsigned32>(palette, index));
		return index;
	}

	inline float AverageBitSize(unsigned8 pointerBits, size_t blockSize)
	{
		return
			((float)(64 +                                                  // Bits for the palette and block pointers (32 + 32)
			(GetBitsPerFullMaterial() * BitHelper::Exp2(pointerBits)) +    // Bits for the palette
			BitHelper::RoundToBytes(blockSize * pointerBits))) /           // Bits for the palette pointers (round to bytes since pointers should be available)
			((float)blockSize);	                                           // Divide by the total number of elements in the block
	}

	inline unsigned8 OptimalBitCount(const std::vector<size_t>& blockSizes)
	{
		// Optimal palette size is calculated greedily, by whatever size requires the least bits per material
		unsigned8 bitsPerFullMaterial = GetBitsPerFullMaterial();
		float minBitsPerMaterial = (float)bitsPerFullMaterial; // Set it to the bits per full material first. Note that in this case, the size for the palette and block pointers are ignored.
		unsigned8 optimalSize = bitsPerFullMaterial;
		for (unsigned8 i = 0; i < (unsigned8)blockSizes.size(); i++)
		{
			unsigned8 bitCount = i + 1;
			float requiredBits = AverageBitSize(bitCount, blockSizes[i]);
			if (requiredBits < minBitsPerMaterial)
			{
				optimalSize = bitCount;
				minBitsPerMaterial = requiredBits;
			}
		}
		return optimalSize;

	}

	inline unsigned8 GetBitsPerFullMaterial() const
	{
		return std::max((unsigned8)1, BitHelper::GetSet(mPaletteMask));
	}

	inline size_t GetBlockPalettesSize() const
	{
		size_t res = 0;
		for (auto it = mBlockPalettes.begin() + 1; it != mBlockPalettes.end(); ++it)
			if (it->size() > 1)
				res += (*it).size();
		return res;
	}

	inline size_t GetPalettesByteSize() const
	{
		size_t palettesByteSize = GetBlockPalettesSize() * GetBitsPerFullMaterial();
		palettesByteSize = BitHelper::RoundToBytes(palettesByteSize);
		return palettesByteSize / 8;
	}

	inline unsigned8 GetPaletteBitSize(const size_t& i) const
	{
		// The first palette is always the "full color" palette. Calculate the actual bitsize for all other palettes
		return (i == 0) ? GetBitsPerFullMaterial() : BitHelper::Log2Ceil(mBlockPalettes[i].size());
	}

	inline unsigned8 GetBlockPalettePointerBitSize(const size_t& i) const
	{
		return GetPaletteBitSize(mPalettePointers[i]);
	}

	// Returns the index of the block in which the material at materialIndex is
	inline unsigned64 GetBlockIndex(const unsigned64& materialIndex) const
	{
		if (mBlockPointers.empty()) return 0;
		//if (materialIndex == 0) return 0;
		auto curBlockIt = std::upper_bound(mBlockPointers.begin(), mBlockPointers.end(), materialIndex);
		return (curBlockIt - mBlockPointers.begin()) - 1;
	}

	// Returns the end material index of block i (e.g. the index where the next block starts)
	inline unsigned64 GetBlockEndIndex(const unsigned64& i) const
	{
		return ((i + 1 >= mBlockPointers.size()) ? size() : mBlockPointers[i + 1]);
	}

	// Returns the size of block i
	inline unsigned64 GetBlockSize(const unsigned64& i) const
	{
		return GetBlockEndIndex(i) - mBlockPointers[i];
	}

	// Returns the size of block i in terms of bytes (e.g. blockSize * blockBitSize)
	inline unsigned64 GetBlockByteSize(const unsigned64& i) const
	{
		unsigned64 res = GetBlockSize(i) * GetBlockPalettePointerBitSize(i);
		res = BitHelper::RoundToBytes(res);
		return res >> 3;
	}

	// Returns the sum of all block sizes in bytes
	inline unsigned64 GetBlocksByteSize() const
	{
		size_t res = 0;
		for (size_t i = 0; i < mBlockPointers.size(); i++)
			res += GetBlockByteSize(i);
		return res;
	}

	inline unsigned8 GetOriginalPointerSize() const
	{
		return BitHelper::RoundToBytes(BitHelper::Log2Ceil(mOriginalSize)) / 8;
	}

	inline unsigned8 GetPalettePointerSize() const
	{
		size_t palettesSize = GetBlockPalettesSize();
		return BitHelper::RoundToBytes(BitHelper::Log2Ceil(palettesSize) + 3) / 8;
	}

	inline unsigned8 GetBlockPointerSize() const
	{
		size_t paletteByteSize = GetPalettesByteSize();
		size_t blocksByteSize = GetBlocksByteSize();
		return BitHelper::RoundToBytes(BitHelper::Log2Ceil(paletteByteSize + blocksByteSize)) / 8;
	}

	void RebuildFirstBlock(const std::vector<T>& materialPointers, size_t fromIndex = 0)
	{
		// The first block should contain all colors, so that it can be picked whenever storing a palette takes up too much memory.
		// First copy the existing first block
		std::vector<T> firstBlock;
		if (!mBlockPalettes.empty() && fromIndex > 0)
		{
			const std::vector<T>& curFirstBlock = mBlockPalettes[0].GetData();
			firstBlock.resize(curFirstBlock.size());
			std::copy(curFirstBlock.begin(), curFirstBlock.end(), firstBlock.begin());
			// Remove existing value from palette indices
			mPaletteIndices.erase(mBlockPalettes[0]);
		}
		// TODO: Do something faster then another full scan of the material pointers...
		// Add the new material pointers to the new first block
		size_t existingFirstBlockSize = firstBlock.size();
		firstBlock.resize(firstBlock.size() + materialPointers.size());
		std::copy(materialPointers.begin(), materialPointers.end(), firstBlock.begin() + existingFirstBlockSize);

		// Build the new first block by taking all unique materials
		CollectionHelper::Unique(firstBlock, [](const T& x, const T& y)	{ return ((unsigned)x) < ((unsigned)y); });
		mFirstBlockPointerSize = BitHelper::Log2Ceil(firstBlock.size());

		// Build a dictionary to quickly find materials in the new first block
		mFirstBlockPointers.clear();
		for (unsigned32 i = 0; i < firstBlock.size(); i++)
			mFirstBlockPointers.insert(std::make_pair((unsigned32)firstBlock[i], i));

		// Add the firstblock as the first entry in the blockpalettes
		Block<T> firstBlockPalette = Block<T>(firstBlock);
		if (mBlockPalettes.empty())
			mBlockPalettes.push_back(firstBlockPalette);
		else
		{
			// If the firstblock changes, all blocks that used the first block (before fromIndex) should be updated accordingly
			// To do this, we build a dictionary that maps the firstBlock indices from the old to the new firstblock
			std::unordered_map<unsigned32, unsigned32> oldToNewFirstBlock;
			const Block<T>& oldFirstBlock = mBlockPalettes[0];
			for (unsigned32 i = 0; i < (unsigned32)oldFirstBlock.size(); i++)
				oldToNewFirstBlock.insert(std::make_pair(i, mFirstBlockPointers[oldFirstBlock[i]]));

			unsigned32 oldFirstBlockPointerMask = BitHelper::GetLSMask<unsigned32>(0, BitHelper::Log2Ceil(oldFirstBlock.size()));
			std::vector<unsigned8> oldFirstBlockBitMap = BitHelper::GetBitMapHS(oldFirstBlockPointerMask);

			unsigned32 newFirstBlockPointerMask = BitHelper::GetLSMask<unsigned32>(0, mFirstBlockPointerSize);
			std::vector<unsigned8> newFirstBlockBitMap = BitHelper::GetBitMapHS(newFirstBlockPointerMask);

			// And then replace all pointers to the old first block to pointers to the new first block
			for (size_t i = 0; i < mBlockPointers.size(); i++)
			{
				if (mPalettePointers[i] != 0) continue;
				if (mBlockPointers[i] >= fromIndex) break;
				unsigned64 blockSize = std::min(GetBlockEndIndex(i), (unsigned64)fromIndex) - mBlockPointers[i];
				std::vector<unsigned32> unpackedOldBlock(blockSize);
				// Unpack the block:
				for (size_t j = 0; j < blockSize; j++)
					unpackedOldBlock[j] = oldToNewFirstBlock[BitHelper::UnpackTightAt<unsigned8, unsigned32>(mPackedBlockPalettePointers[i], j, oldFirstBlockBitMap)];
				mPackedBlockPalettePointers[i] = BitHelper::PackTight<unsigned8, unsigned32>(unpackedOldBlock, newFirstBlockBitMap);
			}

			// Replace the first block palette by the new block palette
			mBlockPalettes[0] = firstBlockPalette;
		}
		// Update the palette indices
		mPaletteIndices.insert(std::pair<Block<T>, unsigned32>(firstBlockPalette, 0));

		// Calculate the first block mask
		unsigned32 firstBlockMask = 0;
		for (T v : firstBlock)
			firstBlockMask |= (unsigned32)v;
		SetPaletteMask(firstBlockMask);
	}

	struct TextureTreeNode
	{
		unsigned64 blockStart;
		unsigned64 blockEnd;
		std::vector<T> palette;

		TextureTreeNode(unsigned64 start, unsigned64 end, T material)
		{
			blockStart = start;
			blockEnd = end;
			palette.push_back(material);
		}

		TextureTreeNode(unsigned64 start, unsigned64 end, std::vector<TextureTreeNode*> children)
		{
			blockStart = start;
			blockEnd = end;
			//this->children = children;
			for (auto child : children)
				palette.insert(palette.end(), child->palette.begin(), child->palette.end());
			CollectionHelper::UniqueSerial(palette, [](const T& a, const T& b) { return (unsigned32)a < (unsigned32)b; });
		}

		TextureTreeNode(unsigned64 start, unsigned64 end, std::unordered_set<T> palette)
		{
			blockStart = start;
			blockEnd = end;
			//this->children = children;
			for (auto mat : palette)
				this->palette.push_back(mat);
		}

		unsigned64 Size() const { return blockEnd - blockStart; }

		void AddChild(TextureTreeNode& child)
		{
			palette.insert(palette.end(), child.palette.begin(), child.palette.end());
			CollectionHelper::UniqueSerial(palette);
		}
	};

	struct PaletteBlock
	{
		unsigned64 blockStart;
		unsigned64 blockEnd;
		Block<T> palette;
		bool fullPalette;

		PaletteBlock() {}

		PaletteBlock(unsigned64 start, unsigned64 end, const Block<T>& palette)
		{
			blockStart = start;
			blockEnd = end;
			this->palette = palette;
			fullPalette = false;
		}

		PaletteBlock(unsigned64 start, unsigned64 end)
		{
			blockStart = start;
			blockEnd = end;
			fullPalette = true;
		}

		unsigned64 Size() const { return blockEnd - blockStart; }
	};

	void AddBlocks(const std::vector<PaletteBlock>& blocks, const std::vector<T>& fullMaterialPointers, unsigned64 fromIndex)
	{
		unsigned64 coveredBlocks = 0, coveredCells = 0, uncoveredBlocks = 0, uncoveredCells = 0;

		// Analyze the blocks:
		for (const PaletteBlock& block : blocks)
		{
			if (block.fullPalette)
			{
				uncoveredBlocks++;
				uncoveredCells += block.Size();
			}
			else
			{
				coveredBlocks++;
				coveredCells += block.Size();
			}
		}
		printf("Covered blocks: %llu, cells: %llu, uncovered blocks: %llu, cells: %llu\n", coveredBlocks, coveredCells, uncoveredBlocks, uncoveredCells);


		std::unordered_map<unsigned32, unsigned32> blockInPalettePointers;
		for (auto block : blocks)
		{
			unsigned8 blockBitCount = block.fullPalette ? BitHelper::Log2Ceil(mBlockPalettes[0].size()) : BitHelper::Log2Ceil(block.palette.size());
			// Create the block palette (if necessary)
			unsigned32 paletteIdx = 0;
			if (!block.fullPalette)
			{
				Block<T> palette = block.palette;
				palette.Sort([](const T& x, const T& y)	{ return ((unsigned)x) < ((unsigned)y); });
				paletteIdx = AddPalette(palette);
			}
			// Add the block and palette pointers
			mBlockPointers.push_back(block.blockStart + fromIndex);
			mPalettePointers.push_back((unsigned32)paletteIdx);

			// Add the in-palette pointers to the materials to mPackedBlockPalettePointers
			// First build a dictionary that maps materials to in-palette pointers
			std::unordered_map<unsigned32, unsigned32>* inPalettePointers = &mFirstBlockPointers;
			if (!block.fullPalette)
			{
				const Block<T>& palette = mBlockPalettes[paletteIdx];
				blockInPalettePointers.clear();
				for (unsigned32 j = 0; j < palette.size(); j++)
					blockInPalettePointers.insert(std::pair<unsigned32, unsigned32>((unsigned32)palette[j], j));
				inPalettePointers = &blockInPalettePointers;
			}

			// Then find the correct in-palette pointer for each material and build an (unpacked) vector containing it
			if (blockBitCount > 0)
			{
				std::vector<unsigned32> unpackedBlockPointers(block.blockEnd - block.blockStart);
				for (size_t j = block.blockStart; j < block.blockEnd; j++)
					unpackedBlockPointers[j - block.blockStart] = inPalettePointers->operator[]((unsigned32)fullMaterialPointers[j]);

				unsigned32 paletteMask = BitHelper::GetLSMask<unsigned32>(0, blockBitCount);
				mPackedBlockPalettePointers.push_back(BitHelper::PackTight<unsigned8, unsigned32>(unpackedBlockPointers, paletteMask));
			}
			else
			{
				// If the palette pointers take 0 bits, just push an empty array.
				mPackedBlockPalettePointers.push_back(std::vector<unsigned8>());
			}
		}
	}

	std::vector<PaletteBlock> BuildBlocks(std::vector<T>& fullMaterials, unsigned8 greedyUntilBits = 3)
	{
		// Build the texture tree leaf nodes, which are all blocks that can be made with a single material
		std::vector<TextureTreeNode*> nodesToProcess;
		unsigned64 curStartIndex = 0;
		T curMaterial = fullMaterials[0];
		for (size_t i = 0; i < fullMaterials.size(); i++)
		{
			if (fullMaterials[i] != curMaterial)
			{
				nodesToProcess.push_back(new TextureTreeNode(curStartIndex, i, curMaterial));
				curMaterial = fullMaterials[i];
				curStartIndex = i;
			}
		}
		nodesToProcess.push_back(new TextureTreeNode(curStartIndex, fullMaterials.size(), curMaterial));

		std::vector<PaletteBlock> blocks;

		BoolArray covered(fullMaterials.size());

		unsigned8 bitsPerFullMaterial = GetBitsPerFullMaterial();
		unsigned8 maxSizeToTry = std::min<unsigned8>(greedyUntilBits, bitsPerFullMaterial - 1); // bits
		std::vector<TextureTreeNode*> curBitNodes;
		std::vector<size_t> blocksPerBitsize(maxSizeToTry + 1);

		size_t steps = (nodesToProcess.size() / MAX_STEP_SIZE) + (nodesToProcess.size() % MAX_STEP_SIZE != 0 ? 1 : 0);
		size_t stepSize = nodesToProcess.size() / steps;
		// Calculate the sizes of the blocks to process
		for (size_t step = 0; step < steps; step++)
		{
			size_t stepStart = step * stepSize;
			size_t stepEnd = std::min(nodesToProcess.size(), (step + 1) * stepSize);

			std::vector<TextureTreeNode*> nodesLeft(nodesToProcess.begin() + stepStart, nodesToProcess.begin() + stepEnd);
			printf("stepSize: %llu, stepStart: %llu, stepEnd: %llu, totalSize: %llu.\n", (unsigned64)stepSize, (unsigned64)stepStart, (unsigned64)stepEnd, (unsigned64)nodesToProcess.size());
			for (unsigned8 bitSize = 0; bitSize <= maxSizeToTry; bitSize++)
			{
				if (nodesLeft.empty()) break;
				unsigned64 bucketSize = BitHelper::Exp2(bitSize);
				// Result of solving the average size so that it is smaller than bitSize + 1 bits.
				// This means solving the following equation for z
				// (96 + (x * 2^y) + (z * y)) / z = y + 1
				// where x = bitsPerFullMaterial, y = bitSize, z = blockSize
				// which gives:
				// z = 96 + (x * 2^y)
				// This makes sense: since each element in a block will take 1 bit more with the next bitsize, the critical size is reached when 
				// the number of block pointers is equal to the header of the block in bits.
				unsigned64 criticalSize = 96 + bitsPerFullMaterial * bucketSize;
				curBitNodes.clear();

				// Build all possible sets of nodes that fit in the current bit size
				if (bitSize > 0)
				{
					unsigned64 checkUntil = nodesLeft.size();
					unsigned64 lastBlockStart = ~((unsigned64)0);
					// Find all (unique) block starts:
					std::vector<unsigned64> blockStarts;
					for (unsigned64 i = 0; i < checkUntil; i++)
					{
						unsigned64 blockStart = nodesLeft[i]->blockStart;
						if (lastBlockStart == blockStart) continue;
						lastBlockStart = blockStart;
						blockStarts.push_back(i);
					}

					nodesLeft.resize(nodesLeft.size() + blockStarts.size());
					tbb::parallel_for((unsigned64)0, (unsigned64)blockStarts.size(), [&](const unsigned64& i)
						//for (size_t i = 0; i < blockStarts.size(); i++)
					{
						TextureTreeNode* curNode = nodesLeft[blockStarts[i]];
						std::unordered_set<T> curPalette;

						for (T material : curNode->palette)
							curPalette.insert(material);
						unsigned64 curBlockEnd = curNode->blockEnd;
						unsigned64 blockStartIdx = i + 1;
						while (blockStartIdx < blockStarts.size())
						{
							unsigned64 j = blockStarts[blockStartIdx];
							unsigned64 nextBlockStart = blockStartIdx + 1 < blockStarts.size() ? blockStarts[blockStartIdx + 1] : checkUntil;
							// If there is a gap between this block and the next block, break;
							if (nodesLeft[j]->blockStart > curBlockEnd) break;
							while (j < nextBlockStart)
							{
								TextureTreeNode* potential = nodesLeft[j];
								std::vector<T> newMaterials;
								for (T mat : potential->palette)
									if (curPalette.find(mat) == curPalette.end())
										newMaterials.push_back(mat);
								if (curPalette.size() + newMaterials.size() <= bucketSize)
								{
									for (auto newMaterial : newMaterials)
										curPalette.insert(newMaterial);
									curBlockEnd = potential->blockEnd;
									blockStartIdx++;
									break;
								}
								// If the palette of this block is too big, try the next one, it has a smaller palette and might fit
								j++;
							}
							blockStartIdx++;
						}
						nodesLeft[checkUntil + i] = new TextureTreeNode(curNode->blockStart, curBlockEnd, /*children,*/ curPalette);
					});
				}

				// Now that we've build all possible blocks, find all blocks that are bigger than the criticalSize, and create them
				// Sort blocks, first on palette size, then on block size
				tbb::parallel_sort(nodesLeft.begin(), nodesLeft.end(), [](TextureTreeNode* a, TextureTreeNode* b)
				{
					unsigned64 sizeA = a->Size();
					unsigned64 sizeB = b->Size();
					if (sizeA != sizeB) return sizeA > sizeB;
					return a->blockStart < b->blockStart;
				});
				TextureTreeNode* curNode = nodesLeft[0];
				size_t i = 0;
				while (curNode->Size() >= criticalSize && i < nodesLeft.size())
				{
					curNode = nodesLeft[i];
					// If the curNode isn't covered yet, build a block out of it and add it to the palette blocks
					if (!covered.Any(curNode->blockStart, curNode->blockEnd))
					{
						blocksPerBitsize[bitSize]++;
						blocks.push_back(PaletteBlock(curNode->blockStart, curNode->blockEnd, Block<T>(curNode->palette)));
						covered.SetRange(curNode->blockStart, curNode->blockEnd, true);
					}
					i++;
				}

				// Remove all covered (and partly covered) nodes
				std::vector<TextureTreeNode*> newNodesLeft;
				for (size_t i = 0; i < nodesLeft.size(); i++)
				{
					curNode = nodesLeft[i];
					if (covered.Any(curNode->blockStart, curNode->blockEnd))
						delete curNode;
					else
						newNodesLeft.push_back(curNode);
				}
				nodesLeft = newNodesLeft;

				// Sort the nodes, first on block start (ascending), then on palette size (descending)
				tbb::parallel_sort(nodesLeft.begin(), nodesLeft.end(), [](TextureTreeNode* a, TextureTreeNode* b)
				{
					if (a->blockStart != b->blockStart) return a->blockStart < b->blockStart;
					return a->Size() > b->Size();
				});
			}
			CollectionHelper::PrintContents(blocksPerBitsize);

			tbb::parallel_for_each(nodesLeft.begin(), nodesLeft.end(), [](TextureTreeNode* node) { delete node; });
			nodesLeft.clear();
			nodesLeft.shrink_to_fit();
		}

		// Create blocks for the uncovered parts of the scene using the serial block builder algorithm
		size_t blockStart = 0;
		bool inBlock = false;
		for (size_t i = 0; i < fullMaterials.size(); i++)
		{
			if (!covered.Get(i))
			{
				if (!inBlock)
				{
					blockStart = i;
					inBlock = true;
				}
			}
			else
			{
				if (inBlock)
				{
					std::vector<PaletteBlock> fillerBlocks = SerialBuildBlocks(fullMaterials, blockStart, i, false);
					unsigned64 fillerBlocksIndex = blocks.size();
					blocks.resize(fillerBlocksIndex + fillerBlocks.size());
					std::move(fillerBlocks.begin(), fillerBlocks.end(), blocks.begin() + fillerBlocksIndex);

					inBlock = false;
				}
			}
		}
		if (inBlock)
		{
			std::vector<PaletteBlock> fillerBlocks = SerialBuildBlocks(fullMaterials, blockStart, fullMaterials.size(), false);
			unsigned64 fillerBlocksIndex = blocks.size();
			blocks.resize(fillerBlocksIndex + fillerBlocks.size());
			std::move(fillerBlocks.begin(), fillerBlocks.end(), blocks.begin() + fillerBlocksIndex);
		}


		// Sort the blocks
		tbb::parallel_sort(blocks.begin(), blocks.end(), [](const PaletteBlock& a, const PaletteBlock& b) { return a.blockStart < b.blockStart; });

		return blocks;
	}

	std::vector<PaletteBlock> SerialBuildBlocks(std::vector<T>& fullMaterialPointers, unsigned64 start = 0, unsigned64 end = ~(unsigned64)0, bool verbose = true)
	{
		// Algorithm:
		// Check for palette sizes of 2, 4, 8, 16, 32, 64, 128, 256 colors how big the block can be.
		// Somehow determine the most efficient size pointer (1 - 8 bits)
		// Create a block of that size.
		if (end > fullMaterialPointers.size()) end = fullMaterialPointers.size();
		std::vector<PaletteBlock> blocks;
		unsigned64 curBlockStart = start;
		unsigned64 lastBlockStart = curBlockStart;
		unsigned8 bitsPerFullMaterial = GetBitsPerFullMaterial();
		unsigned8 maxSizeToTry = std::min(8, bitsPerFullMaterial - 1); // bits
		if (maxSizeToTry == 0) maxSizeToTry = 1;
		size_t maxMaterialsPerBlock = size_t(1) << maxSizeToTry;

		std::vector<size_t> blocksPerBitsize(maxSizeToTry);
		while (curBlockStart < end)
		{
			if (verbose)
			{
				size_t curPercentage = (size_t)((curBlockStart - start) * 100) / (end - start);
				size_t lastPercentage = (size_t)((lastBlockStart - start) * 100) / (end - start);
				if (curPercentage > lastPercentage) printf("%u %%\n", (unsigned32)curPercentage);
			}
			lastBlockStart = curBlockStart;
			std::vector<size_t> blockSizes; // Store the block sizes per amount of bits in the palette pointer
			std::vector<T> materialsInBlock;
			std::unordered_set<unsigned32> materialsInBlockSet;
			std::vector<Block<T>> blockSizePalettes; // The palettes for each of the block counts
			size_t i = curBlockStart;
			unsigned8 lastFilledBitCount = 0;
			// Keep moving through the material pointers until the maximum number of materials has been reached
			while (materialsInBlock.size() <= maxMaterialsPerBlock && i < end)
			{
				T& curMaterial = fullMaterialPointers[i];
				unsigned32 curMaterialInt = (unsigned32)curMaterial;
				// Check if this material is already in the materials
				bool materialAlreadyInBlock = materialsInBlockSet.find(curMaterialInt) != materialsInBlockSet.end();
				if (!materialAlreadyInBlock)
				{
					// Check if the limit for a certain number of bits is reached when the new material is added
					for (unsigned8 bitCount = 1; bitCount <= maxSizeToTry; bitCount++)
						if (materialsInBlock.size() == (((size_t)1) << bitCount))
						{
							lastFilledBitCount = bitCount;
							blockSizes.push_back(i - curBlockStart);
							blockSizePalettes.push_back(Block<T>(materialsInBlock));
						}
					// Add the new material
					materialsInBlock.push_back(curMaterial);
					materialsInBlockSet.insert((unsigned32)curMaterial);
				}
				i++;
			}
			// If we're at the end of the material, fill the rest of the blockSizes with the current block size
			if (lastFilledBitCount < maxSizeToTry)
			{
				// Add some dummy materials to the block so that it is big enough
				materialsInBlock.resize(((size_t)1) << maxSizeToTry);
				for (unsigned8 bitCount = lastFilledBitCount + 1; bitCount <= maxSizeToTry; bitCount++)
				{
					blockSizes.push_back(i - curBlockStart);
					blockSizePalettes.push_back(Block<T>(materialsInBlock, 0, ((size_t)1) << bitCount));
				}
			}
			unsigned8 optimalBitCount = OptimalBitCount(blockSizes);
			// If the optimalBitCount gives no compression, pick the halfway blocksize (so as not to skip to much indexes than can potentially be compressed better)
			size_t blockSize;
			if (optimalBitCount == bitsPerFullMaterial)
				blockSize = blockSizes[blockSizes.size() / 2];
			else
				blockSize = blockSizes[optimalBitCount - 1];

			if (optimalBitCount == bitsPerFullMaterial)
				blocks.push_back(PaletteBlock(curBlockStart, curBlockStart + blockSize));
			else
				blocks.push_back(PaletteBlock(curBlockStart, curBlockStart + blockSize, blockSizePalettes[optimalBitCount - 1]));

			// Update the block starts
			curBlockStart += blockSize;
		}
		
		return blocks;
	}
public:
	PaletteBlockTexture() :
			mOriginalSize(0),
			mFirstBlockPointers(std::unordered_map<unsigned32, unsigned32>()),
			mFirstBlockPointerSize(0),
			mBlockPointers(std::vector<unsigned64>()),
			mPalettePointers(std::vector<unsigned32>()),
			mBlockPalettes(std::vector<Block<T>>()),
			mPaletteIndices(std::unordered_map<Block<T>, unsigned32>()),
			mPackedBlockPalettePointers(std::vector<std::vector<unsigned8>>()),
			mPaletteMask(0),
			mPaletteBitMap(std::vector<unsigned8>())
	{}

	~PaletteBlockTexture() override {}

	T operator[](size_t i) const override
	{
		unsigned64 blockId = GetBlockIndex(i);
		const Block<T>& palette = mBlockPalettes[mPalettePointers[blockId]];
		unsigned8 bitSize = BitHelper::Log2Ceil(palette.size());
		unsigned32 blockPaletteId = 0;
		if (bitSize > 0) // if bitSize == 0, then the palette is only one color big, so the blockPaletteId is always 0.
		{
			unsigned32 paletteMask = BitHelper::GetLSMask<unsigned32>(0, bitSize);
			blockPaletteId = BitHelper::UnpackTightAt<unsigned8, unsigned32>(mPackedBlockPalettePointers[blockId], i - mBlockPointers[blockId], paletteMask);
		}
		return palette[blockPaletteId];
	}

	unsigned64 size() const override { return mOriginalSize; }

	// Returns the (unpacked) material indices from the given index
	std::vector<T> GetTexture(size_t fromIndex = 0) override
	{
		//printf("fromIndex=%u, size=%u \n", fromIndex, size());
		assert(fromIndex <= size());
		std::vector<T> res(size() - fromIndex);
		tbb::parallel_for(fromIndex, (size_t)size(), [&](size_t i)
		{
			res[i - fromIndex] = this->operator[](i);
		});
		return res;
	}

	void SetTexture(const std::vector<T>& materialPointers, size_t fromIndex = 0) override
	{
		//// Debug: just repack everything if fromIndex != 0
		//if (fromIndex != 0)
		//{
		//	std::vector<T> existingMaterials = GetTexture();
		//	std::vector<T> allMaterialPointers(fromIndex + materialPointers.size());
		//	std::move(existingMaterials.begin(), existingMaterials.begin() + fromIndex, allMaterialPointers.begin());
		//	std::copy(materialPointers.begin(), materialPointers.end(), allMaterialPointers.begin() + fromIndex);
		//	SetTexture(allMaterialPointers);
		//	return;
		//}

		// Rebuild the first block to also contain all new materials
		RebuildFirstBlock(materialPointers, fromIndex);

		// Remove all data after fromIndex
		if (!mBlockPointers.empty())
		{
			unsigned64 switchBlockId = GetBlockIndex(fromIndex);
			size_t switchBlockWantedSize = fromIndex - mBlockPointers[switchBlockId];
			if (switchBlockWantedSize == 0)
				switchBlockId--;
			else
			{
				size_t switchBlockBitSize = BitHelper::Log2Ceil(mBlockPalettes[mPalettePointers[switchBlockId]].size());
				mPackedBlockPalettePointers[switchBlockId].resize(BitHelper::RoundToBytes(switchBlockWantedSize * switchBlockBitSize) / 8);
			}

			// Remove the block pointers and block palette pointers that point to locations after fromIndex
			unsigned64 blockPointerEndIndex = switchBlockId + 1;
			if (blockPointerEndIndex < mBlockPointers.size())
			{
				mBlockPointers.resize(blockPointerEndIndex);
				mPalettePointers.resize(blockPointerEndIndex);
				mPackedBlockPalettePointers.resize(blockPointerEndIndex);

				// Remove all palettes that aren't used anymore
				if (blockPointerEndIndex == 0)
				{
					mBlockPalettes.resize(1); // Only keep the first block, because it should always exist
				}
				else
				{
					unsigned32 finalPalette = *std::max_element(mPalettePointers.begin(), mPalettePointers.end());
					mBlockPalettes.resize(finalPalette + 1);
				}
				RebuildPaletteIndices();
			}
		}
		std::vector<T> fullMaterialPointers = materialPointers;
		std::vector<PaletteBlock> blocks = BuildBlocks(fullMaterialPointers);
		//std::vector<PaletteBlock> blocks = SerialBuildBlocks(fullMaterialPointers);
		AddBlocks(blocks, fullMaterialPointers, fromIndex);
		mOriginalSize = fromIndex + fullMaterialPointers.size();
	}

	void ReplaceMaterials(const std::unordered_map<T, T>& replacementMap) override
	{
		bool replacersEqual = true;
		for (auto replacer : replacementMap)
			if (replacer.first != replacer.second)
			{
				replacersEqual = false;
				break;
			}
		if (replacersEqual)
			return;

		// Replace the colors in all palettes:
		tbb::parallel_for_each(mBlockPalettes.begin(), mBlockPalettes.end(), [&](Block<T>& block)
		{
			for (unsigned i = 0; i < block.size(); i++)
			{
				auto replacer = replacementMap.find(block[i]);
				if (replacer != replacementMap.end())
					block.Set(i, replacer->second);
			}
		});

		RebuildPaletteIndices();
	}

	void Recompress() override
	{
		SetTexture(GetTexture()); 
	}

	void ReadFromFile(std::istream& file)
	{
		// Read the original size
		Serializer<unsigned64>::Deserialize(mOriginalSize, file);
		Serializer<unsigned8>::Deserialize(mFirstBlockPointerSize, file);

		unsigned32 paletteMask;
		Serializer<unsigned32>::Deserialize(paletteMask, file);
		SetPaletteMask(paletteMask);

		// Read the block pointers
		unsigned64 blockPointerSize = 0;
		Serializer<unsigned64>::Deserialize(blockPointerSize, file);
		mBlockPointers.resize(blockPointerSize);
		Serializer<unsigned64*>::Deserialize(&mBlockPointers[0], blockPointerSize, file);
		// Also read the palette pointers
		mPalettePointers.resize(blockPointerSize);
		Serializer<unsigned32*>::Deserialize(&mPalettePointers[0], blockPointerSize, file);

		// Read the block palettes
		unsigned64 blockPalettesSize = 0;
		Serializer<unsigned64>::Deserialize(blockPalettesSize, file);
		mBlockPalettes.resize(blockPalettesSize);
		std::vector<T> blockCache;
		for (size_t i = 0; i < (size_t)blockPalettesSize; i++)
		{
			Serializer<std::vector<T>, unsigned32>::Deserialize(blockCache, file);
			mBlockPalettes[i] = Block<T>(blockCache);
		}
		
		// Read the block palette pointers
		Serializer<std::vector<std::vector<unsigned8>>, unsigned32>::Deserialize(mPackedBlockPalettePointers, file);
	}

	void WriteToFile(std::ostream& file) const override
	{
		// Write the original size
		Serializer<unsigned64>::Serialize(mOriginalSize, file);
		Serializer<unsigned8>::Serialize(mFirstBlockPointerSize, file);

		// Write the palette mask
		Serializer<unsigned32>::Serialize(mPaletteMask, file);

		// Write the block pointers
		unsigned64 blockPointerSize = (unsigned64)mBlockPointers.size();
		Serializer<unsigned64>::Serialize(blockPointerSize, file);
		Serializer<unsigned64*>::Serialize(&mBlockPointers[0], blockPointerSize, file);
		Serializer<unsigned32*>::Serialize(&mPalettePointers[0], blockPointerSize, file);

		// Write the block palettes
		unsigned64 blockPalettesSize = mBlockPalettes.size();
		Serializer<unsigned64>::Serialize(blockPalettesSize, file);
		for (unsigned64 i = 0; i < blockPalettesSize; i++)
			Serializer<std::vector<T>, unsigned32>::Serialize(mBlockPalettes[i].GetData(), file);

		// Write the block palette pointers
		Serializer<std::vector<std::vector<unsigned8>>, unsigned32>::Serialize(mPackedBlockPalettePointers, file);
	}

	// Use this method to output the compressed texture as a vector of one byte characters
	std::vector<unsigned8> GetTexturePool() const override {
		// Texture pool should first contain all palettes, followed by the paletteBlockPointers
		size_t poolSize = GetTexturePoolSize();
		std::vector<unsigned8> pool(poolSize);
		size_t curByteIndex = 0;
		// Write all palettes (tightly packed), except the first one!
		{
			std::vector<unsigned32> palettes(GetBlockPalettesSize());
			size_t i = 0;
			for (auto palette = mBlockPalettes.begin() + 1; palette != mBlockPalettes.end(); ++palette)
				// Only write palettes with more than 1 color, palettes with one color will be replaced by the material directly.
				if (palette->size() > 1)
					for (size_t j = 0; j < palette->size(); j++)
						palettes[i++] = (unsigned32)palette->Get(j);
			std::vector<unsigned8> tightPalettes = BitHelper::PackTight<unsigned8, unsigned32>(palettes, mPaletteBitMap);
			std::move(tightPalettes.begin(), tightPalettes.end(), pool.begin());
			curByteIndex += tightPalettes.size();
		}

		std::vector<unsigned8> firstBlockBitMap = BitHelper::GetBitMapHS(BitHelper::GetLSMask<unsigned32>(0, mFirstBlockPointerSize));

		// Now write all mPaletteBlockPointers. These need to be packed tightly according to the size of the palette.
		for (size_t i = 0; i < mBlockPointers.size(); i++) {
			size_t blockSize = GetBlockSize(i);
			std::vector<unsigned8> packedPalettePointers;
			if (mPalettePointers[i] == 0)
			{
				// If we should use the "zero block", then pack the materials directly (the zero block is not an actual palette as it just contains all materials)
				std::vector<unsigned32> blockContents(blockSize);
				for (size_t j = 0; j < blockSize; j++)
					blockContents[j] = (unsigned32)(mBlockPalettes[0].Get(BitHelper::UnpackTightAt<unsigned8, unsigned32>(mPackedBlockPalettePointers[i], j, firstBlockBitMap)));
				packedPalettePointers = BitHelper::PackTight<unsigned8, unsigned32>(blockContents, mPaletteBitMap);
			}
			else
			{
				packedPalettePointers = mPackedBlockPalettePointers[i];
			}
			std::move(packedPalettePointers.begin(), packedPalettePointers.end(), pool.begin() + curByteIndex);
			curByteIndex += packedPalettePointers.size();
			//assert(packedPalettePointers.size() == GetBlockByteSize(i));
		}
		
		return pool;
	}

	size_t GetTexturePoolSize() const override {
		// Calculate the size required for all palettes. These will be packed tightly
		return GetPalettesByteSize() + GetBlocksByteSize();
	}

	// Use this method to output any addition information that might be needed.
	std::vector<unsigned8> GetAdditionalTexturePool() const
	{
		unsigned8 originalPointerSize = GetOriginalPointerSize();
		unsigned8 palettePointerSize = GetPalettePointerSize();
		unsigned8 blockPointerSize = GetBlockPointerSize();

		size_t textureSize = GetAdditionalTexturePoolSize();
		std::vector<unsigned8> additionalPool(textureSize);
		additionalPool[0] = originalPointerSize;
		additionalPool[1] = palettePointerSize;
		additionalPool[2] = blockPointerSize;

		// Build the actual palette pointers, that point to the start of the palette
		std::vector<unsigned64> actualPalettePointers(mBlockPalettes.size());
		actualPalettePointers[0] = -1; // Palette 0 contains all materials, so we store them directly in this case. This is indicating by a 0xFFFFFFFF mask (==-1)
		unsigned64 curActualPalettePointer = 0;
		for (size_t i = 1; i < mBlockPalettes.size(); i++)
		{
			unsigned64 actualPalettePointer = curActualPalettePointer;
			unsigned64 paletteBitSize = GetPaletteBitSize(i);
			if (paletteBitSize == 0) // Meaning the palette is one color, so palette is not stored. This is signaled with mask 0xFFFFFFFE.
				actualPalettePointer = -2;
			else
				actualPalettePointer |= (paletteBitSize - 1) << (palettePointerSize * 8 - 3); // Add mask indicating the size of the palette

			actualPalettePointers[i] = actualPalettePointer;

			// Skip palettes of 0 bits (0 or 1 materials)
			if (paletteBitSize > 0)
				curActualPalettePointer += mBlockPalettes[i].size();
		}

		assert(curActualPalettePointer == GetBlockPalettesSize());
		unsigned64 curActualBlockPointer = GetPalettesByteSize();
		// Write them to the pool
		for (size_t i = 0; i < mBlockPointers.size(); i++)
		{
			size_t curByte = 3 + i * (originalPointerSize + palettePointerSize + blockPointerSize);
			size_t bitSize = GetBlockPalettePointerBitSize(i);

			unsigned64 blockByteSize = GetBlockByteSize(i);
			unsigned64 blockPointer = curActualBlockPointer;
			curActualBlockPointer += blockByteSize;
			if (blockByteSize == 0)
			{ // If blockByteSize = 0, just write the actual color of all nodes in this block as the block pointer
				T mat = mBlockPalettes[mPalettePointers[i]][0];
				blockPointer = (unsigned32)mat;
			}
			BitHelper::SplitInBytesAndMove(mBlockPointers[i], additionalPool, curByte, originalPointerSize);
			BitHelper::SplitInBytesAndMove(actualPalettePointers[mPalettePointers[i]], additionalPool, curByte + originalPointerSize, palettePointerSize);
			BitHelper::SplitInBytesAndMove(blockPointer, additionalPool, curByte + originalPointerSize + palettePointerSize, blockPointerSize);
		}
		assert(curActualBlockPointer == GetPalettesByteSize() + GetBlocksByteSize());
		return additionalPool;
	}

	size_t GetAdditionalTexturePoolSize() const
	{
		// Assume palette pointers and block pointers are both 32 bits
		size_t bytesPerBlock = GetOriginalPointerSize() + GetPalettePointerSize() + GetBlockPointerSize();

		// Calculate the total number of required byte
		return 3 + mBlockPointers.size() * bytesPerBlock;
	}

	std::map<std::string, std::string> GetAdditionalProperties() const override
	{
		std::map<std::string, std::string> res;
		// Write the mask used for the palette pointers
		res.insert(std::pair<std::string, std::string>("mask", std::to_string(mPaletteMask)));
		// Write the blockCount
		size_t blockCount = mBlockPointers.size();
		res.insert(std::pair<std::string, std::string>("blockCount", std::to_string(blockCount)));

		return res;
	};
};