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

#include "/Engine/Shared/RayTracingTypes.h"
#include "../Common.ush"
#include "../MonteCarlo.ush"
#include "../SceneTextureParameters.ush"
#include "../StochasticLighting/StochasticLightingUpsample.ush"

// Ensure that DiffuseAlbedo is not overridden on SSS material (as we don't split lighting with the Lumen/RT integrator)
#define SUBSTRATE_SSS_MATERIAL_OVERRIDE 0
#define SUBSTRATE_COMPLEXSPECIALPATH 0
#define SUBSTRATE_GLINTS_ALLOWED 0

// When tracing from compute, SUBSTRATE_MATERIALCONTAINER_IS_VIEWRESOURCE=0 is not automatically detected, so we notify the use of raytracing here.
#define SUBSTRATE_MATERIALCONTAINER_IS_VIEWRESOURCE 0 

// Additional rewiring to make DeferredShadingCommon happy
#define PreIntegratedGF				ReflectionStruct.PreIntegratedGF
#define PreIntegratedGFSampler		GlobalBilinearClampedSampler

#include "LumenMaterial.ush"
#include "LumenCardCommon.ush"
#include "LumenTracingCommon.ush"
#include "LumenReflectionCommon.ush"
#include "LumenVisualizeTraces.ush"
#include "LumenRadianceCacheCommon.ush"
#include "LumenScreenTracing.ush"
#define USE_HAIRSTRANDS_VOXEL DIM_HAIRSTRANDS_VOXEL
#include "LumenHairTracing.ush"

#ifndef THREADGROUP_SIZE_2D
#define THREADGROUP_SIZE_2D 8
#endif

#ifndef THREADGROUP_SIZE_1D
#define THREADGROUP_SIZE_1D THREADGROUP_SIZE_2D * THREADGROUP_SIZE_2D
#endif

#define RAY_TRACING_PASS_DEFAULT		0
#define RAY_TRACING_PASS_FAR_FIELD		1
#define RAY_TRACING_PASS_HIT_LIGHTING	2

uint2 OutputThreadGroupSize;
RWBuffer<uint> RWHardwareRayTracingIndirectArgs;

#ifdef FLumenReflectionHardwareRayTracingIndirectArgsCS
[numthreads(1, 1, 1)]
void FLumenReflectionHardwareRayTracingIndirectArgsCS()
{
	WriteDispatchIndirectArgs(RWHardwareRayTracingIndirectArgs, 0, (CompactedTraceTexelAllocator[0] + OutputThreadGroupSize.x - 1) / OutputThreadGroupSize.x, 1, 1);
}
#endif

#if LUMEN_HARDWARE_RAYTRACING || LUMEN_HARDWARE_INLINE_RAYTRACING

#include "LumenHardwareRayTracingCommon.ush"

RaytracingAccelerationStructure TLAS;
RaytracingAccelerationStructure FarFieldTLAS;

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

uint HitLightingForceOpaque;
uint HitLightingShadowMode;
uint HitLightingShadowTranslucencyMode;
uint HitLightingDirectLighting;
uint HitLightingSkylight;
uint UseReflectionCaptures;

float RayTracingBias;
float RayTracingNormalBias;
float NearFieldMaxTraceDistance;
float FarFieldBias;
float FarFieldMaxTraceDistance;
float NearFieldMaxTraceDistanceDitherScale;
float NearFieldSceneRadius;
float PullbackBias;
float DistantScreenTracesStartDistance;
uint MaxTraversalIterations;
uint MeshSectionVisibilityTest;
int ApplySkyLight;
int HitLightingForceEnabled;
int SampleSceneColor;
float RelativeDepthThickness;
uint MaxReflectionBounces;
uint MaxRefractionBounces;

Texture2DArray<uint2> TraceBookmark;

RWTexture2DArray<float3> RWTraceRadiance;
RWTexture2DArray<float>  RWTraceHit;
RWTexture2DArray<uint2>  RWTraceBookmark;
RWTexture2DArray<uint>   RWTraceMaterialId;

LUMEN_HARDWARE_RAY_TRACING_ENTRY(LumenReflectionHardwareRayTracing)
{
	uint ThreadIndex = DispatchThreadIndex.x;
	uint GroupIndex = DispatchThreadIndex.y;
	uint DispatchThreadId = GroupIndex * THREADGROUP_SIZE_1D + ThreadIndex;
	
	if (DispatchThreadId < CompactedTraceTexelAllocator[0])
	{
		const FReflectionTracingCoord ReflectionTracingCoord = DecodeTraceTexel(CompactedTraceTexelData[DispatchThreadId]);

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

		float2 ScreenUV = GetScreenUVFromReflectionTracingCoord(ReflectionTracingCoord.Coord);
		float2 ScreenCoord = ScreenUV * View.BufferSizeAndInvSize.xy;
		uint LinearCoord = ScreenCoord.y * View.BufferSizeAndInvSize.x + ScreenCoord.x;

		float SceneDepth = DownsampledDepth.Load(int4(ReflectionTracingCoord.CoordFlatten, 0)).x;
		float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);
		const float3 WorldNormal = ReadMaterialData(uint2(ScreenCoord)).WorldNormal;
		FRayData RayData = GetRayData(ReflectionTracingCoord.CoordFlatten);
		const bool bIsFirstPerson = RayData.bIsFirstPersonPixel;

		float ClippedNearFieldMaxTraceDistance = ClipAndDitherNearFieldMaxTraceDistance(
			TranslatedWorldPosition,
			RayData.Direction,
			ReflectionTracingCoord.Coord,
			NearFieldSceneRadius,
			NearFieldMaxTraceDistance,
			NearFieldMaxTraceDistanceDitherScale);

		float TraceMaxDistance = ClippedNearFieldMaxTraceDistance;
		#if DIM_RADIANCE_CACHE
		{
			TraceMaxDistance = min(TraceMaxDistance, RayData.RadianceCacheMaxTraceDistance);
		}
		#endif
		#if RAY_TRACING_PASS == RAY_TRACING_PASS_FAR_FIELD
		{
			TraceMaxDistance = max(TraceMaxDistance, FarFieldMaxTraceDistance);
		}
		#endif
		#if RAY_TRACING_PASS == RAY_TRACING_PASS_HIT_LIGHTING
		{
			if (TraceHitDistance > ClippedNearFieldMaxTraceDistance)
			{
				TraceMaxDistance = max(TraceMaxDistance, FarFieldMaxTraceDistance);
			}

			// Make sure that the hit-lighting re-trace will hit something
			TraceHitDistance = max(TraceHitDistance - 0.001f, 0.0f);
		}
		#endif

		float RayBias = (RAY_TRACING_PASS == RAY_TRACING_PASS_FAR_FIELD) ? FarFieldBias : RayTracingBias;

		FRayDesc Ray;
		Ray.Origin = TranslatedWorldPosition;
		Ray.Direction = RayData.Direction;
		Ray.TMin = max(TraceHitDistance, RayBias);
		Ray.TMax = max(TraceMaxDistance, RayBias);

		// Apply PullbackBias after screen traces only if ray didn't reach the end, so that we don't retrace it again
		if (Ray.TMin < Ray.TMax)
		{
			Ray.TMin = max(Ray.TMin - PullbackBias, RayBias);
		}

		// Normal bias
		if (RayTracingNormalBias > 0.0f)
		{
			Ray.Origin += WorldNormal * saturate(dot(WorldNormal, Ray.Direction)) * RayTracingNormalBias;
		}

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

		FRayTracedLightingContext Context = CreateRayTracedLightingContext(
			RayCone,
			ReflectionTracingCoord.Coord,
			LinearCoord,
			/*CullingMode*/ RAY_FLAG_CULL_BACK_FACING_TRIANGLES,
			MaxTraversalIterations,
			MeshSectionVisibilityTest != 0);

		Context.bHiResSurface = UseHighResSurface != 0;
		Context.MaxReflectionBounces = MaxReflectionBounces;
		Context.bUseBookmark = true;

	#if RECURSIVE_REFRACTION_TRACES
		Context.MaxRefractionBounces = MaxRefractionBounces;
		Context.InstanceMask |= RAY_TRACING_MASK_TRANSLUCENT;  
	#endif

		// The first person GBuffer bit is not available when static lighting is enabled, so we skip intersections with the first person world space representation to avoid artifacts in that case.
#if HAS_FIRST_PERSON_GBUFFER_BIT
		// When tracing a ray for a first person pixel, we do not want to intersect the world space representation of that first person primitive.
		// Doing so would lead to "self"-intersection artifacts between the two representations of that mesh.
		if (!bIsFirstPerson)
		{
			Context.InstanceMask |= RAY_TRACING_MASK_OPAQUE_FP_WORLD_SPACE;
		}
#endif

		FRayTracedLightingResult Result = CreateRayTracedLightingResult(Ray);
		if (Ray.TMin < Ray.TMax)
		{
			#if LUMEN_HARDWARE_INLINE_RAYTRACING
			{
				Context.HitGroupData = HitGroupData;
				Context.RayTracingSceneMetadata = RayTracingSceneMetadata;
				Context.RWInstanceHitCountBuffer = RWInstanceHitCountBuffer;				
			}
			#endif

			#if RAY_TRACING_PASS == RAY_TRACING_PASS_HIT_LIGHTING
			{
				FLumenRayHitBookmark RayHitBookmark;
				RayHitBookmark.PackedData = TraceBookmark[ReflectionTracingCoord.CoordFlatten].xy;

				if (TraceHitDistance > ClippedNearFieldMaxTraceDistance)
				{
					Context.bIsFarFieldRay = true;
				}

				Context.bHitLightingDirectLighting = HitLightingDirectLighting != 0;
				Context.bHitLightingSkylight = HitLightingSkylight != 0;
				Context.bUseReflectionCaptures = UseReflectionCaptures != 0;
				Context.bForceOpaque = HitLightingForceOpaque;
				Context.HitLightingShadowMode = HitLightingShadowMode;
				Context.HitLightingShadowTranslucencyMode = HitLightingShadowTranslucencyMode;
				Context.HitLightingShadowMaxTraceDistance = NearFieldMaxTraceDistance;

				Result = TraceAndCalculateRayTracedLighting(TLAS, FarFieldTLAS, Ray, Context, RayHitBookmark);
			}
			#elif RAY_TRACING_PASS == RAY_TRACING_PASS_DEFAULT
			{
				Result = TraceSurfaceCacheRay(TLAS, Ray, Context);
			}
			#elif RAY_TRACING_PASS == RAY_TRACING_PASS_FAR_FIELD
			{
				Result = TraceSurfaceCacheFarFieldRay(FarFieldTLAS, Ray, Context);
			}
			#endif
		}

		FConeTraceResult TraceResult;
		TraceResult.Lighting = Result.Radiance;
		TraceResult.Transparency = 1;
		TraceResult.OpaqueHitDistance = Result.TraceHitDistance;
		TraceResult.GeometryWorldNormal = Result.GeometryWorldNormal;

		// Trace against hair voxel structure, if enabled
		#if DIM_HAIRSTRANDS_VOXEL && RAY_TRACING_PASS != RAY_TRACING_PASS_FAR_FIELD
		if (Ray.TMin < Ray.TMax)
		{
			float HairTraceDistance = min(Ray.TMax, 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 bNeedsHitLightingPass = false;
		#if RAY_TRACING_PASS == RAY_TRACING_PASS_DEFAULT || RAY_TRACING_PASS == RAY_TRACING_PASS_FAR_FIELD
		{
			if (Result.bIsHit && (HitLightingForceEnabled != 0 || !Result.bIsRadianceCompleted))
			{
				bNeedsHitLightingPass = true;
			}

			if (Result.bIsHit && SampleSceneColor > 0)
			{
				float3 HitTranslatedWorldPosition = Ray.Origin + Ray.Direction * TraceResult.OpaqueHitDistance;
				if (SampleSceneColorAtHit(HitTranslatedWorldPosition, TraceResult.GeometryWorldNormal, ReflectionTracingCoord.Coord, RelativeDepthThickness, TraceResult.Lighting))
				{
					bNeedsHitLightingPass = false;
				}
			}
		}
		#endif

		#if DISTANT_SCREEN_TRACES
		{
			if (!Result.bIsHit)
			{
				float3 TranslatedDistantRayOrigin = Ray.Origin;
				float MaxTraceDistanceForDistantScreenTrace = DistantScreenTraceMaxTraceDistance;

				// Start the distant linear screen traces where we have no fallback for HZB screen traces
				float2 RaySphereHit = RayIntersectSphere(TranslatedDistantRayOrigin, RayData.Direction, float4(View.TranslatedWorldCameraOrigin, DistantScreenTracesStartDistance));
				if (RaySphereHit.x < 0 && RaySphereHit.y > 0)
				{
					// Pull the origin inside the sphere slightly to reduce the gap created by jittering ray step offset
					// t + (t - 1) / (n * 2 - 1) where n is the number of distant screen trace steps
					RaySphereHit.y = max((32.f / 31.f) * RaySphereHit.y - (1.f / 31.f) * MaxTraceDistanceForDistantScreenTrace, 0);

					TranslatedDistantRayOrigin += RaySphereHit.y * RayData.Direction;
					MaxTraceDistanceForDistantScreenTrace -= RaySphereHit.y;	
				}

				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,
							Result.bIsHit,
							TraceResult.Lighting);
					}
				}
			}
		}
		#endif

		#if DIM_RADIANCE_CACHE
		{
			// If we're tracing a near field ray that missed, which has valid Radiance Cache coverage, use that for final lighting
			if (!Result.bIsHit && RayData.bUseRadianceCache)
			{
				float RandomScalarForStochasticIterpolation = BlueNoiseScalar(ReflectionTracingCoord.Coord, ReflectionsStateFrameIndex);

				float3 WorldPosition = TranslatedWorldPosition - DFHackToFloat(PrimaryView.PreViewTranslation); // LUMEN_LWC_TODO
				FRadianceCacheCoverage Coverage = GetRadianceCacheCoverage(WorldPosition, RayData.Direction, InterleavedGradientNoise(ReflectionTracingCoord.Coord, ReflectionsStateFrameIndexMod8));

				SampleRadianceCacheAndApply(Coverage, WorldPosition, RayData.Direction, RayData.ConeHalfAngle, RandomScalarForStochasticIterpolation, TraceResult.Lighting, TraceResult.Transparency);
				Result.bIsHit = true;
				bNeedsHitLightingPass = false;
			}
		}
		#endif

		if ((ApplySkyLight != 0) && !Result.bIsHit)
		{
			ApplySkylightToTraceResult(RayData.Direction, TraceResult);
		}

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

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

		if (SampleHeightFog > 0)
		{
			float CoverageForFog = 1.0; // There is always something of the sky fallback.
			float3 OriginToCollider = Ray.Direction * TraceResult.OpaqueHitDistance;
			TraceResult.Lighting.rgb = GetFogOnLuminance(TraceResult.Lighting.rgb, CoverageForFog, TranslatedWorldPosition, Ray.Direction, TraceResult.OpaqueHitDistance);
		}

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

		RWTraceRadiance[ReflectionTracingCoord.CoordFlatten] = MakeFinite(TraceResult.Lighting);
		RWTraceHit[ReflectionTracingCoord.CoordFlatten] = EncodeRayDistance(TraceResult.OpaqueHitDistance, Result.bIsHit);

		#if WRITE_DATA_FOR_HIT_LIGHTING_PASS
		{
			RWTraceMaterialId[ReflectionTracingCoord.CoordFlatten] = PackTraceMaterialId(bNeedsHitLightingPass, Result.MaterialShaderIndex);
			RWTraceBookmark[ReflectionTracingCoord.CoordFlatten] = Result.Bookmark.PackedData;
		}
		#endif

		#define DEBUG_SUPPORT_VISUALIZE_TRACE_COHERENCY 0
		#if DEBUG_SUPPORT_VISUALIZE_TRACE_COHERENCY
		{
			if (VisualizeTraceCoherency != 0)
			{
				uint DebugGroupIndex = 10240;
				// UE_RAY_TRACING_DISPATCH_1D
				int DebugTraceIndex = (int)RayIndex - (int)DebugGroupIndex * 32;

				if (DebugTraceIndex >= 0 && DebugTraceIndex < 32)
				{
					float3 WorldPosition = TranslatedWorldPosition - DFHackToFloat(PrimaryView.PreViewTranslation); // LUMEN_LWC_TODO
					WriteTraceForVisualization(DebugTraceIndex, WorldPosition, RayData.Direction, TraceResult.OpaqueHitDistance, / *TraceResult.Lighting* /float3(1, 0, 0));
				}
			}
		}
		#endif
	}
}

#endif // LUMEN_HARDWARE_RAYTRACING