#pragma once

#include "CompressedTexture.h"
#include "../Octree/MultiRootTree.h"
#include "../../core/Serializer.h"
#include "../../core/MathHelper.h"
#include "../PoolBuilder/StandardPoolBuilder.h"

template<typename T>
class MultiRootBasedTexture : public CompressedTexture<T>
{
private:
	MultiRootTree<>* mData;
	std::vector<unsigned8> mBitMap;
	unsigned32 mMask;
	unsigned64 mOriginalSize;

	glm::uvec3 GetTextureCoord(size_t i) const
	{
		size_t nodeIndex = i / 8;
		unsigned8 depth = mData->GetMaxLevel();
		size_t mask = BitHelper::GetLSMask<size_t>(0, depth - 1);
		glm::uvec3 leafIndex(
			(nodeIndex & mask) << 1,
			((nodeIndex >> (depth - 1)) & mask) << 1,
			((nodeIndex >> ((depth - 1) * 2)) & mask) << 1);
		if (BitHelper::GetLS(i, 2)) leafIndex.z++;
		if (BitHelper::GetLS(i, 1)) leafIndex.y++;
		if (BitHelper::GetLS(i, 0)) leafIndex.x++;
		return leafIndex;
	}

	static unsigned8 GetRequiredDepth(size_t nodeCount)
	{
		unsigned8 indexBits = BitHelper::Log2Ceil(nodeCount);
		return indexBits / 3 + ((indexBits % 3 == 0) ? 0 : 1);
	}
public:
	MultiRootBasedTexture() :
			mData(NULL),
			mBitMap(std::vector<unsigned8>()),
			mMask(0),
			mOriginalSize(0)
	{}

	~MultiRootBasedTexture() { delete mData; }

	T operator[](size_t i) const override {
		assert(i < mOriginalSize);
		glm::uvec3 textureCoord = GetTextureCoord(i);
		unsigned32 res = 0;
		for (unsigned32 bit = 0; bit < (unsigned32)mBitMap.size(); bit++)
		{
			bool bitSet;
			if (bit == 0)
				bitSet = mData->HasLeaf(textureCoord);
			else
				bitSet = mData->SlaveHasLeaf(textureCoord, bit - 1);
			if (bitSet) BitHelper::SetHS(res, mBitMap[bit]);
		}
		T resVal;
		resVal = res;
		return resVal;
	}
	unsigned64 size() const override { return mOriginalSize; }

	std::vector<T> GetTexture(size_t fromIndex = 0) override {
		std::vector<T> data(mOriginalSize - fromIndex);
		for (size_t i = fromIndex; i < mOriginalSize; i++) { data[i - fromIndex] = operator[](i); }
		return data;
	}
	// Compresses the texture, replacing everything after fromIndex by whatever is in texture
	void SetTexture(const std::vector<T>& texture, size_t fromIndex = 0) override
	{
		// TODO: correctly handle the fromIndex
		if (fromIndex != 0)
		{
			// Hack to handle fromindex: just unpack and repack the whole texture
			std::vector<T>& oldTexture = GetTexture();
			// Copy the new texture
			oldTexture.resize(fromIndex + texture.size());
			std::copy(texture.begin(), texture.end(), oldTexture.begin() + fromIndex);
			// Build the DAG
			SetTexture(oldTexture);
			return;
		}
		else if (mData != NULL)
		{
			delete mData;
			mData = NULL;
		}

		assert(mData == NULL); // Make sure we are working with a fresh tree, since appending does not work correctly yet

		// Figure out which bits are set:
		unsigned32 mask = 0;
		for (auto mat : texture)
			mask |= (unsigned32)mat;
		
		mMask = mask;
		unsigned8 bitCount = BitHelper::GetSet(mask);
		mBitMap = BitHelper::GetBitMapHS(mask);

		unsigned8 depth = GetRequiredDepth(fromIndex + texture.size());

		if (mData == NULL)
		{
			mData = new MultiRootTree<>(depth, bitCount - 1);
		}
		
		for (size_t i = 0; i < texture.size(); i++)
		{
			glm::uvec3 texCoord = GetTextureCoord(i);
			unsigned32 texel = (unsigned32)texture[i];
			for (unsigned32 bit = 0; bit < (unsigned32)mBitMap.size(); bit++)
			{
				if (BitHelper::GetHS(texel, mBitMap[bit]))
				{
					if (bit == 0) mData->AddLeafNode(texCoord);
					else mData->AddLeafNode(texCoord, bit - 1);
				}
			}
		}
		mData->ToDAG();

		mOriginalSize = fromIndex + texture.size();
	}

	// Replace all materials by other colors. This could be useful as some methods don't need to fully decompress and compress the whole texture to do this.
	void ReplaceMaterials(const std::unordered_map<T, T>& replacers) override
	{
		// TODO: implement this...
		std::vector<T> texture = GetTexture();
		delete mData;
		mData = NULL;
		mOriginalSize = 0;
		mMask = 0;
		mBitMap.clear();
		SetTexture(texture);
	}
	void Recompress() override { return; }

	void ReadFromFile(std::istream& file) override
	{
		Serializer<unsigned64>::Deserialize(mOriginalSize, file);
		Serializer<unsigned32>::Deserialize(mMask, file);
		unsigned8 depth = GetRequiredDepth(mOriginalSize);
		unsigned8 slaveRootCount = BitHelper::GetSet(mMask) - 1;
		if (mData != NULL) delete mData;
		mData = new MultiRootTree<>(depth, slaveRootCount);
		mData->Deserialize(file);
	}
	virtual void WriteToFile(std::ostream& file) const override
	{
		Serializer<unsigned64>::Serialize(mOriginalSize, file);
		Serializer<unsigned32>::Serialize(mMask, file);
		mData->Serialize(file);
	}

	// Use this method to output the compressed texture as a vector of one byte characters
	std::vector<unsigned8> GetTexturePool() const override
	{
		StandardPoolBuilder builder;
		std::vector<unsigned8> res;
		builder.BuildPool(mData, res);
		return res;
	}
	size_t GetTexturePoolSize() const override 
	{
		StandardPoolBuilder builder;
		return builder.GetPoolSize(mData);
	}

	virtual std::map<std::string, std::string> GetAdditionalProperties() const override { return std::map<std::string, std::string>(); }
};