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

#include "/Engine/Shared/RayTracingTypes.h"
#include "../Common.ush"
#include "../MonteCarlo.ush"
#include "../MortonCode.ush"
#include "../SceneTextureParameters.ush"
#if USE_STOCHASTIC
#include "../BlueNoise.ush"
#include "LumenSceneDirectLightingStochastic.ush"
#include "../MegaLights/MegaLightsRayTracing.ush"
#endif
#include "LumenCardCommon.ush"
#include "LumenTracingCommon.ush"
#include "LumenReflectionCommon.ush"

#ifndef LUMEN_HARDWARE_RAYTRACING
#define LUMEN_HARDWARE_RAYTRACING 0
#endif // LUMEN_HARDWARE_RAYTRACING

#ifndef LUMEN_HARDWARE_INLINE_RAYTRACING
#define LUMEN_HARDWARE_INLINE_RAYTRACING 0
#endif // LUMEN_HARDWARE_INLINE_RAYTRACING

#if LUMEN_HARDWARE_RAYTRACING || LUMEN_HARDWARE_INLINE_RAYTRACING
#include "LumenHardwareRayTracingCommon.ush"
#endif // LUMEN_HARDWARE_RAYTRACING || LUMEN_HARDWARE_INLINE_RAYTRACING

#include "LumenCardTile.ush"
#define SUPPORT_CONTACT_SHADOWS 0
#include "../DeferredLightingCommon.ush"
#include "LumenSceneDirectLighting.ush"

Buffer<uint> DispatchLightTilesIndirectArgs;
RWBuffer<uint> RWHardwareRayTracingIndirectArgs;
uint2 OutputThreadGroupSize;
uint bStochastic;

#ifdef LumenDirectLightingHardwareRayTracingIndirectArgsCS
[numthreads(1, 1, 1)]
void LumenDirectLightingHardwareRayTracingIndirectArgsCS()
{
	uint GroupOf64Traces = DispatchLightTilesIndirectArgs[0];
	if (bStochastic)
	{
		GroupOf64Traces = DivideAndRoundUp64(GroupOf64Traces);
	}

	WriteDispatchIndirectArgs(RWHardwareRayTracingIndirectArgs, 0,
		(CARD_TILE_SIZE * CARD_TILE_SIZE + OutputThreadGroupSize.x - 1) / OutputThreadGroupSize.x,
		(GroupOf64Traces + OutputThreadGroupSize.y - 1) / OutputThreadGroupSize.y,
		1);
}
#endif

#if LUMEN_HARDWARE_RAYTRACING || LUMEN_HARDWARE_INLINE_RAYTRACING

RaytracingAccelerationStructure TLAS;
RaytracingAccelerationStructure FarFieldTLAS;

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

Buffer<uint2> ShadowTraceTileData;
Buffer<int> VirtualShadowMapIds;

uint MaxTraversalIterations;
uint MeshSectionVisibilityTest;
uint ViewIndex;

uint LumenLightType;
float MaxTraceDistance;
float FarFieldMaxTraceDistance;

float HardwareRayTracingEndBias;
float HardwareRayTracingShadowRayBias;
float HeightfieldShadowReceiverBias;
float HeightfieldProjectionBiasSearchRadius;

StructuredBuffer<uint> LightTileAllocator;
StructuredBuffer<uint2> LightTiles;

bool IsRayOccluded(FLumenMinimalRayResult RayResult)
{
	return RayResult.bHit || !RayResult.bCompleted;
}

bool IsRayOccluded(FRayTracedLightingResult RayResult)
{
	return RayResult.bIsHit || !RayResult.bIsCompleted;
}

StructuredBuffer<uint> ShadowTraceAllocator;
StructuredBuffer<uint> ShadowTraces;

// Only used by the stochastic lighting
StructuredBuffer<uint> CompactedLightSampleData;
StructuredBuffer<uint> CompactedLightSampleAllocator;
Texture2D<float4> LumenSceneData;
RWTexture2DArray<uint> RWLightSamples;

FLumenLight GetLumenLightData(uint LightIndex, uint ViewIndex)
{
	const FDFVector3 PreViewTranslation = GetPreViewTranslation(ViewIndex);
	return LoadLumenLight(LightIndex, DFHackToFloat(PreViewTranslation), ViewExposure[ViewIndex]);
}

/**
* Compute shadow ray direction and distance for a given reservoir.
*/
#if USE_STOCHASTIC 
FLightSampleTrace GetLightSampleTrace(float3 TranslatedWorldPosition, uint LocalLightIndex, uint2 SampleCoord, uint ViewIndex)
{
	const float2 LightSampleRandom = BlueNoiseVec2(SampleCoord, MegaLightsStateFrameIndex);

	FLightSampleTrace LightSampleTrace;
	LightSampleTrace.Direction = 0.0f;
	LightSampleTrace.Distance = 0.0f;

	if (LocalLightIndex != MAX_LOCAL_LIGHT_INDEX)
	{
		const FDeferredLightData LightData = GetLumenLightData(LocalLightIndex, ViewIndex).DeferredLightData;
		const FLightShaderParameters LightParameters = ConvertToLightShaderParameters(LightData);

		float3 Unused;
		float Unused2;

		bool bValid = GenerateShadowRay(
			LightParameters,
			LightData.bSpotLight,
			LightData.bRectLight,
			!LightData.bRadialLight,
			TranslatedWorldPosition,
			float3(0, 0, 0),
			LightSampleRandom,
			Unused,
			LightSampleTrace.Direction,
			Unused2,
			LightSampleTrace.Distance);
	}

	return LightSampleTrace;
}
#endif

LUMEN_HARDWARE_RAY_TRACING_ENTRY(LumenSceneDirectLightingHardwareRayTracing)
{
	const uint ThreadIndex = DispatchThreadIndex.x;
	const uint GroupIndex = DispatchThreadIndex.y;
	const uint ShadowTraceIndex = GroupIndex * 64 + ThreadIndex;

	#if USE_STOCHASTIC 
	if (ShadowTraceIndex < CompactedLightSampleAllocator[0])
	{
		const FLumenSampleCoord SampleCoord = UnpackLumenSampleCoord(CompactedLightSampleData[ShadowTraceIndex]);
		const FLumenSampleSceneData SceneData = UnpackLumenSampleSceneData(LumenSceneData[SampleCoord.Coord]);
		FLightSample LightSample = UnpackLightSample(RWLightSamples[uint3(SampleCoord.Coord, SampleCoord.LayerIndex)]);
		const FLumenLight LumenLight = GetLumenLightData(LightSample.LocalLightIndex, SceneData.ViewIndex);

		// Trace visibility ray
		{
			float3 L = LumenLight.DeferredLightData.Direction;
			float3 ToLight = L;
			float NearFieldTMax = MaxTraceDistance;
			float FarFieldTMax = FarFieldMaxTraceDistance;

			if (LumenLight.Type != LIGHT_TYPE_DIRECTIONAL)
			{
				ToLight = LumenLight.DeferredLightData.TranslatedWorldPosition - SceneData.TranslatedWorldPosition;
				float LengthToLight = max(length(ToLight) - HardwareRayTracingEndBias, 0.0f);
				NearFieldTMax = min(NearFieldTMax, LengthToLight);
				FarFieldTMax = min(FarFieldTMax, LengthToLight);
				L = normalize(ToLight);
			}

			FRayDesc Ray;
			const float2 RandSample = 0.5;

			float ReceiverBias = 0.0f;
			#if !ENABLE_HEIGHTFIELD_PROJECTION_BIAS
			if (SceneData.bHeightfield)
			{
				ReceiverBias = CalculateDistanceBasedHeightfieldBias(HeightfieldShadowReceiverBias, SceneData.TranslatedWorldPosition, PrimaryView.TranslatedWorldCameraOrigin);
			}
			#endif

			Ray.Origin = GetCardWorldPositionForShadowing(SceneData.TranslatedWorldPosition, L, SceneData.WorldNormal, HardwareRayTracingShadowRayBias + ReceiverBias);
			Ray.Direction = L;
			Ray.TMin = 0;
			Ray.TMax = NearFieldTMax;

			FRayCone RayCone = (FRayCone)0;

			FRayTracedLightingContext Context = CreateRayTracedLightingContext(
				RayCone,
				0, //TODO - CoordInCardTile,
				0, //TODO - CoordInCardTile.x, // dummy coordinate
				/*CullingMode*/ FORCE_TWO_SIDED ? RAY_FLAG_NONE : RAY_FLAG_CULL_FRONT_FACING_TRIANGLES,
				MaxTraversalIterations,
				MeshSectionVisibilityTest != 0);

			// Shadows don't need closest hit distance
			Context.bAcceptFirstHitAndEndSearch = true;

			#if LUMEN_HARDWARE_INLINE_RAYTRACING
			{
				Context.HitGroupData = HitGroupData;
				Context.RayTracingSceneMetadata = RayTracingSceneMetadata;
				Context.RWInstanceHitCountBuffer = RWInstanceHitCountBuffer;
			}
			#endif

			bool bRayOccluded = false;

			#if ENABLE_FAR_FIELD_TRACING
			{
				FRayDesc FarFieldRay = Ray;
				FarFieldRay.TMax = FarFieldTMax;
						
				FRayTracedLightingResult RayResult = TraceSurfaceCacheFarFieldRay(FarFieldTLAS, FarFieldRay, Context);
				bRayOccluded = IsRayOccluded(RayResult);
			}
			#endif

			if (!bRayOccluded)
			{
				Context.InstanceMask = RAY_TRACING_MASK_OPAQUE_SHADOW;
				Context.bIsShadowRay = true;

				FRayTracedLightingResult RayResult = TraceSurfaceCacheRay(TLAS, Ray, Context);
				bRayOccluded = IsRayOccluded(RayResult);
			}

			if (bRayOccluded)
			{
				LightSample.bCompleted = true;
				LightSample.bVisible = false;
				RWLightSamples[uint3(SampleCoord.Coord, SampleCoord.LayerIndex)] = PackLightSample(LightSample);
			}
		}
	}
	#else
	if (ShadowTraceIndex < ShadowTraceAllocator[0])
	{
		FShadowTrace ShadowTrace = UnpackShadowTrace(ShadowTraces[ShadowTraceIndex]);
		uint2 CoordInCardTile = ShadowTrace.LightTileCoord;	
		const FLightTileForShadowMaskPass LightTile = UnpackLightTileForShadowMaskPass(LightTiles[ShadowTrace.LightTileIndex]);

		if (LightTile.ViewIndex == ViewIndex)
		{
			uint2 TexelInCardPageCoord = LightTile.TileCoord * CARD_TILE_SIZE + CoordInCardTile;

			const FLumenLight LumenLight = LoadLumenLight(LightTile.LightIndex, DFHackToFloat(PrimaryView.PreViewTranslation), PrimaryView.PreExposure);

			if (all(CoordInCardTile < CARD_TILE_SIZE))
			{
				FShadowMaskRay ShadowMaskRay;
				ReadShadowMaskRayRW(ShadowTrace.LightTileIndex, CoordInCardTile, ShadowMaskRay);

				// Trace visibility ray
				if (!ShadowMaskRay.bShadowFactorComplete)
				{
					FLumenCardPageData CardPage = GetLumenCardPageData(LightTile.CardPageIndex);
					FLumenCardData Card = GetLumenCardData(CardPage.CardIndex);
					float2 AtlasUV = CardPage.PhysicalAtlasUVRect.xy + CardPage.PhysicalAtlasUVTexelScale * (TexelInCardPageCoord + 0.5f);
					float2 CardUV = CardPage.CardUVRect.xy + CardPage.CardUVTexelScale * (TexelInCardPageCoord + 0.5f);

					FLumenSurfaceCacheData SurfaceCacheData = GetSurfaceCacheData(Card, CardUV, AtlasUV);

					float3 WorldPosition = SurfaceCacheData.WorldPosition;
					float3 WorldNormal = SurfaceCacheData.WorldNormal;
					float3 TranslatedWorldPosition = WorldPosition + DFHackToFloat(PrimaryView.PreViewTranslation); // LUMEN_LWC_TODO

					float3 L = LumenLight.DeferredLightData.Direction;
					float3 ToLight = L;
					float NearFieldTMax = MaxTraceDistance;
					float FarFieldTMax = FarFieldMaxTraceDistance;

					if (LumenLight.Type != LIGHT_TYPE_DIRECTIONAL)
					{
						ToLight = LumenLight.DeferredLightData.TranslatedWorldPosition - TranslatedWorldPosition;
						float LengthToLight = max(length(ToLight) - HardwareRayTracingEndBias, 0.0f);
						NearFieldTMax = min(NearFieldTMax, LengthToLight);
						FarFieldTMax = min(FarFieldTMax, LengthToLight);
						L = normalize(ToLight);
					}

					FRayDesc Ray;
					const float2 RandSample = 0.5;

					float ReceiverBias = 0.0f;
					#if !ENABLE_HEIGHTFIELD_PROJECTION_BIAS
					if (Card.bHeightfield)
					{
						ReceiverBias = CalculateDistanceBasedHeightfieldBias(HeightfieldShadowReceiverBias, TranslatedWorldPosition, PrimaryView.TranslatedWorldCameraOrigin);
					}
					#endif

					Ray.Origin = GetCardWorldPositionForShadowing(TranslatedWorldPosition, L, WorldNormal, HardwareRayTracingShadowRayBias + ReceiverBias);
					Ray.Direction = L;
					Ray.TMin = 0;
					Ray.TMax = NearFieldTMax;

					FRayCone RayCone = (FRayCone)0;

					FRayTracedLightingContext Context = CreateRayTracedLightingContext(
						RayCone,
						CoordInCardTile,
						CoordInCardTile.x, // dummy coordinate
						/*CullingMode*/ FORCE_TWO_SIDED ? RAY_FLAG_NONE : RAY_FLAG_CULL_FRONT_FACING_TRIANGLES,
						MaxTraversalIterations,
						MeshSectionVisibilityTest != 0);

					// Shadows don't need closest hit distance
					Context.bAcceptFirstHitAndEndSearch = true;

					#if LUMEN_HARDWARE_INLINE_RAYTRACING
					{
						Context.HitGroupData = HitGroupData;
						Context.RayTracingSceneMetadata = RayTracingSceneMetadata;
						Context.RWInstanceHitCountBuffer = RWInstanceHitCountBuffer;
					}
					#endif

					bool bRayOccluded = false;

					#if ENABLE_FAR_FIELD_TRACING
					{
						FRayDesc FarFieldRay = Ray;
						FarFieldRay.TMax = FarFieldTMax;
						
						FRayTracedLightingResult RayResult = TraceSurfaceCacheFarFieldRay(FarFieldTLAS, FarFieldRay, Context);
						bRayOccluded = IsRayOccluded(RayResult);
					}
					#endif

					// Find the heightfield intersection that corresponds to the given card position.
					#if ENABLE_HEIGHTFIELD_PROJECTION_BIAS
					if (Card.bHeightfield && !bRayOccluded)
					{
						float SearchRadius = HeightfieldProjectionBiasSearchRadius;
						float3 SearchDirection = float3(0.0, 0.0, 1.0);

						FRayDesc ProjectedRay;
						ProjectedRay.Origin = Ray.Origin - SearchDirection * SearchRadius;
						ProjectedRay.Direction = SearchDirection;
						ProjectedRay.TMin = 0.0f;
						ProjectedRay.TMax = 2.0f * SearchRadius;
						Context.CullingMode = RAY_FLAG_CULL_FRONT_FACING_TRIANGLES;

						FLumenMinimalRayResult SearchResult = TraceLumenMinimalRay(TLAS, ProjectedRay, Context);

						if (IsRayOccluded(SearchResult))
						{
							float Epsilon = 0.01;
							Ray.Origin = ProjectedRay.Origin + ProjectedRay.Direction * SearchResult.HitT + SearchResult.HitNormal * Epsilon;
						}
					}
					#endif

					if (!bRayOccluded)
					{						
						Context.InstanceMask = RAY_TRACING_MASK_OPAQUE_SHADOW;
						Context.bIsShadowRay = true;

						FRayTracedLightingResult RayResult = TraceSurfaceCacheRay(TLAS, Ray, Context);
						bRayOccluded = IsRayOccluded(RayResult);
					}

					ShadowMaskRay.ShadowFactor *= bRayOccluded ? 0.0f : 1.0f;
				}

				ShadowMaskRay.bShadowFactorComplete = true;
				ShadowTrace = UnpackShadowTrace(ShadowTraces[ShadowTraceIndex]);

				WriteShadowMaskRay(ShadowMaskRay, ShadowTrace.LightTileIndex, CoordInCardTile, true);

				if (ShadowTrace.DownSampleLevel > 0)
				{
					const uint2 QuadBaseCoord = CoordInCardTile & ~0x1;
					const uint2 SubQuadCoord = CoordInCardTile - QuadBaseCoord;
					const uint SubQuadIndex = SubQuadCoord.y * 2 + SubQuadCoord.x;
					const bool bRayOccluded = all(ShadowMaskRay.ShadowFactor == 0.0f);
					FShadowMaskRay NeighborShadowMaskRays[3];

					for (uint Index = 0, NeighborIndex = 0; Index < 4; ++Index)
					{
						if (Index != SubQuadIndex)
						{
							const uint2 Offset = uint2(Index % 2, Index / 2);
							ReadShadowMaskRayRW(ShadowTrace.LightTileIndex, QuadBaseCoord + Offset, NeighborShadowMaskRays[NeighborIndex]);
							++NeighborIndex;
						}
					}
					
					for (uint Index = 0, NeighborIndex = 0; Index < 4; ++Index)
					{
						if (Index != SubQuadIndex)
						{
							const uint2 Offset = uint2(Index % 2, Index / 2);
							FShadowMaskRay NeighborShadowMaskRay = NeighborShadowMaskRays[NeighborIndex];
							NeighborShadowMaskRay.ShadowFactor *= bRayOccluded ? 0.0f : 1.0f;
							NeighborShadowMaskRay.bShadowFactorComplete = true;
							WriteShadowMaskRay(NeighborShadowMaskRay, ShadowTrace.LightTileIndex, QuadBaseCoord + Offset, true);
							++NeighborIndex;
						}
					}
				}
			}
		}
	}
	#endif
}

#endif // LUMEN_HARDWARE_RAYTRACING