#include "RayTracer.h"
#include "WhittedRenderer.h"
#include "PathRenderer.h"
#include "VRUT_NsightCustomEvents.h"
using namespace VRUT;

#ifdef USE_TEMPORAL
void RayTracer::resize() {
	vector<OglTexTarget*> textars;
	for (int i = 0; i < 2; i++) {
		textars.clear();

		texReproj[i].intFormat = GL_RGBA16F;
		texReproj[i].format = GL_RGBA;
		texReproj[i].type = GL_FLOAT;
		texReproj[i].magFilter = GL_NEAREST;
		texReproj[i].minFilter = GL_NEAREST;
		texReproj[i].clampBorder = true;

		texAlbedo[i].intFormat = GL_RGBA16F;
		texAlbedo[i].format = GL_RGBA;
		texAlbedo[i].type = GL_FLOAT;
		texAlbedo[i].magFilter = GL_NEAREST;
		texAlbedo[i].minFilter = GL_NEAREST;
		texAlbedo[i].clampBorder = true;

		texND[i].intFormat = GL_RGBA16F;
		texND[i].format = GL_RGBA;
		texND[i].type = GL_FLOAT;
		texND[i].magFilter = GL_NEAREST;
		texND[i].minFilter = GL_NEAREST;
		texND[i].clampBorder = true;

		texVar[i].intFormat = GL_RGBA16F;
		texVar[i].format = GL_RGBA;
		texVar[i].type = GL_FLOAT;
		texVar[i].magFilter = GL_NEAREST;
		texVar[i].minFilter = GL_NEAREST;
		texVar[i].clampBorder = true;

		textars.push_back(&texReproj[i]);
		textars.push_back(&texAlbedo[i]);
		textars.push_back(&texND[i]);
		textars.push_back(&texVar[i]);

		bufferReproj[i].init(textars, 0, texWidth, texHeight);
	}

	int coeff[] = { 3,9,27,1 };
	for (int i = 0; i < 4; i++) {
		for (int ii = 0; ii < 2; ii++) {
			if (ii && i > 1)
				continue;

			texNDF[ii][i].intFormat = GL_RGBA16F;
			texNDF[ii][i].format = GL_RGBA;
			texNDF[ii][i].type = GL_FLOAT;
			texNDF[ii][i].magFilter = GL_NEAREST;
			texNDF[ii][i].minFilter = GL_NEAREST;
			if (i == 2 || ii) {
				texNDF[ii][i].magFilter = GL_LINEAR;
				texNDF[ii][i].minFilter = GL_LINEAR;
			}
			texNDF[ii][i].clampBorder = true;

			texAlbedoF[ii][i].intFormat = GL_RGBA16F;
			texAlbedoF[ii][i].format = GL_RGBA;
			texAlbedoF[ii][i].type = GL_FLOAT;
			texAlbedoF[ii][i].magFilter = GL_NEAREST;
			texAlbedoF[ii][i].minFilter = GL_NEAREST;
			if (i == 2 || ii) {
				texAlbedoF[ii][i].magFilter = GL_LINEAR;
				texAlbedoF[ii][i].minFilter = GL_LINEAR;
			}
			texAlbedoF[ii][i].clampBorder = true;

			textars.clear();
			textars.push_back(&texNDF[ii][i]);
			textars.push_back(&texAlbedoF[ii][i]);
			bufferND[ii][i].init(textars, 0, texWidth / coeff[i], texHeight / coeff[i]);
		}
	}

	for (int i = 0; i < 2; i++) {
		textars.clear();
		texVarBlur[i].intFormat = GL_RGBA16F;
		texVarBlur[i].format = GL_RGBA;
		texVarBlur[i].type = GL_FLOAT;
		texVarBlur[i].magFilter = GL_NEAREST;
		texVarBlur[i].minFilter = GL_NEAREST;
		texVarBlur[i].clampBorder = true;

		textars.push_back(&texVarBlur[i]);
		bufferVarBlur[i].init(textars, 0, texWidth, texHeight);
	}

	textars.clear();
	texWaveletMemory.intFormat = GL_RGBA16F;
	texWaveletMemory.format = GL_RGBA;
	texWaveletMemory.type = GL_FLOAT;
	texWaveletMemory.magFilter = GL_NEAREST;
	texWaveletMemory.minFilter = GL_NEAREST;
	texWaveletMemory.clampBorder = true;

	textars.push_back(&texWaveletMemory);
	bufferWaveletMemory.init(textars, 0, texWidth, texHeight);

	for (int i = 0; i < WAVELEVELS; i++) {
		textars.clear();
		texWavelet[i].intFormat = GL_RGBA16F;
		texWavelet[i].format = GL_RGBA;
		texWavelet[i].type = GL_FLOAT;
		texWavelet[i].magFilter = GL_NEAREST;
		texWavelet[i].minFilter = GL_NEAREST;
		//if (!i) {
		//	texWavelet[i].magFilter = GL_NEAREST;
		//	texWavelet[i].minFilter = GL_NEAREST;
		//}
		texWavelet[i].clampBorder = true;

		textars.push_back(&texWavelet[i]);
		bufferWavelet[i].init(textars, 0, texWidth, texHeight);
	}

	for (int i = 0; i < 2; i++) {
		textars.clear();
		texTemp[i].intFormat = GL_RGBA16F;
		texTemp[i].format = GL_RGBA;
		texTemp[i].type = GL_FLOAT;
		texTemp[i].magFilter = GL_NEAREST;
		texTemp[i].minFilter = GL_NEAREST;
		texTemp[i].clampBorder = true;

		textars.push_back(&texTemp[i]);
		bufferTemp[i].init(textars, 0, texWidth, texHeight);
	}

	GLCHECKERROR(wxT("RayTracer::buffer "));
}

void RayTracer::initGl() {
	//hack to speedup glDrawPixels from http://www.cs.rit.edu/~ncs/Courses/570/UserGuide/OpenGLonWin-31.html
	glDisable(GL_BLEND);
	glDisable(GL_DEPTH_TEST);
	glDisable(GL_DITHER);
	glDisable(GL_STENCIL_TEST);

	resize();
	nowTex = 0; frame = 0;

	//prepare for QUAD drawing
	for (int tex = 0; tex < TEXTURES; tex++) {
		texCPU[tex].init(texWidth, texHeight);
	}

	glPixelStorei(GL_UNPACK_ALIGNMENT, 1);

#if PBO_COUNT > 0
	for (int tex = 0; tex < TEXTURES; tex++)
		glGenBuffers(PBO_COUNT, pboIds[tex]);
#endif
#if USE_GL_FILTER
	wxString shaderPrefix = KERNEL->GetAppPath() + wxT("modules/render/");
	//wxString shaderPrefix = wxT("../base_modules/RayTracer/shaders/");

	//Initialize VBO, VAO, IBO
	float quad[] = { //pos, tcoord
		1, -1, 1, 0,
		1,  1, 1, 1,
		-1, -1, 0, 0,
		-1,  1, 0, 1
	};

	if (glDrawArraysIndirect)
		useIndirectDraw = true;

	if (glBindVertexArray)
		useVertexArray = true;

	glGenBuffers(1, &fsquadVBO);
	glBindBuffer(GL_ARRAY_BUFFER, fsquadVBO);
	glBufferData(GL_ARRAY_BUFFER, 16 * sizeof(float), quad, GL_STATIC_DRAW);
	glBindBuffer(GL_ARRAY_BUFFER, 0);

	if (useVertexArray)
	{
		glGenVertexArrays(1, &fsquadVAO);
		glBindVertexArray(fsquadVAO);
		glEnableVertexAttribArray(0);
		glEnableVertexAttribArray(1);
		glBindBuffer(GL_ARRAY_BUFFER, fsquadVBO);
		glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), 0);
		glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)(2 * sizeof(float)));
		glBindVertexArray(0);
	}

	if (useIndirectDraw)
	{
		typedef  struct {
			GLuint  count;
			GLuint  instanceCount;
			GLuint  first;
			GLuint  baseInstance;
		} DrawArraysIndirectCommand;

		DrawArraysIndirectCommand c;
		c.baseInstance = 0;
		c.count = 4;
		c.first = 0;
		c.instanceCount = 1;

		glGenBuffers(1, &fsquadIBO);
		glBindBuffer(GL_DRAW_INDIRECT_BUFFER, fsquadIBO);
		glBufferData(GL_DRAW_INDIRECT_BUFFER, sizeof(DrawArraysIndirectCommand), &c, GL_STATIC_READ);
		glBindBuffer(GL_DRAW_INDIRECT_BUFFER, 0);
	}


	shV.compile(shaderPrefix + "rt_filter.vert");
	shF.compile(shaderPrefix + "rt_filter.frag");
	shProg.link(shV, shF);

	shVQuad.compile(shaderPrefix + "Quad.vsh");
	shFQuad.compile(shaderPrefix + "Quad.fsh");
	shProgQuad.link(shVQuad, shFQuad);

	shFReproject.compile(shaderPrefix + "Reproject.fsh");
	shProgReproject.link(shVQuad, shFReproject);

	shFReduceND.compile(shaderPrefix + "ReduceND.fsh");
	shProgReduceND.link(shVQuad, shFReduceND);

	shFVarBlur.compile(shaderPrefix + "VarBlur.fsh");
	shProgVarBlur.link(shVQuad, shFVarBlur);

	shFExpandND.compile(shaderPrefix + "ExpandND.fsh");
	shProgExpandND.link(shVQuad, shFExpandND);

	shFWavelet.compile(shaderPrefix + "WaveletBlur.fsh");
	shProgWavelet.link(shVQuad, shFWavelet);

	shFTemp.compile(shaderPrefix + "TempAA.fsh");
	shProgTemp.link(shVQuad, shFTemp);

	shProg.enable();
	shProg.importValue("filterRadius", (float)filterRadius);
	shProg.importValue("frameId", (unsigned int)0);
	shProg.importValue("keyValue", (float)((float)toneMappingKeyValue) * 1.0e-2f);
	shProg.importValue("whitePoint", (float)((float)toneMappingWhitePoint) * 1.0e-2f);
	shProg.importValue("texDownscaleRatio", 1.0f);
	shProg.importValue("imageTexture", &texCPU[0]);
	shProg.importValue("resolution", Vector2(1, 1));

#else
	shader = nullptr;
#endif
	GLCHECKERROR(wxT("RayTracer::Initialize "));

	wxString oglVendor(wxString::FromAscii((char*)glGetString(GL_VENDOR)));
	if (oglVendor.Find(wxT("ATI")) != wxNOT_FOUND || oglVendor.Find(wxT("AMD")) != wxNOT_FOUND)
		openGLVendor = AMD;

}

void RayTracer::renderGl() {
	//wxULongLong lastTime = VRUTGetTime();

	if (useTemporalPathTracing) {
		bindCPUTexture(0);
		subImage2D();

		bindCPUTexture(1);
		subImage2D();

		bindCPUTexture(2);
		subImage2D();



		bufferReproj[nowTex].enable();
		glViewport(0, 0, texWidth, texHeight);
		prepReproject();
		if (filterRadius == 4)
			shProgReproject.importValue("rejectNow", 1);
		drawQuad();



		//reduce
		bufferND[0][0].enable();
		glViewport(0, 0, texWidth / 3, texHeight / 3);
		prepReduceND(&texND[nowTex], &texAlbedo[nowTex], 1);
		drawQuad();

		bufferND[0][1].enable();
		glViewport(0, 0, texWidth / 9, texHeight / 9);
		prepReduceND(&texNDF[0][0], &texAlbedoF[0][0], 3);
		drawQuad();

		bufferND[0][2].enable();
		glViewport(0, 0, texWidth / 27, texHeight / 27);
		prepReduceND(&texNDF[0][1], &texAlbedoF[0][1], 9);
		drawQuad();

		//expand
		bufferND[1][1].enable();
		glViewport(0, 0, texWidth / 9, texHeight / 9);
		prepExpandND(&texNDF[0][2], &texNDF[0][1], &texAlbedoF[0][2], &texAlbedoF[0][1]);
		drawQuad();

		bufferND[1][0].enable();
		glViewport(0, 0, texWidth / 3, texHeight / 3);
		prepExpandND(&texNDF[1][1], &texNDF[0][0], &texAlbedoF[1][1], &texAlbedoF[0][0]);
		drawQuad();

		bufferND[0][3].enable();
		glViewport(0, 0, texWidth, texHeight);
		prepExpandND(&texNDF[1][0], &texND[nowTex], &texAlbedoF[1][0], &texAlbedo[nowTex]);
		drawQuad();



		bufferVarBlur[0].enable();
		glViewport(0, 0, texWidth, texHeight);
		prepVarBlur(&texVar[nowTex],1);
		drawQuad();

		/*bufferVarBlur[1].enable();
		glViewport(0, 0, texWidth, texHeight);
		prepVarBlur(&texVarBlur[0],2);
		drawQuad();*/



		float time = wxGetLocalTimeMillis().ToDouble();
		bufferWavelet[0].enable();
		glViewport(0, 0, texWidth, texHeight);
		prepWavelet(&texVarBlur[0]/*&texReproj[nowTex]*/, &texNDF[0][3], 0, time);
		drawQuad();

		for (int i = 0; i < WAVELEVELS - 1; i++) {
			bufferWavelet[i + 1].enable();
			glViewport(0, 0, texWidth, texHeight);
			prepWavelet(&texWavelet[i], &texNDF[0][3], WAVELEVELS - 1 - i, time);
			drawQuad();
		}

		bufferTemp[nowTex].enable();
		glViewport(0, 0, texWidth, texHeight);
		prepTemp();
		drawQuad();

		glBindFramebuffer(GL_FRAMEBUFFER, 0);
		glViewport(0, 0, texWidth, texHeight);
		//prepQuad(&texWavelet[WAVELEVELS-1]);
		//prepQuad(&texNDF[0][3]);
		//prepQuad(&texCPU[0]);
		prepQuad(&texTemp[nowTex]);

		if (filterRadius == 4)
			prepQuad(&texReproj[nowTex]);
		if (filterRadius == 5)
			prepQuad(&texCPU[1]);
		if (filterRadius == 6)
			prepQuad(&texCPU[2]);
		if (filterRadius == 7)
			prepQuad(&texWavelet[0]);
		if (filterRadius == 8)
			prepQuad(&texWavelet[1]);
		if (filterRadius == 9)
			prepQuad(&texWavelet[2]);
		if (filterRadius == 10)
			prepQuad(&texWavelet[3]);
		if (filterRadius == 11)
			prepQuad(&texWavelet[4]);
		if (filterRadius == 12)
			prepQuad(&texNDF[0][3]);
		if (filterRadius == 13)
			prepQuad(&texAlbedoF[0][3]);
		if (filterRadius == 14)
			prepQuad(&texReproj[nowTex]);
		if (filterRadius == 15)
			prepQuad(&texTemp[nowTex]);
		if (filterRadius == 16)
			prepQuad(&texReproj[nowTex]);
		if (filterRadius == 17)
			prepQuad(&texWavelet[0]);
		if (filterRadius == 18)
			prepQuad(&texWavelet[1]);
		if (filterRadius == 19)
			prepQuad(&texWavelet[2]);
		if (filterRadius == 20)
			prepQuad(&texWavelet[3]);
		if (filterRadius == 21)
			prepQuad(&texWavelet[4]);
		if (filterRadius == 22)
			prepQuad(&texReproj[nowTex]);
		if (filterRadius == 23)
			prepQuad(&texVarBlur[0]);
		//prepQuad(&texVarBlur[1]);
		drawQuad();

		nowTex = !nowTex;
		frame++;

		/*bindCPUTexture(0);
		subImage2D();

		glBindFramebuffer(GL_FRAMEBUFFER, 0);
		prepFilter(0);
		drawQuad();*/
	}
	else {
		glBindFramebuffer(GL_FRAMEBUFFER,0);
		bindCPUTexture(0);
		subImage2D();

		prepFilter(0);
		drawQuad();

		//nowTex = !nowTex;
		//frame++;
	}

	// Render from FBO to HMD
	if (openVRWrapper != nullptr)
		drawFromFBO();

#if PBO_COUNT == 2
	for (int i = 0; i < texturesUsed; i++) {
		if (/*!reset &&*/ texData[i])
		{
			PUSH_RANGE("DrawPixels glMapBuffer", 4);
			if (!glMapBufferRange)
				texData[i] = (float*)glMapBuffer(GL_PIXEL_UNPACK_BUFFER, GL_WRITE_ONLY);
			POP_RANGE;
		}
	}

#endif //PBO_COUNT == 2
}


void RayTracer::prepQuad(int tex) {
	shProgQuad.enable();
	shProgQuad.importValue("tex", &texCPU[tex]);
	shProgQuad.importValue("showAlbedo", showAlbedo);
	//shProgQuad.importValue("leaveProb", 0.0f);
}

void RayTracer::prepQuad(OglTex *tex) {
	shProgQuad.enable();
	shProgQuad.importValue("tex", tex);
	shProgQuad.importValue("texAlbedo", &texAlbedoF[0][3]);
	shProgQuad.importValue("showAlbedo", showAlbedo);

	float avgLuminance = renderer->GetAvgLuminance();
	shProgQuad.importValue("keyValue", ((float)toneMappingKeyValue) * 1.0e-2f);
	shProgQuad.importValue("whitePoint", ((float)toneMappingWhitePoint) * 1.0e-2f);
	shProgQuad.importValue("avgLuminance", avgLuminance);

	//shProgQuad.importValue("leaveProb", 0.0f);
}

void RayTracer::prepReproject() {
	shProgReproject.enable();
	shProgReproject.importValue("texDiff", &texCPU[0]);
	shProgReproject.importValue("texAlbedo", &texCPU[1]);
	shProgReproject.importValue("texNormalDepth", &texCPU[2]);

	shProgReproject.importValue("texDiffPrev", &texReproj[!nowTex]);
	if (!(frame % (TILE_SIZE*TILE_SIZE / 16)) && frame && (filterRadius < 16 || filterRadius > 21) ) {
		shProgReproject.importValue("texDiffFilteredPrev", /*&texWaveletMemory*/&texWavelet[0]);
		shProgReproject.importValue("fullNew", 1);
	}
	else {
		shProgReproject.importValue("texDiffFilteredPrev", &texReproj[!nowTex]);
		shProgReproject.importValue("fullNew", 0);
	}
	shProgReproject.importValue("texAlbedoPrev", &texAlbedo[!nowTex]);
	shProgReproject.importValue("texNDPrev", &texND[!nowTex]);
	shProgReproject.importValue("texVarPrev", &texVar[!nowTex]);

	shProgReproject.importValue("rejectNow", rerender ? 1 : 0);
	//shProgReproject.importValue("rejectNow", 1);

	shProgReproject.importValue("frameId", (unsigned int)((renderer->frameId >> 1) & 0xFFFF));
	shProgReproject.importValue("resolution", Vector2(w, h));
}

void RayTracer::prepReduceND(OglTex *from, OglTex *fromAlbedo, int step) {
	shProgReduceND.enable();
	shProgReduceND.importValue("texNormalDepth", from);
	shProgReduceND.importValue("texAlbedo", fromAlbedo);
	shProgReduceND.importValue("resolution", Vector2(w / step, h / step));
}

void RayTracer::prepExpandND(OglTex *prev, OglTex *now, OglTex *prevAlbedo, OglTex *nowAlbedo) {
	shProgExpandND.enable();
	shProgExpandND.importValue("texNormalDepthPrev", prev);
	shProgExpandND.importValue("texNormalDepthNow", now);

	shProgExpandND.importValue("texAlbedoPrev", prevAlbedo);
	shProgExpandND.importValue("texAlbedoNow", nowAlbedo);
}

void RayTracer::prepVarBlur(OglTex *src, int level) {
	shProgVarBlur.enable();
	shProgVarBlur.importValue("texVar", src);
	shProgVarBlur.importValue("tex", &texReproj[nowTex]);
	shProgVarBlur.importValue("level", level);

	shProgVarBlur.importValue("resolution", Vector2(w, h));
}

void RayTracer::prepWavelet(OglTex *src, OglTex *srcND, int l, float time) {
	shProgWavelet.enable();
	shProgWavelet.importValue("tex", src);
	//shProgWavelet.importValue("firstVarTex", &texVar);
	shProgWavelet.importValue("texNormalDepth", srcND);
	shProgWavelet.importValue("texVar", &texVar[nowTex]);
	shProgWavelet.importValue("res", Vector2(w, h));
	shProgWavelet.importValue("time", time);
	shProgWavelet.importValue("iter", (float)pow(2, l));

	float debug = filterRadius;
	shProgWavelet.importValue("debug", debug);

	if (filterRadius == 22)
		shProgWavelet.importValue("debug", 1.0f);
	else
		shProgWavelet.importValue("debug", 0.0f);
}

void RayTracer::prepTemp() {
	shProgTemp.enable();
	shProgTemp.importValue("tex", &texWavelet[WAVELEVELS - 1]);
	shProgTemp.importValue("prevTex", &texTemp[!nowTex]);
	shProgTemp.importValue("texAlbedo", &texAlbedoF[0][3]);
	shProgTemp.importValue("showAlbedo", showAlbedo);
	shProgTemp.importValue("res", Vector2(w, h));
	shProgTemp.importValue("time", (float)wxGetLocalTimeMillis().ToDouble());
	shProgTemp.importValue("rejectNow", rerender ? 1 : 0);

	float avgLuminance = renderer->GetAvgLuminance();
	shProgTemp.importValue("keyValue", ((float)toneMappingKeyValue) * 1.0e-2f);
	shProgTemp.importValue("whitePoint", ((float)toneMappingWhitePoint) * 1.0e-2f);
	shProgTemp.importValue("avgLuminance", avgLuminance);
}

void RayTracer::prepFilter(int tex) {
#if USE_GL_FILTER
	shProg.enable();
	if (filterRadius < 0)
	{
		// automatic filter radius
		float avgSamples = renderer->GetAvgSamples();
		float radius = 3;
		if ((avgSamples >= 1))
		{
			radius = avgSamples * std::abs(filterRadius) + 2;
			LOG(wxString::Format(wxT("radius %f"), radius));
		}
		shProg.importValue("filterRadius", radius);
	}
	else
		shProg.importValue("filterRadius", (float)filterRadius);
	float avgLuminance = renderer->GetAvgLuminance();

	shProg.importValue("avgLuminance", avgLuminance);
	shProg.importValue("frameId", (unsigned int)((renderer->frameId >> 1) & 0xFFFF));

	shProg.importValue("keyValue", ((float)toneMappingKeyValue) * 1.0e-2f);
	shProg.importValue("whitePoint", ((float)toneMappingWhitePoint) * 1.0e-2f);
	shProg.importValue("texDownscaleRatio", float(texDownscaleFactor));
	shProg.importValue("imageTexture", &texCPU[tex]);
	shProg.importValue("resolution", Vector2(w, h));
#endif // USE_GL_FILTER
}

#endif