#include "FluidContext.h"

#include "globals.h"
#include "intrin.h"

void FluidContext::BuildUp(int screenX, int screenY, D3DMULTISAMPLE_TYPE multiSampleType)
{
	currentFluid = NULL;

	D3DVERTEXELEMENT9 decl[] =
	{
		{0, 0, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION, 0},
		{0, 12, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_NORMAL, 0},
		{1, 0, D3DDECLTYPE_FLOAT4, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 0},
		{1, 16, D3DDECLTYPE_FLOAT1, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 1},
		D3DDECL_END()
	};

	g_d3d_device->CreateVertexDeclaration(decl, &vertexDecl);

	D3DVERTEXELEMENT9 quadDecl[] =
	{
		{0, 0, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION, 0},
		{0, 12, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 0},
		D3DDECL_END()
	};
	
	g_d3d_device->CreateVertexDeclaration(quadDecl, &quadVertexDecl);
	float* quadData;
	g_d3d_device->CreateVertexBuffer(sizeof(float) * 20, D3DUSAGE_WRITEONLY, 0, D3DPOOL_MANAGED, &quadVertexBuffer, NULL);
	quadVertexBuffer->Lock(0, sizeof(float) * 20, (void**)&quadData, 0);
	::memcpy(quadData, particleQuadVerts, sizeof(float) * 20);
	quadVertexBuffer->Unlock();

	int* quadIndexData;
	g_d3d_device->CreateIndexBuffer(6 * sizeof(int), D3DUSAGE_WRITEONLY, D3DFMT_INDEX32, D3DPOOL_MANAGED, &quadIndexBuffer, NULL);
	quadIndexBuffer->Lock(0, 0, (void**)&quadIndexData, 0);
	quadIndexData[0] = 0;
	quadIndexData[1] = 1;
	quadIndexData[2] = 2;
	quadIndexData[3] = 2;
	quadIndexData[4] = 0;
	quadIndexData[5] = 3;
	quadIndexBuffer->Unlock();

	// build up the render targets
	renderTargetWidth = screenX;
	renderTargetHeight = screenY;
	renderTargetDDX = D3DXVECTOR4(1.0f / renderTargetWidth, 1.0f / renderTargetHeight, 0.00f, 0.0f );

	g_d3d_device->CreateRenderTarget( screenX,
			screenY,
			D3DFMT_A16B16G16R16F,
			multiSampleType,
			0,
			FALSE,
			&renderSurface,
			NULL);
}

void FluidContext::RenderParticles(const FluidParticleSystem& particleInfo)
{
	g_d3d_device->SetStreamSource(0, quadVertexBuffer, 0, sizeof(float) * 5);
	g_d3d_device->SetStreamSourceFreq(0, D3DSTREAMSOURCE_INDEXEDDATA | particleInfo.InstanceCount);

	g_d3d_device->SetStreamSource(1, particleInfo.InstanceDataBuffer, 0, sizeof(InstanceData));
	g_d3d_device->SetStreamSourceFreq(1, D3DSTREAMSOURCE_INSTANCEDATA | 1ul);

	g_d3d_device->SetVertexDeclaration(vertexDecl);
	g_d3d_device->SetIndices(quadIndexBuffer);

	g_d3d_device->SetTexture(0, NULL); // pos
	g_d3d_device->SetTexture(1, NULL); // normal

	// Begin

	{
		g_Shaders[ FluidPosition ].Activate();
		g_d3d_device->SetRenderState( D3DRS_ZWRITEENABLE, TRUE );

		g_d3d_device->SetRenderTarget(0, renderSurface); // pos
		g_d3d_device->Clear(0, NULL, D3DCLEAR_TARGET, 0x0, 1.0f, 0);

		g_d3d_device->DrawIndexedPrimitive(D3DPT_TRIANGLELIST, 0, 0, 4, 0, 2);

		g_d3d_device->StretchRect( renderSurface, NULL,
			currentFluid->renderSurfacePingPong[0], NULL, D3DTEXF_LINEAR );
	}

	{
		g_Shaders[ FluidNormal ].Activate();
		g_d3d_device->SetRenderState( D3DRS_ZWRITEENABLE, FALSE );

		g_d3d_device->SetRenderTarget(0, renderSurface); // normal
		g_d3d_device->Clear(0, NULL, D3DCLEAR_TARGET, 0x0, 1.0f, 0);

		g_d3d_device->DrawIndexedPrimitive(D3DPT_TRIANGLELIST, 0, 0, 4, 0, 2);

		g_d3d_device->StretchRect( renderSurface, NULL,
			currentFluid->renderSurfacePingPong[2], NULL, D3DTEXF_LINEAR );
	}
	
	g_d3d_device->SetStreamSourceFreq(0, 1);
	g_d3d_device->SetStreamSourceFreq(1, 1);
	g_d3d_device->SetStreamSource(1, NULL, 0, 0);
}

void FluidContext::RenderParticleDepth(const FluidParticleSystem& particleInfo)
{
	g_d3d_device->SetRenderState(D3DRS_CULLMODE, D3DCULL_NONE);

	D3DXMATRIX mat;
	D3DXMatrixMultiply( &mat, &g_matView, &g_matProjection );
	D3DXMatrixTranspose( &mat, &mat );
	g_d3d_device->SetVertexShaderConstantF( 31, (float*)&mat, 4 ); 
	g_d3d_device->SetPixelShaderConstantF( 31, (float*)&mat, 4 ); 

	g_d3d_device->SetVertexShaderConstantF( 35, (float*)&g_matView, 4 ); 
	g_d3d_device->SetPixelShaderConstantF( 35, (float*)&g_matView, 4 ); 

	g_d3d_device->SetStreamSource(0, quadVertexBuffer, 0, sizeof(float) * 5);
	g_d3d_device->SetStreamSourceFreq(0, D3DSTREAMSOURCE_INDEXEDDATA | particleInfo.InstanceCount);

	g_d3d_device->SetStreamSource(1, particleInfo.InstanceDataBuffer, 0, sizeof(InstanceData));
	g_d3d_device->SetStreamSourceFreq(1, D3DSTREAMSOURCE_INSTANCEDATA | 1ul);

	g_d3d_device->SetVertexDeclaration(vertexDecl);
	g_d3d_device->SetIndices(quadIndexBuffer);

	g_Shaders[ FluidDepth ].Activate();
	g_d3d_device->DrawIndexedPrimitive(D3DPT_TRIANGLELIST, 0, 0, 4, 0, 2);

	g_d3d_device->SetStreamSourceFreq(0, 1);
	g_d3d_device->SetStreamSourceFreq(1, 1);
	g_d3d_device->SetStreamSource(1, NULL, 0, 0);
}


void FluidContext::Begin(const Fluid& fluid)
{
	g_d3d_device->SetRenderState(D3DRS_CULLMODE, D3DCULL_NONE);

	currentFluid = &fluid;

	g_d3d_device->SetPixelShaderConstantF(20, fluid.Settings.Softness, 1);
	g_d3d_device->SetPixelShaderConstantF(21, fluid.Settings.LightTransferAmmount, 1);
	g_d3d_device->SetPixelShaderConstantF(22, fluid.Settings.SpecularFactor, 1);
	g_d3d_device->SetPixelShaderConstantF(23, fluid.Settings.SpecularPower, 1);
	g_d3d_device->SetPixelShaderConstantF(24, fluid.Settings.TranslucencyAmmount, 1);
	g_d3d_device->SetPixelShaderConstantF(25, fluid.Settings.TextureScale, 1);
	g_d3d_device->SetPixelShaderConstantF(26, fluid.Settings.SurfaceTension, 1);
	g_d3d_device->SetPixelShaderConstantF(27, fluid.Settings.RefractAmmount, 1);
	g_d3d_device->SetPixelShaderConstantF(28, fluid.Settings.DepthDifferenceBlur, 1);

	D3DXMATRIX mat;
	D3DXMatrixMultiply( &mat, &g_matView, &g_matProjection );
	D3DXMatrixTranspose( &mat, &mat );
	g_d3d_device->SetVertexShaderConstantF( 31, (float*)&mat, 4 ); 
	g_d3d_device->SetPixelShaderConstantF( 31, (float*)&mat, 4 ); 

	g_d3d_device->SetVertexShaderConstantF( 35, (float*)&g_matView, 4 ); 
	g_d3d_device->SetPixelShaderConstantF( 35, (float*)&g_matView, 4 ); 
}


void FluidContext::End()
{
	g_d3d_device->SetStreamSourceFreq(0, 1);
	g_d3d_device->SetStreamSourceFreq(1, 1);
	g_d3d_device->SetStreamSource(1, NULL, 0, 0);

	g_d3d_device->SetVertexDeclaration(quadVertexDecl);

	//quadddx
	g_d3d_device->SetVertexShaderConstantF(30, (float*)&currentFluid->renderTargetPingPongDDX, 1);
	g_d3d_device->SetDepthStencilSurface(NULL);

	// blur
	{
		g_Shaders[ FluidBlur ].Activate();

		for (int i = 0; i < 3; ++i)
		{
			g_d3d_device->SetSamplerState( i, D3DSAMP_MINFILTER, D3DTEXF_LINEAR );
			g_d3d_device->SetSamplerState( i, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR );
			g_d3d_device->SetSamplerState( i, D3DSAMP_MIPFILTER, D3DTEXF_LINEAR );
			g_d3d_device->SetSamplerState( i, D3DSAMP_ADDRESSU, D3DTADDRESS_CLAMP );
			g_d3d_device->SetSamplerState( i, D3DSAMP_ADDRESSV, D3DTADDRESS_CLAMP );
		}

		for (int k = 0; k < *currentFluid->Settings.BlurSteps * 2; ++k)
		{
			g_d3d_device->SetRenderTarget(0, currentFluid->renderSurfacePingPong[0 + ((k+1)%2)]);
			g_d3d_device->SetRenderTarget(1, currentFluid->renderSurfacePingPong[2 + ((k+1)%2)]);

			// ddx
			g_d3d_device->SetPixelShaderConstantF(29, (float*)&(k == 0 ? renderTargetDDX : currentFluid->renderTargetPingPongDDX), 1);

			for (int i = 0; i < renderTextureCount; ++i)
				g_d3d_device->SetTexture( i, currentFluid->renderTexturePingPong[i*2 + k%2]);

			HRESULT hr = g_d3d_device->DrawPrimitiveUP(D3DPT_TRIANGLELIST, 2, screenQuadVerts, 20);
		}
	}

	currentFluid = NULL;
}


void FluidContext::Display(const Fluid& fluid)
{
	g_Shaders[ FluidCompose ].Activate();
	g_d3d_device->SetRenderState( D3DRS_ZENABLE, FALSE );

	g_d3d_device->SetSamplerState( 3, D3DSAMP_ADDRESSU, D3DTADDRESS_WRAP );
	g_d3d_device->SetSamplerState( 3, D3DSAMP_ADDRESSV, D3DTADDRESS_WRAP );

	g_d3d_device->SetTexture(0, fluid.renderTexturePingPong[1]); // pos
	g_d3d_device->SetTexture(1, fluid.renderTexturePingPong[3]); // normal
	g_d3d_device->SetTexture(2, fluid.Settings.DiffuseTexture); // diffuse, herp
	g_d3d_device->SetTexture(3, g_pFullScreenRT[ RT_COLOR ] /*backgroundTexture*/); // background, derp

	g_d3d_device->SetRenderTarget(1, NULL);
	g_d3d_device->SetRenderTarget(2, NULL);

	g_d3d_device->DrawPrimitiveUP(D3DPT_TRIANGLELIST, 2, screenQuadVerts, 20);
}


FluidParticleSystem FluidContext::GenerateFluidParticleSystem(unsigned int instanceCount, InstanceData* initialData)
{
	FluidParticleSystem result;
	result.InstanceDataBuffer = NULL;
	result.InstanceCount = instanceCount;
	result.IsVisible = false;

	if(instanceCount > 0)
	{
		InstanceData* instanceData;
		g_d3d_device->CreateVertexBuffer(sizeof(InstanceData) * instanceCount, D3DUSAGE_WRITEONLY, 0, D3DPOOL_MANAGED, &result.InstanceDataBuffer, NULL);
		if (initialData != NULL)
		{
			result.InstanceDataBuffer->Lock(0, sizeof(InstanceData) * instanceCount, (void**)&instanceData, 0);
			::memcpy(instanceData, initialData, sizeof(InstanceData) * instanceCount);
			result.InstanceDataBuffer->Unlock();
		}
	}

	return result;
}


Fluid FluidContext::GenerateFluid(const FluidSettings& settings)
{
	Fluid result;
	result.Settings = settings;
	result.renderTargetPingPongDDX = D3DXVECTOR4( 1.0f / (renderTargetWidth * *settings.RenderTextureSizeScale), 1.0f / (renderTargetHeight * *settings.RenderTextureSizeScale), 0.0f, 0.0f );

	for (int i = 0; i < renderTextureCount; ++i)
	{
		for (int k = 0; k < 2; ++k)
		{
			g_d3d_device->CreateTexture((UINT)(renderTargetWidth * *settings.RenderTextureSizeScale),
				(UINT)(renderTargetHeight * *settings.RenderTextureSizeScale),
				1,
				D3DUSAGE_RENDERTARGET,
				D3DFMT_A16B16G16R16F,
				D3DPOOL_DEFAULT,
				&result.renderTexturePingPong[i * 2 + k],
				NULL);
			result.renderTexturePingPong[i * 2 + k]->GetSurfaceLevel(0, &result.renderSurfacePingPong[i * 2 + k]);
		}
	}

	return result;
}


void FluidContext::MeltObject(int objectIndex, int fluidIndex, int pointMapIndex, float time)
{
	if (g_FluidParticleSystem[fluidIndex].InstanceCount == 0 || !g_FluidParticleSystem[fluidIndex].IsVisible)
		return;

	float thickness = 4.0f;
	// initial fill the particle system
	InstanceData* instanceData;
	g_FluidParticleSystem[fluidIndex].InstanceDataBuffer->Lock(0, sizeof(InstanceData) * g_FluidParticleSystem[fluidIndex].InstanceCount, (void**)&instanceData, 0);//D3DLOCK_DISCARD);
	g_Random.setSeed(0);

	D3DXMATRIX mat = g_Objects[ objectIndex ].GetTransformation();
	for (unsigned int k = 0; k < g_FluidParticleSystem[fluidIndex].InstanceCount / 4; ++k)
		for (unsigned int m = 0; m < 4; ++m)
		{
			unsigned int i = k*4+m;
			D3DXVECTOR3 pmPos = D3DXVECTOR3(g_FluidPointMap[pointMapIndex][k].x, g_FluidPointMap[pointMapIndex][k].y, g_FluidPointMap[pointMapIndex][k].z);
			D3DXVec3TransformCoord(&pmPos, &pmPos, &mat);

			float drop = time - g_StartMeltTime;
			drop -= (1.0f - g_FluidPointMap[pointMapIndex][k].w) * (g_EndMeltTime - g_StartMeltTime);
			drop *= max(0.5f, 1.0f - (1.0f - g_FluidPointMap[pointMapIndex][k].w));
			drop = max(0.0f, drop);
			drop = mypow(drop, 2.0f);
			if (drop > 1.0f)
				drop += mypow(drop - 1.0f, 4.0f);

			instanceData[i].pos[0] = (pmPos.x  + g_PointMapDistance[pointMapIndex] * thickness * g_Random.genFloat(-0.25f, 0.25f) );
			instanceData[i].pos[1] = (pmPos.y  + g_PointMapDistance[pointMapIndex] * thickness * g_Random.genFloat(-0.25f, 0.25f) ) - drop;
			instanceData[i].pos[2] = (pmPos.z  + g_PointMapDistance[pointMapIndex] * thickness * g_Random.genFloat(-0.25f, 0.25f) );

			instanceData[i].size = g_FluidParticleSize[fluidIndex];
			instanceData[i].intensity = min(5.0f, drop) / 5.0f;
		}
	g_FluidParticleSystem[fluidIndex].InstanceDataBuffer->Unlock();
}