// Copyright Epic Games, Inc. All Rights Reserved.

/*=============================================================================
	HeightFogPixelShader.usf: Scene fogging pixel shader.
=============================================================================*/

#include "Common.ush"
#include "SceneTexturesCommon.ush"
#include "SHCommon.ush"
#include "HeightFogCommon.ush"
#include "Random.ush"

#if PERMUTATION_SUPPORT_LOCAL_FOG_VOLUME
#include "LocalFogVolumes/LocalFogVolumeCommon.ush"
#endif

float UpsampleJitterMultiplier;

float bOnlyOnRenderedOpaque;
uint bUseWaterDepthTexture;
uint bPropagateAlpha;

float4 WaterDepthTextureMinMaxUV;
float4 OcclusionTextureMinMaxUV;

Texture2D OcclusionTexture;
SamplerState OcclusionSampler;

Texture2D WaterDepthTexture;
SamplerState WaterDepthSampler;

#if PERMUTATION_SAMPLE_FOG_ON_CLOUDS
Texture2D<float> SrcCloudDepthTexture;
SamplerState SrcCloudDepthSampler;
Texture2D<float4> SrcCloudViewTexture;
SamplerState SrcCloudViewSampler;
#endif

void ExponentialPixelMain(
	float2 TexCoord : TEXCOORD0,
	float3 ScreenVector : TEXCOORD1,
	float4 SVPos : SV_POSITION,
	out float4 OutColor : SV_Target0
	)
{
	float SceneDepth;
	float DeviceZ;
#if PERMUTATION_SAMPLE_FOG_ON_CLOUDS

	float2 PixPos = SVPos.xy;
	const float4 CloudLuminanceTransmittance = SrcCloudViewTexture.Load(int3(PixPos, 0));
	const float CloudCoverage = 1.0f - CloudLuminanceTransmittance.a;
	if (CloudLuminanceTransmittance.a > 0.999)
	{
		OutColor = float4(0.0f, 0.0f, 0.0f, 1.0f);
		return;
	}

	const float CloudDepthKm = SrcCloudDepthTexture.Load(int3(PixPos, 0)).r;
	SceneDepth = CloudDepthKm * KILOMETER_TO_CENTIMETER;
	DeviceZ = ConvertToDeviceZ(SceneDepth); // Warning: for simplicity, we use DeviceZ as world distance in kilometer when SAMPLE_ATMOSPHERE_ON_CLOUDS. See special case in IntegrateSingleScatteredLuminance.

#else
	if (bUseWaterDepthTexture)
	{
		float2 SampleUV = clamp(TexCoord, WaterDepthTextureMinMaxUV.xy, WaterDepthTextureMinMaxUV.zw);
		DeviceZ = Texture2DSampleLevel(WaterDepthTexture, WaterDepthSampler, SampleUV, 0).r;
		SceneDepth = ConvertFromDeviceZ(DeviceZ);
	}
	else
	{
	#if SCENE_TEXTURES_DISABLED
		float SceneDepth = SCENE_TEXTURES_DISABLED_SCENE_DEPTH_VALUE;
		DeviceZ = ConvertToDeviceZ(SceneDepth);
	#else
		DeviceZ = Texture2DSampleLevel(SceneTexturesStruct.SceneDepthTexture, SceneTexturesStruct_SceneDepthTextureSampler, TexCoord, 0).r;
		SceneDepth = ConvertFromDeviceZ(DeviceZ); // Fetch the depth buffer Z / W value, solve for W
	#endif
	}
#endif

	bool bIsRendered = (DeviceZ != 0.0);

	float3 WorldPositionRelativeToCamera = ScreenVector.xyz * SceneDepth;

	float ZSlice = log2(SceneDepth * View.VolumetricFogGridZParams.x + View.VolumetricFogGridZParams.y) * View.VolumetricFogGridZParams.z * View.VolumetricFogInvGridSize.z;
	
	uint3 Rand16Bits = Rand3DPCG16(int3(SVPos.xy, View.StateFrameIndexMod8));	
	float3 Rand3D = (float3(Rand16Bits) / float(uint(0xffff))) * 2.0f - 1.0f;
	float3 CellOffset = UpsampleJitterMultiplier * Rand3D;

	float3 VolumeUV = float3(((SVPos.xy - View.ViewRectMin.xy) + CellOffset.xy) * View.VolumetricFogSVPosToVolumeUV, ZSlice);
	VolumeUV.xy = min(VolumeUV.xy, View.VolumetricFogUVMax);

	float4 HeightFogInscatteringAndOpacity = CalculateHeightFog(WorldPositionRelativeToCamera);

#if PERMUTATION_SUPPORT_LOCAL_FOG_VOLUME
	uint2 TilePos = uint2(SVPos.xy - View.ViewRectMin.xy) / LFVTilePixelSize.xx;
	float4 LFVContribution = GetLFVContribution(PrimaryView, TilePos, WorldPositionRelativeToCamera);
	HeightFogInscatteringAndOpacity = float4(LFVContribution.rgb + HeightFogInscatteringAndOpacity.rgb * LFVContribution.a, LFVContribution.a * HeightFogInscatteringAndOpacity.a);
#endif

	float4 FogInscatteringAndOpacity = CombineVolumetricFog(HeightFogInscatteringAndOpacity, VolumeUV, 0, SceneDepth);
	float2 OcclusionTexCoord = clamp(TexCoord, OcclusionTextureMinMaxUV.xy, OcclusionTextureMinMaxUV.zw);
	float LightShaftMask = Texture2DSample(OcclusionTexture, OcclusionSampler, OcclusionTexCoord).x;
	FogInscatteringAndOpacity.rgb *= LightShaftMask;

	FogInscatteringAndOpacity.rgb *= View.PreExposure;


#if PERMUTATION_SAMPLE_FOG_ON_CLOUDS

	// Reduce cloud luminance according to the fog transmittance and add the fog in scattering luminance according to the cloud coverage.
	OutColor.rgb = CloudLuminanceTransmittance.rgb* FogInscatteringAndOpacity.a + FogInscatteringAndOpacity.rgb * CloudCoverage;
	// Coverage of the cloud layer itself does not change.
	OutColor.a   = CloudLuminanceTransmittance.a;

#else // PERMUTATION_SAMPLE_FOG_ON_CLOUDS

	if (bOnlyOnRenderedOpaque > 0.0 && !bIsRendered)
	{
		FogInscatteringAndOpacity.rgb = 0;
		FogInscatteringAndOpacity.a = 1;
	}
	
	FLATTEN
	if (bPropagateAlpha)
	{
		// We need to output coverage so that we can potentially apply alpha hold out blending. See RenderViewFog.
		OutColor = RETURN_COLOR(float4(FogInscatteringAndOpacity.rgb, 1.0 - FogInscatteringAndOpacity.w));
	}
	else
	{
		OutColor = RETURN_COLOR(float4(FogInscatteringAndOpacity.rgb, FogInscatteringAndOpacity.w));
	}
	
#endif // PERMUTATION_SAMPLE_FOG_ON_CLOUDS
}