#include "PVMVoxelizer.h"
#include <algorithm>
#include <numeric>
#include <fstream>
#include "../../PropertyLoader.h"
#include "../OctreeBuilder/SettingsParser.h"
#include "../OctreeBuilder/BaseOctreeBuilder.h"
#include "../../inc/pvm/ddsbase.h"
#include "../../inc/glm/common.hpp"
#include "../ColorHelper.h"
#include "../MathHelper.h"
#include "../StringHelper.h"
#include "../Serializer.h"
#include"../CollectionHelper.h"

struct ColorInterpolator
{
	glm::u8vec3 operator()(const glm::u8vec3& color1, glm::u8vec3& color2, float t)
	{
		//glm::u8vec3 hsv1 = ColorHelper::RGBtoHSV(color1);
		//glm::u8vec3 hsv2 = ColorHelper::RGBtoHSV(color2);
		//glm::u8vec3 hsvRes;
		//hsvRes.r = (unsigned8)MathHelper::lerp(t, (float)hsv1.r, (float)hsv2.r);
		//hsvRes.g = (unsigned8)MathHelper::lerp(t, (float)hsv1.g, (float)hsv2.g);
		//hsvRes.b = (unsigned8)MathHelper::lerp(t, (float)hsv1.b, (float)hsv2.b);
		//return ColorHelper::HSVtoRGB(hsvRes);
		return glm::u8vec3(
			unsigned8(MathHelper::lerp(t, float(color1.r), float(color2.r))),
			unsigned8(MathHelper::lerp(t, float(color1.g), float(color2.g))),
			unsigned8(MathHelper::lerp(t, float(color1.b), float(color2.b)))
			);
	}
};

PVMVoxelizer::PVMVoxelizer() :
	mWidth(0), mHeight(0), mDepth(0), mComponents(0), mScale(0), mVolume(std::vector<unsigned8>()), mMinOpacity(0.3f)
{
	mTransferFunction = TransferFunction<glm::u8vec3>();
	//==============================================
	// Bonsai tree TF:
	//mTransferFunction.AddNode(0, 0, glm::u8vec3(0));
	//mTransferFunction.AddNode(21, 0.f, glm::u8vec3(0, 78, 0));
	//mTransferFunction.AddNode(23, 0.7f, glm::u8vec3(0, 78, 0));
	//mTransferFunction.AddNode(29, 0.7f, glm::u8vec3(118, 97, 71));
	//mTransferFunction.AddNode(255, 1.f, glm::u8vec3(160, 160, 160));
	//===============================================
	// Baby skull TF:
	//mTransferFunction.AddNode(0, 0, glm::u8vec3(0));
	//mTransferFunction.AddNode(50, 0.f, glm::u8vec3(255, 0, 0));
	//mTransferFunction.AddNode(255, 1.f, glm::u8vec3(160, 160, 160));
	//===============================================
	// Christmas tree (lossless):
	// Weird transfer function to make sure that each value has a unique color
	//for (unsigned32 i = 0; i < 256; i++)
	//{
	//	mTransferFunction.AddNode(i * 256, 1.f, glm::u8vec3((i % 2) * 255, (i * 4) % 255, i));
	//}
	//mTransferFunction.AddNode(0, 0.f, glm::u8vec3(0, 0, 0));
	//mTransferFunction.AddNode(4096 * 16, 1.f, glm::u8vec3(255, 255, 255));

	// Christmas tree (beautiful)
	mTransferFunction.AddNode(0, 0, glm::u8vec3(0));
	//mTransferFunction.AddNode(255 * 256, 1.f, glm::u8vec3(255, 255, 255));
	 
	mTransferFunction.AddNode(60 * 256, 0.f, glm::u8vec3(0, 50, 0));
	mTransferFunction.AddNode(90 * 256, 1.f, glm::u8vec3(0, 50, 0));
	mTransferFunction.AddNode(255 * 256, 1.f, glm::u8vec3(255, 255, 0));
	//mTransferFunction.AddNode(100 * 256, 1.f, glm::u8vec3(0, 70, 0));
	//mTransferFunction.AddNode(130 * 256, 1.f, glm::u8vec3(118, 97, 71));
	//mTransferFunction.AddNode(160 * 256, 1.f, glm::u8vec3(255, 255, 0));
	//mTransferFunction.AddNode(255 * 256, 1.f, glm::u8vec3(255, 255, 0));
}
PVMVoxelizer::~PVMVoxelizer() {}

bool PVMVoxelizer::Initialize() { return true; }

bool PVMVoxelizer::LoadScene(const std::string& filename)
{
	
	ReadVolumeData(filename, mWidth, mHeight, mDepth, mComponents, mVolume);
	mScale = BitHelper::Log2Ceil(std::max(std::max(mWidth, mHeight), mDepth));
	// TODO: Error handling
	return true;
}

bool PVMVoxelizer::UnloadScene()
{
	mWidth = mHeight = mDepth = mComponents = mScale = 0;
	mVolume = std::vector<unsigned8>();
	return true;
}


std::vector<glm::uvec3> PVMVoxelizer::GetValidCoords(unsigned8 scale)
{
	if (scale == 0) return std::vector<glm::uvec3>(1, glm::uvec3(0));
	std::vector<glm::uvec3> res;
	Voxelize(scale, scale, glm::uvec3(0), 
		[&](const VoxelInfo& info, bool best) { res.push_back(info.position); },
		[](const std::vector<float>& opacities)->float { return *std::max_element(opacities.begin(), opacities.end()); },
		[](const std::vector<glm::u8vec3>& colors, const std::vector<float>& opacities)->glm::u8vec3 { return glm::u8vec3(0); });
	return res;
}
void PVMVoxelizer::Voxelize(unsigned8 scale, unsigned8 partScale, glm::uvec3 partCoord, const std::function<void(const VoxelInfo&, bool best)>& nodeAdder,
	bool colors, bool normals, bool reflectivity)
{
	Voxelize(scale, partScale, partCoord, nodeAdder, 
		[](const std::vector<float>& opacities)->float { return (float)CollectionHelper::CalculateMean(opacities); },
		[](const std::vector<glm::u8vec3>& colors, const std::vector<float>& opacities)->glm::u8vec3
	{
		float totalOpacity = CollectionHelper::Sum(opacities);
		glm::vec3 avgColor(0);
		for (size_t i = 0; i < colors.size(); i++)
			avgColor += glm::vec3(colors[i]) * (opacities[i] / totalOpacity);
		return glm::u8vec3((unsigned8)avgColor.x, (unsigned8)avgColor.y, (unsigned8)avgColor.z);
	});
}
void PVMVoxelizer::Voxelize(unsigned8 scale, unsigned8 partScale, glm::uvec3 partCoord, 
	const std::function<void(const VoxelInfo&, bool best)>& nodeAdder, 
	const std::function<float(const std::vector<float> values)>& opacityFilter,
	const std::function<glm::u8vec3(const std::vector<glm::u8vec3>, const std::vector<float>)>& colorFilter)
{
	if (scale > mScale)
		printf("Warning: volume data at a lower resolution than the requested scale.");

	float opacity;
	glm::u8vec3 col;
	unsigned valueCount = (unsigned32)BitHelper::Exp2(mComponents * 8);
	std::vector<float> opacityCache(valueCount);
	std::vector<glm::u8vec3> colorCache(valueCount);
	for (unsigned i = 0; i < valueCount; i++)
	{
		mTransferFunction.Evaluate(i, opacity, col, ColorInterpolator());
		opacityCache[i] = opacity;
		colorCache[i] = col;
	}


	unsigned32 resolution = (unsigned32)BitHelper::Exp2(partScale);
	unsigned32 filterSize = mScale >= partScale ? (unsigned32)BitHelper::Exp2(mScale - partScale) : 1;
	unsigned32 maxValue = 0;
	unsigned32 filledCells = 0;
	for (unsigned32 x = 0; x < resolution; x++)
	{
		if (x * filterSize >= mWidth) continue;
		for (unsigned32 y = 0; y < resolution; y++)
		{
			if (y * filterSize >= mHeight) continue;
			for (unsigned32 z = 0; z < resolution; z++)
			{
				if (z * filterSize >= mDepth) continue;
				if (partScale < mScale)
				{
					// Call the filter over the values in the scene part that is enclosed by a single voxel
					std::vector<float> opacities;
					std::vector<glm::u8vec3> colors;
					for (unsigned32 volX = x * filterSize; volX < (x + 1) * filterSize; volX++)
						for (unsigned32 volY = y * filterSize; volY < (y + 1) * filterSize; volY++)
							for (unsigned32 volZ = z * filterSize; volZ < (z + 1) * filterSize; volZ++)
							{
								size_t volIdx = (volX + volY * mWidth + volZ * (mWidth * mHeight)) * mComponents;
								if (volIdx >= mVolume.size()) continue;
								unsigned32 value = 0;
								BitHelper::JoinBytesLittleEndian(mVolume, value, volIdx, mComponents);
								if (value > maxValue) maxValue = value;
								colors.push_back(colorCache[value]);
								opacities.push_back(opacityCache[value]);
							}
					opacity = opacityFilter(opacities);
					if (opacity > mMinOpacity)
						col = colorFilter(colors, opacities);
				}
				else
				{
					glm::uvec3 volCoord(
						x + partCoord.x * resolution,
						y + partCoord.y * resolution,
						z + partCoord.z * resolution);
					size_t volIdx= (volCoord.x + volCoord.y * mWidth + volCoord.z * (mWidth * mHeight)) * mComponents;
					unsigned32 value = 0;
					BitHelper::JoinBytesLittleEndian(mVolume, value, volIdx, mComponents);
					if (value > maxValue) maxValue = value;
					opacity = opacityCache[value];
					col = colorCache[value];
				}
				if (opacity > mMinOpacity)
				{
					filledCells++;
					VoxelInfo info(glm::uvec3(x, z, y), col, glm::vec3(0), 1.f, opacity);
					nodeAdder(info, true);
				}
				
			}
		}
	}
	printf("Max value: %u, filled: %u / %u (%f %%)\n", maxValue, filledCells, mWidth * mHeight * mDepth, (float(filledCells) / float(mWidth * mHeight * mDepth)) * 100.f);
}

void PVMVoxelizer::ReadVolumeData(const std::string& filename, unsigned32& width, unsigned32& height, unsigned32& depth, unsigned32& components, std::vector<unsigned8>& data)
{
	std::string filetype = filename.substr(filename.length() - 3, 3);
	StringHelper::ToUpper(filetype);
	if (filetype == "PVM")
	{
		unsigned char* volumeData = readPVMvolume(filename.c_str(), &width, &height, &depth, &components);
		data = std::vector<unsigned8>(width * height * depth * components);
		for (size_t i = 0; i < data.size(); i++) data[i] = volumeData[i];
		delete volumeData;
	}
	else if (filetype == "DAT")
	{
		std::ifstream file(filename, std::ios::binary);
		unsigned16 width16;
		unsigned16 height16;
		unsigned16 depth16;
		Serializer<unsigned16>::Deserialize(width16, file);
		Serializer<unsigned16>::Deserialize(height16, file);
		Serializer<unsigned16>::Deserialize(depth16, file);
		width = width16;
		height = height16;
		depth = depth16;
		components = 2;
		data = std::vector<unsigned8>(width * height * depth * components);
		Serializer<unsigned8*>::Deserialize(&data[0], width * height * depth * components, file);
		file.close();
	}
}