#pragma once
#include "../../core/Defines.h"
#include "../../core/Serializer.h"
#include "../../core/CollectionHelper.h"
#include "MaterialLibraryPointer.h"
#include "../../inc/tbb/parallel_sort.h"
#include <vector>
#include <iterator>

// T should be some type with a method called "Serialize()". 
// This method should return some type that can be iterated over using operator [], and each item of the iteration should be an unsigned char (unsigned8) to put in the texture colors.
// If sizeof(result) for the value is used, it should give the number of bytes needed to store in pixels.
// An example of a working return type is std::vector<unsigned8>, glm::u8vec3, and unsigned char[].
template<typename T, typename Comparer = std::less<T>, unsigned8 channelsPerPixel = 3>
class MaterialLibrary
{
private:
	std::map<T, MaterialLibraryPointer, Comparer> mMaterialPointers;
	bool finalized = false;
	unsigned short mTextureSize;

protected:
	const size_t SIZE_OF_MATERIAL = sizeof(T);
	const unsigned32 PIXELS_PER_MATERIAL = (unsigned32)(SIZE_OF_MATERIAL / channelsPerPixel + (SIZE_OF_MATERIAL % channelsPerPixel == 0 ? 0 : 1));

	std::vector<T>* mMaterials;

	inline MaterialLibraryPointer WrapSetIndex(size_t setIndex) const
	{
		assert(finalized);
		setIndex *= PIXELS_PER_MATERIAL;
		return MaterialLibraryPointer(
			(unsigned16)(setIndex % (size_t)mTextureSize),
			(unsigned16)(setIndex / (size_t)mTextureSize)
			);
	}

	inline size_t GetSetIndex(MaterialLibraryPointer textureIndex) const { assert(finalized); return (textureIndex.x + textureIndex.y * mTextureSize) / PIXELS_PER_MATERIAL; }

	inline void RebuildMaterialPointers()
	{
		assert(finalized);
		mMaterialPointers.clear();
		// Build the result material pointers map:
		for (size_t i = 0; i < mMaterials->size(); i++)
			mMaterialPointers.insert(std::pair<T, MaterialLibraryPointer>(mMaterials->at(i), WrapSetIndex(i)));
	}
public:
	MaterialLibrary() :
			mMaterialPointers(std::map<T, MaterialLibraryPointer, Comparer>()),
			finalized(false),
			mTextureSize(0),
			mMaterials(new std::vector<T>())
	{}

	MaterialLibrary(std::vector<unsigned char> texture, unsigned short textureSize, MaterialLibraryPointer highestMaterialIndex) : MaterialLibrary()
	{
		ReadTexture(texture, textureSize, highestMaterialIndex);
	}

	// Copy constructor
	MaterialLibrary(const MaterialLibrary& other) :
			mMaterialPointers(std::map<T, MaterialLibraryPointer, Comparer>(other.mMaterialPointers)),
			finalized(other.finalized),
			mTextureSize(other.mTextureSize),
			mMaterials(new std::vector<T>(*other.mMaterials))
	{}

	~MaterialLibrary() { delete mMaterials; }
	
	void Serialize(std::ostream& file)
	{
		unsigned16 materialTextureSize = GetTextureSize();
		MaterialLibraryPointer maxTextureIndex = GetMaxTextureIndex();
		std::vector<unsigned8> texture = GetTexture();

		// The material texture size and max texture index used to be stored in one 32 bit unsigned integer.
		// However, this caused problems if the material texture contained more then 1024 colors. 
		// Therefore, we now use 48 bits for this, precluded by an empty header in the old style, to ensure that the new style can be recognized correctly.
		Serializer<unsigned32>::Serialize(0, file); 
		Serializer<unsigned16>::Serialize(materialTextureSize, file);
		Serializer<MaterialLibraryPointer>::Serialize(maxTextureIndex, file);
		Serializer<unsigned8*>::Serialize(&texture[0], (size_t)materialTextureSize * (size_t)materialTextureSize * (size_t)channelsPerPixel, file);
	}

	void Deserialize(std::istream& file)
	{
		unsigned16 materialTextureSize = 0;
		MaterialLibraryPointer maxTextureIndex = MaterialLibraryPointer(0);

		// If the first unsigned32 of a material library is 0, then the new style information is placed after it. 
		// If it isn't 0, then use the old style
		unsigned materialTextureSizeSummary;
		Serializer<unsigned>::Deserialize(materialTextureSizeSummary, file);
		if (materialTextureSizeSummary == 0)
		{
			// New style (supports bigger libraries)
			Serializer<unsigned16>::Deserialize(materialTextureSize, file);
			Serializer<MaterialLibraryPointer>::Deserialize(maxTextureIndex, file);
		}
		else
		{
			// Old style
			unsigned mask1 = BitHelper::GetLSMask<unsigned32>(20, 30);
			unsigned mask2 = BitHelper::GetLSMask<unsigned32>(10, 20);
			unsigned mask3 = BitHelper::GetLSMask<unsigned32>(0, 10);
			maxTextureIndex.x = (mask1 & materialTextureSizeSummary) >> 20;
			maxTextureIndex.y = (mask2 & materialTextureSizeSummary) >> 10;
			materialTextureSize = BitHelper::CeilToNearestPowerOfTwo(mask3 & materialTextureSizeSummary);
		}

		size_t textureArraySize = (size_t)materialTextureSize * (size_t)materialTextureSize * (size_t)channelsPerPixel;
		std::vector<unsigned8> texture(textureArraySize);
		Serializer<unsigned8*>::Deserialize(&texture[0], textureArraySize, file);
		ReadTexture(texture, materialTextureSize, maxTextureIndex);
	}



	void AddMaterial(const T& material) { 
		if(!finalized) mMaterials->push_back(material);
	}
	void RemoveMaterial(const T& material) { if (!finalized) mMaterials->erase(material); }
	void Finalize(bool filterUnique = true) 
	{ 
		if (finalized) return;

		if (filterUnique)
			CollectionHelper::Unique(*mMaterials, Comparer());

		mTextureSize = GetTextureSize();
		finalized = true;

		RebuildMaterialPointers();
	}
	bool IsFinalized() const { return finalized; }
	std::vector<T> GetMaterials() const
	{
		std::vector<T> materials(mMaterials->size());
		tbb::parallel_for(size_t(0), mMaterials->size(), [&](size_t i)
		{
			materials[i] = mMaterials->at(i);
		});
		return materials;
	}

	T GetNearestMaterial(T material) const { 
		return NearestFinder<T>()(material, *mMaterials);
	}
	T GetMaterial(MaterialLibraryPointer materialPointer)
	{
		assert(finalized);
		return mMaterials->at(GetSetIndex(materialPointer));
	}
	std::vector<std::pair<T, MaterialLibraryPointer>> GetMaterialTextureIndices() const
	{
		assert(finalized);
		std::vector<std::pair<T, MaterialLibraryPointer>> materials;
		size_t setIndex = 0;
		for (T material : (*mMaterials))
		{
			auto textureIndex = WrapSetIndex(setIndex++);
			materials.push_back(std::pair<T, MaterialLibraryPointer>(material, textureIndex));
		}
		return materials;
	}
	MaterialLibraryPointer GetTextureIndex(const T& material) const
	{
		assert(finalized);
		auto materialIt = mMaterialPointers.find(material);
		if (materialIt == mMaterialPointers.end())
			return GetTextureIndex(GetNearestMaterial(material));
		return materialIt->second;
	}

	bool Contains(const T& material) const
	{
		auto materialIt = mMaterialPointers.find(material);
		return materialIt != mMaterialPointers.end();
	}

	MaterialLibraryPointer GetMaxTextureIndex() const
	{
		return WrapSetIndex(mMaterials->size());
	}

	unsigned short GetTextureSize() const {
		if (finalized) return mTextureSize;
		if (mMaterials->empty())
			return 0;
		size_t requiredPixels = PIXELS_PER_MATERIAL * mMaterials->size();
		if (requiredPixels < 4) requiredPixels = 4; // Make sure the texture is always 2x2
		unsigned16 v = (unsigned16)std::ceil(std::sqrt(double(requiredPixels)));
		return BitHelper::CeilToNearestPowerOfTwo(v);
	}

	std::vector<unsigned8> GetTexture() const
	{
		std::vector<unsigned8> texture(mTextureSize * mTextureSize * channelsPerPixel);
		size_t i = 0;
		for (T material : (*mMaterials))
		{
			auto materialProperties = material.Serialize();
			unsigned props = (unsigned)materialProperties.size();
			unsigned prop = 0;
			assert(i + channelsPerPixel * PIXELS_PER_MATERIAL <= texture.size());

			for (unsigned pixel = 0; pixel < PIXELS_PER_MATERIAL; pixel++)
			{
				for (unsigned8 channel = 0; channel < channelsPerPixel; channel++)
				{
					if (prop < props)
						texture[i + channel] = materialProperties[prop];
					prop++;
				}
				i += channelsPerPixel;
			}
		}
		return texture;
	}

	void ReadTexture(std::vector<unsigned char> texture, unsigned textureSize, MaterialLibraryPointer maxTextureIndex)
	{
		mTextureSize = textureSize;
		unsigned props = PIXELS_PER_MATERIAL * channelsPerPixel;
		std::vector<unsigned8> curMatProps(SIZE_OF_MATERIAL);
		unsigned maxIndex = maxTextureIndex.x + maxTextureIndex.y * textureSize;
		bool endFound = maxIndex != 0;
		bool firstEmptyMaterialFound = false;
		size_t i = 0;
		while (!(endFound && i >= maxIndex * channelsPerPixel))
		{
			// Read the materials from the current pixels
			for (size_t j = 0; j < SIZE_OF_MATERIAL; j++)
				curMatProps[j] = texture[i + j];
			if (!endFound)
			{
				bool allZero = true;
				for (unsigned char prop : curMatProps)
					if (prop != 0)
					{
						allZero = false;
						break;
					}
				if (firstEmptyMaterialFound && allZero) endFound = true;
				firstEmptyMaterialFound |= allZero;
			}

			T curMat;
			curMat.Deserialize(curMatProps);
			AddMaterial(curMat);

			i += props;
		}
		finalized = true;
		RebuildMaterialPointers();
	}

	bool operator==(const MaterialLibrary& other) const
	{
		// Check if the material count and finalized state are equal
		if (this->IsFinalized() != other.IsFinalized() 
			|| this->mMaterials->size() != other.mMaterials->size())
			return false;

		// Check if the materials are equal
		for (size_t i = 0; i < mMaterials->size(); i++)
			if (this->mMaterials->at(i) != other.mMaterials->at(i))
				return false;
		return true;
	}

	bool operator!=(const MaterialLibrary& other) const
	{
		return !(*this == other);
	}
};