#include "App.h"


BaseApp *app = new App();

void App::moveCamera(const float3 &dir){
	float3 newPos = camPos + dir * (speed * frameTime);

	float3 point;
	const BTri *tri;
	if (bsp.intersects(camPos, newPos, &point, &tri)){
		newPos = point + tri->plane.xyz();
	}
	bsp.pushSphere(newPos, 20);

	camPos = newPos;
}

void App::resetCamera(){
	camPos = float3(-670, -30, 0);
	wx = 0;
	wy = -PI / 2;
}

bool App::onMouseButton(const int x, const int y, const MouseButton button, const bool pressed){
	if (D3DApp::onMouseButton(x, y, button, pressed)) return true;

	if (button == MOUSE_LEFT){
		if (pressed){
			lastMouseTime = time;
		} else {
			if (lastMouseTime > 0){
				lightPos = camPos;
				lightRadius = 50 * (time - lastMouseTime);

				drawLightmap();
				lastMouseTime = 0;
			}
		}

		return true;
	}

	return false;
}

bool App::init(){
	lightPos = float3(0, 400, 0);
	lightRadius = 50;
	lightRange = 2500;
	lastMouseTime = 0;

	map = new Model();
	if (!map->loadT3d("../Models/Corridor/Map.t3d")) return false;
	map->scale(0, float3(-1, 1, 1));
	map->removeStream(4);
	map->removeStream(3);
	map->removeStream(2);
	map->computeTangentSpace(true);

	TexturePacker texPacker;

	uint nIndices = map->getIndexCount();
	float3 *src = (float3 *) map->getStream(0).vertices;
	uint *inds = map->getStream(0).indices;

	for (uint i = 0; i < nIndices; i += 6){
		float3 v0 = src[inds[i]];
		float3 v1 = src[inds[i + 1]];
		float3 v3 = src[inds[i + 5]];

		float w = length(v1 - v0) / 8;
		float h = length(v3 - v0) / 8;

		uint tw = (int) w;
		uint th = (int) h;

		if (tw < 8) tw = 8;
		if (th < 8) th = 8;

		tw = (tw + 4) & ~7;
		th = (th + 4) & ~7;

		texPacker.addRectangle(tw, th);
	}

	lmWidth  = 512;
	lmHeight = 512;

	if (!texPacker.assignCoords(&lmWidth, &lmHeight, widthComp)){
		ErrorMsg("Lightmap too small");
		return false;
	}

//	tw = getUpperPowerOfTwo(tw);
//	th = getUpperPowerOfTwo(th);



	uint nLmCoords = (nIndices / 6) * 4;
	float2 *lmCoords = new float2[nLmCoords];
	uint *lmIndices = new uint[nIndices];

	uint index = 0;
	for (uint i = 0; i < nIndices; i += 6){
		TextureRectangle *rect = texPacker.getRectangle(index);

		float x0 = float(rect->x + 0.5f) / lmWidth;
		float y0 = float(rect->y + 0.5f) / lmHeight;
		float x1 = float(rect->x + rect->width  - 0.5f) / lmWidth;
		float y1 = float(rect->y + rect->height - 0.5f) / lmHeight;

		lmCoords[4 * index + 0] = float2(x0, y0);
		lmCoords[4 * index + 1] = float2(x1, y0);
		lmCoords[4 * index + 2] = float2(x1, y1);
		lmCoords[4 * index + 3] = float2(x0, y1);

		lmIndices[i + 0] = 4 * index;
		lmIndices[i + 1] = 4 * index + 1;
		lmIndices[i + 2] = 4 * index + 2;
		lmIndices[i + 3] = 4 * index;
		lmIndices[i + 4] = 4 * index + 2;
		lmIndices[i + 5] = 4 * index + 3;

		float3 v0 = src[inds[i]];
		float3 v1 = src[inds[i + 1]];
		float3 v2 = src[inds[i + 2]];
		float3 v3 = src[inds[i + 5]];

		bsp.addTriangle(v0, v1, v2);
		bsp.addTriangle(v0, v2, v3);

		index++;
	}
	bsp.build();

	map->addStream(TYPE_TEXCOORD, 2, nLmCoords, (float *) lmCoords, lmIndices, true);


	Image positionMapImage[SAMPLE_COUNT];
	if (!positionMapImage[0].loadDDS("PositionMap0.dds")){
		float2 positions[SAMPLE_COUNT] = {
			float2( 0.375,  0.125),
			float2( 0.125, -0.375),
			float2(-0.375, -0.125),
			float2(-0.125,  0.375),
		};

		float4 *pMap[SAMPLE_COUNT];
		for (int i = 0; i < SAMPLE_COUNT; i++){
			pMap[i] = (float4 *) positionMapImage[i].create(FORMAT_RGBA32F, lmWidth, lmHeight, 1, 1);
		}

		index = 0;
		for (uint i = 0; i < nIndices; i += 6){
			TextureRectangle *rect = texPacker.getRectangle(index);

			float3 dirS = (src[inds[i + 1]] - src[inds[i]]);
			float3 dirT = (src[inds[i + 5]] - src[inds[i]]);
			float3 pos = src[inds[i]];
			float3 normal = cross(dirS, dirT);

			for (int t = 0; t < int(rect->height); t++){
				for (int s = 0; s < int(rect->width); s++){
					int index = ((rect->y + t) * lmWidth + (rect->x + s));

					for (int i = 0; i < SAMPLE_COUNT; i++){
						float3 samplePos = pos +
							(saturate((float(s) / (rect->width  - 1))) + positions[i].x / (rect->width  - 1)) * dirS +
							(saturate((float(t) / (rect->height - 1))) + positions[i].y / (rect->height - 1)) * dirT;
						bsp.pushSphere(samplePos, 5);
						bsp.pushSphere(samplePos, 4.9f);
						bsp.pushSphere(samplePos, 4.8f);

						pMap[i][index] = float4(samplePos, 0);
					}
				}
			}
			index++;
		}

		for (int i = 0; i < SAMPLE_COUNT; i++){
			positionMapImage[i].convert(FORMAT_RGBA16F);


			char str[256];
			sprintf(str, "PositionMap%d.dds", i);
			positionMapImage[i].saveDDS(str);
		}
	}

	map->cleanUp();

	return true;
}

void App::exit(){
	delete map;
}

bool App::load(){
	if (caps.PixelShaderVersion < D3DPS_VERSION(2,0) || caps.VertexShaderVersion < D3DVS_VERSION(1,1)){
		ErrorMsg("This demo requires ps2.0 and vs1.1.");
		return false;
	}

	// Sampler states
	if ((nearestClamp = renderer->addSamplerState(NEAREST, CLAMP, CLAMP, CLAMP)) == SS_NONE) return false;
	if ((bilinearClamp = renderer->addSamplerState(BILINEAR, CLAMP, CLAMP, CLAMP)) == SS_NONE) return false;
	if ((trilinearAniso = renderer->addSamplerState(TRILINEAR_ANISO, WRAP, WRAP, WRAP)) == SS_NONE) return false;



	// Render targets
	if (SUCCEEDED(d3d->CheckDeviceFormat(D3DADAPTER_DEFAULT, D3DDEVTYPE_HAL, D3DFMT_X8R8G8B8, 0, D3DRTYPE_TEXTURE, D3DFMT_R16F))){
		if ((shadowMap = renderer->addRenderTarget(256, 256, FORMAT_R16F, nearestClamp, CUBEMAP)) == TEXTURE_NONE) return false;
	} else {
		// Nvidia doesn't support R16F, so we'll use RG16F
		if ((shadowMap = renderer->addRenderTarget(256, 256, FORMAT_RG16F, nearestClamp, CUBEMAP)) == TEXTURE_NONE) return false;
	}
	if ((shadowMapDepth = renderer->addRenderDepth(256, 256, 16)) == TEXTURE_NONE) return false;

	if ((rtLightMap[0] = renderer->addRenderTarget(lmWidth, lmHeight, FORMAT_RGBA16, linearClamp)) == TEXTURE_NONE) return false;
	if ((rtLightMap[1] = renderer->addRenderTarget(lmWidth, lmHeight, FORMAT_RGBA16, linearClamp)) == TEXTURE_NONE) return false;


	// Shaders
	if ((lighting = renderer->addShader("lighting.shd")) == SHADER_NONE) return false;
	if ((light = renderer->addShader("light.shd")) == SHADER_NONE) return false;
	if ((depth = renderer->addShader("depth.shd")) == SHADER_NONE) return false;
	if ((lightMapping = renderer->addShader("lightMapping.shd")) == SHADER_NONE) return false;
	if ((lightMappingBlend = renderer->addShader("lightMapping.shd", "#define BLEND\n")) == SHADER_NONE) return false;

	// Textures
	if ((base[0] = renderer->addTexture("../Textures/floor_wood_3.dds", true, trilinearAniso)) == TEXTURE_NONE) return false;
	if ((base[1] = renderer->addTexture("../Textures/stone34.dds",      true, trilinearAniso)) == TEXTURE_NONE) return false;
	if ((base[2] = renderer->addTexture("../Textures/floor_wood_4.dds", true, trilinearAniso)) == TEXTURE_NONE) return false;

	if ((bump[0] = renderer->addNormalMap("../Textures/floor_wood_3Bump.dds", FORMAT_RGBA8, true, trilinearAniso)) == TEXTURE_NONE) return false;
	if ((bump[1] = renderer->addNormalMap("../Textures/stone34Bump.dds",      FORMAT_RGBA8, true, trilinearAniso)) == TEXTURE_NONE) return false;
	if ((bump[2] = renderer->addNormalMap("../Textures/floor_wood_4Bump.dds", FORMAT_RGBA8, true, trilinearAniso)) == TEXTURE_NONE) return false;

	if ((lightTex = renderer->addTexture("../Textures/Corona.dds", true, bilinearClamp)) == TEXTURE_NONE) return false;

	for (int i = 0; i < SAMPLE_COUNT; i++){
		char str[256];
		sprintf(str, "PositionMap%d.dds", i);
		if ((positionMap[i] = renderer->addTexture(str, false, nearestClamp)) == TEXTURE_NONE) return false;
	}

	// Vertex declaration
	FormatDesc quadFormat[] = { 0, TYPE_VERTEX, FORMAT_FLOAT, 2 };
	if ((quadVF = renderer->addVertexFormat(quadFormat, elementsOf(quadFormat))) == VF_NONE) return false;

	FormatDesc lightFormat[] = {
		0, TYPE_VERTEX,   FORMAT_FLOAT, 3,
		0, TYPE_TEXCOORD, FORMAT_FLOAT, 2,
	};
	if ((lightVF = renderer->addVertexFormat(lightFormat, elementsOf(lightFormat))) == VF_NONE) return false;

	if ((blendAdd = renderer->addBlendState(ONE, ONE)) == BS_NONE) return false;

	// Upload map to vertex/index buffer
	if (!map->makeDrawable(renderer, true)) return false;

	currRTLightMap = 0;
	drawLightmap();

	return true;
}

void App::unload(){

}

bool App::onReset(){
	currRTLightMap = 0;
	drawLightmap();
	
	return true;
}

void App::drawLight(const float3 &dx, const float3 &dy){
	renderer->reset();
	renderer->setShader(light);
	renderer->setTexture("Tex", lightTex, bilinearClamp);
	renderer->setVertexFormat(lightVF);
	renderer->setBlendState(blendAdd);
	renderer->apply();

	struct Vertex {
		float3 position;
		float2 texCoord;
	};

	Vertex vertices[4];

	vertices[0].position = lightPos + lightRadius * (-dx + dy);
	vertices[1].position = lightPos + lightRadius * ( dx + dy);
	vertices[2].position = lightPos + lightRadius * ( dx - dy);
	vertices[3].position = lightPos + lightRadius * (-dx - dy);

	vertices[0].texCoord = float2(0, 0);
	vertices[1].texCoord = float2(1, 0);
	vertices[2].texCoord = float2(1, 1);
	vertices[3].texCoord = float2(0, 1);

	dev->DrawPrimitiveUP(D3DPT_TRIANGLEFAN, 2, vertices, sizeof(Vertex));
}

void App::drawLightmap(){
	dev->BeginScene();

	mat4 proj = cubeProjectionMatrixD3D(1, 5000);

	float3 samples[LS_COUNT];
	samples[0] = lightPos;
	for (int i = 1; i < LS_COUNT; i++){
		float3 v;
		do {
			do {
				v.x = float(rand()) / RAND_MAX * 2.0f - 1.0f;
				v.y = float(rand()) / RAND_MAX * 2.0f - 1.0f;
				v.z = float(rand()) / RAND_MAX * 2.0f - 1.0f;
			} while (dot(v, v) > 1.0f);

			samples[i] = lightPos + lightRadius * v;
		} while (!bsp.isInOpenSpace(samples[i]));
		bsp.pushSphere(samples[i], 5.0f);
		bsp.pushSphere(samples[i], 4.9f);
		bsp.pushSphere(samples[i], 4.8f);
	}



	for (int i = 0; i < LS_COUNT; i++){
		/*
			Shadow map pass
		*/
		for (int face = 0; face < 6; face++){
			renderer->changeRenderTarget(shadowMap, shadowMapDepth, face);

			mat4 mvp = proj * cubeViewMatrix(face);

			renderer->reset();
			renderer->setShader(depth);
			renderer->setRasterizerState(cullBack);
			renderer->setShaderConstant4x4f("mvp", mvp * translate(-samples[i]));
			renderer->setShaderConstant3f("lightPos", samples[i]);
			renderer->setShaderConstant1f("invRadius", 1.0f / lightRange);
			renderer->apply();

			dev->Clear(0, NULL, D3DCLEAR_ZBUFFER, 0, 1.0f, 0);

			map->draw(renderer);
		}

		int otherRT = currRTLightMap;
		currRTLightMap = 1 - otherRT;


		/*
			Lightmap pass using the shadow map
		*/
		renderer->changeRenderTarget(rtLightMap[currRTLightMap]);

		renderer->reset();
		if (i == 0){
			renderer->setShader(lightMapping);
		} else {
			renderer->setShader(lightMappingBlend);
			renderer->setTexture("Target", rtLightMap[otherRT], nearestClamp);
		}
		renderer->setVertexFormat(quadVF);
		renderer->setTexture("PosMap0", positionMap[0], nearestClamp);
		renderer->setTexture("PosMap1", positionMap[1], nearestClamp);
		renderer->setTexture("PosMap2", positionMap[2], nearestClamp);
		renderer->setTexture("PosMap3", positionMap[3], nearestClamp);
		renderer->setTexture("ShadowMap", shadowMap, nearestClamp);
		renderer->setDepthState(noDepthTest);
		renderer->setShaderConstant3f("lightPos", samples[i]);
		renderer->setShaderConstant2f("halfPixel", float2(0.5f / lmWidth, 0.5f / lmHeight));
		renderer->setShaderConstant1f("invRadius", 1.0f / lightRange);
		renderer->setShaderConstant1f("weight", 1.0f / LS_COUNT);
		renderer->apply();

		const float2 vertices[] = { float2(0, 2), float2(3, -1), float2(-3, -1) };
		dev->DrawPrimitiveUP(D3DPT_TRIANGLELIST, 1, vertices, sizeof(float2));
	}

	renderer->changeToMainFramebuffer();

	dev->EndScene();
}

void App::drawFrame(){
	mat4 projection = toD3DProjection(perspectiveMatrixX(1.5f, width, height, 10, 7000));
	mat4 modelview = rotateXY(-wx, -wy) * translate(-camPos);
	mat4 mvp = projection * modelview;

    dev->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, 0, 1.0f, 0);

	for (uint i = 0; i < map->getBatchCount(); i++){
		renderer->reset();
		renderer->setRasterizerState(cullBack);
		renderer->setShader(lighting);
		renderer->setTexture("Base", base[i], trilinearAniso);
		renderer->setTexture("Bump", bump[i], trilinearAniso);
		renderer->setTexture("LightMap", rtLightMap[currRTLightMap], linearClamp);
		renderer->setShaderConstant4x4f("mvp", mvp);
		renderer->setShaderConstant3f("camPos", camPos);
		renderer->setShaderConstant3f("lightPos", lightPos);
		renderer->setShaderConstant1f("invRadius", 1.0f / lightRange);
		renderer->apply();
	
		map->drawBatch(renderer, i);
	}


	drawLight(modelview.rows[0].xyz(), modelview.rows[1].xyz());

	if (lastMouseTime > 0.0){
		renderer->setup2DMode(0, (float) width, 0, (float) height);

		char str[64];
		sprintf(str, "Light radius: %.0f", 50 * (time - lastMouseTime));
		renderer->drawText(str, 10, height - 42.0f, 30, 36, defaultFont, linearClamp, blendSrcAlpha, noDepthTest);
	}
}
