/**
 * @ Version: SCREEN SPACE SHADERS - UPDATE 18
 * @ Description: Water implementation
 * @ Modified time: 2023-09-30 10:46
 * @ Author: https://www.moddb.com/members/ascii1457
 * @ Mod: https://www.moddb.com/mods/stalker-anomaly/addons/screen-space-shaders
 */

#include "common.h"
#include "anomaly_shaders.h"
#include "reflections.h"
#include "lmodel.h"

// Screen space functions
#include "check_screenspace.h"
#include "screenspace_water.h"

#ifdef SSFX_FOG
	#include "screenspace_fog.h"
#endif

struct vf
{
	float2	tbase		: TEXCOORD0;	// base
	float4	tnorm0		: TEXCOORD1;	// nm0
	float3	position_w	: TEXCOORD2;	// nm1
	float3	M1			: TEXCOORD3;
	float3	M2			: TEXCOORD4;
	float3	M3			: TEXCOORD5;
	float3	v2point_w	: TEXCOORD6;
	float4	tctexgen	: TEXCOORD7;
	float4	c0			: COLOR0;
	float	fog			: FOG;
	float4	hpos		: SV_Position;
};

Texture2D	s_nmap;
Texture2D	s_bluenoise;
Texture2D	s_rainsplash;
Texture2D	s_watercaustics;

uniform float4 ssfx_is_underground;
//float3	water_intensity; // Unused?

float4 main( vf I ) : SV_Target
{
	// 3d view space pos reconstruction math
	float3 Pv = float3(I.tctexgen.z * (I.hpos.xy * pos_decompression_params.zw - pos_decompression_params.xy), I.tctexgen.z);

	// Normal Textures
	float3	n0	= s_nmap.Sample(smp_base, I.tnorm0.xy);
	float3	n1	= s_nmap.Sample(smp_base, I.tnorm0.zw);
	float3	basenormal = ((n0 + n1) * 0.5f) * 2.0f - 1.0f;
	basenormal += float3(0.1f, 0.1f, 0.1f); // Offset the normal to "center" the textures again
	
	// Load Data
	float2 PosTc = I.tctexgen.xy / I.tctexgen.z;
	gbuffer_data gbd = gbuffer_load_data( PosTc, I.hpos.xy );
	float4 _P = float4( gbd.P, 1.0f );

	// Water Surface to world space
	float3 w_s = mul(m_inv_V, Pv );

	// Trick to check if rain can fall on the water, gloss include the rain_apply_gloss
	float rain_cover = step(0.035f * (-0.4f + rain_params.x), gbd.gloss - 0.04f );
	
	// Rain intensity. Fadeout and rain cover
	half is_raining = rain_params.x > 0 ? 1.0f : 0.0f;
	float RainInt = clamp(rain_params.x * 1.6f, 0.65f * is_raining, 1.0f);
	
	// Ripples normal
	float2 Ripples = ssfx_rain_ripples( s_rainsplash, (w_s.xz + eye_position.xz) * 0.13f, float3(RainInt, G_SSR_WATER_RAIN, 15.0f), _P.z);
	basenormal += float3(Ripples, 0) * is_raining;
	
	// Water wave intensity
	float3 Navg = normalize(float3(basenormal.x * G_SSR_WATER_WAVES, basenormal.y * G_SSR_WATER_WAVES, basenormal.z));
		
	// Refractions normal ( Used by shadows and refraction )
	float2 N_Scr = normalize(float3(basenormal.x * G_SSR_WATER_REFRACTION, basenormal.y * G_SSR_WATER_REFRACTION, basenormal.z));
	
	// Discard refractions from things above the water ( Like weapons )
	float Refraction_Discard = Pv.z < SSFX_get_depth(PosTc + N_Scr, 0);

	// Screen UV + Discard
	float2 Refraction_UV = N_Scr * Refraction_Discard;

	// Shadows to do some tricks
#ifndef USE_MSAA
	float water_shadows = saturate(s_accumulator.Sample(smp_nofilter, PosTc + Refraction_UV).r * 2000);
#else
	float water_shadows = saturate(s_accumulator.Load(int3((PosTc + Refraction_UV) * screen_res.xy, 0), 0).r * 2000);
#endif

	float3	Nw	= mul (float3x3(I.M1, I.M2, I.M3), Navg);
			Nw	= normalize (Nw);
	
	float3 v2point	= normalize (I.v2point_w);
	float3 vreflect = reflect(v2point, Nw.xyz);

	float3 eyedir = normalize(Pv);
	float3 Nv = normalize(mul(m_V, Nw));

	// Some vars to put our reflection
	float3 refl_ray, refl_skybox;
	float3 reflection;
	float3 ssr_hit_uv = 0.0f;
		
#ifdef NEED_REFLECTIONS
	
	// Blue Noise & Normal for noise
	float3 NN = normalize(float3(basenormal.x * 0.15f, basenormal.y * 0.15f, basenormal.z));
	float blue_noise = s_bluenoise.Sample(smp_linear, I.tbase * float2(0.05f, 0.05f) + NN).b * 1.5f;
	
	// Compute reflection bounce
	float3 wreflect = reflect(eyedir, Nv);

	// Don't trace rays which face the camera. Still worth to avoid the rays mess when you look down.
	float edgeFade = step(dot(-eyedir, wreflect), 0);
	
	// Trace a ray
	if (edgeFade > 0.02f)
		ssr_hit_uv = SSFX_ssr_water_ray(Pv, wreflect, blue_noise, 0);

	// Get current Skybox
	refl_skybox = SSFX_calc_sky(vreflect) * G_SSR_WATER_SKY_REFLECTION;

	if (all(ssr_hit_uv.xy))
	{
		// Get reflection pixel from scene screen
		refl_ray = SSFX_get_image(ssr_hit_uv.xy, 0);

		// Adjust reflection intensity using ssr_hit_uv.y and edgeFade
		float ray_fade = ssr_hit_uv.y * 5.0f;
		
		// Reflection fog fadeout
		float fog = 1.0 - saturate(( length(float3(Pv.x,Pv.y,ssr_hit_uv.z)) * fog_params.w + fog_params.x) * 1.4f);
		
		float refl_power = saturate(ray_fade * edgeFade * fog);
		
		// Fallback to Skybox
		reflection = lerp(refl_skybox * !ssfx_is_underground.x, refl_ray * G_SSR_WATER_MAP_REFLECTION, refl_power);
	}
	else
	{
		// No reflection data, we use only refl_skybox
		reflection = refl_skybox * !ssfx_is_underground.x;
	}

#else
	reflection = L_hemi_color.rgb;
#endif

	// Fresnel
#ifndef G_SSR_WATER_CHEAPFRESNEL
	float fresnel		 = SSFX_calc_fresnel(eyedir, Nv, 1.33f); 
	float fresnel_amount = saturate(fresnel * 1.5);
#else
	float fresnel		 = saturate (dot(vreflect, v2point));
	float fresnel_amount = pow(fresnel, 3.0);
#endif

	// Normal for the base texture...
	float2 N_Tex = normalize(float3(basenormal.x * G_SSR_WATER_TEX_DISTORTION, basenormal.y * G_SSR_WATER_TEX_DISTORTION, basenormal.z));

	// Get Position with Refraction discard
	float3 _P2 = SSFX_get_position(PosTc + Refraction_UV, 0);

	// 3d view space pos reconstruction math
	_P2 = float3( _P2.z * ( I.hpos.xy * pos_decompression_params.zw - pos_decompression_params.xy ), _P2.z);

	// Bottom of the water to world space ( Project the caustics and water fog )
	float3 w_b = mul(m_inv_V, float4(_P2 , 1));

	// Caustics offset anim
	float4 CausticsT = float4( timers.x * 0.1f, 0.0f, -timers.x * 0.07f, 0.2f );

	// Fake Caustics
	float3 Ca0 = saturate(s_watercaustics.Sample(smp_base, w_b.xz * (0.19f / G_SSR_WATER_CAUSTICS_SCALE) + CausticsT.xy + N_Scr ) - 0.1f );
	float3 Ca1 = saturate(s_watercaustics.Sample(smp_base, w_b.xz * (0.11f / G_SSR_WATER_CAUSTICS_SCALE) + CausticsT.zw + N_Scr ) - 0.1f );
	float3 CA = min(Ca0, Ca1) * water_shadows * G_SSR_WATER_CAUSTICS; // Caustics if light hit the water

	// Water Depth
	float waterDepth = w_s.y - w_b.y + eye_position.y;

	// Water Fog
	float waterFog = exp(waterDepth) - 1.0f;

	// Screen buffer
	float3 screen = SSFX_get_image(PosTc + Refraction_UV, 0);

	// Base texture * base color
	float3 base_tex = s_base.Sample(smp_base, (w_s.xz + eye_position.xz) * 0.1f + N_Tex).rgb;

	// Texture * base color
	base_tex *= I.c0.rgb;

	// Mix refraction ( 100% clear water ) & base texture ( Turbidity )
	float3 turbidity = lerp(base_tex, screen, smoothstep(G_SSR_WATER_FOG_MAXDEPTH, -G_SSR_WATER_TURBIDITY, waterFog));
	
	// Let's start the accumulation... First the water result and reflection.
	float3 acc = lerp(turbidity, reflection, fresnel_amount * G_SSR_WATER_REFLECTION);
	
	// Caustics. Fade through water fog
	acc = acc + CA * smoothstep(G_SSR_WATER_FOG_MAXDEPTH + 0.5f, 0.0f, waterFog) * saturate(waterFog * 3.0f);
		
	// Specular
#ifdef NEED_SPECULARS
	float3 Nf = normalize(float3(Nw.x, Nw.y * G_SSR_WATER_SPECULAR_NORMAL, Nw.z)); // flatten the water normal to get better specular
	float3 spec = L_sun_color.rgb * pow(abs(dot(normalize(v2point + L_sun_dir_w), Nf)), 512) * G_SSR_WATER_SPECULAR;
	acc += spec * water_shadows;
#endif

	// Fogging
#ifdef SSFX_FOG
	float fogging = SSFX_FOGGING(1.0 - I.fog, w_s.y);
#else
	float fogging = I.fog;
#endif
	acc = lerp(fog_color, acc, fogging);
	
	// Soft border
	float border_alpha = smoothstep( 0.0 , G_SSR_WATER_SOFTBORDER, waterDepth);
	
	// Done
	return  float4(acc, fogging * fogging * border_alpha);
}