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

#include "/Engine/Shared/RayTracingTypes.h"
#include "../Common.ush"
#include "../MonteCarlo.ush"
#include "../SceneTextureParameters.ush"
#include "LumenCardCommon.ush"
#include "LumenTracingCommon.ush"
#define RADIANCE_CACHE_DEPTH_TEST_SPHERE_PARALLAX 1
#include "LumenRadianceCacheCommon.ush"
#include "LumenScreenProbeCommon.ush"
#include "LumenScreenProbeTracingCommon.ush"
#include "LumenFloatQuantization.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

RWBuffer<uint> RWHardwareRayTracingIndirectArgs;
uint2 OutputThreadGroupSize;

#ifdef FLumenScreenProbeHardwareRayTracingIndirectArgsCS
[numthreads(1, 1, 1)]
void FLumenScreenProbeHardwareRayTracingIndirectArgsCS()
{
	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;
float MaxTraceDistance;
float NearFieldSceneRadius;
float NearFieldMaxTraceDistance;
float FarFieldBias;
float FarFieldMaxTraceDistance;
float PullbackBias;
float NormalBias;
float BiasForTracesFromHairPixels;
float AvoidSelfIntersectionTraceDistance;
float SkipFirstTwoSidedHitDistance;
float MinTraceDistanceToSampleSurfaceCache;
uint MaxTraversalIterations;
uint MeshSectionVisibilityTest;

LUMEN_HARDWARE_RAY_TRACING_ENTRY(LumenScreenProbeGatherHardwareRayTracing)
{
	uint ThreadIndex = DispatchThreadIndex.x;
	uint GroupIndex = DispatchThreadIndex.y;
	uint CompactedTraceIndex = GroupIndex * THREADGROUP_SIZE_1D + ThreadIndex;

	if (CompactedTraceIndex < CompactedTraceTexelAllocator[0])
	{
		uint ScreenProbeIndex;
		uint2 TraceTexelCoord;
		DecodeScreenProbeTraceTexel(CompactedTraceTexelData[CompactedTraceIndex], ScreenProbeIndex, TraceTexelCoord);

		uint2 ScreenProbeScreenPosition = GetScreenProbeScreenPosition(ScreenProbeIndex);
		uint2 ScreenTileCoord = GetScreenTileCoord(ScreenProbeScreenPosition);
		uint LinearCoord = ScreenTileCoord.y * ScreenProbeViewSize.x + ScreenTileCoord.x;

		float2 ScreenUV = GetScreenUVFromScreenProbePosition(ScreenProbeScreenPosition);
		uint2 ScreenProbeAtlasCoord = uint2(ScreenProbeIndex % ScreenProbeAtlasViewSize.x, ScreenProbeIndex / ScreenProbeAtlasViewSize.x);
		uint2 ScreenProbeTraceCoord = GetTraceBufferCoord(ScreenProbeAtlasCoord, TraceTexelCoord);

		float TraceHitDistance = DecodeProbeRayDistance(RWTraceHit[ScreenProbeTraceCoord]).HitDistance;

		float SceneDepth = GetScreenProbeDepth(ScreenProbeAtlasCoord);
		float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);

		float2 ProbeTexelCenter = GetProbeTexelCenter(ScreenTileCoord);
		const float Noise = ProbeTexelCenter.y;

		float2 ProbeUV;
		float ConeHalfAngle;
		GetProbeTracingUV(ScreenProbeTraceCoord, TraceTexelCoord, ProbeTexelCenter, 1, ProbeUV, ConeHalfAngle);

		float RayBias = 0.0f;

		if (GetScreenProbeIsHair(ScreenProbeAtlasCoord))
		{
			RayBias = BiasForTracesFromHairPixels;
		}

		#if ENABLE_FAR_FIELD_TRACING
		{
			RayBias = max(RayBias, FarFieldBias);
		}
		#endif

		float3 RayDirection = EquiAreaSphericalMapping(ProbeUV);

		float ClippedNearFieldMaxTraceDistance = ClipAndDitherNearFieldMaxTraceDistance(
			TranslatedWorldPosition,
			RayDirection,
			ScreenTileCoord,
			NearFieldSceneRadius,
			NearFieldMaxTraceDistance,
			/*NearFieldMaxTraceDistanceDitherScale*/ 0.0f);

		float TraceMaxDistance = ClippedNearFieldMaxTraceDistance;

		bool bReachedRadianceCache = false;
		#if DIM_RADIANCE_CACHE
		{
			float3 RayOriginWorldPos = TranslatedWorldPosition - DFHackToFloat(PrimaryView.PreViewTranslation);
			FRadianceCacheCoverage Coverage = GetRadianceCacheCoverage(RayOriginWorldPos, RayDirection, 0); // LUMEN_LWC_TODO
			if (Coverage.bValid)
			{
				TraceMaxDistance = min(TraceMaxDistance, Coverage.MinTraceDistanceBeforeInterpolation);
				bReachedRadianceCache = true;
			}
		}
		#endif
		#if RAY_TRACING_PASS == RAY_TRACING_PASS_FAR_FIELD
		{
			TraceMaxDistance = max(TraceMaxDistance, FarFieldMaxTraceDistance);
		}
		#endif

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

		// Apply PullbackBias 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);
		}

		float3 ProbeWorldNormalForBiasing = GetScreenProbeNormalForBiasing(ScreenProbeAtlasCoord, Ray.Direction);
		Ray.Origin += NormalBias * ProbeWorldNormalForBiasing;

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

		FRayTracedLightingContext Context = CreateRayTracedLightingContext(
			RayCone,
			ScreenTileCoord,
			ScreenTileCoord.y * ScreenProbeViewSize.x + ScreenTileCoord.x,
			/*CullingMode*/ RAY_FLAG_CULL_BACK_FACING_TRIANGLES,
			MaxTraversalIterations,
			MeshSectionVisibilityTest != 0);

		Context.MinTraceDistanceToSampleSurfaceCache = MinTraceDistanceToSampleSurfaceCache;

		#if LUMEN_HARDWARE_INLINE_RAYTRACING
		{
			Context.HitGroupData = HitGroupData;
			Context.RayTracingSceneMetadata = RayTracingSceneMetadata;
			Context.RWInstanceHitCountBuffer = RWInstanceHitCountBuffer;
		}
		#endif
		
		FRayTracedLightingResult Result = CreateRayTracedLightingResult(Ray);
		if (Ray.TMin < Ray.TMax)
		{
			#if RAY_TRACING_PASS == RAY_TRACING_PASS_HIT_LIGHTING
			{
				Context.bHitLightingDirectLighting = HitLightingDirectLighting != 0;
				Context.bHitLightingSkylight = HitLightingSkylight != 0;
				Context.bUseReflectionCaptures = false;
				Context.bForceOpaque = HitLightingForceOpaque;
				Context.HitLightingShadowMode = HitLightingShadowMode;
				Context.HitLightingShadowTranslucencyMode = HitLightingShadowTranslucencyMode;
				Context.HitLightingShadowMaxTraceDistance = NearFieldMaxTraceDistance;
			
				Result = TraceAndCalculateRayTracedLighting(TLAS, FarFieldTLAS, Ray, Context);
			}
			#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
		}

		bool bMoving = false;
		if (Result.bIsHit)
		{
			float3 HitTranslatedWorldPosition = Ray.Origin + Ray.Direction * Result.TraceHitDistance;
			float3 HitWorldVelocity = HitTranslatedWorldPosition - CalcPrevTranslatedWorldPositionFromGPUSceneInstanceIndex(HitTranslatedWorldPosition, Result.SceneInstanceIndex);
			bMoving = IsTraceMoving(TranslatedWorldPosition, SceneDepth, ScreenProbeAtlasCoord, HitTranslatedWorldPosition, HitWorldVelocity);
		}

		Result.Radiance += GetSkylightLeaking(Ray.Direction, Result.TraceHitDistance);
		Result.Radiance *= View.PreExposure;

		if (SampleHeightFog > 0)
		{
			float3 OriginToCollider = Ray.Direction * Result.TraceHitDistance;
			float CoverageForFog = 1.0; // There is always something or the sky fallback.
			Result.Radiance.rgb = GetFogOnLuminance(Result.Radiance.rgb, CoverageForFog, Ray.Origin, Ray.Direction, Result.TraceHitDistance);
		}

		RWTraceRadiance[ScreenProbeTraceCoord] = QuantizeForFloatRenderTarget(Result.Radiance, Noise);
		RWTraceHit[ScreenProbeTraceCoord] = EncodeProbeRayDistance(min(Result.TraceHitDistance, MaxTraceDistance), Result.bIsHit, bMoving, bReachedRadianceCache);
	}
}

#endif // LUMEN_HARDWARE_RAYTRACING || LUMEN_HARDWARE_INLINE_RAYTRACING