float3 lightDir : register(c1);
float3 viewPos : register(c2);
float3 viewDir : register(c4);
float3 passID2expIsLast : register(c23);

float4 g_fGlow : register(c61);

sampler2D colorSampler : register(s0);
sampler2D depthSampler : register(s1);
sampler2D waveSampler : register(s3);
sampler2D shadowSampler : register(s4);

// Tweakables
static float4 g_fRayColor = float4(1.75f * 0.45f, 1.75f * 0.4f, 1.75f * 0.35f, 1.0f);
static float g_fRayNearPlane = 10.0f;
static float g_fRayFarPlane = 500.0f;
static float g_fRayDensity = 0.25f;
static float g_fRayDecay = 0.9125f;
static float g_fRayAnglePersistence = 4.f;

float4 ps_ray_mask(float2 TexCoord : TEXCOORD0, float2 LightPos : TEXCOORD2) : COLOR0
{
	float2 d = LightPos - TexCoord;
	float r = saturate( 1.f - dot(d, d) );
	r *= saturate( g_fRayAnglePersistence * dot(-lightDir, viewDir) );

	float m = tex2D(depthSampler, TexCoord).x;
	m = saturate( (m - g_fRayNearPlane) / (g_fRayFarPlane - g_fRayNearPlane) );
//	m = saturate(1.f - m);
	
//	float2 p = TexCoord * 2.0f - 1.0f;
//	float b = saturate( 1.0f - 0.7071f * dot(p, p) );

	return /* b */ m * r * g_fRayColor;
}

// Pixel shader
float4 ps_ray_extrude(float2 TexCoord : TEXCOORD0, float2 LightPos : TEXCOORD2) : COLOR0
{	
	float2 s = g_fRayDensity / (passID2expIsLast.y) * (LightPos - TexCoord) / 8.f;
	
	float4 r = 0.f, ro = 0.f, rm = 0.f, rd = .125f;
	float2 c = TexCoord;
	
	for(int i = 0; i < 8; i++)
	{
		float4 rs = tex2D(colorSampler, c);
		ro = (i == 0) ? rs : ro;
		
		rm = max(rs, rm);
		rs.xyz = rm.xyz;
		r += rs * rd;
		
		rd.xyz *= (float3)g_fRayDecay;
		c += s;
	}
	
	r.xyz *= 1.0f + 0.25f * g_fGlow.z;
	r.a = min(ro.a, r.a);

	return r * saturate(1.f - passID2expIsLast.z * r.a);
}

float3 g_fTime : register(c3);

// Tweakables
static float g_fWaterLevel = 5.0f;

static float4 g_fLightFilterColor = float4(0.4f, 0.75f, 0.8f, 0.99f); // float4(0.8f, 0.85f, 0.9f, 0.8f);
static float g_fFilterOffset = 0.0f;
static float g_fLightFilterDensity = 0.02f;

static float g_fCausticScaling = 0.004f;
static float g_fCausticDistortion = 0.1f;
static float4 g_fCausticPatternScaling = 0.25f;
static float2 g_vCausticSpeed = float2(0.05f, 0.1f);
static float g_fCausticBrightness = 0.35f;

float4 ps_underwater_filter(float2 TexCoord : TEXCOORD0, float3 EyeDir : TEXCOORD1,
							float2 LightPos : TEXCOORD2) : COLOR0
{
	// Compute eye space position
	float fEyeDepth = tex2D(depthSampler, TexCoord).x;
	float3 vEyePoint = viewPos + fEyeDepth * EyeDir;
	float3 vEyeNormal = normalize( cross( ddx(vEyePoint), ddy(vEyePoint) ) );
	
	// Compute distance to water level
	float fUnderWater = g_fWaterLevel - vEyePoint.y;
	float fLightTravel = max( fUnderWater / max(-lightDir.y, 0.0001f), 0.0f );
	
	// Intersect water plane
	float3 vSubmerged = vEyePoint - viewPos;
	float fSubmergedBelow = (fUnderWater >= 0.0f)
		? abs(vSubmerged.y)
		: (g_fWaterLevel - viewPos.y);
	float fSubmerged = (viewPos.y > g_fWaterLevel)
		? max(fUnderWater, 0.0f)
		: fSubmergedBelow;
	vSubmerged *= fSubmerged / max( abs(vSubmerged.y), 0.0001f );

	// Compute submerged length
	fSubmerged = length(vSubmerged);
	
	// Compute light color
	float4 fFilterDist= fLightTravel + float2(fSubmerged, 0.0f).xxxy;
	float4 fLightColor = pow(g_fLightFilterColor, g_fLightFilterDensity * fFilterDist );
	
	// Color hack
	fLightColor.xyz *= pow(g_fLightFilterColor.xyz, max(fLightTravel - 30, 0) * 0.08f);
	
	// Compute light coordinate space
	float3 vLightU = normalize( cross( lightDir, float3(1.0f, 0.0f, 0.0f) ) );
	float3 vLightV = normalize( cross( vLightU, lightDir ) );
	
	// Scale time
	float fTime = g_fTime.z * 0.000095f;
	
	// Animate caustics
	float2 fCausticCoord = float2( dot(vLightU, vEyePoint), dot(vLightV, vEyePoint) ) * g_fCausticScaling;
	float2 fCausticCoordDelta = fTime * g_vCausticSpeed * g_fCausticScaling;
	
	float3 bx = float3(2.0f, 1.0f, 0.0f);
	
	// Get caustic normal
	float3 vCausticNormal1 = normalize( tex2D(waveSampler, fCausticCoord - fCausticCoordDelta).xzy * bx.xyx - bx.yzy );
	float3 vCausticNormal2 = normalize( tex2D(waveSampler, fCausticCoord.yx + fCausticCoordDelta).xzy * bx.xyx - bx.yzy );
	float3 vCausticNormal3 = normalize( tex2D(waveSampler, 5 * (fCausticCoord - 2 * fCausticCoordDelta)).xzy * bx.xyx - bx.yzy );
	float3 vCausticNormal4 = normalize( tex2D(waveSampler, 5 * (fCausticCoord.yx + 2 * fCausticCoordDelta)).xzy * bx.xyx - bx.yzy );
	
	// Compute normals used for distortion & lighting
	float3 vCausticNormalX1 = vCausticNormal1 + vCausticNormal4 / 3.0f;
	float3 vCausticNormalX2 = vCausticNormal2 + vCausticNormal3 / 3.0f;
	float3 vCausticNormal = vCausticNormalX1 + vCausticNormalX2;
	vCausticNormal = normalize(vCausticNormal);
	
	// Compute caustic pattern
	float4 fCausticPattern = float2( dot(vEyePoint, vLightU), dot(vEyePoint, vLightV) ).xyxy;
	fCausticPattern.xy += g_fCausticDistortion * vCausticNormalX1.xz * (1.0f + fEyeDepth); // + 1000 * dc.yx;
	fCausticPattern.zw += g_fCausticDistortion * vCausticNormalX2.xz * (1.0f + fEyeDepth); // - 1000 * dc.xy;
	fCausticPattern *= g_fCausticPatternScaling;
	fCausticPattern = 2.0f * (fCausticPattern - round(fCausticPattern));
	float fCaustics = dot(fCausticPattern, fCausticPattern) * g_fCausticBrightness;

	// Sharpen
	fCaustics *= fCaustics;
	fCaustics *= fCaustics;
	
	// Cut at surface
	fCaustics *= saturate(fUnderWater);
	
	// Remove underside lighting
	fCaustics *= saturate( dot(vEyeNormal, -lightDir) ) * tex2D(shadowSampler, TexCoord).x;
	
	float3 vSurfacePoint = vEyePoint - vSubmerged;
	vSurfacePoint.y = 0.0f;
	float fFadeOut = saturate( max(dot(vSurfacePoint, vSurfacePoint) - 2000000.0f, 0.0f) / 10000000.0f )
		* saturate(fUnderWater);
	
	// Filter light
	return lerp(float4(fLightColor.xyz, fLightColor.w * fCaustics), float2(1, 0).xxxy, fFadeOut);
}

// Tweakables
static float4 g_fFogColor = float4(0.1f, 0.18f, 0.17f, 0.0f); // float4(0.6f, 0.7f, 0.9f, 0.25f);
static float g_fFogDensity = 0.7f;
static float g_fSceneFog = 0.6f;
static float g_fRayScaling = 0.024f;
static float2 g_vRaySpeed = float2(0.03f, 0.02f);

float4 ps_underwater(float2 TexCoord : TEXCOORD0, float3 EyeDir : TEXCOORD1,
					 float2 LightPos : TEXCOORD2) : COLOR0
{
	// Compute eye space position
	float fEyeDepth = tex2D(depthSampler, TexCoord).x;
	float3 vEyePoint = viewPos + fEyeDepth * EyeDir;
	
	// Compute distance to water level
	float fUnderWater = g_fWaterLevel - vEyePoint.y;
	
	// Intersect water plane
	float3 vSubmerged = vEyePoint - viewPos;
	float fSubmergedBelow = (fUnderWater >= 0.0f)
		? abs(vSubmerged.y)
		: (g_fWaterLevel - viewPos.y);
	float fSubmerged = (viewPos.y > g_fWaterLevel)
		? max(fUnderWater, 0.0f)
		: fSubmergedBelow;
	vSubmerged *= fSubmerged / max( abs(vSubmerged.y), 0.0001f );

	// Compute submerged length & normalize
	fSubmerged = length(vSubmerged);
	float3 vSubmergedDir = vSubmerged / max(fSubmerged, 0.0001f);
	
	// Pre-compute light travel quantities
	float fEyeLightTravel = max( (g_fWaterLevel - viewPos.y) / max(-lightDir.y, 0.0001f), 0.0f );
	float fDeltaLightTravel = vSubmergedDir.y / min(lightDir.y, -0.0001f);
	float fSubmergedLightTravel = fEyeLightTravel + fDeltaLightTravel * fSubmerged;
	
	float4 fLogLightFilterColor = log(g_fLightFilterColor);
	float4 fFog = g_fFogColor / (fLogLightFilterColor * g_fLightFilterDensity * fDeltaLightTravel - g_fFogDensity);
	fFog *= exp(fLogLightFilterColor * (g_fFilterOffset + g_fLightFilterDensity * fSubmergedLightTravel) - g_fFogDensity * fSubmerged)
		- exp(fLogLightFilterColor * (g_fFilterOffset + g_fLightFilterDensity) * fEyeLightTravel);
	
	// "Absorb" more scene light
	fFog.w = g_fSceneFog * (1.0f - exp(-g_fFogDensity * fSubmerged));
//	fFog.w = dot(fFog.xyz, 0.1f);
//	fFog.w = 1.0f - fFog.w;
//	fFog.w = 0;

	// Compute light coordinate space
	float3 vLightU = normalize( cross( lightDir, float3(1.0f, 0.0f, 0.0f) ) );
	float3 vLightV = normalize( cross( vLightU, lightDir ) );

	//  Parallel ray direction vector
	float3 vRayDir = EyeDir - lightDir * dot(lightDir, EyeDir);
	vRayDir = normalize(vRayDir);
	float3 vFullRayDir = EyeDir / dot(EyeDir, vRayDir);
	
	// 1 / submerged flat
	float3 vSubmergedFlatInvScale = vSubmerged - lightDir * dot(lightDir, vSubmerged);
	vSubmergedFlatInvScale /= max( dot(vSubmergedFlatInvScale, vSubmergedFlatInvScale), 0.0001f );
	
	// 4 ray Planes
	float4 fPlanes = float4(5.0f, 10.0f, 20.0f, 40.0f);
	float4 fPlaneIntersect = fPlanes * dot(vRayDir, vSubmergedFlatInvScale);
	float4 fPlaneSubmerged = vSubmerged.y * fPlaneIntersect;
//	float fWeightScaling = saturate( 1.0f - 0.1f * dot(viewDir, -lightDir) ); // polar
	float4 fPlaneWeight = 1.0f - saturate( 0.5f * (fPlaneIntersect - 1.0f) );
	
	// Compute light travel
	float4 fPlanesY = g_fWaterLevel - (viewPos.y + vFullRayDir.y * fPlanes);
	float4 fPlanesLightTravel = max( fPlanesY / max(-lightDir.y, 0.0001f), 0.0f );
	float4 fPlanesLightInt = 1.0f / (1.0f + 0.01f * fPlanesLightTravel * fPlanesLightTravel);
	
//	return float4(fPlanesLightInt.yyy, 1.0f);

	// Snap view pos
//	float3 vSnappedViewPos = vLightU * floor( dot(vLightU, viewPos) / 25.0f ) * 25.0f
//		+ vLightV * floor( dot(vLightV, viewPos) / 25.0f ) * 25.0f;
//	fPlanes += dot(viewPos - vSnappedViewPos, vRayDir);
	
	// Scale time
	float fTime = g_fTime.z * 0.000075f;
	
	// Animate rays
	float4 fRayCoordsX = dot(vLightU, viewPos) + dot(vLightU, vRayDir) * fPlanes;
	float4 fRayCoordsY = dot(vLightV, viewPos) + dot(vLightV, vRayDir) * fPlanes;
//	fPlanes = dot(vLightU, vRayDir) * tx - dot(vLightV, vRayDir) * ty - 2.0f * dot(viewPos, vRayDir);
	fRayCoordsX *= g_fRayScaling; fRayCoordsY *= g_fRayScaling;
	float2 fRayCoordDelta = fTime * g_vRaySpeed * g_fRayScaling;
	
	// Sample rays
	float4 fRay1 = tex2D(waveSampler, float2(fRayCoordsX.x, fRayCoordsY.x) + fPlanes.x * fRayCoordDelta).xzyw;
	float4 fRay2 = tex2D(waveSampler, float2(fRayCoordsX.y, fRayCoordsY.y) - fPlanes.y * fRayCoordDelta).xzyw;
	float4 fRay3 = tex2D(waveSampler, float2(fRayCoordsX.z, fRayCoordsY.z) + fPlanes.z * fRayCoordDelta).xzyw;
	float4 fRay4 = tex2D(waveSampler, float2(fRayCoordsX.w, fRayCoordsY.w) - fPlanes.w * fRayCoordDelta).xzyw;
	
	// Decode Normal
	float3 vRayNormal = (fRay1.xyz + fRay2.xyz + fRay3.xyz + fRay4.xyz) / 4.0f;
	vRayNormal.xz = vRayNormal.xz * 2.0f - 1.0f;
	vRayNormal = normalize(vRayNormal);
	
	float4 fLight;
	
	// Light
	fLight.x = fRay1.w;
	fLight.y = fRay2.w;
	fLight.z = fRay3.w;
	fLight.w = fRay4.w;

//	fLight *= fLight * fLight;

	float fLightSubmerged = 0.05f * max( fPlaneIntersect * dot(vSubmerged, lightDir), 0.0f);
	
	// Fade rays
	fLight *= fPlanesLightInt * fPlaneWeight * saturate(fSubmerged) / (1.0f + fLightSubmerged * fLightSubmerged);
	fLight *= saturate(g_fWaterLevel - viewPos.y + 1.0f);
	
	float fRayAngle = saturate( 1.0f - dot(vRayNormal, -lightDir) );
	float fRayColor = 0.1f * saturate( dot(EyeDir, vRayNormal - lightDir) );
	fRayAngle = frac( fRayAngle );
	
	float3 fLightColor = 1.0f;
	fLightColor.r += saturate( 1.0f - pow(3.0f * fRayAngle - 3.0f * (fRayAngle > 1.0f / 3.0f), 2) ) * fRayColor;
	fLightColor.g += saturate( 1.0f - pow(3.0f * fRayAngle - 1.0f, 2) ) * fRayColor;
	fLightColor.b += saturate( 1.0f - pow(3.0f * fRayAngle - 2.0f, 2) ) * fRayColor;
	
	// God rays
	fFog.xyz += dot(fLight, 0.04f) * fLightColor;
	
	float3 vSurfacePoint = vEyePoint - vSubmerged;
	vSurfacePoint.y = 0.0f;
	float fFadeOut = 1.0f - saturate( max(dot(vSurfacePoint, vSurfacePoint) - 2000000.0f, 0.0f) / 10000000.0f )
		* saturate(fUnderWater);

	return fFog * fFadeOut;
}