#include "shaders/splat_btf_shader.h"
#include "shaderc.h"
#include <glm/glm.hpp>
#include <glm/ext.hpp>

using namespace glm;
using namespace std;

SplatBtfShader::SplatBtfShader(PointCloud& pointcloud, Camera& camera, glm::mat4& modelT, float& pSizeT, ColorShade colorShade):
	pointcloud(pointcloud), camera(camera), modelT(modelT), pSizeT(pSizeT), colorShade(colorShade)
{
	generateFootprint();
	splatProgram = ShaderC::LoadShaderProgram("splat_layer.vert", "splat_layer.frag");
	blendProgram = ShaderC::LoadShaderProgram("splat_blend.vert", "splat_blend.frag");

	glGenVertexArrays(1, &vao);
	glGenBuffers(1, &vbo);
	glGenBuffers(1, &ebo);
	glGenTextures(1, &textureId);
	glBindTexture(GL_TEXTURE_2D, textureId);
	glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_R32F, SAMPLE_RES, SAMPLE_RES, 0, GL_RED, GL_FLOAT, &footprint[0]);
	glGenerateMipmap(GL_TEXTURE_2D);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

	glUseProgram(splatProgram);
	glUniform1i(glGetUniformLocation(splatProgram, "fpSample"), 0);

	glUniform1f(glGetUniformLocation(splatProgram, "ambientI"), AMBIENT_I);
	glUniform1f(glGetUniformLocation(splatProgram, "diffuseI"), DIFFUSE_I);
	glUniform1f(glGetUniformLocation(splatProgram, "specularI"), SPECULAR_I);
	glUniform1f(glGetUniformLocation(splatProgram, "specularPower"), SPECULAR_POW);
	vec3 lightDir = glm::normalize(vec3(0, -1, -1));
	glUniform3f(glGetUniformLocation(splatProgram, "lightDir"),
		lightDir.x, lightDir.y, lightDir.z);
	glUseProgram(0);
	loadVertexData(false);
	
	float quadVertices[] = { 
		// positions   // texCoords
		-1.0f,  1.0f,  0.0f, 1.0f,
		-1.0f, -1.0f,  0.0f, 0.0f,
		 1.0f, -1.0f,  1.0f, 0.0f,

		-1.0f,  1.0f,  0.0f, 1.0f,
		 1.0f, -1.0f,  1.0f, 0.0f,
		 1.0f,  1.0f,  1.0f, 1.0f
	};
	glGenVertexArrays(1, &quadVAO);
	glGenBuffers(1, &quadVBO);
	glBindVertexArray(quadVAO);
	glBindBuffer(GL_ARRAY_BUFFER, quadVBO);
	glBufferData(GL_ARRAY_BUFFER, sizeof(quadVertices), &quadVertices, GL_STATIC_DRAW);
	glEnableVertexAttribArray(0);
	glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)0);
	glEnableVertexAttribArray(1);
	glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)(2 * sizeof(float)));

	glGenFramebuffers(1, &framebuffer);
	glGenTextures(1, &textureColor);
	glGenTextures(1, &textureAlphaWeight);
	int viewport[4];
	glGetIntegerv(GL_VIEWPORT, viewport);
	int viewportWidth = viewport[2];
	int viewportHeight = viewport[3];
	
	glBindFramebuffer(GL_FRAMEBUFFER, framebuffer);

	glBindTexture(GL_TEXTURE_2D, textureColor);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB32F, viewportWidth, viewportHeight, 0, GL_RGBA, GL_UNSIGNED_BYTE, NULL);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
	glBindTexture(GL_TEXTURE_2D, 0);

	glBindTexture(GL_TEXTURE_2D, textureAlphaWeight);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RG32F, viewportWidth, viewportHeight, 0, GL_RG, GL_UNSIGNED_BYTE, NULL);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
	glBindTexture(GL_TEXTURE_2D, 0);

	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, textureColor, 0);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1, GL_TEXTURE_2D, textureAlphaWeight, 0);

	const GLenum buffers[] = { GL_COLOR_ATTACHMENT0, GL_COLOR_ATTACHMENT1 };
	glDrawBuffers(2, buffers);
	glBindFramebuffer(GL_FRAMEBUFFER, 0);

	glUseProgram(blendProgram);
	glUniform1i(glGetUniformLocation(blendProgram, "ColorSum"), 0);
	glUniform1i(glGetUniformLocation(blendProgram, "AlphaWeightSum"), 1);
	glUseProgram(0);
}

void SplatBtfShader::generateFootprint()
{
	for (int i = 0; i < SAMPLE_RES; i++) {
		for (int j = 0; j < SAMPLE_RES; j++) {
			int ind = i * SAMPLE_RES + j;
			float xoff = (float)j / (float)(SAMPLE_RES - 1);
			float yoff = (float)i / (float)(SAMPLE_RES - 1);
			vec2 v = vec2((-0.5f + xoff) * 2.0f, (-0.5f + yoff) * 2.0f);
			float x = glm::dot(v, v);
			footprint[ind] = glm::exp(-(x*x) / 0.2f);
		}
	}
}

void SplatBtfShader::loadVertexData(bool realloc)
{
	const vector<Point>& pc = pointcloud.Points();
	vertexBuffer.clear();
	if (!realloc)vertexBuffer.reserve(pc.size());
	for (uint32_t i = 0; i < pc.size(); i++) {
		Point p = pc[i];
		VertexData vd;
		vd.position = p.position;
		vd.size = p.pointSize;
		switch (colorShade.color) {
		case ColorMode::WHITE:
			vd.color = vec3(1, 1, 1);
			break;
		case ColorMode::COLOR:
			vd.color = p.color;
			break;
		case ColorMode::LABEL:
			vec3 lc(0.2f, 0.2f, 0.2f);
			lc = COLOR_LABEL[p.label % MAX_LABELS];
			if (p.label < MAX_LABELS) {
			}
			vd.color = lc;
			break;
		case ColorMode::PRED:
			vec3 pc(0.2f, 0.2f, 0.2f);
			pc = COLOR_LABEL[p.part % MAX_LABELS];
			if (p.part < MAX_LABELS) {
			}
			vd.color = pc;
			break;
		case ColorMode::NORMAL:
			vd.color = glm::abs(p.normal);
			break;
		case ColorMode::CURVATURE:
			vd.color = vec3(1, 1, 1) * p.curvature / 10.0f;
			break;
		}
		if (colorShade.shade == ShadeMode::SHADE) {
			vd.normal = p.normal;
		}
		vertexBuffer.push_back(vd);
	}

	if (realloc) {
		glBindVertexArray(vao);
		glBindBuffer(GL_ARRAY_BUFFER, vbo);
		glBufferSubData(GL_ARRAY_BUFFER, 0, sizeof(VertexData) * vertexBuffer.size(), &vertexBuffer[0]);
	}
	else {
		glBindVertexArray(vao);
		glBindBuffer(GL_ARRAY_BUFFER, vbo);
		glBufferData(GL_ARRAY_BUFFER, sizeof(VertexData) * vertexBuffer.size(), &vertexBuffer[0], GL_STATIC_DRAW);
		glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo);
		glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(unsigned int) * vertexBuffer.size(), nullptr, GL_STREAM_DRAW);

		glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(VertexData), (void*)0);
		glEnableVertexAttribArray(0);
		glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, sizeof(VertexData), (void*)(3 * sizeof(float)));
		glEnableVertexAttribArray(1);
		glVertexAttribPointer(2, 3, GL_FLOAT, GL_FALSE, sizeof(VertexData), (void*)(6 * sizeof(float)));
		glEnableVertexAttribArray(2);
		glVertexAttribPointer(3, 1, GL_FLOAT, GL_FALSE, sizeof(VertexData), (void*)(9 * sizeof(float)));
		glEnableVertexAttribArray(3);

		glBindVertexArray(0);
	}
}

void SplatBtfShader::ChangeColorShade(ColorShade colorShade)
{
	this->colorShade = colorShade;
	loadVertexData(true);
}

void SplatBtfShader::ReloadPointCloud()
{
	loadVertexData(true);
}


void SplatBtfShader::Begin()
{
	glEnable(GL_VERTEX_PROGRAM_POINT_SIZE);
	glPointParameteri(GL_POINT_SPRITE_COORD_ORIGIN, GL_LOWER_LEFT);
	glDisable(GL_DEPTH_TEST);
	glBindVertexArray(vao);
}


void SplatBtfShader::Draw()
{
	glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
	updateUniforms();
	if (accurateSort) {
		pointcloud.Sort(cameraDir, pointcloud.AverageDist());
	}
	else {
		//pointcloud.ApproxSort(cameraDir, pointcloud.AverageDist());
		pointcloud.Sort(cameraDir, pointcloud.AverageDist());
	}
	
	const vector<PointCloudSlice>& slices = pointcloud.Slices();
	
	for (auto s : slices) {
		setLayerElements(s.begin, s.count);
		drawLayer();
	}
	
	/*const vector<Point>& pc = pointcloud.Points();
	vector<uint32_t> elements;
	for (uint32_t i = 0; i < pc.size(); i++) {
		elements.push_back(i);
	}
	setLayerElements(&elements[0], elements.size());
	drawLayer();*/
}

void SplatBtfShader::setLayerElements(unsigned int* elements, unsigned int n)
{
	elementCount = n;
	glBindVertexArray(vao);
	glBufferSubData(GL_ELEMENT_ARRAY_BUFFER, 0, sizeof(unsigned int) * n, elements);
	glBindVertexArray(0);
}

void SplatBtfShader::drawLayer()
{
	glEnable(GL_BLEND);

	glBindFramebuffer(GL_FRAMEBUFFER, framebuffer);
	glBlendFunc(GL_ONE, GL_ONE);
	glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
	glClear(GL_COLOR_BUFFER_BIT);
	glUseProgram(splatProgram);
	glBindVertexArray(vao);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_2D, textureId);
	glDrawElements(GL_POINTS, elementCount, GL_UNSIGNED_INT, 0);

	glBindFramebuffer(GL_FRAMEBUFFER, 0);
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
	glUseProgram(blendProgram);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_2D, textureColor);
	glActiveTexture(GL_TEXTURE1);
	glBindTexture(GL_TEXTURE_2D, textureAlphaWeight);
	glBindVertexArray(quadVAO);
	glDrawArrays(GL_TRIANGLES, 0, 6);
	glDisable(GL_BLEND);
}


void inline SplatBtfShader::updateUniforms()
{
	mat4 view = camera.GetViewMatrix();
	mat4 projection = camera.GetProjMatrix();
	mat4 viewDir = camera.GetViewDirMatrix();
	glUseProgram(splatProgram);
	glUniformMatrix4fv(glGetUniformLocation(splatProgram, "model"), 1, GL_FALSE, value_ptr(modelT));
	glUniformMatrix4fv(glGetUniformLocation(splatProgram, "view"), 1, GL_FALSE, value_ptr(view));
	glUniformMatrix4fv(glGetUniformLocation(splatProgram, "projection"), 1, GL_FALSE, value_ptr(projection));
	glUniformMatrix4fv(glGetUniformLocation(splatProgram, "viewDir"), 1, GL_FALSE, value_ptr(viewDir));
	glUniform1f(glGetUniformLocation(splatProgram, "psizet"), pSizeT);
	if (colorShade.shade == ShadeMode::SHADE) {
		glUniform1i(glGetUniformLocation(splatProgram, "shading"), true);
	}
	else {
		glUniform1i(glGetUniformLocation(splatProgram, "shading"), false);
	}
	cameraDir.x = view[0][2];
	cameraDir.y = view[1][2];
	cameraDir.z = view[2][2];
	glUseProgram(0);
}

void SplatBtfShader::ChangeSortMode(bool accurate)
{
	accurateSort = accurate;
}

SplatBtfShader::~SplatBtfShader()
{
	glDeleteBuffers(1, &vbo);
	glDeleteVertexArrays(1, &vao);
	glDeleteProgram(splatProgram);
	glDeleteTextures(1, &textureId);
}