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

/*=============================================================================
	VirtualShadowMapProjection.usf: 
=============================================================================*/

#define SUPPORT_CONTACT_SHADOWS 1
#include "../HairStrands/HairStrandsVisibilityCommonStruct.ush"
#include "../HairStrands/HairStrandsVoxelPageCommonStruct.ush"
#include "../Common.ush"
#include "../SceneTexturesCommon.ush"
#include "../DeferredShadingCommon.ush"
#include "../DeferredLightingCommon.ush"
#include "../LightShaderParameters.ush"
#include "../LightGridCommon.ush"
#include "../HairStrands/HairStrandsVisibilityCommon.ush"
#include "../Visualization.ush"
#include "PageAccessCommon.ush"
#include "ProjectionCommon.ush"
#include "ProjectionDirectional.ush"
#include "ProjectionSpot.ush"
#include "PageCacheCommon.ush"

#if HAS_HAIR_STRANDS
#define VOXEL_TRAVERSAL_TYPE VOXEL_TRAVERSAL_LINEAR_MIPMAP
#define VOXEL_TRAVERSAL_DEBUG 0
#include "../HairStrands/HairStrandsVoxelPageCommon.ush"
#include "../HairStrands/HairStrandsVoxelPageTraversal.ush"
#endif

#ifndef DIRECTIONAL_LIGHT
#define DIRECTIONAL_LIGHT 0
#endif

#ifndef DEBUG_OUTPUT
#define DEBUG_OUTPUT 0
#endif

#define CONTACT_SHADOW_SAMPLES 8

float4 EncodeLightAttenuationFromMask(float ShadowMask)
{
	const float ShadowFadeFraction = 1;
	float SSSTransmission = ShadowMask;

	// 0 is shadowed, 1 is unshadowed
	// RETURN_COLOR not needed unless writing to SceneColor;
	//float FadedShadow = lerp(1.0f, Square(ShadowMask), ShadowFadeFraction);
	//float FadedSSSShadow = lerp(1.0f, Square(SSSTransmission), ShadowFadeFraction);
	float FadedShadow = lerp(1.0f, ShadowMask, ShadowFadeFraction);
	float FadedSSSShadow = lerp(1.0f, SSSTransmission, ShadowFadeFraction);

	// the channel assignment is documented in ShadowRendering.cpp (look for Light Attenuation channel assignment)
	return EncodeLightAttenuation(half4(FadedShadow, FadedSSSShadow, FadedShadow, FadedSSSShadow));
}

Texture2D<float> InputShadowFactor;

void VirtualShadowMapCompositePS(
	in float4 SvPosition : SV_Position,
	out float4 OutShadowMask : SV_Target)
{
	float ShadowFactor = InputShadowFactor.Load(int3(SvPosition.xy, 0));	
	OutShadowMask = EncodeLightAttenuationFromMask(ShadowFactor);
}


SCHEDULER_MIN_PRESSURE
MAX_OCCUPENCY

// Returns hit distance or negative if miss
float VirtualShadowMapScreenRayCast(
	float3 RayOriginTranslatedWorld, float3 RayDirection, float RayLength,
	int NumSteps, float StepOffset)
{
	bool bHitCastContactShadow = true;
	float HitDistance = ShadowRayCast(
		RayOriginTranslatedWorld, RayDirection, RayLength,
		NumSteps, StepOffset, bHitCastContactShadow);

	return (bHitCastContactShadow) ? HitDistance : -1.0f;
}

float3 GetEstimatedGeoWorldNormal(float3 TranslatedWorldPosition, float2 ScreenUV)
{
	// Figure out slope, we do world space since that is the space where we might override using the shading normal...
	float3 TranslatedWorldPositionDDX = DDX(TranslatedWorldPosition);
	float3 TranslatedWorldPositionDDY = DDY(TranslatedWorldPosition);
	float3 EstimatedGeoWorldNormal = cross(TranslatedWorldPositionDDX, TranslatedWorldPositionDDY);

	// Handle NaNs; will cause it to revert to shading normal below
	float LengthSq = dot(EstimatedGeoWorldNormal, EstimatedGeoWorldNormal);
	EstimatedGeoWorldNormal = LengthSq > 1e-8f ? normalize(EstimatedGeoWorldNormal) : float3(0, 0, 0);

#if 1
	// NOTE: Gbuffer normal is not the surface normal for hair; hair lighting takes a different path but possibly
	// necessary to do some sort of special casing here (disable normal offset and biasing entirely?).
	FGBufferData GBufferData = GetGBufferData(ScreenUV);
	// If the estimated geo normal is too far out we assume it's broken (derivative includes other surfaces or background) and fall back to the shading normal
	if (dot(GBufferData.WorldNormal, EstimatedGeoWorldNormal) < 0.25f)
	{
		EstimatedGeoWorldNormal = GBufferData.WorldNormal;
	}
#endif

	return EstimatedGeoWorldNormal;
}

FLightShaderParameters ConvertFromLocal( const FLocalLightData LightData )
{
	FLightShaderParameters Light = (FLightShaderParameters)0;
	Light.Position			= LightData.LightPositionAndInvRadius.xyz;
	Light.InvRadius			= LightData.LightPositionAndInvRadius.w;
	Light.Color				= LightData.LightColorAndFalloffExponent.xyz;
	Light.FalloffExponent	= LightData.LightColorAndFalloffExponent.w;
	Light.Direction			= LightData.LightDirectionAndShadowMask.xyz;
	Light.Tangent			= LightData.LightTangentAndSoftSourceRadius.xyz;
	Light.SpotAngles		= LightData.SpotAnglesAndSourceRadiusPacked.xy;
	Light.SpecularScale		= 1;
	Light.SourceRadius		= LightData.SpotAnglesAndSourceRadiusPacked.z;
	Light.SoftSourceRadius	= LightData.LightTangentAndSoftSourceRadius.w;
	Light.SourceLength		= f16tof32(asuint(LightData.SpotAnglesAndSourceRadiusPacked.w));
	return Light;
}

float4 ComputeRandom4(uint2 PixelPosition)
{
	const uint InSeed = View.StateFrameIndexMod8;
	const uint2 Seed0 = Rand3DPCG16(int3(PixelPosition, InSeed)).xy;
	const uint2 Seed1 = Rand3DPCG16(int3(PixelPosition + 17, InSeed)).xy;
	return float4(
		Hammersley16(InSeed, 8, Seed0),
		Hammersley16(InSeed, 8, Seed1));
}

int4 ProjectionRect;
float ContactShadowLength;
float NormalBias;
int SMRTRayCount;					// 0 = off
int SMRTSamplesPerRay;
float SMRTRayLengthScale;			// Directional lights
float SMRTCotMaxRayAngleFromLight;	// Spot/point lights

// One pass projection
StructuredBuffer< int > VirtualShadowMapIdRemap;
uint NumDirectionalLightSmInds;
RWTexture2D< uint > RWShadowMaskBits;

// Single light per pass
// Light parameters loaded via GetRootLightShaderParameters();
int LightUniformVirtualShadowMapId;
RWTexture2D< float4 > RWShadowFactor;

// Debug output pass
StructuredBuffer< FPhysicalPageMetaData > PhysicalPageMetaData;
RWTexture2D< float4 > RWDebug;
int DebugOutputType;
int DebugVirtualShadowMapId;

// Type of input data consume by the page allocation (i.e., data read from the source buffer: Gbuffer, HairStrands data, ...)
#define INPUT_TYPE_GBUFFER 0
#define INPUT_TYPE_HAIRSTRANDS 1
uint InputType;

FVirtualShadowMapSampleResult ProjectLight(
	int VirtualShadowMapId,
	FLightShaderParameters Light,
	uint2 PixelPos,
	float SceneDepth,
	float ContactShadowLengthWorld,
	float3 TranslatedWorldPosition,
	half3 GBufferNormal,
	const half Noise)
{
	const float3 BaseWorldPosition = TranslatedWorldPosition - View.PreViewTranslation;
	// NOTE: If View.PreViewTranslation is guaranteed to be the camera position we can simplify
	const float DistanceToCamera = length(TranslatedWorldPosition - View.TranslatedWorldCameraOrigin);
	const float NormalBiasLength = max(0.01f, NormalBias * DistanceToCamera / View.ViewToClip._11);

#if DIRECTIONAL_LIGHT
	uint LightType = LIGHT_TYPE_DIRECTIONAL;
	float3 L = Light.Direction;
	// TODO: Find a way to branch out backfacing surfaces again that doesn't make assumptions
	// about the Gbuffer normal that break with hair
	bool bInLightRegion = true;
#else
	uint LightType = Light.SpotAngles.x > -2.0f ? LIGHT_TYPE_SPOT : LIGHT_TYPE_POINT;
	float3 ToLight = Light.Position - BaseWorldPosition;
	float d2 = dot( ToLight, ToLight );
	float InvDist = rsqrt( d2 );
	half3 L = ToLight * InvDist;
	
	bool bInLightRadius = InvDist >= Light.InvRadius;
	bool bInLightCone = dot( L, Light.Direction ) >= Light.SpotAngles.x;
	bool bInLightRegion = bInLightRadius && bInLightCone;
	
#endif // DIRECTIONAL_LIGHT

	FVirtualShadowMapSampleResult Result = InitVirtualShadowMapSampleResult();

	BRANCH
	if (bInLightRegion)
	{
		const bool bIsHair = InputType == INPUT_TYPE_HAIRSTRANDS;
		const bool bBackfaceCull = false;		// TODO
		const half3 EstimatedGeoWorldNormal = bIsHair ? L : GBufferNormal;
		
		TranslatedWorldPosition += EstimatedGeoWorldNormal * NormalBiasLength;

		if (SMRTRayCount > 0)
		{
		#if DIRECTIONAL_LIGHT
			Result = TraceDirectional(
				VirtualShadowMapId,
				Light,
				PixelPos,
				SceneDepth,
				TranslatedWorldPosition,
				ContactShadowLengthWorld,
				SMRTRayCount,
				SMRTSamplesPerRay,
				SMRTRayLengthScale,
				Noise,
				bBackfaceCull);
		#else
			Result = TraceLocalLight(
				VirtualShadowMapId,
				Light,
				LightType,
				PixelPos,
				SceneDepth,
				TranslatedWorldPosition,
				ContactShadowLengthWorld,
				SMRTRayCount,
				SMRTSamplesPerRay,
				SMRTCotMaxRayAngleFromLight,
				Noise,
				bBackfaceCull);
		#endif // DIRECTIONAL_LIGHT
		}
		else
		{
			Result = SampleVirtualShadowMap(
				VirtualShadowMapId,
				TranslatedWorldPosition - View.PreViewTranslation,
				ContactShadowLengthWorld,
				EstimatedGeoWorldNormal);
		}

		// Do not run contact shadow when computing the hair shadow mask (i.e. shadow mask applied on hair, has the scene 
		// depth buffer contains fully opaque depth, which create false positive intersection resulting in wrong self shadowing)
		if (ContactShadowLengthWorld > 0.0f && !bIsHair)
		{
			half ContactShadowHitDistance = VirtualShadowMapScreenRayCast(TranslatedWorldPosition, L, ContactShadowLengthWorld, CONTACT_SHADOW_SAMPLES, Noise);
			if (ContactShadowHitDistance >= 0.0f)
			{
				Result.ShadowFactor = 0.0f;
			}
		}

		// Hair strands voxel traversal to apply hair shadow on top of opaque geometry
		#if HAS_HAIR_STRANDS
		if (!bIsHair)
		{
			float3 RayStart = 0;
			float3 RayEnd = 0;
			bool bCastHairRay = false;
			float4 Random = 0;
			if (Result.ShadowFactor > 0)
			{
				Random = ComputeRandom4(PixelPos);
				uint RayIndex = min(Random.w * SMRTRayCount, SMRTRayCount - 1);

			#if DIRECTIONAL_LIGHT
				bCastHairRay = GenerateRayDirectional(
					Light,
					PixelPos,
					TranslatedWorldPosition,
					ContactShadowLengthWorld,
					RayIndex,
					SMRTRayCount,
					RayStart,
					RayEnd);
			#else
				bCastHairRay = GenerateRayLocalLight(
					Light,
					PixelPos,
					TranslatedWorldPosition,
					SMRTCotMaxRayAngleFromLight,
					RayIndex,
					SMRTRayCount,
					RayStart,
					RayEnd);
			#endif
			}

			if (bCastHairRay)
			{
				// Jitter start position to mask/compensate coarse voxel resolution
				{
					const float PositionBiasScale = 0.5f;
					const float3 DepthBias = VirtualVoxel.VoxelWorldSize * (VirtualVoxel.DepthBiasScale_Shadow * L + PositionBiasScale * (Random.xyz * 2 - 1));
					RayStart += DepthBias;
				}

				const float DistanceThreshold = 100000.0f;
				const float CoverageThreshold = 0.995f; // When Coverage is high, we do not trace shadow on opaque since hair/fur is covering the background.

				FVirtualVoxelCommonDesc CommonDesc;
				CommonDesc.PageCountResolution		= VirtualVoxel.PageCountResolution;
				CommonDesc.VoxelWorldSize			= VirtualVoxel.VoxelWorldSize;
				CommonDesc.PageTextureResolution	= VirtualVoxel.PageTextureResolution;
				CommonDesc.PageResolution			= VirtualVoxel.PageResolution;

				FHairTraversalSettings TraversalSettings = InitHairTraversalSettings();
				TraversalSettings.DensityScale		= VirtualVoxel.DensityScale_Shadow;
				TraversalSettings.CountThreshold	= 1;
				TraversalSettings.DistanceThreshold = DistanceThreshold;
				TraversalSettings.bDebugEnabled		= false;
				TraversalSettings.SteppingScale		= VirtualVoxel.SteppingScale_Shadow;
				TraversalSettings.Random			= Random.xyz;
				TraversalSettings.PixelRadius		= SceneDepth * VirtualVoxel.HairCoveragePixelRadiusAtDepth1;
				TraversalSettings.bUseOpaqueVisibility = false;

				const uint VoxelDescCount = VirtualVoxel.NodeDescCount;
				for (uint VoxelDescIt=0; VoxelDescIt<VoxelDescCount; ++VoxelDescIt)
				{
					const FPackedVirtualVoxelNodeDesc PackedNode = VirtualVoxel.NodeDescBuffer[VoxelDescIt];
					const FVirtualVoxelNodeDesc NodeDesc = UnpackVoxelNode(PackedNode, VirtualVoxel.PageResolution);

					FHairTraversalResult HairResult = InitHairTraversalResult();
					HairResult = ComputeHairCountVirtualVoxel(
						RayStart,
						RayEnd,
						CommonDesc,
						NodeDesc,
						VirtualVoxel.PageIndexBuffer,
						VirtualVoxel.PageIndexOccupancyBuffer,
						VirtualVoxel.PageTexture,
						TraversalSettings);

					Result.ShadowFactor = min(Result.ShadowFactor, saturate(1 - HairResult.HairCount));
				}
			}
		}
		#endif
	}

	return Result;
}

void OutputDebug( uint2 PixelPos, FVirtualShadowMapSampleResult Result )
{
	float3 Output = float3( 1, 0, 1 );
		
	if ( Result.bValid )
	{
		uint2 vPage = Result.VirtualTexelAddress.xy >> VSM_LOG2_PAGE_SIZE;
		uint2 pPage = Result.PhysicalTexelAddress.xy >> VSM_LOG2_PAGE_SIZE;
		bool bValidPageAddress = all( pPage < uint2(VirtualShadowMap.PhysicalPoolSizePages) );

		if ( DebugOutputType == 1 )			// ShadowFactor
		{
			Output = Result.ShadowFactor.xxx;
		}
		if ( DebugOutputType == 2 )			// ClipmapIndexOrMip (mix in a bit of shadow for clarity)
		{
			float3 Color = IntToColor( Result.ClipmapIndexOrMipLevel );
			Output = lerp( Color, Result.ShadowFactor.xxx, 0.15f );
		}
		else if ( DebugOutputType == 3 && bValidPageAddress )	// Virtual page (mix in a bit of shadow for clarity)
		{
			float3 Color = IntToColor( vPage.x + ( vPage.y << 10U ) + (Result.ClipmapIndexOrMipLevel << 20U) );
			Output = lerp( Color, Result.ShadowFactor.xxx, 0.15f );
		}
		else if ( DebugOutputType == 4 && bValidPageAddress )	// Physical page age (mix in a bit of shadow for clarity)
		{
			uint Age = PhysicalPageMetaData[ PhysPageAddressToIndex(pPage) ].Age;
			float3 Color = GreenToRedTurbo( 1.0f - 0.1f * float(Age) );
			Output = lerp( Color, Result.ShadowFactor.xxx, 0.15f );
		}
		else if ( DebugOutputType == 5 )	// RayCount
		{
			Output = GreenToRedTurbo( float( Result.RayCount ) / float( SMRTRayCount ) );
		}
		else if ( DebugOutputType == 6 )	// OccluderDistance
		{
			Output = ColorMapViridis( 0.001f * Result.OccluderDistance );
		}
	}

	RWDebug[ PixelPos ] = float4( Output, 1.0f );
}

[numthreads(8, 8, 1)]
void VirtualShadowMapProjection(
	uint3	GroupId				: SV_GroupID,
	uint	GroupIndex			: SV_GroupIndex,
	uint3	DispatchThreadId	: SV_DispatchThreadID )
{
	// Morton order within a group so page access/atomics are more coherent and wave-swizzled gradients are possible.
	uint2 PixelCoord = DispatchThreadId.xy;
	uint2 LocalPixelPos = 8 * GroupId.xy + MortonDecode( GroupIndex );
	uint2 PixelPos = LocalPixelPos + uint2( ProjectionRect.xy );
	if ( any( PixelPos >= uint2( ProjectionRect.zw ) ) )
	{
		return;
	}
	
	float DeviceZ = SceneTexturesStruct.SceneDepthTexture.Load( int3( PixelPos, 0 ) ).r;
	#if HAS_HAIR_STRANDS
	if (InputType == INPUT_TYPE_HAIRSTRANDS)
	{
		FCategorizationData CatData = DecodeCategorizationData(HairStrands.HairCategorizationTexture.Load( int3(PixelPos, 0 ) ) );
		DeviceZ = CatData.ClosestDepth;
		if (CatData.PixelCoverage == 0)
		{
			return;
		}
	}
	#endif
	const float SceneDepth = ConvertFromDeviceZ( DeviceZ );

	const float4 SvPosition = float4( float2( PixelPos ) + 0.5, DeviceZ, 1.0f );
	const float3 TranslatedWorldPosition = SvPositionToTranslatedWorld( SvPosition );

	const float ContactShadowLengthWorld = ContactShadowLength * View.ClipToView[1][1] * SceneDepth;
	const half Noise = InterleavedGradientNoise( SvPosition.xy, View.StateFrameIndexMod8 );

	const half3 GBufferNormal = GetGBufferDataUint( PixelPos, true ).WorldNormal;

	FVirtualShadowMapSampleResult DebugResult = InitVirtualShadowMapSampleResult();
	bool bWriteDebug = false;

#if ONE_PASS_PROJECTION
	uint EyeIndex = 0; // TODO: Instanced stereo
	uint GridLinearIndex = ComputeLightGridCellIndex( LocalPixelPos, SceneDepth, EyeIndex );
	const FCulledLightsGridData CulledLightGridData = GetCulledLightsGrid( GridLinearIndex, EyeIndex );

	// We can only handle 10 lights in our output encoding right now, so no purpose in computing more
	const uint LightCount = min(GetShadowMaskMaxLightCount(), CulledLightGridData.NumLocalLights );

	uint ShadowMask = 0xffffffff;

	LOOP
	for (uint Index = 0; Index < LightCount; ++Index)
	{
		// The Virtual Shadow Remap stores directional lights first
		const uint LocalLightIndex = ForwardLightData.CulledLightDataGrid[CulledLightGridData.DataStartIndex + Index];
		int VirtualShadowMapId = VirtualShadowMapIdRemap[VirtualShadowMap.NumDirectionalLights + LocalLightIndex];

		if (VirtualShadowMapId != INDEX_NONE)
		{
			const FLocalLightData LightData = GetLocalLightData( CulledLightGridData.DataStartIndex + Index, EyeIndex );
			FLightShaderParameters Light = ConvertFromLocal( LightData );

			FVirtualShadowMapSampleResult Result = ProjectLight(
				VirtualShadowMapId,
				Light,
				PixelPos,
				SceneDepth,
				ContactShadowLengthWorld,
				TranslatedWorldPosition,
				GBufferNormal,
				Noise);

			PackShadowMask(ShadowMask, Result.ShadowFactor, Index);

			if (DebugVirtualShadowMapId == VirtualShadowMapId)
			{
				DebugResult = Result;
				bWriteDebug = true;
			}
		}
	}

	RWShadowMaskBits[ PixelPos ] = ~ShadowMask;

#else // !ONE_PASS_PROJECTION
	{
		int VirtualShadowMapId = LightUniformVirtualShadowMapId;
		FLightShaderParameters Light = GetRootLightShaderParameters();

		FVirtualShadowMapSampleResult Result = ProjectLight(
			VirtualShadowMapId,
			Light,
			PixelPos,
			SceneDepth,
			ContactShadowLengthWorld,
			TranslatedWorldPosition,
			GBufferNormal,
			Noise);

		RWShadowFactor[ PixelPos ] = Result.ShadowFactor;

		if (DebugVirtualShadowMapId == VirtualShadowMapId)
		{
			DebugResult = Result;
			bWriteDebug = true;
		}
	}
#endif // ONE_PASS_PROJECTION

#if DEBUG_OUTPUT
	if (bWriteDebug && InputType == INPUT_TYPE_GBUFFER)
	{
		OutputDebug( PixelPos, DebugResult );
	}
#endif // DEBUG_OUTPUT
}