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

using namespace glm;
using namespace std;

SplatFtbShader::SplatFtbShader(PointCloud& pointcloud, Camera& camera, glm::mat4& modelT, float& pSizeT, ColorShade colorShade) :
	pointcloud(pointcloud), camera(camera), modelT(modelT), pSizeT(pSizeT), colorShade(colorShade)
{
	generateFootprint();
	generateShaderProgram();
	generateFramebuffers();
	

	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)));

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

void SplatFtbShader::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 SplatFtbShader::generateShaderProgram()
{
	splatProgram = ShaderC::LoadShaderProgram("splat_layer.vert", "splat_layer.frag");
	blendProgram = ShaderC::LoadShaderProgram("splat_blend.vert", "splat_blend2.frag");
	stencilProgram = ShaderC::LoadShaderProgram("splat_stencil.vert", "splat_stencil.frag");
	backgroundProgram = ShaderC::LoadShaderProgram("texture_blend.vert", "texture_blend.frag");
}

void SplatFtbShader::generateFramebuffers()
{
	int viewport[4];
	glGetIntegerv(GL_VIEWPORT, viewport);
	int viewportWidth = viewport[2];
	int viewportHeight = viewport[3];

	//layer frame buffer
	glGenFramebuffers(1, &layerFramebuffer);
	glGenTextures(1, &textureColor);
	glGenTextures(1, &textureAlphaWeight);
	glBindFramebuffer(GL_FRAMEBUFFER, layerFramebuffer);

	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);

	glGenRenderbuffers(1, &layerStencil);
	glBindRenderbuffer(GL_RENDERBUFFER, layerStencil);
	glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, viewportWidth, viewportHeight);
	glFramebufferRenderbuffer(
		GL_FRAMEBUFFER, GL_STENCIL_ATTACHMENT, GL_RENDERBUFFER, layerStencil
	);

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

	//composite frame buffer
	glGenFramebuffers(1, &compositeFramebuffer);
	glGenTextures(1, &compositeColor);
	glBindFramebuffer(GL_FRAMEBUFFER, compositeFramebuffer);

	glBindTexture(GL_TEXTURE_2D, compositeColor);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, 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);

	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, compositeColor, 0);

	glGenRenderbuffers(1, &compositeStencil);
	glBindRenderbuffer(GL_RENDERBUFFER, compositeStencil);
	glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, viewportWidth, viewportHeight);
	glFramebufferRenderbuffer(
		GL_FRAMEBUFFER, GL_STENCIL_ATTACHMENT, GL_RENDERBUFFER, compositeStencil
	);

	const GLenum buffers2[] = { GL_COLOR_ATTACHMENT0 };
	glDrawBuffers(1, buffers2);

	glBindFramebuffer(GL_FRAMEBUFFER, 0);
}

void SplatFtbShader::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 SplatFtbShader::ChangeColorShade(ColorShade colorShade)
{
	this->colorShade = colorShade;
	loadVertexData(true);
}

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


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


void SplatFtbShader::Draw()
{
	glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
	glBindFramebuffer(GL_FRAMEBUFFER, compositeFramebuffer);
	glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
	glClear(GL_COLOR_BUFFER_BIT);
	glStencilMask(0xFF);
	glClearStencil(0);
	glClear(GL_STENCIL_BUFFER_BIT);

	glBindFramebuffer(GL_FRAMEBUFFER, layerFramebuffer);
	glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
	glClear(GL_COLOR_BUFFER_BIT);
	glStencilMask(0xFF);
	glClearStencil(0);
	glClear(GL_STENCIL_BUFFER_BIT);

	glEnable(GL_BLEND);
	glEnable(GL_STENCIL_TEST);
	glDisable(GL_DEPTH_TEST);

	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();
	}

	glDisable(GL_STENCIL_TEST);

	glBindFramebuffer(GL_FRAMEBUFFER, 0);
	glBlendFunc(GL_ONE, GL_SRC_ALPHA);
	glUseProgram(backgroundProgram);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_2D, compositeColor);
	glBindVertexArray(quadVAO);
	glDrawArrays(GL_TRIANGLES, 0, 6);

	glDisable(GL_BLEND);
}

void SplatFtbShader::setLayerElements(unsigned int* elements, unsigned int n)
{
	elementCount = n;
	glBindVertexArray(vao);
	glBufferSubData(GL_ELEMENT_ARRAY_BUFFER, 0, sizeof(unsigned int) * n, elements);
	glBindVertexArray(0);
	//std::cout << "ftb";
}

void SplatFtbShader::drawLayer()
{
	//glDisable(GL_STENCIL_TEST);
	//step 1: draw individual layer & write stencil
	glBindFramebuffer(GL_FRAMEBUFFER, layerFramebuffer);
	glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
	glClear(GL_COLOR_BUFFER_BIT);
	glStencilFunc(GL_NOTEQUAL, 0xFF, 0xF0);
	glStencilOp(GL_KEEP, GL_KEEP, GL_REPLACE);
	glStencilMask(0x0F);
	glBlendFunc(GL_ONE, GL_ONE);
	glUseProgram(splatProgram);
	glBindVertexArray(vao);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_2D, textureId);
	glDrawElements(GL_POINTS, elementCount, GL_UNSIGNED_INT, 0);

	//step 2: sync stencil
	int viewport[4];
	glGetIntegerv(GL_VIEWPORT, viewport);
	int viewportWidth = viewport[2];
	int viewportHeight = viewport[3];
	glBindFramebuffer(GL_DRAW_FRAMEBUFFER, compositeFramebuffer);
	glBindFramebuffer(GL_READ_FRAMEBUFFER, layerFramebuffer);
	glBlitFramebuffer(0, 0, viewportWidth, viewportHeight,
		0, 0, viewportWidth, viewportHeight,
		GL_STENCIL_BUFFER_BIT, GL_NEAREST);

	//step 3: compose
	glBindFramebuffer(GL_FRAMEBUFFER, compositeFramebuffer);
	glStencilFunc(GL_EQUAL, 0xFF, 0x0F);
	glStencilOp(GL_KEEP, GL_KEEP, GL_KEEP);
	glStencilMask(0x00);
	glBlendEquation(GL_FUNC_ADD);
	glBlendFuncSeparate(GL_DST_ALPHA, GL_ONE, GL_ZERO, 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);

	//step 3: write stencil
	glStencilFunc(GL_EQUAL, 0xFF, 0x0F);
	glStencilOp(GL_KEEP, GL_KEEP, GL_REPLACE);
	glStencilMask(0xF0);
	glUseProgram(stencilProgram);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_2D, compositeColor);
	glBindVertexArray(quadVAO);
	glDrawArrays(GL_TRIANGLES, 0, 6);

	//step 4: clear stencil
	glStencilMask(0x0F);
	glClearStencil(0);
	glClear(GL_STENCIL_BUFFER_BIT);

	//step 5: sync stencil
	glBindFramebuffer(GL_DRAW_FRAMEBUFFER, layerFramebuffer);
	glBindFramebuffer(GL_READ_FRAMEBUFFER, compositeFramebuffer);
	glBlitFramebuffer(0, 0, viewportWidth, viewportHeight,
		0, 0, viewportWidth, viewportHeight,
		GL_STENCIL_BUFFER_BIT, GL_NEAREST);
}


void inline SplatFtbShader::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 SplatFtbShader::ChangeSortMode(bool accurate)
{
	accurateSort = accurate;
}

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