#include "../../Direct3D/Shader/RuShader_D3D_Water.h"

#pragma unmanaged

// ************************************************************************************************************************************************************

CRuShader_Water_PS20_Base::CRuShader_Water_PS20_Base(CRuEngineD3D *engine)
:	CRuShader_Base(engine)
{
	// Setup vertex elements
	CRuVertexElement vertexElements[] =
	{
		{ 0,	0,	ruELEMENTTYPE_FLOAT3,		ruELEMENTMETHOD_DEFAULT,	ruELEMENTUSAGE_POSITION,	0 },
		{ 0,	12,	ruELEMENTTYPE_D3DCOLOR,		ruELEMENTMETHOD_DEFAULT,	ruELEMENTUSAGE_COLOR,		0 },
		RUVERTEXDECL_END()
	};

	// Construct vertex declaration
	m_vertexDecl = m_engine->CreateVertexDeclaration(vertexElements);
}

CRuShader_Water_PS20_Base::~CRuShader_Water_PS20_Base()
{
}

void CRuShader_Water_PS20_Base::SetParameters(const CRuParameterCache *paramCache)
{
	DWORD param;

	CRuShader_Base::SetParameters(paramCache);

	if(paramCache->GetParameter(ruPARAMTYPE_TRANSFORM_PROJECTION, param))
	{
		// Setup texture thingy
		CRuMatrix4x4 texTransform = *((CRuMatrix4x4 *) param);

		texTransform._11 *= 0.5f;
		texTransform._22 *= -0.5f;
		texTransform._31 += 0.5f;
		texTransform._32 += 0.51f;	// HACK, param)) offset a bit to hide the water line

		INT32 uvTransformOffset = 27 * 4;

		for(UINT32 i = 0, k = 0; i < 4; ++i)
		{
			for(UINT32 j = 0; j < 4; ++j, ++k)
			{
				m_shaderConstants[k + uvTransformOffset] = texTransform.m_elements[j][i];
			}
		}
	}

	if(paramCache->GetParameter(ruPARAMTYPE_WATER_BASECOLOR, param))
	{
		INT32 waterColorOffset = 31 * 4;

		m_shaderConstants[waterColorOffset++] = ((param >> 16) & 0xFF) / 255.0f;
		m_shaderConstants[waterColorOffset++] = ((param >> 8) & 0xFF) / 255.0f;
		m_shaderConstants[waterColorOffset++] = ((param) & 0xFF) / 255.0f;
	}

	if(paramCache->GetParameter(ruPARAMTYPE_WATER_REFLECTFACTOR, param))
	{
		INT32 factorOffset = 2 * 4;
		float factor = RuCastDWORDToFloat(param);

		m_pixelShaderConstants[factorOffset] = factor;
	}

	if(paramCache->GetParameter(ruPARAMTYPE_WATER_REFRACTFACTOR, param))
	{
		INT32 factorOffset = 2 * 4 + 1;
		float factor = RuCastDWORDToFloat(param);

		m_pixelShaderConstants[factorOffset] = factor;
	}

	if(paramCache->GetParameter(ruPARAMTYPE_WATER_ALPHAFACTOR, param))
	{
		INT32 factorOffset = 31 * 4 + 3;
		float factor = RuCastDWORDToFloat(param);

		m_shaderConstants[factorOffset] = factor;
	}

	if(paramCache->GetParameter(ruPARAMTYPE_WATER_BASEFACTOR, param))
	{
		INT32 factorOffset = 2 * 4 + 2;
		float factor = RuCastDWORDToFloat(param);

		m_pixelShaderConstants[factorOffset] = factor;
	}

	if(paramCache->GetParameter(ruPARAMTYPE_WATER_SPECULARSTRENGTH, param))
	{
		INT32 factorOffset = 2 * 4 + 3;
		float factor = RuCastDWORDToFloat(param);

		m_pixelShaderConstants[factorOffset] = factor;
	}
}

void CRuShader_Water_PS20_Base::Instance(INT32 streamIndex, IRuMesh *mesh, UINT32 morphTargetIndex, void *buffer) const
{
	if(streamIndex == 0)
	{
		IRuMorphTarget *morphTarget = mesh->GetMorphTarget(morphTargetIndex);
		CRuVertex_Bare_Colored *vertices = (CRuVertex_Bare_Colored *) buffer;
		const DWORD *diffuse = morphTarget->GetDiffuse();

		if(morphTarget)
		{
			for(INT32 i = 0, j = 0; i < mesh->GetNumVertices(); i++)
			{
				vertices[i].x = morphTarget->GetPosition()[i].x;
				vertices[i].y = morphTarget->GetPosition()[i].y;
				vertices[i].z = morphTarget->GetPosition()[i].z;

				if(diffuse)
				{
					vertices[i].diffuse = diffuse[i];
				}
			}
		}
	}
}

// ************************************************************************************************************************************************************

CRuShader_Water_PS20_Full::CRuShader_Water_PS20_Full(CRuEngineD3D *engine)
:	CRuShader_Water_PS20_Base(engine)
{
	char *vsSource = "\
																													\n\
	vs_1_1															; Vertex shader version 1.1						\n\
																													\n\
	dcl_position0		v0											; Declare position								\n\
	dcl_color			v1											; Declare diffuse								\n\
																													\n\
	; Transform position into homogeneous space																		\n\
	mov					r0,			c20.yyyy						; Move 1 into r0								\n\
	mov					r1,			c20.yyyy						; Move 1 into r1								\n\
	m4x3				r0,			v0,			c0					; Transform vertex to world space				\n\
	m4x3				r1,			r0,			c3					; Transform vertex to view space				\n\
	m4x4				r0,			r1,			c6					; Transform vertex to homogeneous space			\n\
	mov					oPos,		r0																				\n\
																													\n\
	; Fog factor																									\n\
	sub					r2.x,		c25.z,		c25.y				; Fog End - Fog Start into r2.x					\n\
	rcp					r2.x,		r2.x							; Reciprocal of fog range into r0.x				\n\
	sub					r3.x,		r1.z,		c25.y				; Offset vertex z by fog start into r3.x		\n\
	mul					r2.x,		r3.x,		r2.x				; Divide z by fog range into r0					\n\
	sub					r2.x,		c20.y,		r2.x				; Invert fog factor								\n\
	max					r2.x,		r2.x,		c20.w				; Clamp to [0, 1]								\n\
	min					r2.x,		r2.x,		c20.y																\n\
	mov					oFog,		r2.x							; Output fog factor								\n"
	"																												\n\
	; Ambient lighting into r4																						\n\
	mov					r4,			c10																				\n\
																													\n\
	mov					r3,			c20.w																			\n\
	mov					r3.y,		c20.y																			\n\
																													\n\
	; Object-level spherical harmonics																				\n\
	add					r4,			c11,		r4					; Constant term									\n\
	mad					r4,			c12,		r3.x,		r4		; 1st-order terms								\n\
	mad					r4,			c13,		r3.y,		r4														\n\
	mad					r4,			c14,		r3.z,		r4														\n\
	mul					r5.x,		r3.x,		r3.z				; X * Z											\n\
	mul					r5.y,		r3.z,		r3.y				; Z * Y											\n\
	mul					r5.z,		r3.y,		r3.x				; Y * X											\n\
	mad					r4,			c15,		r5.x,		r4		; 2nd-order terms								\n\
	mad					r4,			c16,		r5.y,		r4														\n\
	mad					r4,			c17,		r5.z,		r4														\n\
	mul					r5.xyz,		r3,			r3																	\n\
	mul					r5.z,		r5.z,		c20.x																\n\
	sub					r5.z,		r5.z,		c20.y																\n\
	mad					r4,			c18,		r5.z,		r4														\n\
	sub					r5.x,		r5.x,		r5.y																\n\
	mad					r4,			c19,		r5.x,		r4														\n\
																													\n\
	; Diffuse color																									\n\
	mul					r4,			r4,			c26					; Multiply diffuse by global diffuse			\n\
	mul					r4,			r4,			c31					; Modulate light color by water color			\n\
	mov					oD0,		r4								; Set diffuse									\n\
																													\n\
	; Compute uv																									\n\
	m4x4				oT0.xyzw,	r1,			c27																	\n\
	m4x4				oT1.xyzw,	r1,			c27																	\n\
	mul					r2,			c20.z,		c20.z																\n\
	mul					r2,			r2,			c20.z																\n\
	mul					r2,			r2,			c20.z																\n\
	mul					r1.x,		v0.x,		r2.x																\n\
	mul					r1.y,		v0.z,		r2.y																\n\
	mov					oT2,		r1																				\n\
	";

	char *psSource = "\
																													\n\
	ps_2_0																											\n\
																													\n\
	; Declare texture samplers																						\n\
	dcl_2d				s0											; REFLECT										\n\
	dcl_2d				s1											; REFRACT										\n\
	dcl_2d				s2											; BASE											\n\
	dcl_2d				s3											; CAUSTIC										\n\
																													\n\
	; Declare texture coordinate registers																			\n\
	dcl					t0.xyzw																						\n\
	dcl					t1.xyzw																						\n\
	dcl					t2.xyzw																						\n\
																													\n\
	dcl					v0.xyzw																						\n\
																													\n\
	; Divide texture coordinates by their z component																\n\
	rcp					r3,			t0.w																			\n\
	mul					r2,			t0,			r3																	\n\
																													\n\
	; Load caustic map																								\n\
	texld				r0,			t2,			s3					; Sample CAUSTIC								\n\
	sub					r0,			r0,			c0.y																\n\
	mul					r0,			r0,			c0.y																\n\
	mul					r0,			r0,			c0.y																\n\
	mul					r0,			r0,			c0.y																\n\
	mul					r4,			r0,			c0.y																\n\
	add					r2.x,		r2.x,		r4.x				; Perturb projection UV along u					\n\
																													\n\
	; Texture sampling																								\n\
	texld				r0,			r2,			s0					; Sample REFLECT								\n\
	texld				r1,			r2,			s1					; Sample REFRACT								\n\
	texld				r2,			t2,			s2					; Sample BASE									\n\
																													\n\
	mul					r3,			r0,			c2.x				; Accumulate reflection color					\n\
	mad					r3,			r1,			c2.y,		r3		; Accumulate refraction color					\n\
	mad					r3,			r2,			c2.z,		r3		; Accumulate base color							\n\
	mul					r3,			r3,			v0					; Multiply by diffuse color						\n\
	mul					r4,			v0,			r4					; Add specular?									\n\
	mad					r3,			r4,			c2.w,		r3														\n\
	mad					r3,			v0,			c2.w,		r3														\n\
																													\n\
	mov					oC0,		r3								; Move result into output						\n\
	";

	// Assemble vertex shader
	m_vertexShader = RuD3DAssembleVertexShader(vsSource);

	// Assemble pixel shader
	m_pixelShader = RuD3DAssemblePixelShader(psSource);

	// Overrides
	m_override_ZWrite = TRUE;
}

CRuShader_Water_PS20_Full::~CRuShader_Water_PS20_Full()
{
}

void CRuShader_Water_PS20_Full::BeginPass(INT32 passNumber)
{
	// Call base BeginPass to perform any additional setup
	CRuShader_Water_PS20_Base::BeginPass(passNumber);

	if(passNumber == 0)
	{
		m_vertexDecl->Select();
		m_engine->SetVertexShader(m_vertexShader);
		m_engine->SetPixelShader(m_pixelShader);
		m_engine->GetD3DDevice()->SetVertexShaderConstantF(0, m_shaderConstants, m_numBaseConstants + 2);
		m_engine->GetD3DDevice()->SetPixelShaderConstantF(0, m_pixelShaderConstants, m_numBasePixelShaderConstants + 1);

		m_engine->SetRenderState(D3DRS_ALPHATESTENABLE, FALSE);
		m_engine->SetRenderState(D3DRS_FILLMODE, D3DFILL_SOLID);

		m_engine->SetRenderState(D3DRS_ALPHABLENDENABLE, FALSE);

		m_engine->SetRenderState(D3DRS_ZWRITEENABLE, TRUE);
		m_engine->SetRenderState(D3DRS_ZFUNC, D3DCMP_LESSEQUAL);

		m_engine->SetSamplerState(0, D3DSAMP_MINFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(0, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(0, D3DSAMP_MIPFILTER, D3DTEXF_LINEAR);

		m_engine->SetSamplerState(0, D3DSAMP_ADDRESSU, D3DTADDRESS_CLAMP);
		m_engine->SetSamplerState(0, D3DSAMP_ADDRESSV, D3DTADDRESS_CLAMP);

		m_engine->SetSamplerState(1, D3DSAMP_MINFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(1, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(1, D3DSAMP_MIPFILTER, D3DTEXF_LINEAR);

		m_engine->SetSamplerState(1, D3DSAMP_ADDRESSU, D3DTADDRESS_CLAMP);
		m_engine->SetSamplerState(1, D3DSAMP_ADDRESSV, D3DTADDRESS_CLAMP);

		m_engine->SetSamplerState(2, D3DSAMP_MINFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(2, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(2, D3DSAMP_MIPFILTER, D3DTEXF_LINEAR);

		m_engine->SetSamplerState(2, D3DSAMP_ADDRESSU, D3DTADDRESS_WRAP);
		m_engine->SetSamplerState(2, D3DSAMP_ADDRESSV, D3DTADDRESS_WRAP);

		m_engine->SetSamplerState(3, D3DSAMP_MINFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(3, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(3, D3DSAMP_MIPFILTER, D3DTEXF_LINEAR);

		m_engine->SetSamplerState(3, D3DSAMP_ADDRESSU, D3DTADDRESS_WRAP);
		m_engine->SetSamplerState(3, D3DSAMP_ADDRESSV, D3DTADDRESS_WRAP);
	}
}

void CRuShader_Water_PS20_Full::EndPass(INT32 passNumber)
{
}

// ************************************************************************************************************************************************************

CRuShader_Water_PS20_Diffuse::CRuShader_Water_PS20_Diffuse(CRuEngineD3D *engine)
:	CRuShader_Water_PS20_Base(engine)
{
	char *vsSource = "\
																													\n\
	vs_1_1															; Vertex shader version 1.1						\n\
																													\n\
	dcl_position0		v0											; Declare position								\n\
	dcl_color			v1											; Declare diffuse								\n\
																													\n\
	; Transform position into homogeneous space																		\n\
	mov					r0,			c20.yyyy						; Move 1 into r0								\n\
	mov					r1,			c20.yyyy						; Move 1 into r1								\n\
	m4x3				r0,			v0,			c0					; Transform vertex to world space				\n\
	m4x3				r1,			r0,			c3					; Transform vertex to view space				\n\
	m4x4				r0,			r1,			c6					; Transform vertex to homogeneous space			\n\
	mov					oPos,		r0																				\n\
																													\n\
	; Fog factor																									\n\
	sub					r2.x,		c25.z,		c25.y				; Fog End - Fog Start into r2.x					\n\
	rcp					r2.x,		r2.x							; Reciprocal of fog range into r0.x				\n\
	sub					r3.x,		r1.z,		c25.y				; Offset vertex z by fog start into r3.x		\n\
	mul					r2.x,		r3.x,		r2.x				; Divide z by fog range into r0					\n\
	sub					r2.x,		c20.y,		r2.x				; Invert fog factor								\n\
	max					r2.x,		r2.x,		c20.w				; Clamp to [0, 1]								\n\
	min					r2.x,		r2.x,		c20.y																\n\
	mov					oFog,		r2.x							; Output fog factor								\n"
	"																												\n\
	; Ambient lighting into r4																						\n\
	mov					r4,			c10																				\n\
																													\n\
	mov					r3,			c20.w																			\n\
	mov					r3.y,		c20.y																			\n\
																													\n\
	; Object-level spherical harmonics																				\n\
	add					r4,			c11,		r4					; Constant term									\n\
	mad					r4,			c12,		r3.x,		r4		; 1st-order terms								\n\
	mad					r4,			c13,		r3.y,		r4														\n\
	mad					r4,			c14,		r3.z,		r4														\n\
	mul					r5.x,		r3.x,		r3.z				; X * Z											\n\
	mul					r5.y,		r3.z,		r3.y				; Z * Y											\n\
	mul					r5.z,		r3.y,		r3.x				; Y * X											\n\
	mad					r4,			c15,		r5.x,		r4		; 2nd-order terms								\n\
	mad					r4,			c16,		r5.y,		r4														\n\
	mad					r4,			c17,		r5.z,		r4														\n\
	mul					r5.xyz,		r3,			r3																	\n\
	mul					r5.z,		r5.z,		c20.x																\n\
	sub					r5.z,		r5.z,		c20.y																\n\
	mad					r4,			c18,		r5.z,		r4														\n\
	sub					r5.x,		r5.x,		r5.y																\n\
	mad					r4,			c19,		r5.x,		r4														\n\
																													\n\
	; Diffuse color																									\n\
	mul					r4,			r4,			c26					; Multiply diffuse by global diffuse			\n\
	mul					r4,			r4,			c31					; Modulate light color by water color			\n\
	mul					r4,			r4,			v1					; Modulate by vertex color						\n\
	mov					oD0,		r4								; Set diffuse									\n\
																													\n\
	; Compute uv																									\n\
	mul					r2,			c20.z,		c20.z																\n\
	mul					r2,			r2,			c20.z																\n\
	mul					r2,			r2,			c20.z																\n\
	mul					r1.x,		v0.x,		r2.x																\n\
	mul					r1.y,		v0.z,		r2.y																\n\
	mov					oT0,		r1																				\n\
	";

	char *psSource = "\
																													\n\
	ps_2_0																											\n\
																													\n\
	; Declare texture samplers																						\n\
	dcl_2d				s0											; BASE											\n\
	dcl_2d				s1											; CAUSTIC										\n\
																													\n\
	; Declare texture coordinate registers																			\n\
	dcl					t0.xyzw																						\n\
																													\n\
	dcl					v0.xyzw																						\n\
																													\n\
	; Texture sampling																								\n\
	texld				r0,			t0,			s0					; Sample BASE									\n\
																													\n\
	mul					r0,			r0,			v0					; Multiply by diffuse color						\n\
																													\n\
	mov					oC0,		r0								; Move result into output						\n\
	";

	// Assemble vertex shader
	m_vertexShader = RuD3DAssembleVertexShader(vsSource);

	// Assemble pixel shader
	m_pixelShader = RuD3DAssemblePixelShader(psSource);

	// Overrides
	m_override_ZWrite = TRUE;
}

CRuShader_Water_PS20_Diffuse::~CRuShader_Water_PS20_Diffuse()
{
}

void CRuShader_Water_PS20_Diffuse::BeginPass(INT32 passNumber)
{
	// Call base BeginPass to perform any additional setup
	CRuShader_Water_PS20_Base::BeginPass(passNumber);

	if(passNumber == 0)
	{
		m_vertexDecl->Select();
		m_engine->SetVertexShader(m_vertexShader);
		m_engine->SetPixelShader(m_pixelShader);
		m_engine->GetD3DDevice()->SetVertexShaderConstantF(0, m_shaderConstants, m_numBaseConstants + 2);
		m_engine->GetD3DDevice()->SetPixelShaderConstantF(0, m_pixelShaderConstants, m_numBasePixelShaderConstants + 1);

		m_engine->SetRenderState(D3DRS_FILLMODE, D3DFILL_SOLID);

		m_engine->SetRenderState(D3DRS_ALPHATESTENABLE, FALSE);

		m_engine->SetRenderState(D3DRS_ALPHABLENDENABLE, TRUE);
		m_engine->SetRenderState(D3DRS_SRCBLEND, D3DBLEND_SRCALPHA);
		m_engine->SetRenderState(D3DRS_DESTBLEND, D3DBLEND_INVSRCALPHA);

		m_engine->SetRenderState(D3DRS_ZWRITEENABLE, TRUE);
		m_engine->SetRenderState(D3DRS_ZFUNC, D3DCMP_LESSEQUAL);

		m_engine->SetRenderState(D3DRS_CULLMODE, D3DCULL_NONE);

		m_engine->SetSamplerState(0, D3DSAMP_MINFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(0, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(0, D3DSAMP_MIPFILTER, D3DTEXF_LINEAR);

		m_engine->SetSamplerState(0, D3DSAMP_ADDRESSU, D3DTADDRESS_WRAP);
		m_engine->SetSamplerState(0, D3DSAMP_ADDRESSV, D3DTADDRESS_WRAP);

		m_engine->SetSamplerState(1, D3DSAMP_MINFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(1, D3DSAMP_MAGFILTER, D3DTEXF_LINEAR);
		m_engine->SetSamplerState(1, D3DSAMP_MIPFILTER, D3DTEXF_LINEAR);

		m_engine->SetSamplerState(1, D3DSAMP_ADDRESSU, D3DTADDRESS_WRAP);
		m_engine->SetSamplerState(1, D3DSAMP_ADDRESSV, D3DTADDRESS_WRAP);
	}
}

void CRuShader_Water_PS20_Diffuse::EndPass(INT32 passNumber)
{
	m_engine->SetRenderState(D3DRS_CULLMODE, D3DCULL_CCW);
}

// ************************************************************************************************************************************************************

#pragma managed
