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

/*=============================================================================
	LumenTranslucencyVolumeHardwareRayTracing.usf
=============================================================================*/

// TODO: Remove hair dependency
#include "../HairStrands/HairStrandsVoxelPageCommonStruct.ush"

#include "/Engine/Shared/RayTracingTypes.h"
#include "../Common.ush"
#include "../SHCommon.ush"
#include "../MonteCarlo.ush"
#include "../OctahedralCommon.ush"
#include "../SceneTextureParameters.ush"

#include "LumenCardCommon.ush"
#include "LumenTracingCommon.ush"
#include "LumenRadianceCacheCommon.ush"
#include "LumenTranslucencyVolumeLightingShared.ush"

#if LUMEN_HARDWARE_RAYTRACING
#include "LumenHardwareRayTracingCommon.ush"

#ifndef DIM_LIGHTING_MODE
#define DIM_LIGHTING_MODE LIGHTING_FROM_SURFACE_CACHE
#endif

RaytracingAccelerationStructure TLAS;
RaytracingAccelerationStructure FarFieldTLAS;

#if LUMEN_HARDWARE_INLINE_RAYTRACING
	StructuredBuffer<FHitGroupRootConstants> HitGroupData;
	StructuredBuffer<FRayTracingSceneMetadataRecord> RayTracingSceneMetadata;
	RWStructuredBuffer<uint> RWInstanceHitCountBuffer;
#endif // LUMEN_HARDWARE_INLINE_RAYTRACING

RWTexture3D<float3> RWVolumeTraceRadiance;
RWTexture3D<float> RWVolumeTraceHitDistance;
uint MaxTraversalIterations;
uint MeshSectionVisibilityTest;

// Applying an offset from the depth buffer helps with reducing self intersection with global distance field tracing.
void ApplyDepthconstraintsToOffset(float3 GridCoordinate, inout float3 LocalFrameJitterOffset)
{
	if (GridCenterOffsetFromDepthBuffer >= 0.0)
	{
		float4 GridSampleClip	= mul(float4(ComputeCellTranslatedWorldPosition(GridCoordinate, LocalFrameJitterOffset), 1), View.TranslatedWorldToClip);
		float3 GridSampleScreen	= GridSampleClip.xyz / GridSampleClip.w;
		float3 GridSampleUVz	= float3(GridSampleScreen.xy * View.ScreenPositionScaleBias.xy + View.ScreenPositionScaleBias.wz, GridSampleScreen.z);
		float GridSampleSceneDepth = ConvertFromDeviceZ(SceneTexturesStruct.SceneDepthTexture.SampleLevel(SceneTexturesStruct_SceneDepthTextureSampler, GridSampleUVz.xy, 0).r);

		float DepthGridCoordinateWithZOffset = ComputeZSliceFromDepth(GridSampleSceneDepth) - GridCenterOffsetFromDepthBuffer;

		float3 GridCoordinateOffset = float3(GridCoordinate) + LocalFrameJitterOffset;
		
		const float DepthToSampleGridOffset = DepthGridCoordinateWithZOffset - GridCoordinateOffset.z;	// <0 if samples is in front of the depth buffer
		if (DepthToSampleGridOffset < 0.0																// Move the sample forward only if the depth buffer grid coordinate is closer,
			&& (-DepthToSampleGridOffset < GridCenterOffsetThresholdToAcceptDepthBufferOffset)			// But only if it is less than N slice away (to not bring all the sample behind forward)
			)
		{
			// If the grid sample is further than the grid z for the depth buffer, adjust the offset 
			LocalFrameJitterOffset.z += DepthToSampleGridOffset;
		}
	}
}

LUMEN_HARDWARE_RAY_TRACING_ENTRY(LumenTranslucencyVolumeHardwareRayTracing)
{
	uint OctahedralAtlasSizeX = TranslucencyGIGridSize.x * TranslucencyVolumeTracingOctahedronResolution;
	// DispatchThreadId to match compute version
	uint3 DispatchThreadId = uint3(DispatchThreadIndex.x % OctahedralAtlasSizeX, DispatchThreadIndex.y, DispatchThreadIndex.x / OctahedralAtlasSizeX);

#define DEINTERLEAVED_VOLUME_TRACING 0
#if DEINTERLEAVED_VOLUME_TRACING
	uint3 GridCoordinate = uint3(DispatchThreadId.xy % TranslucencyGIGridSize.xy, DispatchThreadId.z);
	uint2 TraceTexelCoord = DispatchThreadId.xy / TranslucencyGIGridSize.xy;
#else
	uint3 GridCoordinate = uint3(DispatchThreadId.xy / TranslucencyVolumeTracingOctahedronResolution, DispatchThreadId.z);
	uint2 TraceTexelCoord = DispatchThreadId.xy - GridCoordinate.xy * TranslucencyVolumeTracingOctahedronResolution;
#endif

	if (all(GridCoordinate < TranslucencyGIGridSize) && all(TraceTexelCoord < TranslucencyVolumeTracingOctahedronResolution))
	{
		FRayTracedLightingResult TraceLightingResult;
		TraceLightingResult.Radiance = 0;
		TraceLightingResult.TraceHitDistance = MaxTraceDistance;

		if (IsFroxelVisible(GridCoordinate))
		{
			float3 LocalFrameJitterOffset = FrameJitterOffset.xyz;
			ApplyDepthconstraintsToOffset(GridCoordinate, LocalFrameJitterOffset);
			float3 WorldPosition = ComputeCellWorldPosition(GridCoordinate, LocalFrameJitterOffset);

			float2 ProbeUV;
			float ConeHalfAngle;
			GetProbeTracingUV(TraceTexelCoord, GetProbeTexelCenter(GridCoordinate.xy), ProbeUV, ConeHalfAngle);

			float3 WorldConeDirection = EquiAreaSphericalMapping(ProbeUV);

			float TraceDistance = MaxTraceDistance;

			FRadianceCacheCoverage Coverage;
			Coverage.bValid = false;

			#if PROBE_SOURCE_MODE == PROBE_SOURCE_MODE_RADIANCE_CACHE
				Coverage = GetRadianceCacheCoverage(WorldPosition, WorldConeDirection, .5f);
				
				if (Coverage.bValid)
				{
					TraceDistance = min(TraceDistance, Coverage.MinTraceDistanceBeforeInterpolation);
				}
			#endif

			FRayDesc Ray;
			Ray.Origin = WorldPosition + DFHackToFloat(PrimaryView.PreViewTranslation); // LUMEN_LWC_TODO
			Ray.Direction = WorldConeDirection;
			Ray.TMin = 0;
			Ray.TMax = TraceDistance;

			FRayCone RayCone = (FRayCone)0;
			RayCone = PropagateRayCone(RayCone, ConeHalfAngle, 0.0);

			const uint LinearCoord = GridCoordinate.y * TranslucencyGIGridSize.x + GridCoordinate.x;

			FRayTracedLightingContext Context = CreateRayTracedLightingContext(
				RayCone,
				GridCoordinate.xy,
				LinearCoord,
				/*CullingMode*/ RAY_FLAG_NONE, // Disable culling in order to minimize leaking when rays start inside walls
				MaxTraversalIterations,
				MeshSectionVisibilityTest != 0);

			#if DIM_LIGHTING_MODE == LIGHTING_FROM_SURFACE_CACHE
			{
#if LUMEN_HARDWARE_INLINE_RAYTRACING
				Context.HitGroupData = HitGroupData;
				Context.RayTracingSceneMetadata = RayTracingSceneMetadata;
				Context.RWInstanceHitCountBuffer = RWInstanceHitCountBuffer;
#endif
				TraceLightingResult = TraceSurfaceCacheRay(TLAS, Ray, Context);
			}
			#else
			{
				TraceLightingResult = TraceAndCalculateRayTracedLighting(TLAS, FarFieldTLAS, Ray, Context, DIM_LIGHTING_MODE, false);
			}
			#endif

			if (!TraceLightingResult.bIsHit)
			{
				float Transparency = 1.0f;

			#if PROBE_SOURCE_MODE == PROBE_SOURCE_MODE_RADIANCE_CACHE
				if (Coverage.bValid)
				{
					SampleRadianceCacheAndApply(Coverage, WorldPosition, WorldConeDirection, ConeHalfAngle, /*RandomScalarForStochasticIterpolation*/ 0.5f, TraceLightingResult.Radiance, Transparency);
				}
				else
			#endif
				{
					FConeTraceResult TraceResult;
					TraceResult.Lighting = TraceLightingResult.Radiance;
					TraceResult.Transparency = Transparency;
					ApplySkylightToTraceResult(WorldConeDirection, TraceResult);
					TraceLightingResult.Radiance = TraceResult.Lighting;
				}
			}

			TraceLightingResult.Radiance *= View.PreExposure;

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

			if (MaxLighting > MaxRayIntensity)
			{
				TraceLightingResult.Radiance *= MaxRayIntensity / MaxLighting;
			}
		}

		uint3 WriteCoord = uint3(GridCoordinate.xy * TranslucencyVolumeTracingOctahedronResolution + TraceTexelCoord, GridCoordinate.z);
		RWVolumeTraceRadiance[WriteCoord] = TraceLightingResult.Radiance;
		RWVolumeTraceHitDistance[WriteCoord] = min(TraceLightingResult.TraceHitDistance, MaxHalfFloat);
	}
}

#endif // LUMEN_HARDWARE_RAYTRACING