#include "data.h"
#include "shaders.h"

struct TEXTURE_2D_SRV_UAV_RTV {
	ID3D11Texture2D* TEX;
	ID3D11ShaderResourceView* SRV;
	ID3D11UnorderedAccessView* UAV;
	ID3D11RenderTargetView* RTV;
};

void setNullUAVSRV() {
		ID3D11UnorderedAccessView* nullUAV[] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
		ID3D11ShaderResourceView* nullSRV[] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
		context->CSSetUnorderedAccessViews(0, 8, nullUAV, NULL);
		context->CSSetShaderResources(0, 8, nullSRV);
}

void createSRVUAVRTV(TEXTURE_2D_SRV_UAV_RTV* tex, int width, int height) {
	D3D11_TEXTURE2D_DESC texDesc = {
		width, height, 0, 1, DXGI_FORMAT_R16G16B16A16_FLOAT,
		{1, 0}, D3D11_USAGE_DEFAULT,
		D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_UNORDERED_ACCESS,
		0, D3D11_RESOURCE_MISC_GENERATE_MIPS
	};
	device->CreateTexture2D(&texDesc, NULL, &tex->TEX);
	device->CreateShaderResourceView(tex->TEX, NULL, &tex->SRV);
	device->CreateUnorderedAccessView(tex->TEX, NULL, &tex->UAV);
	device->CreateRenderTargetView(tex->TEX, NULL, &tex->RTV);
}

int main(int argc, char** argv) {
#ifdef FULLSCREEN
	ShowCursor(0);
	hWnd = CreateWindowExA(0, (LPCSTR)ATOM_STATIC, 0, WS_POPUP | WS_VISIBLE | WS_MAXIMIZE, 0, 0, 0, 0, 0, 0, 0, 0);
	//ShowWindow(hWnd, SW_SHOW);
#else
#ifdef _DEBUG
	WNDCLASSEX wndclass;
	wndclass.cbSize = sizeof (wndclass);
	wndclass.style = CS_HREDRAW | CS_VREDRAW | CS_OWNDC;
	wndclass.lpfnWndProc = (WNDPROC)DefWindowProc;
	wndclass.cbClsExtra = 0;
	wndclass.cbWndExtra = 0;
	wndclass.hInstance = NULL;
	wndclass.hIcon = LoadIcon(NULL, IDI_APPLICATION);
	wndclass.hCursor = LoadCursor(NULL, IDC_ARROW);
	wndclass.hbrBackground = (HBRUSH)GetStockObject(BLACK_BRUSH);
	wndclass.lpszMenuName = NULL;
	wndclass.lpszClassName = "BluFlame intro window class";
	wndclass.hIconSm = LoadIcon(NULL, IDI_APPLICATION);

	RegisterClassEx(&wndclass);

	hWnd = CreateWindowExA(0,
		wndclass.lpszClassName,
		"Intro",
		WS_EX_APPWINDOW,
		GetSystemMetrics(SM_CXSCREEN) - WIDTH - 16,
		0,
		WIDTH + 16,
		HEIGHT + 40,
		NULL,
		NULL,
		GetModuleHandle(NULL),
		NULL
		);
	ShowWindow(hWnd, SW_NORMAL);
	SetWindowPos(hWnd, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE);
#else
	hWnd = CreateWindowExA(0, (LPCSTR)ATOM_STATIC, 0, WS_POPUP | WS_VISIBLE, 0, 0, WIDTH, HEIGHT, 0, 0, 0, 0);
#endif
#endif
	swapChainDesc.OutputWindow = hWnd;
	D3D11CreateDeviceAndSwapChain(NULL, D3D_DRIVER_TYPE_HARDWARE, NULL, D3D11_CREATE_DEVICE_DEBUG, NULL, 0, D3D11_SDK_VERSION, (DXGI_SWAP_CHAIN_DESC*)&swapChainDesc, &swapChain, &device, NULL, &context);
	swapChain->GetBuffer(0, __uuidof(ID3D11Texture2D), (LPVOID*)&backBufferTexture);
	device->CreateUnorderedAccessView(backBufferTexture, NULL, &backBufferUAV);

	device->CreateBuffer(&constantBufferDesc, NULL, &constantBuffer);

	device->CreateBuffer(&particleBufferDesc, NULL, &particleBuffer);
	device->CreateUnorderedAccessView(particleBuffer, &particleUavDesc, &particleBufferUAV);
	device->CreateShaderResourceView(particleBuffer, &particleSrvDesc, &particleBufferSRV);

	device->CreateRenderTargetView(backBufferTexture, &backBufferRtvDesc, &backBufferRtv);

	device->CreateBlendState(&alphaBlendDesc, &alphaBlendState);
	device->CreateRasterizerState(&rasterizerDesc, &rasterizerState);

#ifndef _DEBUG
	CreateShaders();
#else
#ifndef SHADERDEBUG
	InitShader();
	shaderCompileEvent = ::CreateEvent(NULL, FALSE, FALSE, TEXT("WriteEvent"));
	SetThreadPriority((HANDLE)CreateThread(0, 0, &filemon, 0, 0, 0), THREAD_PRIORITY_BELOW_NORMAL);
#endif
	CreateShaders();
#endif

	//---
	// Postprocessing setup
	//---
	D3D11_SAMPLER_DESC samplerDesc = {
		D3D11_FILTER_MIN_MAG_MIP_LINEAR, //D3D11_FILTER_MIN_MAG_MIP_LINEAR, D3D11_FILTER_ANISOTROPIC, D3D11_FILTER_MIN_MAG_MIP_POINT
		D3D11_TEXTURE_ADDRESS_CLAMP,
		D3D11_TEXTURE_ADDRESS_CLAMP,
		D3D11_TEXTURE_ADDRESS_CLAMP,
		0.0f,
		1,
		D3D11_COMPARISON_ALWAYS,
		{0,0,0,0},
		0,
		D3D11_FLOAT32_MAX
	};

	ID3D11SamplerState* sampler;
	device->CreateSamplerState(&samplerDesc, &sampler);
	context->CSSetSamplers(0, 1, &sampler);

	TEXTURE_2D_SRV_UAV_RTV tex0;
	createSRVUAVRTV(&tex0, WIDTH, HEIGHT);

	TEXTURE_2D_SRV_UAV_RTV pingPongTex1;
	createSRVUAVRTV(&pingPongTex1, WIDTH, HEIGHT);

	TEXTURE_2D_SRV_UAV_RTV ghostsTex4;
	createSRVUAVRTV(&ghostsTex4, WIDTH/4, HEIGHT/4);

	TEXTURE_2D_SRV_UAV_RTV streaksTex4;
	createSRVUAVRTV(&streaksTex4, WIDTH/4, HEIGHT/4);

	TEXTURE_2D_SRV_UAV_RTV lensDirt2;
	createSRVUAVRTV(&lensDirt2, WIDTH/2, HEIGHT/2);


	TEXTURE_2D_SRV_UAV_RTV tmpTex1;
	TEXTURE_2D_SRV_UAV_RTV blurTex1;
	createSRVUAVRTV(&tmpTex1, WIDTH, HEIGHT);
	createSRVUAVRTV(&blurTex1, WIDTH, HEIGHT);

	TEXTURE_2D_SRV_UAV_RTV tmpTex8;
	TEXTURE_2D_SRV_UAV_RTV blurTex8;
	createSRVUAVRTV(&tmpTex8, WIDTH/8, HEIGHT/8);
	createSRVUAVRTV(&blurTex8, WIDTH/8, HEIGHT/8);

	TEXTURE_2D_SRV_UAV_RTV tmpTex32;
	TEXTURE_2D_SRV_UAV_RTV blurTex32;
	createSRVUAVRTV(&tmpTex32, WIDTH/32, HEIGHT/32);
	createSRVUAVRTV(&blurTex32, WIDTH/32, HEIGHT/32);



	ID3D11UnorderedAccessView* uavs[] = { particleBufferUAV };
	ID3D11ShaderResourceView* srvs[] = { particleBufferSRV };

	context->RSSetState(rasterizerState);
	context->CSSetUnorderedAccessViews(0, 1, uavs, NULL);
	context->CSSetConstantBuffers(0, 1, &constantBuffer);
	context->CSSetShader(initShader, NULL, 0);
	context->Dispatch(PARTICLE_COUNT / 16, PARTICLE_COUNT / 16, 1);


	context->VSSetConstantBuffers(0, 1, &constantBuffer);
	context->PSSetConstantBuffers(0, 1, &constantBuffer);
	context->OMSetBlendState(alphaBlendState, NULL, 0xffffffff);
	context->RSSetViewports(1, &swapChainViewport);

	float demoTime = 0.0f;
	LARGE_INTEGER startTime;
	QueryPerformanceCounter(&startTime);
	do {
#ifdef _DEBUG
		MSG msg;
		while (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE))
		{
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}
#endif
		//context->ClearRenderTargetView(backBufferRtv, zero);
		// For postprocessing
		context->ClearRenderTargetView(tex0.RTV, zero);

		context->Map(constantBuffer, 0, D3D11_MAP_WRITE_DISCARD, 0, &mappedResource);
		auto constantBufferVertexData = (ConstantBufferType*)mappedResource.pData;

		LARGE_INTEGER time;
		LARGE_INTEGER freq;
		QueryPerformanceCounter(&time);
		QueryPerformanceFrequency(&freq);
		float deltaTime = (float(time.QuadPart - startTime.QuadPart) / freq.QuadPart);
		startTime = time;
		demoTime += deltaTime;
		constantBufferVertexData->demoTime = demoTime;
		constantBufferVertexData->deltaTime = deltaTime;
		constantBufferVertexData->dummy1 = 1.0f;
		constantBufferVertexData->dummy2 = 1.0f;

		context->Unmap(constantBuffer, 0);

		//context->OMSetRenderTargets(1, &backBufferRtv, NULL;
		// For postprocessing
		context->OMSetRenderTargets(1, &tex0.RTV, NULL);

		// particle update
		if (true)
		{
			context->CSSetUnorderedAccessViews(0, 1, uavs, NULL);
			context->CSSetShader(updateShader, NULL, 0);
			context->Dispatch(PARTICLE_COUNT / 16, PARTICLE_COUNT / 16, 1);
			context->CSSetUnorderedAccessViews(0, 1, (ID3D11UnorderedAccessView* const *)zero, NULL);
		}

		// raymarching draw
		context->IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLESTRIP);
		context->VSSetShader(raymarchingVertexShader, NULL, 0);
		context->GSSetShader(NULL, NULL, 0);
		context->PSSetShader(raymarchingPixelShader, NULL, 0);
		context->Draw(4, 0);

		// particle draw
		if (true)
		{
			context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_POINTLIST);
			context->VSSetShader(particleVertexShader, NULL, 0);
			context->GSSetShader(particleGeometryShader, NULL, 0);
			context->PSSetShader(particlePixelShader, NULL, 0);
			context->VSSetShaderResources(0, 1, srvs);
			context->Draw(PARTICLE_COUNT * PARTICLE_COUNT, 0);
			context->VSSetShaderResources(0, 1, (ID3D11ShaderResourceView* const*)zero);
		}

		ID3D11RenderTargetView* nullRTV[] = {NULL};
		context->OMSetRenderTargets(1, nullRTV, NULL);

		// tex0 is the input texture
		context->GenerateMips(tex0.SRV);
		setNullUAVSRV();
		// Distort Chroma
		context->CSSetUnorderedAccessViews(0, 1, &tmpTex1.UAV, NULL);
		context->CSSetShaderResources(0, 1, &tex0.SRV);
		context->CSSetShader(csDistortChroma, NULL, 0);
		context->Dispatch(80, 45, 1);
		setNullUAVSRV();
		
		context->GenerateMips(tmpTex1.SRV);
		
		// Radial Blur
		context->CSSetUnorderedAccessViews(0, 1, &pingPongTex1.UAV, NULL);
		context->CSSetShaderResources(0, 1, &tmpTex1.SRV);
		context->CSSetShader(csRadialBlur, NULL, 0);
		context->Dispatch(80, 45, 1);
		setNullUAVSRV();
		
		context->GenerateMips(pingPongTex1.SRV);
		
		// Circumferential Blur -> (MipMap) -> Streaks
		context->CSSetUnorderedAccessViews(0, 1, &tex0.UAV, NULL);
		context->CSSetShaderResources(0, 1, &pingPongTex1.SRV);
		context->CSSetShader(csCircumferentialBlur, NULL, 0);
		context->Dispatch(80, 45, 1);
		setNullUAVSRV();

		context->GenerateMips(tex0.SRV);

		// Streaks
		context->CSSetUnorderedAccessViews(0, 1, &streaksTex4.UAV, NULL);
		context->CSSetShaderResources(0, 1, &tex0.SRV);
		context->CSSetShader(csStreaks, NULL, 0);
		context->Dispatch(20, 12, 1);

		// Lens dirt
		context->CSSetUnorderedAccessViews(0, 1, &lensDirt2.UAV, NULL);
		context->CSSetShader(csLensDirt, NULL, 0);
		context->Dispatch(40, 23, 1);

		// BlurV (1/1)
		context->CSSetUnorderedAccessViews(0, 1, &tmpTex1.UAV, NULL);
		context->CSSetShaderResources(0, 1, &tex0.SRV);
		context->CSSetShader(csBlurV, NULL, 0);
		context->Dispatch(80, 45, 1);

		// BlurH (1/1) -> blurTex1
		context->CSSetUnorderedAccessViews(0, 1, &blurTex1.UAV, NULL);
		context->CSSetShaderResources(0, 1, &tmpTex1.SRV);
		context->CSSetShader(csBlurH, NULL, 0);
		context->Dispatch(80, 45, 1);

		context->GenerateMips(blurTex1.SRV);

		// BlurV (1/8)
		context->CSSetUnorderedAccessViews(0, 1, &tmpTex8.UAV, NULL);
		context->CSSetShaderResources(0, 1, &blurTex1.SRV);
		context->CSSetShader(csBlurV, NULL, 0);
		context->Dispatch(10, 6, 1);

		// BlurH (1/8) -> blurTex8 -> merge / ghosts
		context->CSSetUnorderedAccessViews(0, 1, &blurTex8.UAV, NULL);
		context->CSSetShaderResources(0, 1, &tmpTex8.SRV);
		context->CSSetShader(csBlurH, NULL, 0);
		context->Dispatch(10, 6, 1);

		context->GenerateMips(blurTex8.SRV);

		// BlurV (1/32)
		context->CSSetUnorderedAccessViews(0, 1, &tmpTex32.UAV, NULL);
		context->CSSetShaderResources(0, 1, &blurTex8.SRV);
		context->CSSetShader(csBlurV, NULL, 0);
		context->Dispatch(3, 2, 1);

		// BlurH (1/32) -> blurTex32 -> merge / lensDirt / smallDirt / ghosts
		context->CSSetUnorderedAccessViews(0, 1, &blurTex32.UAV, NULL);
		context->CSSetShaderResources(0, 1, &tmpTex32.SRV);
		context->CSSetShader(csBlurV, NULL, 0);
		context->Dispatch(3, 2, 1);

		context->GenerateMips(tex0.SRV);

		// Ghosts
		ID3D11ShaderResourceView* ghostsSRV[] = {blurTex1.SRV, blurTex8.SRV};
		context->CSSetUnorderedAccessViews(0, 1, &ghostsTex4.UAV, NULL);
		context->CSSetShaderResources(0, 2, ghostsSRV);
		context->CSSetShader(csGhosts, NULL, 0);
		context->Dispatch(20, 12, 1);

		setNullUAVSRV();

		// Merge
		ID3D11ShaderResourceView* mergeSRV[] = {tex0.SRV, blurTex1.SRV, blurTex8.SRV, blurTex32.SRV, streaksTex4.SRV, lensDirt2.SRV, ghostsTex4.SRV, NULL};
		context->CSSetUnorderedAccessViews(0, 1, &backBufferUAV, NULL);
		context->CSSetShaderResources(0, 8, mergeSRV);
		context->CSSetShader(csMerge, NULL, 0);
		context->Dispatch(80, 45, 1);
		setNullUAVSRV();

		swapChain->Present(1, 0);

#ifdef _DEBUG
#ifndef SHADERDEBUG
		if (::WaitForSingleObject(shaderCompileEvent, 0) == WAIT_OBJECT_0)
		{
			::Sleep(50);
			CreateShaders();
			context->CSSetUnorderedAccessViews(0, 1, uavs, NULL);
			context->CSSetShader(initShader, NULL, 0);
			context->Dispatch(64, 64, 1);
			::Sleep(50);
			QueryPerformanceCounter(&time);
			startTime = time;
		}
#endif
#endif
	} while (!GetAsyncKeyState(VK_ESCAPE));

	return EXIT_SUCCESS;
}