#include "../Framework2/Direct3D/Direct3DApp.h"
#include "../Framework2/Util/Direct3DModel.h"
#include "../Framework2/Math/PerlinNoise.h"
#include "../Framework2/Util/DemoTrack.h"

#define DIST_SCALE 0.0003f
#define RT_SIZE 256

class MainApp : public Direct3DApp {
public:
	void initMenu();
	//void setKey(unsigned int key, bool value);
	void resetCamera();

	void createHorizonMap(Image3D *hMap, Image *image, int wt, int ht, int n, float height);
	void createAngleMap(int size);

	bool init();
	bool exit();
	bool load();
	bool unload();

	void drawSky();
	void drawDepth();
	void drawDistance();
	void drawRoom();
	void drawLight();
	bool drawFrame();
protected:
	ShaderID standard, offset, selfShadow, offsetSelfShadow, distance, zOnly, sky;

	TextureID light, angleMap, noise, shadow, shadowDepth;

	TextureID base[3];
	TextureID bump[3];
	TextureID horizonMap[3];

	vec3 lightPos;

	bool useShadows, useOffset, movingLight;

	DemoTrack demoTrack;

	Model *model;
};

void MainApp::initMenu(){
	Menu *menu = menuSystem->getMainMenu();
	menu->addMenuItem("Selfshadowing: ", &useShadows, INPUT_BOOL);
	menu->addMenuItem("Offset mapping: ", &useOffset, INPUT_BOOL);
	menu->addMenuItem("Moving light: ", &movingLight, INPUT_BOOL);

	App::initMenu();
}

/*void MainApp::setKey(unsigned int key, bool value){
	App::setKey(key, value);

	if (value){
		if (key == KEY_INSERT){
			demoTrack.add(position);
		} else if (key == 'S'){
			demoTrack.writeToFile("lightPath.pth");
		} else if (key == 'O'){
			demoTrack.loadFromFile("lightPath.pth");
		}
	}
}*/

void MainApp::resetCamera(){
	position = vec3(800, -310, -800);

	wx = 0;
	wy = 0.5f;
	wz = 0;
}

#define round(x) (int) (x + 0.5f)

inline unsigned char Pack1(const float x){
    return (unsigned char) (255.0f * fabsf(x));
}

inline unsigned char Pack2(const float x){
    return (unsigned char) (255.0f * x);
}

void MainApp::createHorizonMap(Image3D *hMap, Image *image, int wt, int ht, int n, float height){
	int w = image->getWidth();
	int h = image->getHeight();
	unsigned char *src = image->getImage();

	unsigned char *dest = new unsigned char[wt * ht * n];

	hMap->loadFromMemory(dest, wt, ht, n, FORMAT_I8, true, false);

	for (int z = 0; z < n; z++){
		float xv = cosf(z * 2 * PI / n);
		float yv = sinf(z * 2 * PI / n);

		float invMax = 1.0f / max(fabsf(xv), fabsf(yv));
		xv *= invMax;
		yv *= invMax;

		float len = sqrtf(xv * xv + yv * yv);

		for (int y = 0; y < ht; y++){
			for (int x = 0; x < wt; x++){
				float maxAng = 0;

				float xp = float(x * w) / wt;
				float yp = float(y * h) / ht;

				float bh = float(src[int(yp) * w + int(xp)]);

				float dist = 0;
				//do {
				for (unsigned int i = 0; i < 256; i++){
					xp += xv;
					yp += yv;
					dist += len;

					int xt = round(xp);
					int yt = round(yp);

					while (xt < 0) xt += w;
					while (yt < 0) yt += h;

					xt %= w;
					yt %= h;

					//if (xt < 0 || xt >= w || yt < 0 || yt >= h) break;

					float ang = (float(src[yt * w + xt]) - bh) / dist;
					if (ang > maxAng) maxAng = ang;
				}
				//} while (true);

				*dest++ = Pack2(sinf(atanf(height * maxAng)));
			}
		}
	}
}

void MainApp::createAngleMap(int size){
	unsigned char *imgNeg, *imgPos, *destNeg, *destPos;
	float vx, vy, vz;
	int x, y, z;
	float mid = 0.5f * (size - 1);

	Image *images[6];
	for (unsigned int i = 0; i < 6; i++) images[i] = new Image();

	destPos = imgPos = new unsigned char[size * size * 3];
	destNeg = imgNeg = new unsigned char[size * size * 3];
	for (y = 0; y < size; y++){
		for (z = 0; z < size; z++){
			vx = mid;
			vy = mid - y;
			vz = mid - z;

			float ang0 = atan2f(vy, vx) / PI;
			float ang1 = dot(normalize(vec3(vx, vz, vy)), normalize(vec3(vx, 0, vy)));
			ang1 = sqrtf(fabsf(1 - ang1 * ang1));

			*destPos++ = 0;
			*destPos++ = Pack1(ang0);
			*destPos++ = Pack2(ang1 * (vz > 0));
			*destNeg++ = 0;
			*destNeg++ = Pack1(ang0 - sign(vy));
			*destNeg++ = Pack2(ang1 * (vz < 0));
		}
	}
	images[0]->loadFromMemory(imgPos, size, size, FORMAT_RGB8, true, false);
	images[1]->loadFromMemory(imgNeg, size, size, FORMAT_RGB8, true, false);

	destPos = imgPos = new unsigned char[size * size * 3];
	destNeg = imgNeg = new unsigned char[size * size * 3];
	for (z = 0; z < size; z++){
		for (x = 0; x < size; x++){
			vx = x - mid;
			vy = mid;
			vz = z - mid;

			float ang0 = atan2f(vy, vx) / PI;
			float ang1 = dot(normalize(vec3(vx, vz, vy)), normalize(vec3(vx, 0, vy)));
			ang1 = sqrtf(fabsf(1 - ang1 * ang1));

			*destPos++ = 0;
			*destPos++ = Pack1(ang0);
			*destPos++ = Pack2(ang1 * (vz > 0));
			*destNeg++ = 0;
			*destNeg++ = Pack1(-ang0);
			*destNeg++ = Pack2(ang1 * (vz < 0));
		}
	}
	images[2]->loadFromMemory(imgPos, size, size, FORMAT_RGB8, true, false);
	images[3]->loadFromMemory(imgNeg, size, size, FORMAT_RGB8, true, false);

	destPos = imgPos = new unsigned char[size * size * 3];
	destNeg = imgNeg = new unsigned char[size * size * 3];
	for (y = 0; y < size; y++){
		for (x = 0; x < size; x++){
			vx = x - mid;
			vy = mid - y;
			vz = mid;

			float ang0 = atan2f(vy, vx) / PI;
			float ang1 = dot(normalize(vec3(vx, vz, vy)), normalize(vec3(vx, 0, vy)));
			ang1 = sqrtf(fabsf(1 - ang1 * ang1));

			*destPos++ = 0;
			*destPos++ = Pack1(ang0);
			*destPos++ = Pack2(ang1);
			*destNeg++ = 0;
			*destNeg++ = Pack1(sign(vy) - ang0);
			*destNeg++ = 0;
		}
	}
	images[4]->loadFromMemory(imgPos, size, size, FORMAT_RGB8, true, false);
	images[5]->loadFromMemory(imgNeg, size, size, FORMAT_RGB8, true, false);

	angleMap = renderer->addCubemap(images);
}

bool MainApp::init(){
	useShadows = true;
	useOffset = true;
	movingLight = true;

	model = new Direct3DModel();
	model->loadFromFile("../Models/SimpleCastle/Map.hmdl");

	demoTrack.loadFromFile("lightPath.pth");
	//demoTrack.setNodeSize(3);

	return true;
}

bool MainApp::exit(){
	delete model;
	return true;
}

bool MainApp::load(){
	static char *message =
		"I need to precompute some large volumetric textures now. This process may take a minute or two.\n"
		"The demo will run once this is done. The resulting textures will be stored in the demo's directory,\n"
		"so this only needs to be done once.";
	bool messageShown = false;

	if (caps.PixelShaderVersion < D3DPS_VERSION(2,0) || caps.VertexShaderVersion < D3DVS_VERSION(2,0)){
		String error("This demo requires pixel and vertex shader version 2.0.\n"
					 "Your card and driver only supports pixel shader ");
		error.sprintf("%d.%d", (caps.PixelShaderVersion  >> 8) & 255, caps.PixelShaderVersion  & 255);
		error += " and vertex shader ";
		error.sprintf("%d.%d", (caps.VertexShaderVersion >> 8) & 255, caps.VertexShaderVersion & 255);

		addToLog(error);
		return false;
	}

	setDefaultFont("../Textures/Fonts/Future.font", "../Textures/Fonts/Future.dds");

	shadow = renderer->addRenderCubemap(RT_SIZE, FORMAT_RGBA8);
	shadowDepth = renderer->addRenderDepth(RT_SIZE, RT_SIZE, FORMAT_D16);

	if ((standard = renderer->addShader("standard.shd")) == SHADER_NONE) return false;
	if ((offset = renderer->addShader("offset.shd")) == SHADER_NONE) return false;
	if ((selfShadow = renderer->addShader("selfShadow.shd")) == SHADER_NONE) return false;
	if ((offsetSelfShadow = renderer->addShader("offsetSelfShadow.shd")) == SHADER_NONE) return false;

	if ((distance = renderer->addShader("distance.shd")) == SHADER_NONE) return false;

	if ((zOnly = renderer->addShader("zOnly.shd")) == SHADER_NONE) return false;

	if ((sky = renderer->addShader("sky.shd" )) == SHADER_NONE) return false;

	{
		Image *image = new Image();
		image->loadImage("../Textures/brick01Bump.png");

		Image3D *hMap = new Image3D();

		if (!hMap->loadFromFile("hzMap0.dds")){
			MessageBox(hwnd, message, "Precomputing", MB_OK | MB_ICONINFORMATION);
			messageShown = true;

			createHorizonMap(hMap, image, 64, 64, 16, 0.12f);
			hMap->saveToFile("hzMap0.dds");
		}
		if ((horizonMap[0] = renderer->addTexture3D(hMap, TEX_NOMIPMAPPING)) == TEXTURE_NONE) return false;

		if ((base[0] = renderer->addTexture("../Textures/brick01.dds")) == TEXTURE_NONE) return false;
		if ((bump[0] = renderer->addTexture(image, TEX_NORMALHEIGHTMAP)) == TEXTURE_NONE) return false;
	}

	{
		Image *image = new Image();
		//image->loadImage("../Textures/stone04Bump.png");
		image->loadImage("../Textures/laying_rock7Bump.png");

		Image3D *hMap = new Image3D();

		if (!hMap->loadFromFile("hzMap1.dds")){
			if (!messageShown) MessageBox(hwnd, message, "Precomputing", MB_OK | MB_ICONINFORMATION);
			messageShown = true;

			createHorizonMap(hMap, image, 64, 64, 16, 0.16f);
			hMap->saveToFile("hzMap1.dds");
		}
		if ((horizonMap[1] = renderer->addTexture3D(hMap, TEX_NOMIPMAPPING)) == TEXTURE_NONE) return false;

		//if ((base[1] = renderer->addTexture("../Textures/stone04.dds")) == TEXTURE_NONE) return false;
		if ((base[1] = renderer->addTexture("../Textures/laying_rock7.dds")) == TEXTURE_NONE) return false;
		if ((bump[1] = renderer->addTexture(image, TEX_NORMALHEIGHTMAP)) == TEXTURE_NONE) return false;
	}

	{
		Image *image = new Image();
		image->loadImage("../Textures/StoneWall_1-4Bump.png");

		Image3D *hMap = new Image3D();

		if (!hMap->loadFromFile("hzMap2.dds")){
			if (!messageShown) MessageBox(hwnd, message, "Precomputing", MB_OK | MB_ICONINFORMATION);
			messageShown = true;

			createHorizonMap(hMap, image, 64, 64, 16, 0.16f);
			hMap->saveToFile("hzMap2.dds");
		}
		if ((horizonMap[2] = renderer->addTexture3D(hMap, TEX_NOMIPMAPPING)) == TEXTURE_NONE) return false;

		if ((base[2] = renderer->addTexture("../Textures/StoneWall_1-4.dds")) == TEXTURE_NONE) return false;
		if ((bump[2] = renderer->addTexture(image, TEX_NORMALHEIGHTMAP)) == TEXTURE_NONE) return false;
	}

	if ((light = renderer->addTexture("../Textures/Particle.png")) == TEXTURE_NONE) return false;
	createAngleMap(256);


	initPerlin();

#define SIZE 128
#define RANGE 3.37f

	Image3D *nMap = new Image3D();

	if (!nMap->loadFromFile("nMap.dds")){
		if (!messageShown) MessageBox(hwnd, message, "Precomputing", MB_OK | MB_ICONINFORMATION);
		messageShown = true;

		unsigned int maxVal = 0, minVal = 255;
		unsigned char *noiseImg = new unsigned char[SIZE * SIZE * SIZE];
		unsigned char *dest = noiseImg;
		int x, y, z;
		for (z = 0; z < SIZE; z++){
			for (y = 0; y < SIZE; y++){
				for (x = 0; x < SIZE; x++){
					unsigned char val = (unsigned char) (127.5f + 127.5f * tileableTurbulence3(x * RANGE / SIZE, y * RANGE / SIZE, z * RANGE / SIZE, RANGE, RANGE, RANGE, 12));
					*dest++ = val;
					if (val > maxVal) maxVal = val;
					if (val < minVal) minVal = val;
				}
			}
		}
		dest = noiseImg;
		for (z = 0; z < SIZE; z++)
			for (y = 0; y < SIZE; y++)
				for (x = 0; x < SIZE; x++)
					*dest++ = (255 * (*dest - minVal)) / (maxVal - minVal);

		nMap->loadFromMemory(noiseImg, SIZE, SIZE, SIZE, FORMAT_I8, true, false);
		nMap->saveToFile("nMap.dds");
	}
	noise = renderer->addTexture3D(nMap, TEX_NOMIPMAPPING);

	((Direct3DModel *) model)->uploadToVertexBuffer(dev);
	return true;
}

bool MainApp::unload(){
	((Direct3DModel *) model)->freeVertexBuffer();

	return true;
}

void MainApp::drawSky(){
	//renderer->setDepthFunc(DEPTH_NONE);
	//renderer->setMask(COLOR);
	renderer->setTextures(noise);
	renderer->setShader(sky);
	renderer->apply();

	renderer->changeShaderConstant4x4f("view_proj_matrix", projection * modelView);
	renderer->changeShaderConstant1f("time", time);

	static vec3 vertices[] = {
		vec3(-8000, 1200, -8000),
		vec3( 8000, 1200, -8000),
		vec3( 8000, 1200,  8000),
		vec3(-8000, 1200,  8000),
	};

	dev->SetFVF(D3DFVF_XYZ);
	dev->DrawPrimitiveUP(D3DPT_TRIANGLEFAN, 2, vertices, sizeof(vec3));
}

void MainApp::drawDepth(){
	renderer->setShader(zOnly);
	renderer->setMask(DEPTH);
	renderer->apply();

	renderer->changeShaderConstant4x4f("view_proj_matrix", projection * modelView);

	((Direct3DModel *) model)->draw(dev);
}

void MainApp::drawDistance(){
	renderer->setShader(distance);
	renderer->apply();

	renderer->changeShaderConstant3f("lightPos", lightPos);
	renderer->changeShaderConstant1f("scale", DIST_SCALE);

	mat4 proj = projectionMatrixX(PI / 2, 1, 1, 3000);


	for (int face = 0; face < 6; face++){
		renderer->changeRenderTarget(shadow, shadowDepth, face);
		dev->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_RGBA(255, 255, 255, 255), 1.0f, 0);

		mat4 view = cubemapMatrix(face);
		view.translate(-lightPos);
		renderer->changeShaderConstant4x4f("view_proj_matrix", proj * view);

		((Direct3DModel *) model)->draw(dev);
	}
	renderer->changeToDefaultRenderTarget();
}

void MainApp::drawRoom(){
	for (unsigned int i = 0; i < 3; i++){
		if (useShadows){
			renderer->setTextures(base[i], bump[i], angleMap, horizonMap[i], shadow);
			if (useOffset){
				renderer->setShader(offsetSelfShadow);
			} else {
				renderer->setShader(selfShadow);
			}
		} else {
			renderer->setTextures(base[i], bump[i], shadow);
			if (useOffset){
				renderer->setShader(offset);
			} else {
				renderer->setShader(standard);
			}
		}
		renderer->apply();

		renderer->changeShaderConstant4x4f("view_proj_matrix", projection * modelView);
		renderer->changeShaderConstant3f("lightPos", lightPos);
		renderer->changeShaderConstant3f("camPos", position);
		renderer->changeShaderConstant1f("scale", DIST_SCALE);
		renderer->changeShaderConstant1f("texWorldSize", (i < 2)? 512.0f : 256.0f);
		renderer->changeShaderConstant1f("ambient", 0.14f);
		renderer->changeShaderConstant1f("bias", 0.022f);

		((Direct3DBatch *) model->getBatch(i))->draw(dev);
	}
}

void MainApp::drawLight(){
	renderer->setTextures(light);
	renderer->setBlending(ONE, ONE);
	renderer->apply();

	vec3 dx(modelView.elem[0][0], modelView.elem[0][1], modelView.elem[0][2]);
	vec3 dy(modelView.elem[1][0], modelView.elem[1][1], modelView.elem[1][2]);

	struct Vertex {
		vec3 position;
		vec2 texCoord;
	};

	Vertex vertices[4];

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

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


	dev->SetFVF(D3DFVF_XYZ | D3DFVF_TEXCOORDSIZE2(0) | (1 << D3DFVF_TEXCOUNT_SHIFT));
	dev->DrawPrimitiveUP(D3DPT_TRIANGLEFAN, 2, vertices, sizeof(Vertex));
}

bool MainApp::drawFrame(){
	drawDistance();

	dev->SetRenderState(D3DRS_CULLMODE, D3DCULL_CCW);
    dev->Clear(0, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_RGBA(0, 0, 0, 0), 1.0f, 0);

	projection = projectionMatrixX(1.5f, float(height) / float(width), 1, 10000);
	dev->SetTransform(D3DTS_PROJECTION, (D3DMATRIX *) (const float *) transpose(projection));

	modelView = rotateZXY(-wx, -wy, -wz);
	modelView.translate(-position);
	dev->SetTransform(D3DTS_VIEW, (D3DMATRIX *) (const float *) transpose(modelView));

	//lightPos = vec3(0,-100,100);
	/*if (demoTrack.getNodeCount()) */
	if (movingLight) demoTrack.getAt(0.6f * time, (float *) &lightPos, 0, 3);

	drawDepth();
	drawRoom();

	modelView = rotateZXY(-wx, -wy, -wz);
	dev->SetTransform(D3DTS_VIEW, (D3DMATRIX *) (const float *) transpose(modelView));
	drawSky();

	modelView.translate(-position);
	dev->SetTransform(D3DTS_VIEW, (D3DMATRIX *) (const float *) transpose(modelView));

	drawLight();

	return true;
}

App *app = new MainApp();
