CDAG/Research/scene/TextureCompressor/BlockCompressedTexture.h

#pragma once

#include "CompressedTexture.h"
#include "BasicTexture.h"
#include "../Material/Block.h"
#include "../../inc/tbb/parallel_for.h"
#include "BlockHashers.h"
#include <algorithm>
#include "../../core/Serializer.h"

template<typename T, typename BlockTextureType = BasicTexture<T>>
class BlockCompressedTexture : public CompressedTexture<T>
{
private:
	size_t mActualSize;
	std::vector<unsigned> mBlockPointers;
	BlockTextureType mBlocks;
	unsigned mBlockSize;
public:
	BlockCompressedTexture(unsigned blockSize) :
			mActualSize(0),
			mBlockPointers(std::vector<unsigned>()),
			mBlocks(BlockTextureType()),
			mBlockSize(blockSize)
	{}

	~BlockCompressedTexture() override {}

	T operator[](size_t i) const override
	{
		size_t blockPointerIndex = i / mBlockSize;
		size_t blockIndex = mBlockPointers[blockPointerIndex] * mBlockSize;
		blockIndex += i % mBlockSize;
		return mBlocks[blockIndex];
	}

	unsigned64 size() const override { return mActualSize; }

	// Returns the (unpacked) material indices from the given index
	std::vector<T> GetTexture(size_t fromIndex = 0) override
	{
		assert(fromIndex <= mActualSize);
		std::vector<T> res(mActualSize - fromIndex);
		if (res.size() == 0) return res;
		size_t startBlockIndex = fromIndex / mBlockSize;
		size_t indexInsideStartBlock = fromIndex % mBlockSize;
		size_t indexFromStartBlock = mBlockSize - indexInsideStartBlock;
		// Unpack the part of the first block from the "fromIndex"
		for (size_t j = indexInsideStartBlock; j < mBlockSize && j < res.size(); j++)
			res[j - indexInsideStartBlock] = mBlocks[mBlockPointers[startBlockIndex] * mBlockSize + j];
		// Unpack the rest of the blocks
		for (size_t i = startBlockIndex + 1; i < mBlockPointers.size(); i++)
		{
			size_t blockResIndex = indexFromStartBlock + ((i - startBlockIndex - 1) * mBlockSize);
			size_t blocksIndex = mBlockPointers[i] * mBlockSize;
			for (size_t j = 0; j < mBlockSize; j++)
			{
				if (blockResIndex + j >= res.size()) break;
				res[blockResIndex + j] = mBlocks[blocksIndex + j];
			}
		}
		return res;
	}

	void SetTexture(const std::vector<T>& materialPointers, size_t fromIndex = 0) override
	{
		// TODO: only copy if necessary
		std::vector<T> nodeMaterialPointers;

		mActualSize = fromIndex + materialPointers.size();
		unsigned fromBlockIndex = (unsigned)fromIndex / mBlockSize;
		unsigned highestIndex = mBlockPointers.empty() ? 0 : *std::max_element(mBlockPointers.begin(), mBlockPointers.begin() + fromBlockIndex);

		// Build a map from blocks to indices that already exist (only necessary if fromIndex > 0)
		std::unordered_map<Block<T>, unsigned> blockIndices;
		if (fromIndex >= mBlockSize)
		{
			std::vector<T> uncompressedBlocks = mBlocks.GetTexture();
			for (size_t i = 0; i < highestIndex; i++)
			{
				Block<T> cur(uncompressedBlocks, i * mBlockSize, i * mBlockSize + mBlockSize);
				blockIndices.insert(std::pair<Block<T>, unsigned>(cur, (unsigned)i));
			}
		}
		// Add what was already in the block in which the new nodepointers should be added to nodeMaterialPointers
		unsigned switchIndex = (unsigned)fromIndex % mBlockSize;
		if (switchIndex != 0)
		{
			std::vector<T> switchBlockStart(switchIndex);
			auto switchBlockIndex = mBlockPointers[fromBlockIndex] * mBlockSize;
			for (unsigned j = 0; j < switchIndex; j++)
				switchBlockStart[j] = mBlocks[switchBlockIndex + j];
			nodeMaterialPointers.resize(switchIndex + materialPointers.size());
			std::copy(switchBlockStart.begin(), switchBlockStart.end(), nodeMaterialPointers.begin());
			std::copy(materialPointers.begin(), materialPointers.end(), nodeMaterialPointers.begin() + switchIndex);
		}
		else
		{
			nodeMaterialPointers.resize(materialPointers.size());
			std::copy(materialPointers.begin(), materialPointers.end(), nodeMaterialPointers.begin());
		}
		if (nodeMaterialPointers.size() % mBlockSize != 0)
		{
			unsigned emptyNodesToAdd = mBlockSize - nodeMaterialPointers.size() % mBlockSize;
			nodeMaterialPointers.resize(nodeMaterialPointers.size() + emptyNodesToAdd);
		}
		std::vector<T> newBlocks;
		mBlockPointers.resize(mActualSize / mBlockSize + (((mActualSize % mBlockSize) == 0) ? 0 : 1));

		// Compress the material pointers so that material pointers that are the same won't be reused
		size_t reuseCount = 0;
		unsigned newBlockPointer = highestIndex;
		unsigned blockId = fromBlockIndex;
		for (size_t i = 0; i < nodeMaterialPointers.size(); i += mBlockSize)
		{
			// Build the current block
			Block<T> current(nodeMaterialPointers, i, i + mBlockSize);
			// Check if the current block is already in the texture
			auto existingCurrent = blockIndices.find(current);
			unsigned curBlockPointer;
			if (existingCurrent == blockIndices.end())
			{
				// If it isn't, copy the current block to the block texture.
				for (size_t j = 0; j < mBlockSize; j++)
					newBlocks.push_back(current.Get(j));
				curBlockPointer = newBlockPointer;
				blockIndices.insert(std::pair<Block<T>, unsigned>(current, newBlockPointer));
				newBlockPointer++;
			}
			else
			{
				reuseCount++;
				curBlockPointer = existingCurrent->second;
			}
			mBlockPointers[blockId++] = curBlockPointer;
		}
		mBlocks.SetTexture(newBlocks, highestIndex * mBlockSize);
		printf("%llu Blocks were reused during compression.\n", (unsigned64)reuseCount);
	}

	void ReplaceMaterials(const std::unordered_map<T, T>& replacementMap) override
	{
		mBlocks.ReplaceMaterials(replacementMap);
	}

	void Recompress() override
	{
		// TODO: make this more efficient
		SetTexture(GetTexture());
	}

	void ReadFromFile(std::istream& file)
	{
		// Read the block pointers
		Serializer<std::vector<unsigned32>, unsigned64>::Deserialize(mBlockPointers, file);

		// Read the blocks
		mBlocks.ReadFromFile(file);
	}

	void WriteToFile(std::ostream& file) const override
	{
		// Write the block pointers
		Serializer<std::vector<unsigned32>, unsigned64>::Serialize(mBlockPointers, file);

		// Write the blocks
		mBlocks.WriteToFile(file);
	}

	// Texture pool is taken from how the blocks are stored
	std::vector<unsigned8> GetTexturePool() const override	{ return mBlocks.GetTexturePool(); }
	size_t GetTexturePoolSize() const override 				{ return mBlocks.GetTexturePoolSize(); }

	// Additional texture pool contains for each pixels which block from the blocktexture it is from.
	std::vector<unsigned8> GetAdditionalTexturePool() const
	{
		size_t textureSize = GetAdditionalTexturePoolSize();
		std::vector<unsigned8> additionalPool(textureSize);
		tbb::parallel_for(size_t(0), mBlockPointers.size(), [&](const size_t& i)
		{
			BitHelper::SplitInBytesAndMove(mBlockPointers[i], additionalPool, i * 4);
		});	
		return additionalPool;
	}
	size_t GetAdditionalTexturePoolSize() const
	{
		return mBlockPointers.size() * 4;
	}

	std::map<std::string, std::string> GetAdditionalProperties() const override
	{
		std::map<std::string, std::string> res;
		res.insert(std::pair<std::string, std::string>("blockSize", std::to_string(mBlockSize)));
		
		// Get the additional properties from the block texture
		std::map<std::string, std::string> blockProperties = mBlocks.GetAdditionalProperties();
		for (auto prop : blockProperties)
			res.insert(prop);

		return res;
	};
};