/**
 * @ Version: SCREEN SPACE SHADERS - UPDATE 14.3
 * @ Description: SSR implementation
 * @ Modified time: 2023-01-29 08:40
 * @ Author: https://www.moddb.com/members/ascii1457
 * @ Mod: https://www.moddb.com/mods/stalker-anomaly/addons/screen-space-shaders
 */

#include "common.h"
#include "lmodel.h"
#include "hmodel.h"
#include "settings_screenspace_SSR.h"

TextureCube s_env0;
TextureCube s_env1;

static const int2 q_ssr_steps[6] =
{
	int2(8,200),
	int2(16,35),
	int2(24,18),
	int2(32,5),
	int2(48,1),
	int2(64,1),
};

static const float q_ssr_noise[6] =
{
	float(0.04f),
	float(0.04f),
	float(0.04f),
	float(0.06f),
	float(0.08f),
	float(0.08f),
};

#define MAT_FLORA 0.15f
#define MAT_FLORA_ELIPSON 0.04f
#define SKY_EPS float(0.001)

float SSFX_noise(float2 tc)
{
	float2	noise = frac(tc.xy * 0.5f);
	return noise.x + noise.y * 0.5f;
}

struct RayTrace
{
	float2 r_pos;
	float2 r_step;
	float2 r_start;
	float r_length;
	float z_start;
	float z_end;
};

float2 SSFX_view_to_uv(float3 Pos)
{
	float4 tc = mul(m_P, float4(Pos, 1));
	return (tc.xy / tc.w) * float2(0.5f, -0.5f) + 0.5f;
}

RayTrace SSFX_ray_init(float3 ray_start_vs, float3 ray_dir_vs, float ray_max_dist, int ray_steps, float noise)
{
	RayTrace rt;
	
	float3 ray_end_vs = ray_start_vs + ray_dir_vs * ray_max_dist;
	
	// Ray depth ( from start and end point )
	rt.z_start		= ray_start_vs.z;
	rt.z_end		= ray_end_vs.z;

	// Compute ray start and end (in UV space)
	rt.r_start		= SSFX_view_to_uv(ray_start_vs);
	float2 ray_end	= SSFX_view_to_uv(ray_end_vs);

	// Compute ray step
	float2 ray_diff	= ray_end - rt.r_start;
	rt.r_step		= ray_diff / (float)ray_steps; 
	rt.r_length		= length(ray_diff);
	
	rt.r_pos		= rt.r_start + rt.r_step * noise;
	
	return rt;
}

float SSFX_get_depth(float2 tc, uint iSample : SV_SAMPLEINDEX)
{
	#ifndef USE_MSAA
		return s_position.Sample(smp_nofilter, tc).z;
	#else
		return s_position.Load(int3(tc * pos_decompression_params2.xy, 0), iSample).z;
	#endif
}

float3 SSFX_ray_intersect(RayTrace Ray, uint iSample)
{
	float len = length(Ray.r_pos - Ray.r_start);
	float alpha = len / Ray.r_length;
	float depth_ray = (Ray.z_start * Ray.z_end) / lerp(Ray.z_end, Ray.z_start, alpha);
	float depth_scene = SSFX_get_depth(Ray.r_pos, iSample);
	
	return float3(depth_ray.x - depth_scene, depth_scene, len);
}

float4 SSFX_ssr_fast_ray(float3 ray_start_vs, float3 ray_dir_vs, float2 tc, uint iSample : SV_SAMPLEINDEX)
{
	float2 sky_tc = 0;
	float2 behind_hit = 0;

	// Noise to "improve" consistency between steps
	float3 noise = SSFX_noise(tc * float2(70,35) * 20) * q_ssr_noise[G_SSR_QUALITY];

	// Initialize Ray
	RayTrace ssr_ray = SSFX_ray_init(ray_start_vs, ray_dir_vs, 150, q_ssr_steps[G_SSR_QUALITY].x, 1.0f);

	// Save the original step.x
	float ori_x = ssr_ray.r_step.x;

	// Depth from the start of the ray
	float ray_depthstart = SSFX_get_depth(ssr_ray.r_start, iSample);
	
	// Ray-march
	[unroll (q_ssr_steps[G_SSR_QUALITY].x)]
	for (int i = 0; i < q_ssr_steps[G_SSR_QUALITY].x; i++)
	{
		// Ray out of screen...
		if (ssr_ray.r_pos.y < 0.0f || ssr_ray.r_pos.y > 1.0f)
			return 0;

		// Trick for the horizontal out of bounds. Mirror border of the screen.
		if (ssr_ray.r_pos.x < 0.0f || ssr_ray.r_pos.x > 1.0f)
		{
			ssr_ray.r_pos -= ssr_ray.r_step; // Step back
			ssr_ray.r_step.x = -ssr_ray.r_step.x; // Invert Horizontal
			ssr_ray.r_pos += ssr_ray.r_step; // Step
		}

		// Ray intersect check
		float2 ray_check = SSFX_ray_intersect(ssr_ray, iSample);

		// Sampled depth is not weapon or sky ( SKY_EPS float(0.001) )
		bool NoWpnSky = ray_check.y > 1.3f;

		// Disable weapon and sky
		ray_check.x *= NoWpnSky;

		// Return if ray is not reflecting backward
		if (ray_check.x > 0)
		{
			
			if (ray_check.x <= q_ssr_steps[G_SSR_QUALITY].y)
				return float4(ssr_ray.r_pos, 0, 0);

#if G_SSR_QUALITY > 2 // 1 Binary Search step in higher quality settigns ( Quality 4 & 5 )
			
			// Current ray pos & step to restore later...
			float4 prev_step = 0;
			prev_step.xy = ssr_ray.r_pos;
			prev_step.zw = ssr_ray.r_step;

			// Half and flip
			ssr_ray.r_step *= -0.5f;

			// Step ray
			ssr_ray.r_pos += ssr_ray.r_step;

			// Ray intersect check
			ray_check = SSFX_ray_intersect(ssr_ray, iSample);

			// Depth test... Conditions to use as reflections...
			if (abs(ray_check.x) <= 1.25f)
				return float4(ssr_ray.r_pos, 0, 0);

			// Restore previous ray position & step
			ssr_ray.r_pos = prev_step.xy;
			ssr_ray.r_step = prev_step.zw;

#endif

		}
		else
		{
			// TexCoor for sky ( Used to fade "SSFX_calc_SSR_fade" )
			if (ray_check.y <= SKY_EPS)
				sky_tc = ssr_ray.r_pos;
			
			behind_hit = ssr_ray.r_pos;

			// Reset or keep depending on... ( > 1.3f = no interaction with weapons and sky )
			behind_hit *= (ray_depthstart - 2.0f < ray_check.y) && NoWpnSky;
		}
		
		// Step the ray
		ssr_ray.r_pos += ssr_ray.r_step * (1.0f + noise.x * (1.0f - smoothstep(0, q_ssr_steps[G_SSR_QUALITY].x * 0.33f, i)));

	}

	return float4(behind_hit, sky_tc);
}

float3 SSFX_yaw_vector(float3 Vec, float Rot)
{
	float s, c;
	sincos(Rot, s, c);

	// y-axis rotation matrix
	float3x3 rot_mat = 
	{
		c, 0, s,
		0, 1, 0,
		-s, 0, c
	};
	return mul(rot_mat, Vec);
}

float3 SSFX_calc_sky(float3 dir)
{
	dir = SSFX_yaw_vector(dir, -L_sky_color.w); // Sky rotation
	
	dir.y = (dir.y - max(cos(dir.x) * 0.65f, cos(dir.z) * 0.65f)) * 2.1f; // Fix perspective
	dir.y -= -0.35; // Altitude
	
	float3 sky0 = s_env0.SampleLevel(smp_base, dir, 0).xyz;
	float3 sky1 = s_env1.SampleLevel(smp_base, dir, 0).xyz;
	
	// Use hemi color or real sky color if the modded executable is installed.
#ifndef SSFX_MODEXE
	return saturate(L_hemi_color.rgb * 3.0f) * lerp(sky0, sky1, L_ambient.w);
#else
	return saturate(L_sky_color.bgr * 3.0f) * lerp(sky0, sky1, L_ambient.w);
#endif
}

// float3 SSFX_calc_env(float3 dir)
// {
	// float3 env0 = env_s0.SampleLevel(smp_base, dir, 0).xyz;
	// float3 env1 = env_s1.SampleLevel(smp_base, dir, 0).xyz;
	
	// return env_color.xyz * lerp( env0, env1, env_color.w );
// }

// Full scene lighting
float3 SSFX_get_scene(float2 tc, uint iSample : SV_SAMPLEINDEX)
{
	#ifndef USE_MSAA
		float4 rP = s_position.Sample( smp_nofilter, tc );
	#else
		float4 rP = s_position.Load(int3(tc * pos_decompression_params2.xy, 0), iSample);
	#endif

	if (rP.z <= 0.05f)
		return 0;
	
	#ifndef USE_MSAA
		float4 rD = s_diffuse.Sample( smp_nofilter, tc );
		float4 rL = s_accumulator.Sample(smp_nofilter, tc);
	#else
		float4 rD = s_diffuse.Load(int3(tc * pos_decompression_params2.xy, 0), iSample);
		float4 rL = s_accumulator.Load(int3(tc * pos_decompression_params2.xy, 0), iSample);
	#endif
	
	// Remove emissive materials for now...
	if (length(rL) > 10.0f)
		rL = 0;
	
	float3 rN = gbuf_unpack_normal( rP.xy );
	float rMtl = gbuf_unpack_mtl( rP.w );
	float rHemi = gbuf_unpack_hemi( rP.w );
	
//	float3 nw = mul( m_invV, rN );
	
	// #ifdef SSFX_ENHANCED_SHADERS

		// rL.rgb += rL.a * SRGBToLinear(rD.rgb);

		// hemisphere
		// float3 hdiffuse, hspecular;
		// hmodel(hdiffuse, hspecular, rMtl, rHemi, rD, rP, rN);

		// Final color 
		// float3 rcolor = rL.rgb + hdiffuse.rgb;
		// return LinearTosRGB(rcolor);

	// #else
		// hemisphere
		float3 hdiffuse, hspecular;
		float4 hsun;
		hmodel(hdiffuse, hspecular, hsun, rMtl, rHemi, rD.w, rP, rN);
		
		// Final color
		rL += hsun;
		float4 light = float4(rL.rgb + hdiffuse, rL.w);
		float4 C = rD * light;
		float3 spec = C.www * rL.rgb;
		
		return C.rgb + spec;

	// #endif
	
	return 0;
}

void SSFX_ScreenSpaceReflections(float2 tc, float4 P, float3 N, float gloss, inout float3 color, uint iSample : SV_SAMPLEINDEX)
{
	// Note: Distance falloff on "rain_patch_normal.ps"
	
	// Material conditions ( MAT_FLORA and Terrain for now... )
	bool m_terrain = abs(P.w - 0.95f) <= 0.02f;
	bool m_flora = abs(P.w - MAT_FLORA) <= 0.04f;

	// Let's start with pure gloss.
	float refl_power = gloss;

	// Calc reflection bounce
	float3 inVec = normalize(P.xyz); // Incident
	float3 reVec = reflect(inVec , N); // Reflected

	// Transform space and calc reflection vector ( Skybox & Fresnel )
	float3 nw		 = mul(m_invV, N);
	float3 v2point	 = mul(m_invV, inVec);
	float3 v2reflect = reflect(v2point, nw);

	// Fresnel
	float fresnel = saturate (dot(v2reflect, v2point));
	float fresnel_amount = pow(fresnel, 3);
	refl_power *= fresnel_amount;

	float4 hit_uv = 0;
	
	// Calc SSR ray. Discard low reflective pixels
	if (refl_power > 0.02f)
		hit_uv = SSFX_ssr_fast_ray(P.xyz, reVec, tc, iSample);

	float3 refl_ray;
	float3 reflection = 0;
	float2 uvcoor = 0;
	
	// Sky is the reflection base...
#ifdef G_SSR_CHEAP_SKYBOX
	// reflection = SSFX_calc_env(v2reflect) * G_SSR_SKY_INTENSITY;
#else
	reflection = SSFX_calc_sky(v2reflect) * G_SSR_SKY_INTENSITY;
#endif

	// Valid UV coor? SSFX_trace_ssr_ray return 0.0f if uv is out of bounds or sky.
	if (all(hit_uv.xy))
	{
		// Get scene reflection
		refl_ray = SSFX_get_scene(hit_uv.xy, iSample);

		// Set reflection UV
		uvcoor = hit_uv.xy;

		// Let's fade the reflection based on ray XY coor to avoid abrupt changes and glitches
		float HitFade = saturate(hit_uv.y * G_SSR_VERTICAL_SCREENFADE);

		// Mix base reflection ( skybox ) with ray reflection
		reflection = lerp(reflection, refl_ray, HitFade);
	}
	else
	{
		// Reset gloss.
		refl_power = gloss * fresnel_amount;

		// Set reflection UV
		uvcoor = hit_uv.zw;
	}
	
	// Fade sky if !m_terrain ( Terrain MAT )
	float ray_fade = saturate(saturate(uvcoor.y * G_SSR_VERTICAL_SCREENFADE) + 1.0f * m_terrain);

	// Adjust the intensity of MAT_FLORA
	refl_power *= m_flora ? G_SSR_FLORA_INTENSITY : 1.0f;
	
	// Weapon Attenuation factor.
	float WeaponFactor = smoothstep(G_SSR_WEAPON_MAX_LENGTH - 0.2f, G_SSR_WEAPON_MAX_LENGTH, length(P.xyz));

	// Terrain MAT overwrite WeaponFactor.
	WeaponFactor = saturate(WeaponFactor + 1.0f * m_terrain);
	
	// Global intensity and limit max value.
	float main_clamp = clamp(refl_power * G_SSR_INTENSITY, 0, G_SSR_MAX_INTENSITY);
	
	// Raise reflection intensity and max limit when raining. ( NOTE: Reverted to rain intensity, but improvements are on the way... )
	float rain_extra = G_SSR_WEAPON_RAIN_FACTOR * 1.0f;
	// float rain_extra = G_SSR_WEAPON_RAIN_FACTOR * rain_params.x; // TODO: я заебался править этот шейдер
	
	// Weapon intensity and limit max value.
	float wpn_clamp = clamp((refl_power + rain_extra) * G_SSR_WEAPON_INTENSITY, 0, G_SSR_WEAPON_MAX_INTENSITY + rain_extra);

	#ifdef G_SSR_WEAPON_REFLECT_ONLY_WITH_RAIN
		// wpn_clamp *= rain_params.x; // TODO
		wpn_clamp *= 1.0f;
	#endif
	
	// Lerp between general reflections and weapon reflections.
	refl_power = lerp(wpn_clamp, main_clamp, WeaponFactor);

	// Apply SSR fade to reflection.
	refl_power *= ray_fade;

	// 'Beefs Shader Based NVGs' optional intensity adjustment
// #ifdef G_SSR_BEEFS_NVGs_ADJUSTMENT
	// refl_power *= saturate(1.0f - (1.0f - G_SSR_BEEFS_NVGs_ADJUSTMENT) * (shader_param_8.x > 0.0f));
// #endif

	// Add the reflection to the scene.
	color = lerp(color, reflection, refl_power);
}