// OBGE Custom Water
// by Tomerk

/* User Editable Input Variables, Default Values In The Comments */

static const float causticsStrength = 0.5; //The strength of the caustics (0.35)
static const float choppiness = 1.0; //How choppy the waves are. The higher this value, the choppier the water (1)
static const float turbidity = 1.0; //Density of particles in the water. Lower values make water more transparent (1)
static const float inScattCoeff = 0.18; //Scatter Coefficient of the water (0.18)
static const float3 inExtCoeff = float3(0.3984, 0.2684, 0.2602); //Extinction Coefficient of the water (0.3984, 0.2684, 0.2602)
static const float depthDarkness = 0.1; //Affects how quickly the water volume gets darker with depth underwater (0.1)

/* Plugin Input Variables */

extern float waterHeight = 0; //Height of the water in the cell
extern float4 sunColor = float4(1, 0.71, 0.33, 1);//

extern float fogStart = 4096; //above water fog start
extern float fogEnd = 170000; //above water fog end
extern float4 fogColor = float4(0.39,0.57,0.75,1); //above water fog color

extern float lava = 0; //Whether its currently rendering lava

extern float FOV = 75;
extern float frame = 0;

/* End of Input Variables */

static const float3 extCoeff = inExtCoeff * turbidity;
static const float scattCoeff = inScattCoeff * turbidity;

texture thisframe;
texture Depth;
texture watertexture
<
string filename="effects\\water_NRM.dds";
>;
texture lavatexture
<
string filename="water\\alternatelavaX.dds";
>;
texture reflection;

sampler frameSampler = sampler_state { texture = ; AddressU=Clamp; AddressV=Clamp;};
sampler depthSampler = sampler_state { texture = ; AddressU=Clamp; AddressV=Clamp; };
sampler waterSampler = sampler_state { texture = ; AddressU=Wrap; AddressV=Wrap; MagFilter=Linear; MinFilter = Linear;};
sampler lavaSampler = sampler_state { texture = ; AddressU=Wrap; AddressV=Wrap; MagFilter=Linear; MinFilter = Linear;};
sampler reflectionSampler = sampler_state { texture = ; AddressU=Clamp; AddressV=Clamp;};

matrix m44world;
matrix m44view;
matrix m44proj;
float4 f4SunDir;
float2 rcpres;
bool bHasReflection; //Whether the reflection map is available

static const float animFrames = 32;
static const float3 eyepos = float3( -m44world[3][0], -m44world[3][1], -m44world[3][2] - waterHeight );

static const float nearZ = m44proj._43 / m44proj._33;
static const float farZ = (m44proj._33 * nearZ) / (m44proj._33 - 1.0f);
static const float Zmul = nearZ * farZ;
static const float Zdiff = farZ - nearZ;
static const float depthRange = nearZ - farZ;

struct VSOUT
{
float4 vertPos : POSITION;
float2 UVCoord : TEXCOORD0;
};

struct VSIN
{
float4 vertPos : POSITION0;
float2 UVCoord : TEXCOORD0;
};

VSOUT FrameVS(VSIN IN)
{
VSOUT OUT = (VSOUT)0.0f; // initialize to zero, avoid complaints.
OUT.vertPos = IN.vertPos;
OUT.UVCoord = IN.UVCoord;
return OUT;
}

float3 toWorld(float2 tex)
{
float3 v = float3(m44view[0][2], m44view[1][2], m44view[2][2]);
v += (1/m44proj[0][0] * (2*tex.x-1)).xxx * float3(m44view[0][0], m44view[1][0], m44view[2][0]);
v += (-1/m44proj[1][1] * (2*tex.y-1)).xxx * float3(m44view[0][1], m44view[1][1], m44view[2][1]);
return v;
}

float readDepth(in float2 coord : TEXCOORD0)
{
float posZ = tex2D(depthSampler, coord).x;
posZ = Zmul / ((posZ * Zdiff) - farZ);
return posZ;
}

float correctDepth(in float depth, in float2 coord : TEXCOORD0)
{
coord = (coord * float2(2.0, -2.0) - float2(1.0, -1.0));
float corr = pow(cos(radians((FOV/2) * coord.x)), 1.5) * pow(cos(radians((FOV/2) * coord.y)), 1.5);
return depth * corr;
}

float2 get2DTex( float3 tex )
{
float2 return_tex = float2(0, tex.y);
return_tex.x = frac(tex.x)/animFrames + round( tex.z*(animFrames) )/animFrames;
return return_tex;
}

float3 getWaterNorm( float2 tex, float dist, float camera_vector_z, inout float3 specNorm )
{
float lod = 1;
float2 uvLOD = get2DTex( float3(tex / 1024, frame) );
float2 temp_norm = tex2Dlod( waterSampler, float4(uvLOD,0,0) ).rg * 2 - 1;
float3 norm = normalize(float3(temp_norm * choppiness * lod,1));
specNorm = normalize(float3(temp_norm * choppiness * max(0.5, lod), 1));
return norm;
}

float getFresnelAboveWater( float3 ray, float3 norm )
{
float temp_cos = dot( -ray, norm );
float2 vec = float2(temp_cos, sqrt(1-temp_cos*temp_cos));

float fresnel = vec.x - 1.33 * sqrt(1 - 0.565*vec.y*vec.y);
fresnel /= vec.x + 1.33 * sqrt(1 - 0.565*vec.y*vec.y);
fresnel = saturate(fresnel * fresnel);

return fresnel;
}

float getFresnelBelowWater( float3 ray, float3 norm )
{
float temp_cos = dot( ray, norm );
temp_cos = max(temp_cos, 0.659325);
float2 vec = float2(temp_cos, sqrt(1-temp_cos*temp_cos));

float fresnel = 1.33 * vec.x - sqrt(1 - 1.769*vec.y*vec.y);
fresnel /= 1.33 * vec.x + sqrt(1 - 1.769*vec.y*vec.y);
fresnel = saturate(fresnel * fresnel);

return fresnel;
}

float4 Water( VSOUT IN ) : COLOR0
{
float4 color = tex2D(frameSampler, IN.UVCoord);
float depth = readDepth(IN.UVCoord);
float3 camera_vector = toWorld(IN.UVCoord);
float3 norm_camera_vector = normalize( camera_vector );
float3 world_pos = eyepos + camera_vector*depth;

float falloffStart = correctDepth( fogStart, IN.UVCoord );
float falloffEnd = correctDepth( fogEnd, IN.UVCoord );

/* Render water surface from above water */
if ( (world_pos.z <= 0) && (eyepos.z >= 0) )
{
float4 refract_color = color;
float uw_pos = world_pos.z / camera_vector.z;
float2 surfPos = world_pos.xy - camera_vector.xy * uw_pos;

float3 specNorm = float3(0,0,0);
float3 normal = getWaterNorm( surfPos, depth - uw_pos, -camera_vector.z, specNorm);

float2 refPos = IN.UVCoord + 0.01*normal.yx;
float3 refract_world_pos = eyepos + toWorld( refPos )*readDepth( refPos );

if (refract_world_pos.z <= 0)
refract_color = tex2D( frameSampler, refPos );

else
refract_world_pos = world_pos;

//Render Caustics
float3 dx = ddx(world_pos);
float3 dy = ddy(world_pos);
float3 waterfloorNorm = normalize(cross(dx,dy));

float3 causticsPos = refract_world_pos - f4SunDir.xyz * (refract_world_pos.z / f4SunDir.z);
float caustics = causticsStrength*tex2D(waterSampler, get2DTex(float3(causticsPos.xy / 512,frame)) ).b;
float causticsAngle = saturate( dot(-waterfloorNorm, f4SunDir.xyz) );
refract_color.rgb *= 1 + caustics * causticsAngle * sunColor.xyz - 0.3;

//Calculate Refraction color
float refract_uw_pos = refract_world_pos.z / camera_vector.z;
refract_color.rgb *= exp( -extCoeff * (refract_uw_pos - refract_world_pos.z) / 70 );

float SinBoverSinA = -norm_camera_vector.z;
float3 waterVolColor = scattCoeff * fogColor.xyz / ( extCoeff * (1 + SinBoverSinA) );

waterVolColor *= 1 - exp( -extCoeff * (1 + SinBoverSinA) * refract_uw_pos / 70 );

refract_color.rgb += waterVolColor;

//Calculate reflection color
float4 reflection = fogColor;

if (bHasReflection)
{
refPos = IN.UVCoord + 0.05*normal.yx;
float3 reflect_world_pos = eyepos + toWorld( refPos )*readDepth( refPos );

if (reflect_world_pos.z > 0)
refPos = IN.UVCoord;

reflection = tex2D(reflectionSampler, float2(refPos.x,1-refPos.y) );
}

float4 fresnel = getFresnelAboveWater( norm_camera_vector, normal );

float4 water_result = lerp( refract_color, reflection, fresnel );

float specular = saturate(dot( norm_camera_vector, reflect( f4SunDir.xyz, specNorm ) ));
water_result.xyz += pow(specular, (depth-uw_pos)*0.3 + 750) * sunColor.xyz;

//Add above water fog
float fog = saturate( (depth-uw_pos-falloffStart)/(falloffEnd-falloffStart) );
water_result = lerp(water_result, reflection, fog);

//Smooth shore transitions
color = lerp(water_result, color, saturate( pow(saturate(exp(world_pos.z/800)), 90) ));

if (lava)
{
float lod = saturate( (norm_camera_vector.z*norm_camera_vector.z) * 64 * m44proj[0][0] /(rcpres.x * eyepos.z/-norm_camera_vector.z) );
color = lerp( float4(1,0.33,0.1,1),tex2Dlod(lavaSampler, float4(surfPos / 2048,0,0) ), lod );
color.r += 0.6;
}
}

/* Underwater rendering */
if ( (eyepos.z < 0) )
{
/* Render Underwater Caustics */
float sunDist = max( -world_pos.z / f4SunDir.z, 0 );
if ( world_pos.z < 0 )
{
float3 dx = ddx(world_pos);
float3 dy = ddy(world_pos);
float3 waterfloorNorm = normalize(cross(dx,dy));

float3 causticsPos = world_pos + f4SunDir.xyz * sunDist;
float caustics = causticsStrength*tex2D(waterSampler, get2DTex(float3(causticsPos.xy / 512,frame)) ).b;
float causticsAngle = saturate( dot(-waterfloorNorm, f4SunDir.xyz) );
color.rgb *= 1 + caustics * causticsAngle * sunColor.xyz - 0.3;
}

/* Render Water surface from underwater */
float uw_pos = abs( eyepos.z / camera_vector.z );

float specular = 0;

if ( world_pos.z > 0 )
{
float2 surfPos = eyepos.xy + camera_vector.xy * uw_pos;

float3 specNorm = float3(0,0,0);
float3 normal = getWaterNorm( surfPos, uw_pos, camera_vector.z, specNorm );
float2 refPos = IN.UVCoord + 0.02 * normal.yx;

float refDepth = readDepth( refPos );
float3 refWorldPos = eyepos + toWorld( refPos )*refDepth;
float4 refract_color = color;

if (refWorldPos.z > 0)
refract_color = tex2D(frameSampler, refPos);

else
refDepth = depth;

if (refDepth / depthRange >= 0.999)
refract_color = sqrt(refract_color);

else
{
/* Render Above Water Fog */
float fog = saturate( (refDepth + eyepos.z / camera_vector.z - falloffStart)/(falloffEnd-falloffStart) );
refract_color = lerp(refract_color, fogColor, fog);
}

float fresnel = getFresnelBelowWater( norm_camera_vector, normal );

specular = saturate(dot( float3(1,1,-1)*norm_camera_vector, reflect(f4SunDir.xyz, specNorm) ));
specular = pow(specular, 750)*(1-fresnel);

float4 reflection = float4( scattCoeff * fogColor.xyz / ( extCoeff * (1) ), 1);

color = lerp( refract_color, reflection, fresnel );

if (lava)
{
float lod = saturate( (norm_camera_vector.z*norm_camera_vector.z) * 64 * m44proj[0][0] /(rcpres.x * eyepos.z/-norm_camera_vector.z) );
color = lerp( float4(1,0.33,0.1,1),tex2D(lavaSampler, float4(surfPos / 2048,0,0)), lod );
color.r += 0.6;
}
}

/* Render Underwater Light Scattering and Absorption */
float waterfogDist = depth;
float3 waterVolColor = float3(0,0,0);
float SinBoverSinA = abs(norm_camera_vector.z);

if (camera_vector.z > 0)
{
float SurfRay = -eyepos.z / camera_vector.z;
waterfogDist = min( SurfRay, depth );

float v = saturate(SinBoverSinA) * 0.5;
waterVolColor = scattCoeff * fogColor.xyz / ( extCoeff * (1 - v) );

if (lava)
waterVolColor = float3(1,0.55,0.1);

waterVolColor *= ( exp( extCoeff * (1 - v) * eyepos.z*depthDarkness / 70 ) - exp( -extCoeff * (1 - v) * (waterfogDist) / 70 ) );
}

else
{
waterVolColor = scattCoeff * fogColor.xyz / ( extCoeff * (1 + SinBoverSinA) );

if (lava)
waterVolColor = float3(1,0.55,0.1);

waterVolColor *= ( exp( extCoeff * (1 + SinBoverSinA) * eyepos.z*depthDarkness / 70 ) - exp( -extCoeff * (1 + SinBoverSinA) * (waterfogDist) / 70 ) );
}

color.rgb *= exp( -extCoeff * ( max(-world_pos.z + eyepos.z*(1-depthDarkness*depthDarkness), 0) + waterfogDist ) / 70 );
color.rgb += waterVolColor;
color.rgb += exp( -extCoeff * uw_pos*depthDarkness / 160 )*specular * sunColor.xyz;
}

return color;
}

technique t0
{
pass p0
{
VertexShader = compile vs_3_0 FrameVS();
PixelShader = compile ps_3_0 Water();
}
}