#pragma once
#include "CompressedTexture.h"
#include "../../inc/tbb/parallel_reduce.h"
#include "../../core/BitHelper.h"

template<typename T>
class TightlyPackedTexture : public CompressedTexture<T>
{
private:
	std::vector<unsigned8> mData;
	// Mask containing bits that are set in any of the items
	size_t mOriginalSize;
	unsigned mMask;
	std::vector<unsigned8> mBitMap;

	void SetMask(unsigned32 mask)
	{
		mMask = mask;
		mBitMap = BitHelper::GetBitMapHS(mask);
	}

	inline unsigned32 GetValueAt(size_t i) const { return BitHelper::UnpackTightAt<unsigned8, unsigned32>(mData, i, mBitMap); }
public:
	TightlyPackedTexture() :
			mData(std::vector<unsigned8>()),
			mOriginalSize(0),
			mMask(0),
			mBitMap(std::vector<unsigned8>())
	{
		static_assert(sizeof(T) == 4, "Only types of size 4 can be tightly packed");
	}

	~TightlyPackedTexture() override {}

	T operator[](size_t i) const override
	{
		T value;
		value = GetValueAt(i);
		return value;
	}

	unsigned64 size() const override { return mOriginalSize; }

	std::vector<T> GetTexture(size_t fromIndex = 0) override
	{
		assert(fromIndex <= size());
		//fromIndex = std::min(fromIndex, size());
		std::vector<T> res(mOriginalSize - fromIndex);
		//for (size_t i = fromIndex; i < mOriginalSize; i++)
		tbb::parallel_for(fromIndex, mOriginalSize, [&](size_t i)
		{
			unsigned value = GetValueAt(i);
			res[i - fromIndex] = value;
		});
		return res;
	}

	void SetTexture(const std::vector<T>& nodeMaterialPointers, size_t fromIndex = 0) override 
	{
		if (nodeMaterialPointers.empty()) return;
		if (fromIndex != 0)
		{
			auto curTexture = GetTexture();
			curTexture.resize(fromIndex);
			curTexture.insert(curTexture.end(), nodeMaterialPointers.begin(), nodeMaterialPointers.end());
			SetTexture(curTexture, 0);
			return;
		}

		// Find out which bits are set in both the existing and new texture:
		unsigned mask = 0;
		size_t setBits = BitHelper::GetSet(mMask);
		for (T value : nodeMaterialPointers)
			mask |= static_cast<unsigned32>(value);
		for (size_t i = 0; i < fromIndex; i++)
			mask |= GetValueAt(i);
		// Make sure at least one bit is set (preventing divide by zero and similar problems)
		if (mask == 0) mask = 1;

		// If the mask changed, we need to repack the whole texture
		if (mask != mMask && !mData.empty() && fromIndex != 0)
		{
			auto curTexture = GetTexture();
			curTexture.resize(fromIndex);
			curTexture.insert(curTexture.end(), nodeMaterialPointers.begin(), nodeMaterialPointers.end());
			SetTexture(curTexture, 0);
		}

		// Update the current mask
		SetMask(mask);
		setBits = BitHelper::GetSet(mMask);
		printf("Current texture requires %u bits per pointer\n", (unsigned32)setBits);

		// Calculate how many bits are needed each part
		size_t existingBits = fromIndex * setBits;
		size_t newBits = nodeMaterialPointers.size() * setBits;
		size_t totalBits = existingBits + newBits;

		// If the existing bits don't fit in a fixed number of bytes, we need to repack the first byte of the new set
		unsigned8 bitOffset = (unsigned8)(existingBits & 0x07);
		unsigned8 switchBytes = (bitOffset == 0) ? 0 : 1;

		// The existing- and newbyte counts don't include the switchbyte
		size_t existingBytes = existingBits >> 3;
		size_t totalBytes = (totalBits >> 3) + (((totalBits & 0x07) == 0) ? 0 : 1);

		size_t switchByteIndex = existingBytes;

		// Compress the new data
		std::vector<unsigned8> packed;
		{
			std::vector<unsigned32> unpacked(nodeMaterialPointers.size());
			tbb::parallel_for(size_t(0), nodeMaterialPointers.size(), [&](size_t i) { unpacked[i] = (unsigned32)nodeMaterialPointers[i]; });
			packed = BitHelper::PackTight<unsigned8, unsigned32>(unpacked, mBitMap, bitOffset);
		}

		mData.resize(totalBytes);
		// If the edge between the existing and new materials is not on the edge of a byte,
		// we need to construct a "switchByte", which contains part of the last existing materials, 
		// and part of the first new material:
		if (switchBytes != 0)
			mData[switchByteIndex] = (mData[switchByteIndex] & BitHelper::GetHSMask<unsigned8>(0, bitOffset)) | packed[0];
		// Move the packed data over to mData
		std::move(packed.begin() + switchBytes, packed.end(), mData.begin() + existingBytes + ((size_t)switchBytes));
		mOriginalSize = fromIndex + nodeMaterialPointers.size();
	}

	// Replace all materials. Note that for this to work, all materials currently in the texture need to have a key in the dictionary
	void ReplaceMaterials(const std::unordered_map<T, T>& replacementMap) override
	{
		// Check if the replacements aren't the same as the originals:
		bool replacersEqual = true;
		for (auto replacer : replacementMap)
			if (replacer.first != replacer.second)
			{
				replacersEqual = false; break;
			}
		if (replacersEqual) return;

		 //TODO: Check if the same amount of bits is set
		 //      If that is the case, we can just replace the data. Otherwise, a full repacking is needed
		unsigned newMask = 0;
		for (auto replacer : replacementMap)
			newMask |= (unsigned)replacer.second;
		if (newMask == mMask)
		{ 
			// If the mask doesn't change, we can keep using the same bits, only replacing their values
			unsigned8 setBits = BitHelper::GetSet(mMask);
			std::unordered_map<unsigned, unsigned> bitReplacementMap;
			for (auto value : replacementMap)
			{
				unsigned source = 0;
				unsigned dest = 0;
				for (unsigned8 j = 0; j < setBits; j++)
				{
					if (BitHelper::GetHS((unsigned32)value.first, mBitMap[j])) BitHelper::SetLS(source, setBits - j - 1);
					if (BitHelper::GetHS((unsigned32)value.second, mBitMap[j])) BitHelper::SetLS(dest, setBits - j - 1);
				}
				bitReplacementMap.insert(std::pair<unsigned, unsigned>(source, dest));
			}
			for (size_t i = 0; i < mOriginalSize; i++)
			{
				unsigned source = 0;
				// Find the current value of material i, and replace it with it's replacement map counterpart.
				unsigned8 j = 0;
				while (j < setBits)
				{
					size_t bit = i * setBits + j;
					size_t byte = bit >> 3;
					unsigned8 startBitInByte = bit & 0x07;
					unsigned8 endBitInByte = std::min(startBitInByte + setBits - j, 8);
					unsigned8 bitsInByte = endBitInByte - startBitInByte;
					source |= (unsigned)((mData[byte] & BitHelper::GetHSMask<unsigned8>(startBitInByte, endBitInByte)) >> (8 - endBitInByte)) << (setBits - bitsInByte - j);
					j += bitsInByte;
				}
				// Find the value to replace it with:
				auto replacer = bitReplacementMap.find(source);
				//assert(replacer != bitReplacementMap.end());
				if (replacer != bitReplacementMap.end())
				{
					unsigned dest = replacer->second;
					j = 0;
					while (j < setBits)
					{
						size_t startBit = i * setBits + j;
						size_t byte = startBit >> 3;
						unsigned8 startBitInByte = startBit & 0x07;
						unsigned8 endBitInByte = std::min(startBitInByte + setBits - j, 8);
						unsigned8 bitsInByte = endBitInByte - startBitInByte;

						unsigned8 bitsToCopy = (unsigned8)(dest >> (setBits - (j + bitsInByte)));
						unsigned8 inPlaceBitsToCopy = bitsToCopy << (8 - endBitInByte);

						unsigned8 mask1 = BitHelper::GetHSMask<unsigned8>(0, startBitInByte) | BitHelper::GetHSMask<unsigned8>(endBitInByte, 8);
						unsigned8 mask2 = BitHelper::GetHSMask<unsigned8>(startBitInByte, endBitInByte);
						mData[byte] =
							(mask1 & mData[byte]) |	    // Copy the original bits
							(mask2 & inPlaceBitsToCopy); // Insert the new bits

						j += bitsInByte;
					}
				}
			}
		}
		else
		{
			// If the mask changes, we need to replace the whole texture
			std::vector<T> cur = GetTexture();
			//for (size_t i = 0; i < cur.size(); i++)
			tbb::parallel_for(size_t(0), cur.size(), [&](size_t i)
			{
				auto item = replacementMap.find(cur[i]);
				//assert(item != replacementMap.end());
				if (item != replacementMap.end())
					cur[i] = item->second;
			});
			SetTexture(cur);
		}
	}
	void Recompress() override
	{
		return;
	}

	void ReadFromFile(std::istream& file) 
	{
		mData.clear();
		unsigned64 originalSize;
		Serializer<unsigned64>::Deserialize(originalSize, file);
		mOriginalSize = originalSize;
		unsigned32 mask;
		Serializer<unsigned32>::Deserialize(mask, file);
		SetMask(mask);
		unsigned64 bitCount = ((unsigned64)BitHelper::GetSet(mMask)) * mOriginalSize;
		unsigned64 byteCount = BitHelper::RoundToBytes(bitCount) >> 3;
		mData.resize(byteCount);
		Serializer<unsigned8*>::Deserialize(&mData[0], mData.size(), file);
	}
	void WriteToFile(std::ostream& file) const override
	{
		unsigned64 originalSize = mOriginalSize;
		Serializer<unsigned64>::Serialize(originalSize, file);
		Serializer<unsigned32>::Serialize(mMask, file);
		Serializer<unsigned8*>::Serialize(&mData[0], mData.size(), file);
	}

	// Use this method to output the compressed texture as a vector of one byte characters
	std::vector<unsigned8> GetTexturePool() const override { 
		size_t poolSize = GetTexturePoolSize();
		std::vector<unsigned8> res(poolSize);
		std::copy(mData.begin(), mData.end(), res.begin());
		return res;
	}
	size_t GetTexturePoolSize() const
	{
		return mData.size();
	}

	std::map<std::string, std::string> GetAdditionalProperties() const override
	{
		std::map<std::string, std::string> res;
		res.insert(std::pair<std::string, std::string>("mask", std::to_string(mMask)));
		return res;
	};
};