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

#include "../Common.ush"
#include "LumenMaterial.ush"
#include "../DeferredShadingCommon.ush"
#include "../RayTracing/RayTracingRayCone.ush"
#include "LumenCardCommon.ush"
#include "LumenTracingCommon.ush"
#include "LumenSoftwareRayTracing.ush"
#include "LumenRadianceCacheCommon.ush"
#include "LumenReflectionCommon.ush"
#include "../SceneTextureParameters.ush"
#include "../SHCommon.ush"
#include "../ShaderPrint.ush"
#define HZB_TRACE_INCLUDE_FULL_RES_DEPTH 1
#include "LumenScreenTracing.ush"
#include "LumenHairTracing.ush"

#ifndef THREADGROUP_SIZE
#define THREADGROUP_SIZE 1
#endif	

#define DEBUG_VISUALIZE_TRACE_TYPES 0


RWTexture2DArray<float>  RWTraceHit;
RWTexture2DArray<float3> RWTraceRadiance;
RWTexture2DArray<uint>   RWTraceMaterialId;
RWTexture2DArray<float3> RWTraceBackgroundVisibility;

#ifdef ReflectionClearTracesCS

[numthreads(REFLECTION_THREADGROUP_SIZE_1D, 1, 1)]
void ReflectionClearTracesCS(
	uint GroupId : SV_GroupID,
	uint GroupThreadId : SV_GroupThreadID)
{
	FReflectionTileData TileData;
	uint2 TmpReflectionTracingCoord = GetReflectionTracingScreenCoord(GroupId/*TileCoord*/, GroupThreadId/*InnerTileCoord*/, TileData).xy;
	const FReflectionTracingCoord ReflectionTracingCoord = GetReflectionTracingCoord(TmpReflectionTracingCoord, TileData.ClosureIndex);

	RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = 0.0f;
	RWTraceHit[ReflectionTracingCoord.CoordFlatten] = EncodeRayDistance(0.0f, false);

	#if CLEAT_TRACE_MATERIAL_ID
	{
		RWTraceMaterialId[ReflectionTracingCoord.CoordFlatten] = PackTraceMaterialId(false, 0);
	}
	#endif

	#if CLEAR_BACKGROUND_VISIBILITY
	{
		RWTraceBackgroundVisibility[ReflectionTracingCoord.CoordFlatten] = 1.0f;
	}
	#endif
}
#endif

float MaxTraceDistance;
float MaxHierarchicalScreenTraceIterations;
float RelativeDepthThickness;
float HistoryDepthTestRelativeThickness;
uint MinimumTracingThreadOccupancy;
float4 FirstPersonWorldSpaceRepresentationBounds; // Collective bounding sphere for all the invisible primitives tagged as FirstPersonWorldSpaceRepresentation
float FirstPersonMinimumHitDistanceOnScreenTraceMiss;

#ifdef ReflectionTraceScreenTexturesCS

[numthreads(REFLECTION_THREADGROUP_SIZE_1D, 1, 1)]
void ReflectionTraceScreenTexturesCS(
	uint GroupId : SV_GroupID,
	uint GroupThreadId : SV_GroupThreadID)
{
	FReflectionTileData TileData;
	const uint2 TmpReflectionTracingCoord = GetReflectionTracingScreenCoord(GroupId/*TileCoord*/, GroupThreadId/*InnerTileCoord*/, TileData).xy;
	const FReflectionTracingCoord ReflectionTracingCoord = GetReflectionTracingCoord(TmpReflectionTracingCoord, TileData.ClosureIndex);

	if (all(ReflectionTracingCoord.Coord < ReflectionTracingViewMin + ReflectionTracingViewSize))
	{
		float2 ScreenUV = GetScreenUVFromReflectionTracingCoord(ReflectionTracingCoord.Coord);
		float SceneDepth = DownsampledDepth.Load(int4(ReflectionTracingCoord.CoordFlatten, 0)).x;

		if (SceneDepth > 0.0f)
		{
			float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);
			FRayData RayData = GetRayData(ReflectionTracingCoord.CoordFlatten);

			float3 TranslatedRayOrigin = TranslatedWorldPosition;

			// Bias along the normal to avoid self-intersecting camera facing scene depth texels
			// HZB traversal uses point filtering, so scene depth texels stick out from the plane they are representing
			{
				float2 CornerScreenUV = ScreenUV + .5f * View.BufferSizeAndInvSize.zw;
				float2 CornerScreenPosition = (CornerScreenUV - View.ScreenPositionScaleBias.wz) / View.ScreenPositionScaleBias.xy;
				float3 CornerTranslatedWorldPosition = mul(float4(GetScreenPositionForProjectionType(CornerScreenPosition, SceneDepth), SceneDepth, 1), View.ScreenToTranslatedWorld).xyz;
			#if SUBTRATE_GBUFFER_FORMAT==1
				const float3 WorldNormal = SubstrateUnpackTopLayerData(Substrate.TopLayerTexture.Load(uint3(CornerScreenUV * View.BufferSizeAndInvSize.xy, 0))).WorldNormal;
			#else
				const float3 WorldNormal = GetGBufferDataFromSceneTextures(ScreenUV).WorldNormal;
			#endif
				// Project texel corner onto normal
				float NormalBias = abs(dot(CornerTranslatedWorldPosition - TranslatedWorldPosition, WorldNormal)) * 2.0f;
				TranslatedRayOrigin += NormalBias * WorldNormal;

				float4 RayStartClip = mul(float4(TranslatedRayOrigin, 1.0f), View.TranslatedWorldToClip);

				// Skip screen traces if the normal bias pushed our starting point off-screen, as TraceScreen() doesn't handle this
				if (any(RayStartClip.xy < -RayStartClip.w) || any(RayStartClip.xy > RayStartClip.w))
				{
					return;
				}
			}

			bool bHit;
			bool bUncertain;
			float3 HitUVz;
			float3 LastVisibleHitUVz;
			float HitTileZ;

			TraceScreen(
				TranslatedRayOrigin,
				RayData.Direction,
				RayData.RadianceCacheMaxTraceDistance,
				HZBUvFactorAndInvFactor,
				MaxHierarchicalScreenTraceIterations, 
				RelativeDepthThickness,
				0,
				MinimumTracingThreadOccupancy,
				bHit,
				bUncertain,
				HitUVz,
				LastVisibleHitUVz,
				HitTileZ);

			#if USE_HAIRSTRANDS_SCREEN
			if (!bHit)
			{
				bool   Hair_bHit;
				bool   Hair_bUncertain;
				float3 Hair_HitUVz;
				float3 Hair_LastVisibleHitUVz;
				float Hair_HitTileZ;

				TraceHairScreen(
					TranslatedWorldPosition,
					RayData.Direction,
					RayData.RadianceCacheMaxTraceDistance,
					HZBUvFactorAndInvFactor,
					MaxHierarchicalScreenTraceIterations,
					RelativeDepthThickness,
					0,
					Hair_bHit,
					Hair_bUncertain,
					Hair_HitUVz,
					Hair_LastVisibleHitUVz,
					Hair_HitTileZ);

				if (Hair_bHit && !Hair_bUncertain)
				{
					bHit = Hair_bHit;
					HitUVz = Hair_HitUVz;
					LastVisibleHitUVz = Hair_LastVisibleHitUVz;
					bUncertain = Hair_bUncertain;
					HitTileZ = Hair_HitTileZ;
				}
			}
			#endif

			bHit = bHit && !bUncertain;

			if (bHit)
			{
				float2 HitScreenUV = HitUVz.xy;
				float2 HitScreenPosition = (HitUVz.xy - View.ScreenPositionScaleBias.wz) / View.ScreenPositionScaleBias.xy;
				float HitDeviceZ = HitTileZ;

				float3 HitHistoryScreenPosition = GetHistoryScreenPosition(HitScreenPosition, HitScreenUV, HitDeviceZ);

				float Vignette = min(ComputeHitVignetteFromScreenPos(HitScreenPosition), ComputeHitVignetteFromScreenPos(HitHistoryScreenPosition.xy));
				float Noise = InterleavedGradientNoise(ReflectionTracingCoord.Coord + 0.5f, ReflectionsStateFrameIndexMod8);

				// Skip reporting a hit if near the edge of the screen
				if (Vignette < Noise)
				{
					bHit = false;
					HitUVz = LastVisibleHitUVz;
				}
	
				if (bHit)
				{
					// Calculate the expected depth of the pixel last frame
					float PrevDeviceZ = HitHistoryScreenPosition.z;

					// Lookup the actual depth at the same screen position last frame
					float2 HitHistoryScreenUVForDepth = HitHistoryScreenPosition.xy * PrevScreenPositionScaleBiasForDepth.xy + PrevScreenPositionScaleBiasForDepth.zw;
					float HistoryDeviceZ = Texture2DSampleLevel(HistorySceneDepth, GlobalPointClampedSampler, HitHistoryScreenUVForDepth, 0).x;

					bHit = abs(HistoryDeviceZ - PrevDeviceZ) < HistoryDepthTestRelativeThickness * lerp(.5f, 2.0f, Noise);

					if (!bHit)
					{
						// We found a hit but discarded it, rewind to the last time we know the ray hadn't hit for the next method to resume from
						HitUVz = LastVisibleHitUVz;
					}
				}

				if (bHit)
				{
					float2 HitHistoryScreenUV = HitHistoryScreenPosition.xy * PrevScreenPositionScaleBias.xy + PrevScreenPositionScaleBias.zw;
					HitHistoryScreenUV = clamp(HitHistoryScreenUV, PrevSceneColorBilinearUVMin, PrevSceneColorBilinearUVMax);
					float3 Lighting = SampleScreenColor(PrevSceneColorTexture, GlobalPointClampedSampler, HitHistoryScreenUV).xyz * PrevSceneColorPreExposureCorrection;

					float MaxLighting = max3(Lighting.x, Lighting.y, Lighting.z);
					if (MaxLighting > MaxRayIntensity)
					{
						Lighting *= MaxRayIntensity / MaxLighting;
					}

					#if DEBUG_VISUALIZE_TRACE_TYPES
						RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = float3(.5f, 0, 0) * View.PreExposure;
					#else
						RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = Lighting;
					#endif
				}
			}

			// Make the ray slightly longer as distance reported from screen space traces isn't very reliable.
			// This way we can skip tracing if we reached trace max distance. If not, then the next pass will correct for it using RT pullback bias.
			float HitDistanceOffset = !bHit && !bUncertain ? 2.0f : 0.0f;
			float HitDistance = min(sqrt(ComputeRayHitSqrDistance(TranslatedWorldPosition, HitUVz)) + HitDistanceOffset, MaxTraceDistance);

#if HAS_FIRST_PERSON_GBUFFER_BIT
			// First person world space representation primitives are invisible to the camera but visible in shadows and HWRT reflections.
			// Screen traces can miss intersections with these invisible meshes, so when HWRT picks up, we end up seeing holes and other artifacts
			// in the reflected meshes. To work around this issue, we move the ray position at which we hand off to HWRT outside the bounds of the
			// invisible geometry such that HWRT has a chance to intersect it.
			if (!bHit && FirstPersonWorldSpaceRepresentationBounds.w > 0.0f)
			{
				const float2 Intersections = RayIntersectSphere(TranslatedRayOrigin, RayData.Direction, FirstPersonWorldSpaceRepresentationBounds);
				if (all(Intersections != -1.0f))
				{
					HitDistance = min(HitDistance, max(Intersections.x, FirstPersonMinimumHitDistanceOnScreenTraceMiss));
				}
			}
#endif // HAS_FIRST_PERSON_GBUFFER_BIT

			RWTraceHit[ReflectionTracingCoord.CoordFlatten] = EncodeRayDistance(HitDistance, bHit);
		}
	}
}

#endif

RWBuffer<uint> RWCompactedTraceTexelAllocator;
RWBuffer<uint> RWReflectionCompactionIndirectArgs;
Buffer<uint> ReflectionTracingTileIndirectArgs;
uint CompactionThreadGroupSize;

#ifdef SetupCompactionIndirectArgsCS

[numthreads(64, 1, 1)]
void SetupCompactionIndirectArgsCS(uint GroupThreadId : SV_GroupThreadID)
{
	uint NumTracingThreads = ReflectionTracingTileIndirectArgs[0] * REFLECTION_THREADGROUP_SIZE_1D;
	uint NumCompactionGroups = DivideAndRoundUp(NumTracingThreads, CompactionThreadGroupSize);

	if (GroupThreadId < 3)
	{
		RWReflectionCompactionIndirectArgs[GroupThreadId] = GroupThreadId == 0 ? NumCompactionGroups : 1;
	}

	if (GroupThreadId == 0)
	{
		RWCompactedTraceTexelAllocator[0] = 0;
	}
}

#endif

RWBuffer<uint> RWCompactedTraceTexelData;

// LumenReflections::ETraceCompactionMode
#define TRACE_COMPACTION_MODE_DEFAULT		0
#define TRACE_COMPACTION_MODE_FAR_FIELD		1
#define TRACE_COMPACTION_MODE_HIT_LIGHTING	2

uint CullByDistanceFromCamera;
float CompactionTracingEndDistanceFromCamera;
float CompactionMaxTraceDistance;
float NearFieldMaxTraceDistance;
float FarFieldDitheredStartDistanceFactor;
Texture2DArray<uint> TraceMaterialId;

#ifdef ReflectionCompactTracesCS

groupshared uint SharedGlobalTraceTexelStartOffset;

//@todo - ordered compaction for non-wave ops path
#if WAVE_OPS
	groupshared uint SharedGroupSum;
#else
	groupshared uint SharedTraceTexelAllocator;
	groupshared uint SharedTraceTexels[THREADGROUP_SIZE];
#endif

#if COMPILER_SUPPORTS_WAVE_SIZE && WAVE_OPS
	WAVESIZE(32)
#endif
[numthreads(THREADGROUP_SIZE, 1, 1)]
void ReflectionCompactTracesCS(
	uint GroupId : SV_GroupID,
	uint GroupThreadId : SV_GroupThreadID)
{
	if (GroupThreadId == 0)
#if WAVE_OPS
	{
		SharedGroupSum = 0;
	}
#else
	{
		SharedTraceTexelAllocator = 0;
	}

	GroupMemoryBarrierWithGroupSync();
#endif

	uint ReflectionTileIndex = GroupId * (THREADGROUP_SIZE / REFLECTION_THREADGROUP_SIZE_1D) + GroupThreadId / REFLECTION_THREADGROUP_SIZE_1D;
	uint ReflectionGroupThreadId = GroupThreadId % REFLECTION_THREADGROUP_SIZE_1D;
	uint TraceTexelForThisThread = 0;
	bool bTraceValid = false;

	if (ReflectionTileIndex < ReflectionTracingTileIndirectArgs[0])
	{
		FReflectionTileData TileData;
		const uint2 TmpReflectionTracingCoord = GetReflectionTracingScreenCoordZOrder(ReflectionTileIndex/*TileCoord*/, ReflectionGroupThreadId/*InnerTileCoord*/, TileData);
		const FReflectionTracingCoord ReflectionTracingCoord = GetReflectionTracingCoord(TmpReflectionTracingCoord, TileData.ClosureIndex);

		if (all(ReflectionTracingCoord.Coord < ReflectionTracingViewMin + ReflectionTracingViewSize))
		{
			float2 ScreenUV = GetScreenUVFromReflectionTracingCoord(ReflectionTracingCoord.Coord);
			float SceneDepth = DownsampledDepth.Load(int4(ReflectionTracingCoord.CoordFlatten, 0)).x;

			bool bHit;
			float TraceHitDistance = DecodeRayDistance(TraceHit[ReflectionTracingCoord.CoordFlatten].x, bHit);
			float MaxTraceDistance = CompactionMaxTraceDistance;

			bool bAcceptTrace = false;

			#if TRACE_COMPACTION_MODE == TRACE_COMPACTION_MODE_HIT_LIGHTING
			{
				const FTraceMaterialId MaterialId = UnpackTraceMaterialId(TraceMaterialId[ReflectionTracingCoord.CoordFlatten]);
				if (bHit && MaterialId.bNeedsHitLightingPass)
				{
					bAcceptTrace = true;
				}
			}
			#elif TRACE_COMPACTION_MODE == TRACE_COMPACTION_MODE_FAR_FIELD
			{
				bAcceptTrace = !bHit;
			}
			#else
			{
				bAcceptTrace = !bHit;
			}
			#endif

			if (SceneDepth > 0 
				&& bAcceptTrace
				&& (CullByDistanceFromCamera == 0 || SceneDepth < CompactionTracingEndDistanceFromCamera)
				&& TraceHitDistance <= CompactionMaxTraceDistance)
			{
#if WAVE_OPS
				bTraceValid = true;
				TraceTexelForThisThread = EncodeTraceTexel(ReflectionTracingCoord.Coord, ReflectionTracingCoord.ClosureIndex);
#else
				uint SharedTexelOffset;
				InterlockedAdd(SharedTraceTexelAllocator, 1, SharedTexelOffset);
				SharedTraceTexels[SharedTexelOffset] = EncodeTraceTexel(ReflectionTracingCoord.Coord, ReflectionTracingCoord.ClosureIndex);
#endif
			}
		}
	}

	GroupMemoryBarrierWithGroupSync();

#if WAVE_OPS

	const uint LastLaneIndex = WaveGetLaneCount() - 1;
	const uint LaneIndex = WaveGetLaneIndex();
	const uint OffsetInWave = WavePrefixCountBits(bTraceValid);
	uint OffsetInGroup = 0;

	if (LaneIndex == LastLaneIndex)
	{
		const uint ThisWaveSum = OffsetInWave + (bTraceValid ? 1 : 0);
		InterlockedAdd(SharedGroupSum, ThisWaveSum, OffsetInGroup);
	}
	OffsetInGroup = WaveReadLaneAt(OffsetInGroup, LastLaneIndex) + OffsetInWave;

	GroupMemoryBarrierWithGroupSync();

	// Allocate this group's compacted traces from the global allocator
	if (GroupThreadId == 0)
	{
		InterlockedAdd(RWCompactedTraceTexelAllocator[0], SharedGroupSum, SharedGlobalTraceTexelStartOffset);
	}

	GroupMemoryBarrierWithGroupSync();

	if (bTraceValid)
	{
		RWCompactedTraceTexelData[SharedGlobalTraceTexelStartOffset + OffsetInGroup] = TraceTexelForThisThread;
	}

#else // !WAVE_OPS
	uint ThreadIndex = GroupThreadId;

	if (ThreadIndex == 0)
	{
		InterlockedAdd(RWCompactedTraceTexelAllocator[0], SharedTraceTexelAllocator, SharedGlobalTraceTexelStartOffset);
	}

	GroupMemoryBarrierWithGroupSync();

	if (ThreadIndex < SharedTraceTexelAllocator)
	{
		RWCompactedTraceTexelData[SharedGlobalTraceTexelStartOffset + ThreadIndex] = SharedTraceTexels[ThreadIndex];
	}
#endif
}

#endif

#ifdef SetupCompactedTracesIndirectArgsCS

RWBuffer<uint> RWReflectionCompactTracingIndirectArgs;
RWBuffer<uint> RWReflectionCompactRayTraceDispatchIndirectArgs;

[numthreads(1, 1, 1)]
void SetupCompactedTracesIndirectArgsCS()
{
	// THREADGROUP_SIZE_64
	WriteDispatchIndirectArgs(RWReflectionCompactTracingIndirectArgs, 0, DivideAndRoundUp(CompactedTraceTexelAllocator[0],64u), 1, 1);

	// THREADGROUP_SIZE_32
	WriteDispatchIndirectArgs(RWReflectionCompactTracingIndirectArgs, 1, DivideAndRoundUp(CompactedTraceTexelAllocator[0],32u), 1, 1);

	// THREADGROUP_SIZE_256
	WriteDispatchIndirectArgs(RWReflectionCompactTracingIndirectArgs, 2, DivideAndRoundUp(CompactedTraceTexelAllocator[0],256u), 1, 1);

	// Ray tracing dispatch dimensions are defined simply in terms of threads/rays, not thread groups.
	WriteDispatchIndirectArgs(RWReflectionCompactRayTraceDispatchIndirectArgs, 0, CompactedTraceTexelAllocator[0], 1, 1);
}

#endif

#ifdef ReflectionSortTracesByMaterialCS

#ifndef THREADGROUP_SIZE_1D
	#define THREADGROUP_SIZE_1D 1
#endif
#define NUM_BINS (THREADGROUP_SIZE_1D / 2)
groupshared uint Bins[NUM_BINS];

#if DIM_WAVE_OPS
groupshared uint SharedGroupSum;
#define Offsets Bins // No need for a separate array of offsets as prefix sum is computed in-place
#else // DIM_WAVE_OPS
#if NUM_BINS > 0
groupshared uint Offsets[NUM_BINS];
#else
groupshared uint Offsets[1];
#endif
#endif // DIM_WAVE_OPS

#define ELEMENTS_PER_THREAD 16
#define NUM_ELEMENTS THREADGROUP_SIZE_1D * ELEMENTS_PER_THREAD

#if COMPILER_SUPPORTS_WAVE_SIZE && DIM_WAVE_OPS
	WAVESIZE(32)
#endif
[numthreads(THREADGROUP_SIZE_1D, 1, 1)]
void ReflectionSortTracesByMaterialCS(
	uint GroupId : SV_GroupID,
	uint GroupThreadId : SV_GroupThreadID)
{
	const uint GroupOffset = GroupId * NUM_ELEMENTS;

#if DIM_WAVE_OPS
	if (GroupThreadId == 0)
	{
		SharedGroupSum = 0;
	}
#endif

	if (GroupThreadId < NUM_BINS)
	{
		Bins[GroupThreadId] = 0;
		Offsets[GroupThreadId] = 0;
	}

	GroupMemoryBarrierWithGroupSync();

	uint Hash[NUM_ELEMENTS / THREADGROUP_SIZE_1D];
	uint TraceDataCache[NUM_ELEMENTS / THREADGROUP_SIZE_1D];
	uint TraceBinCache[NUM_ELEMENTS / THREADGROUP_SIZE_1D];

	for (int i = GroupThreadId; i < NUM_ELEMENTS; i += THREADGROUP_SIZE_1D)
	{
		uint RayIndex = GroupOffset + i;
		if (RayIndex < CompactedTraceTexelAllocator[0])
		{
			uint PackedTraceData = CompactedTraceTexelData[RayIndex];
			FReflectionTracingCoord ReflectionTracingCoord = DecodeTraceTexel(PackedTraceData);
			const FTraceMaterialId MaterialId = UnpackTraceMaterialId(TraceMaterialId[ReflectionTracingCoord.CoordFlatten]);
			uint BinIndex = MaterialId.MaterialId % NUM_BINS;

			TraceDataCache[i / THREADGROUP_SIZE_1D] = PackedTraceData;
			TraceBinCache[i / THREADGROUP_SIZE_1D] = BinIndex;

			InterlockedAdd(Bins[BinIndex], 1, Hash[i / THREADGROUP_SIZE_1D]);
		}
	}

	GroupMemoryBarrierWithGroupSync();

#if DIM_WAVE_OPS
	{
		const uint LastLaneIndex = WaveGetLaneCount() - 1;
		const uint LaneIndex = WaveGetLaneIndex();
		const uint BinIndex = GroupThreadId;
		
		uint Value = 0;
		if (BinIndex < NUM_BINS)
		{
			Value = Bins[BinIndex];
		}

		uint BinOffset = WavePrefixSum(Value);
		const uint ThisWaveSum = BinOffset + Value;
		
		uint WaveOffset = 0;
		if (LaneIndex == LastLaneIndex && ThisWaveSum > 0)
		{
			InterlockedAdd(SharedGroupSum, ThisWaveSum, WaveOffset);
		}
		BinOffset += WaveReadLaneAt(WaveOffset, LastLaneIndex);

		GroupMemoryBarrierWithGroupSync();

		if (BinIndex < NUM_BINS)
		{
			Offsets[BinIndex] = BinOffset;
		}
	}
#else // DIM_WAVE_OPS
	if (GroupThreadId < NUM_BINS)
	{
		for (int i = 0; i < GroupThreadId; ++i)
		{
			Offsets[GroupThreadId] += Bins[i];
		}
	}
#endif // DIM_WAVE_OPS

	GroupMemoryBarrierWithGroupSync();

	for (int i = GroupThreadId; i < NUM_ELEMENTS; i += THREADGROUP_SIZE_1D)
	{
		uint RayIndex = GroupOffset + i;
		if (RayIndex < CompactedTraceTexelAllocator[0])
		{
			uint BinIndex = TraceBinCache[i / THREADGROUP_SIZE_1D];
			uint OutputIndex = GroupOffset + Offsets[BinIndex] + Hash[i / THREADGROUP_SIZE_1D];

			RWCompactedTraceTexelData[OutputIndex] = TraceDataCache[i / THREADGROUP_SIZE_1D];
		}
	}
}

#endif

uint CardGridPixelSizeShift;
float3 CardGridZParams;
uint3 CullGridSize;

uint ComputeCardGridCellIndex(uint2 PixelPos, float SceneDepth)
{
	uint ZSlice = (uint)(max(0, log2(SceneDepth * CardGridZParams.x + CardGridZParams.y) * CardGridZParams.z));
	ZSlice = min(ZSlice, (uint)(CullGridSize.z - 1));
	uint3 GridCoordinate = uint3(PixelPos >> CardGridPixelSizeShift, ZSlice);
	uint GridIndex = (GridCoordinate.z * CullGridSize.y + GridCoordinate.y) * CullGridSize.x + GridCoordinate.x;
	return GridIndex;
}

float SurfaceBias;
float CardTraceEndDistanceFromCamera;
float MaxMeshSDFTraceDistance;

void TraceMeshSDFs(FReflectionTracingCoord ReflectionTracingCoord, float TraceHitDistance)
{
	float2 ScreenUV = GetScreenUVFromReflectionTracingCoord(ReflectionTracingCoord.Coord);
	float SceneDepth = DownsampledDepth.Load(int4(ReflectionTracingCoord.CoordFlatten, 0)).x;

	float3 WorldPosition = GetWorldPositionFromScreenUV(ScreenUV, SceneDepth);
	float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);
	FRayData RayData = GetRayData(ReflectionTracingCoord.CoordFlatten);

	float3 SamplePosition = WorldPosition + SurfaceBias * RayData.Direction;
	float3 SampleTranslatedPosition = TranslatedWorldPosition + SurfaceBias * RayData.Direction;
	float PullbackForSurfaceExpand = SurfaceBias;

	FRayCone RayCone = (FRayCone)0;
	RayCone.SpreadAngle = View.EyeToPixelSpreadAngle;
	RayCone = PropagateRayCone(RayCone, RayData.ConeHalfAngle, SceneDepth);

	FConeTraceInput TraceInput;
	TraceInput.Setup(SamplePosition, SampleTranslatedPosition, RayData.Direction, RayCone.SpreadAngle, 0.0f, max(TraceHitDistance - PullbackForSurfaceExpand, 0.0f), RayData.RadianceCacheMaxTraceDistance, 1.0f);
	TraceInput.bHiResSurface = UseHighResSurface != 0;
	TraceInput.VoxelTraceStartDistance = MaxMeshSDFTraceDistance;
	TraceInput.DitherScreenCoord = ReflectionTracingCoord.CoordFlatten.xy;

	uint CardGridCellIndex = ComputeCardGridCellIndex(ReflectionTracingCoord.Coord * ReflectionDownsampleFactorXY, SceneDepth);
	TraceInput.NumMeshSDFs = NumGridCulledMeshSDFObjects[CardGridCellIndex];
	TraceInput.MeshSDFStartOffset = GridCulledMeshSDFObjectStartOffsetArray[CardGridCellIndex];
	TraceInput.CardInterpolateInfluenceRadius = 10;
	// Only expand the SDF surface when the ray traversal is far enough away to not self-intersect
	TraceInput.bExpandSurfaceUsingRayTimeInsteadOfMaxDistance = false;
	TraceInput.InitialMaxDistance = 0;

	FConeTraceResult TraceResult;
	ConeTraceLumenSceneCards(TraceInput, TraceResult);

	TraceInput.NumHeightfields = NumGridCulledHeightfieldObjects[CardGridCellIndex];
	TraceInput.HeightfieldStartOffset = GridCulledHeightfieldObjectStartOffsetArray[CardGridCellIndex];
	ConeTraceLumenSceneHeightfields(TraceInput, TraceResult);

	// Trace against hair voxel structure, if enabled and not already done by other tracing method
	#if USE_HAIRSTRANDS_VOXEL
	{
		float HairTraceDistance = min(RayData.RadianceCacheMaxTraceDistance, TraceResult.OpaqueHitDistance);

		bool bHairHit;
		float HairTransparency;
		float HairHitT;
		
		TraceHairVoxels(
			ReflectionTracingCoord.Coord, 
			SceneDepth, 
			// Use (Translated)WorldPosition instead of SamplePosition, as the bias is too strong otherwise. This is not an issue as 
			// the voxel structure does not cause any self shadowing issue
			TranslatedWorldPosition, 
			RayData.Direction, 
			HairTraceDistance,
			true, 
			bHairHit,
			HairTransparency,
			HairHitT);

		if (bHairHit && HairHitT < HairTraceDistance)
		{
			TraceResult.Lighting *= HairTransparency;
			TraceResult.Transparency *= HairTransparency;
			TraceResult.OpaqueHitDistance = min(HairHitT, TraceResult.OpaqueHitDistance);
		}
	}
	#endif

	float3 Lighting = TraceResult.Lighting * View.PreExposure;
	float Transparency = TraceResult.Transparency;
	float OpaqueHitDistance = TraceResult.OpaqueHitDistance;

	float DistanceFromViewpoint = GetDistanceToCameraFromViewVector(DFHackToFloat(PrimaryView.WorldCameraOrigin) - WorldPosition);
	float DistanceFade = saturate(4 * DistanceFromViewpoint / CardTraceEndDistanceFromCamera - 3);
	DistanceFade = max(DistanceFade, saturate((OpaqueHitDistance - MaxMeshSDFTraceDistance * .7f) / (MaxMeshSDFTraceDistance * .3f)));

	if(SampleHeightFog > 0)
	{
		float CoverageForFog = saturate(1.0 - Transparency);	// LUMEN_TODO single evaluation for now even if we can hit translucent hair above.
		float3 OriginToCollider = RayData.Direction * TraceResult.OpaqueHitDistance;
		Lighting.rgb = GetFogOnLuminance(Lighting.rgb, CoverageForFog, TranslatedWorldPosition, RayData.Direction, TraceResult.OpaqueHitDistance);
	}

	Transparency = lerp(Transparency, 1, DistanceFade);

	Lighting += GetSkylightLeakingForReflections(RayData.Direction, TraceResult.GeometryWorldNormal, OpaqueHitDistance) * View.PreExposure;

	float MaxLighting = max3(Lighting.x, Lighting.y, Lighting.z);

	if (MaxLighting > MaxRayIntensity)
	{
		Lighting *= MaxRayIntensity / MaxLighting;
	}

	#if DEBUG_VISUALIZE_TRACE_TYPES
		RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = float3(0, .5f, 0) * View.PreExposure;
	#else
		RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = Lighting;
	#endif

	bool bHit = false;

	if (Transparency < InterleavedGradientNoise(ReflectionTracingCoord.Coord + 0.5f, 0))
	{
		bHit = true;
	}

	RWTraceHit[ReflectionTracingCoord.CoordFlatten] = EncodeRayDistance(OpaqueHitDistance, bHit);
}

#if THREADGROUP_SIZE_32
	#define REFLECTION_TRACE_THREADGROUP_SIZE_1D 32
#else
	#define REFLECTION_TRACE_THREADGROUP_SIZE_1D 64
#endif

#ifdef ReflectionTraceMeshSDFsCS

[numthreads(REFLECTION_TRACE_THREADGROUP_SIZE_1D, 1, 1)]
void ReflectionTraceMeshSDFsCS(uint DispatchThreadId : SV_DispatchThreadID)
{
	if (DispatchThreadId < CompactedTraceTexelAllocator[0])
	{
		const FReflectionTracingCoord ReflectionTracingCoord = DecodeTraceTexel(CompactedTraceTexelData[DispatchThreadId]);

		bool bUnused;
		float TraceHitDistance = DecodeRayDistance(RWTraceHit[ReflectionTracingCoord.CoordFlatten], bUnused);

		TraceMeshSDFs(ReflectionTracingCoord, TraceHitDistance);
	}
}

#endif

void TraceVoxels(FReflectionTracingCoord ReflectionTracingCoord, float TraceHitDistance)
{
	float2 ScreenUV = GetScreenUVFromReflectionTracingCoord(ReflectionTracingCoord.Coord);
	float SceneDepth = DownsampledDepth.Load(int4(ReflectionTracingCoord.CoordFlatten, 0)).x;

	float3 WorldPosition = GetWorldPositionFromScreenUV(ScreenUV, SceneDepth); // LUMEN_LWC_TODO
	float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);
	FRayData RayData = GetRayData(ReflectionTracingCoord.CoordFlatten);

	float3 SamplePosition = WorldPosition + SurfaceBias * RayData.Direction;
	float3 SampleTranslatedPosition = TranslatedWorldPosition + SurfaceBias * RayData.Direction;

	FRayCone RayCone = (FRayCone)0;
	RayCone.SpreadAngle = View.EyeToPixelSpreadAngle;
	RayCone = PropagateRayCone(RayCone, RayData.ConeHalfAngle, SceneDepth);

	FConeTraceInput TraceInput;
	TraceInput.Setup(SamplePosition, SampleTranslatedPosition, RayData.Direction, RayCone.SpreadAngle, 0.0f, 0.0f, RayData.RadianceCacheMaxTraceDistance, 1.0f);
	TraceInput.DitherScreenCoord = ReflectionTracingCoord.Coord;
	TraceInput.bHiResSurface = UseHighResSurface != 0;

	uint ClipmapForSurfaceExpand = ComputeGlobalDistanceFieldClipmapIndex(TranslatedWorldPosition + TraceHitDistance * RayData.Direction);
	// Pull back from where the previous tracing method left off so that we start outside the surface, so that the SDF expansion works properly and prevents leaking through thin objects
	float PullbackForSurfaceExpand = GlobalVolumeTranslatedCenterAndExtent[ClipmapForSurfaceExpand].w * GlobalVolumeTexelSize * 4.0f;
	TraceInput.VoxelTraceStartDistance = max(TraceHitDistance - PullbackForSurfaceExpand, 0.0f);
	// Only expand the SDF surface when the ray traversal is far enough away to not self-intersect
	TraceInput.bExpandSurfaceUsingRayTimeInsteadOfMaxDistance = false;
	TraceInput.InitialMaxDistance = 0;

	// Dither step factor to smooth out SDF stepping artifacts, but only when near mirror
	float StepFactorNoise = lerp(.95f, 1.0f / .95f, InterleavedGradientNoise(ReflectionTracingCoord.Coord, ReflectionsStateFrameIndexMod8));
	TraceInput.SDFStepFactor = lerp(StepFactorNoise, 1.0f, saturate(RayData.ConeHalfAngle / (PI / 256.0f)));

	FConeTraceResult TraceResult = (FConeTraceResult)0;
	TraceResult.Lighting = 0;
	TraceResult.Transparency = 1;
	TraceResult.OpaqueHitDistance = TraceInput.MaxTraceDistance;
	TraceResult.GeometryWorldNormal = 0.0f;

	#if TRACE_GLOBAL_SDF
	{
		ConeTraceLumenSceneVoxels(TraceInput, TraceResult);
	}
	#endif

	// Trace against hair voxel structure, if enabled and not already done by other tracing method
	#if USE_HAIRSTRANDS_VOXEL
	{
		float HairTraceDistance = min(RayData.RadianceCacheMaxTraceDistance, TraceResult.OpaqueHitDistance);

		bool bHairHit;
		float HairTransparency;
		float HairHitT;

		TraceHairVoxels(
			ReflectionTracingCoord.Coord,
			SceneDepth,
			// Use (Translated)WorldPosition instead of SamplePosition, as the bias is too strong otherwise. This is not an issue as 
			// the voxel structure does not cause any self shadowing issue
			TranslatedWorldPosition,
			RayData.Direction,
			HairTraceDistance,
			true,
			bHairHit,
			HairTransparency,
			HairHitT);

		if (bHairHit && HairHitT < HairTraceDistance)
		{
			TraceResult.Lighting *= HairTransparency;
			TraceResult.Transparency *= HairTransparency;
			TraceResult.OpaqueHitDistance = min(HairHitT, TraceResult.OpaqueHitDistance);
		}
	}
	#endif

	bool bHit = false;

	if (TraceResult.Transparency <= .5f)
	{
		TraceHitDistance = TraceResult.OpaqueHitDistance;
		bHit = true;
	}
	else
	{
		TraceHitDistance = GetMaxHitDistance();
	}

#if SAMPLE_SCENE_COLOR
	if (bHit)
	{
		float3 HitTranslatedWorldPosition = SampleTranslatedPosition + RayData.Direction * TraceResult.OpaqueHitDistance;
		SampleSceneColorAtHit(HitTranslatedWorldPosition, TraceResult.GeometryWorldNormal, ReflectionTracingCoord.Coord, RelativeDepthThickness, TraceResult.Lighting);
	}
#endif

	float CoverageForFog = saturate(1.0 - TraceResult.Transparency);

#if DISTANT_SCREEN_TRACES
	if (!bHit)
	{
		float3 TranslatedDistantRayOrigin = TranslatedWorldPosition;
		float MaxTraceDistanceForDistantScreenTrace = DistantScreenTraceMaxTraceDistance;

		// Start the distant linear screen traces where we have no fallback for HZB screen traces
		ClipRayToStartOutsideGlobalSDF(TranslatedDistantRayOrigin, RayData.Direction, MaxTraceDistanceForDistantScreenTrace);

		if (MaxTraceDistanceForDistantScreenTrace > 0)
		{
			float4 StartRayClip = mul(float4(TranslatedDistantRayOrigin, 1.0), View.TranslatedWorldToClip);

			// If we start tracing from outside the screen (from a position on the screen) then we know we will never hit any valueable screen information.
			// So we skip the tracing if the start position is clipped outside of the view frustum.
			if (StartRayClip.w >= 0.0 && all(-StartRayClip.ww <= StartRayClip.xy) && all(StartRayClip.xy <= StartRayClip.ww))
			{
				// Update SceneDepth according to the new start position moved to the edge of the clipmap. This is important for the SSR HZB tracing to be correct.
				float StartRayDeviceZ = StartRayClip.z / StartRayClip.w;
				float DistantTracingSceneDepth = ConvertFromDeviceZ(StartRayDeviceZ);

				DistantScreenTrace(
					ReflectionTracingCoord.Coord + 0.5f,
					HZBUvFactorAndInvFactor,
					TranslatedDistantRayOrigin,
					RayData.Direction,
					MaxTraceDistanceForDistantScreenTrace,
					DistantTracingSceneDepth,
					bHit,
					TraceResult.Lighting);
			}
		}
	}
#endif

	if (!bHit)
	{
#if RADIANCE_CACHE
		if (RayData.RadianceCacheMaxTraceDistance < MaxTraceDistance * .99f)
		{
			FRadianceCacheCoverage Coverage = GetRadianceCacheCoverage(WorldPosition, RayData.Direction, InterleavedGradientNoise(ReflectionTracingCoord.Coord, ReflectionsStateFrameIndexMod8));
			SampleRadianceCacheAndApply(Coverage, WorldPosition, RayData.Direction, RayData.ConeHalfAngle, /*RandomScalarForStochasticIterpolation*/ 0.5f, TraceResult.Lighting, TraceResult.Transparency);
			TraceHitDistance = RayData.RadianceCacheMaxTraceDistance;
		}
		else
#endif
		{
			TraceHitDistance = MaxTraceDistance;
			ApplySkylightToTraceResult(RayData.Direction, TraceResult);
			CoverageForFog = 1.0f;
		}
	}

	TraceResult.Lighting += GetSkylightLeakingForReflections(TraceResult.GeometryWorldNormal, RayData.Direction, TraceResult.OpaqueHitDistance);

	TraceResult.Lighting *= View.PreExposure;

	if (SampleHeightFog > 0)
	{
		float3 OriginToCollider = RayData.Direction * TraceHitDistance;
		TraceResult.Lighting.rgb = GetFogOnLuminance(TraceResult.Lighting.rgb, CoverageForFog, TranslatedWorldPosition, RayData.Direction, TraceHitDistance);
	}

	float MaxLighting = max3(TraceResult.Lighting.x, TraceResult.Lighting.y, TraceResult.Lighting.z);

	if (MaxLighting > MaxRayIntensity)
	{
		TraceResult.Lighting *= MaxRayIntensity / MaxLighting;
	}

	#if DEBUG_VISUALIZE_TRACE_TYPES
		RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = float3(0, 0, .5f) * View.PreExposure;
	#else
		RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = TraceResult.Lighting;
	#endif

	RWTraceHit[ReflectionTracingCoord.CoordFlatten] = EncodeRayDistance(TraceHitDistance, bHit);
}

#ifdef ReflectionTraceVoxelsCS

[numthreads(REFLECTION_TRACE_THREADGROUP_SIZE_1D, 1, 1)]
void ReflectionTraceVoxelsCS(uint DispatchThreadId : SV_DispatchThreadID)
{
	if (DispatchThreadId < CompactedTraceTexelAllocator[0])
	{
		const FReflectionTracingCoord ReflectionTracingCoord = DecodeTraceTexel(CompactedTraceTexelData[DispatchThreadId]);

		bool bUnused;
		float TraceHitDistance = DecodeRayDistance(RWTraceHit[ReflectionTracingCoord.CoordFlatten], bUnused);

		TraceVoxels(ReflectionTracingCoord, TraceHitDistance);
	}
}

#endif

#ifdef VisualizeReflectionTracesCS

[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void VisualizeReflectionTracesCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	if (all(DispatchThreadId == 0))
	{
		int2 QueryProbeScreenPosition = View.CursorPosition.x >= 0 ? View.CursorPosition * View.ViewResolutionFraction / ReflectionDownsampleFactorXY : -1;

		if (all(QueryProbeScreenPosition >= ReflectionTracingViewMin) && all(QueryProbeScreenPosition < ReflectionTracingViewMin + ReflectionTracingViewSize))
		{
			const FReflectionTracingCoord ReflectionTracingCoord = GetReflectionTracingCoord(QueryProbeScreenPosition, 0 /*TODO*/);
			float2 ScreenUV = GetScreenUVFromReflectionTracingCoord(ReflectionTracingCoord.Coord);
			float SceneDepth = DownsampledDepth.Load(int4(ReflectionTracingCoord.CoordFlatten,0)).x;

			float3 WorldPosition = GetWorldPositionFromScreenUV(ScreenUV, SceneDepth); // LUMEN_LWC_TODO
			float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);
			FRayData RayData = GetRayData(ReflectionTracingCoord.CoordFlatten);

			bool bHit;
			float TraceHitDistance = DecodeRayDistance(TraceHit[ReflectionTracingCoord.CoordFlatten].x, bHit);
			float3 RayHitRadiance = TraceRadiance[ReflectionTracingCoord.CoordFlatten];

			FShaderPrintContext Context = InitShaderPrintContext(true, float2(0.1, 0.1));

			Print(Context, TEXT("Pixel: "));
			Print(Context, QueryProbeScreenPosition.x);
			Print(Context, QueryProbeScreenPosition.y);
			Newline(Context);

			if (SceneDepth > 0.0f)
			{
				Print(Context, TEXT("Distance: "));
				Print(Context, TraceHitDistance);
				Newline(Context);

				Print(Context, TEXT("Radiance: "));
				Print(Context, RayHitRadiance.x);
				Print(Context, RayHitRadiance.y);
				Print(Context, RayHitRadiance.z);

				const float3 HitPoint = TranslatedWorldPosition + RayData.Direction * TraceHitDistance;
				AddLineTWS(Context, TranslatedWorldPosition, HitPoint, float4(RayHitRadiance, 1.0f));
				AddCrossTWS(Context, HitPoint, 2.0f, float4(1, 1, 0, 1));				
			}
		}
	}
}

#endif