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

#ifndef VIRTUAL_SHADOW_MAP
	#define VIRTUAL_SHADOW_MAP 0
#endif
#ifndef DENSE_SHADOW_MAP
	#define DENSE_SHADOW_MAP 0
#endif

#include "../Common.ush"
#include "LumenCardCommon.ush"
#include "LumenCardTile.ush"
#include "LumenCardTileShadowDownsampleFactor.ush"
#define SUPPORT_CONTACT_SHADOWS 0
#include "../DeferredLightingCommon.ush"
#include "../VolumeLightingCommon.ush"
#include "../ForwardShadowingCommon.ush"
#include "../LightDataUniforms.ush"
#include "SurfaceCache/LumenSurfaceCache.ush"
#if LIGHT_FUNCTION
	#include "/Engine/Generated/Material.ush"
	#include "../LightFunctionCommon.ush"
#endif
#include "LumenSceneDirectLightingPerLightShadowCommon.ush"
#if USE_CLOUD_TRANSMITTANCE
	#include "../VolumetricCloudCommon.ush"
#endif

#define ADAPTIVE_SHADOW_TRACING_USE_WAVE_OP (FEATURE_LEVEL >= FEATURE_LEVEL_SM6 || PLATFORM_SUPPORTS_SM6_0_WAVE_OPERATIONS)

RWStructuredBuffer<uint> RWShadowTraces;
RWStructuredBuffer<uint> RWShadowTraceAllocator;

float HeightfieldShadowReceiverBias;
int bAdaptiveShadowTracing;

groupshared uint SharedNumShadowTraces;
groupshared uint SharedShadowTraces[THREADGROUP_SIZE_X * THREADGROUP_SIZE_Y];
groupshared uint SharedGlobalShadowTraceOffset;

groupshared uint bSharedAllTexelsInQuadWriteTrace[(THREADGROUP_SIZE_X / 2) * (THREADGROUP_SIZE_Y / 2)];

uint ComputeTileTraceDownsampleLevel(uint2 GroupThreadId, bool bWriteTrace, bool bTileHasUniformlyVisibility)
{
	uint DownsampleLevel = 0;

	if (bTileHasUniformlyVisibility)
	{
		const uint2 QuadId = GroupThreadId.xy / 2;
		const uint QuadIndex = QuadId.y * (THREADGROUP_SIZE_X / 2) + QuadId.x;

	#if ADAPTIVE_SHADOW_TRACING_USE_WAVE_OP
		uint Ballot = ~((bWriteTrace ? 0u : 1u) << QuadIndex);
		Ballot = WaveActiveBitAnd(Ballot);
		DownsampleLevel = (Ballot & (1u << QuadIndex)) != 0 ? 1 : 0;
	#else
		uint GroupThreadIndex = GroupThreadId.y * THREADGROUP_SIZE_X + GroupThreadId.x;
		
		if (GroupThreadIndex < (THREADGROUP_SIZE_X / 2) * (THREADGROUP_SIZE_Y / 2))
		{
			bSharedAllTexelsInQuadWriteTrace[GroupThreadIndex] = 1;
		}

		GroupMemoryBarrierWithGroupSync();

		if (!bWriteTrace)
		{
			bSharedAllTexelsInQuadWriteTrace[QuadIndex] = 0;
		}

		GroupMemoryBarrierWithGroupSync();

		DownsampleLevel = bSharedAllTexelsInQuadWriteTrace[QuadIndex] != 0 ? 1 : 0;
	#endif
	}

	return DownsampleLevel;
}

[numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, 1)]
void LumenSceneDirectLightingShadowMaskFromLightAttenuationCS( 
	uint3 GroupId : SV_GroupID,
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint3 GroupThreadId : SV_GroupThreadID)
{
	uint LinearGroupThreadId = GroupThreadId.x + GroupThreadId.y * THREADGROUP_SIZE_X;

#if THREADGROUP_SIZE_32
	uint LocalLightTileIndex = GroupId.x / 2;
	uint ShadowMaskSize = SHADOW_MASK_CARD_TILE_DWORDS / 2;
	uint ShadowMaskOffset = GroupId.x % 2 ? 1 : 0;
#else
	uint LocalLightTileIndex = GroupId.x;
	uint ShadowMaskSize = SHADOW_MASK_CARD_TILE_DWORDS;
	uint ShadowMaskOffset = 0;
#endif

	if (LinearGroupThreadId < ShadowMaskSize)
	{
		SharedShadowMask[LinearGroupThreadId] = 0;
	}

	if (LinearGroupThreadId == 0)
	{
		SharedNumShadowTraces = 0;
		SharedGlobalShadowTraceOffset = 0;
		SharedShadowTraces[0] = 0;
	}

	// Need to resolve view for light function computations
	ResolvedView = ResolveView();

	GroupMemoryBarrierWithGroupSync();

	uint SlotIndex = bUseLightTilesPerLightType != 0 && LIGHT_TYPE > 0 ? LIGHT_TYPE - 1 : NUM_BATCHABLE_LIGHT_TYPES + LightIndex;
	uint LightTileIndex = LocalLightTileIndex + LightTileOffsetsPerLight[SlotIndex * NumViews + ViewIndex];
	FLightTileForShadowMaskPass LightTile = UnpackLightTileForShadowMaskPass(LightTiles[LightTileIndex]);

	FLumenCardPageData CardPage = GetLumenCardPageData(LightTile.CardPageIndex + DummyZeroForFixingShaderCompilerBug);

	bool bTileHasUniformlyVisibility = false;
	if (bAdaptiveShadowTracing != 0)
	{
		uint2 TilePhysicalAtlasCoord = CardPage.PhysicalAtlasCoord / CARD_TILE_SIZE + LightTile.TileCoord;
		uint PhysicalAtlasWidthInTiles = LumenCardScene.PhysicalAtlasSize.x / CARD_TILE_SIZE;
		bTileHasUniformlyVisibility = LumenCardTileHasUniformVisibilityToLight(TilePhysicalAtlasCoord, PhysicalAtlasWidthInTiles, LightTile.LightIndex);
	}

	uint2 CoordInCardTile = Unflatten2D(LinearGroupThreadId, uint2(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y)) + uint2(0, ShadowMaskOffset ? 4 : 0);
	uint2 TexelInCardPageCoord = LightTile.TileCoord * CARD_TILE_SIZE + CoordInCardTile;
	float2 AtlasUV = CardPage.PhysicalAtlasUVRect.xy + CardPage.PhysicalAtlasUVTexelScale * (TexelInCardPageCoord + 0.5f);
	float2 CardUV = CardPage.CardUVRect.xy + CardPage.CardUVTexelScale * (TexelInCardPageCoord + 0.5f);

	FShadowMaskRay ShadowMaskRay;
	ShadowMaskRay.ShadowFactor = 0.0f;
	ShadowMaskRay.bShadowFactorComplete = true;

	FLumenCardData Card = GetLumenCardData(CardPage.CardIndex);
	FLumenSurfaceCacheData SurfaceCacheData = GetSurfaceCacheData(Card, CardUV, AtlasUV);

	if (SurfaceCacheData.bValid)
	{
		FDeferredLightData LightData;
		if (bUseLightTilesPerLightType != 0)
		{
			LightData = LoadLumenLight(LightTile.LightIndex, DFHackToFloat(PrimaryView.PreViewTranslation), PrimaryView.PreExposure).DeferredLightData;
		}
		else
		{
			LightData = InitDeferredLightFromUniforms(LIGHT_TYPE);
		}

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

		float3 L = LightData.Direction;
		float3 ToLight = L;
		float CombinedAttenuation = 1.0f;
		float Attenuation = 1.0f;
		float LightMask = 1.0f;

#if LIGHT_TYPE != LIGHT_TYPE_DIRECTIONAL
		if (LightData.bRadialLight)
		{
			LightMask = GetLocalLightAttenuation(TranslatedWorldPosition, LightData, ToLight, L);
		}

		if (LightData.bRectLight)
		{
			FRect Rect = GetRect(ToLight, LightData);
			Attenuation = IntegrateLight(Rect);
		}
		else
		{
			FCapsuleLight Capsule = GetCapsule(ToLight, LightData);
			Capsule.DistBiasSqr = 0;
			Attenuation = IntegrateLight(Capsule, LightData.bInverseSquared);
		}

		CombinedAttenuation *= Attenuation * LightMask;
#endif

		CombinedAttenuation *= saturate(dot(WorldNormal, L));

		if (CombinedAttenuation > 0)
		{
			ShadowMaskRay.ShadowFactor = 1.0f;
			ShadowMaskRay.bShadowFactorComplete = false;

			#if USE_LIGHT_FUNCTION_ATLAS
			if (!ShadowMaskRay.bShadowFactorComplete && LightData.LightFunctionAtlasLightIndex > 0)
			{
				ShadowMaskRay.ShadowFactor *= GetLocalLightFunctionCommon(TranslatedWorldPosition, LightData.LightFunctionAtlasLightIndex);
				ShadowMaskRay.bShadowFactorComplete = !IsVisibleToLight(ShadowMaskRay.ShadowFactor);
			}
			#elif LIGHT_FUNCTION
			if (!ShadowMaskRay.bShadowFactorComplete)
			{
				ShadowMaskRay.ShadowFactor *= GetLumenLightFunction(TranslatedWorldPosition);
				ShadowMaskRay.bShadowFactorComplete = !IsVisibleToLight(ShadowMaskRay.ShadowFactor);
			}
			#endif

			#if USE_CLOUD_TRANSMITTANCE
			if (!ShadowMaskRay.bShadowFactorComplete)
			{
				float OutOpticalDepth = 0.0f;
				ShadowMaskRay.ShadowFactor *= lerp(1.0f, GetCloudVolumetricShadow(TranslatedWorldPosition, CloudShadowmapTranslatedWorldToLightClipMatrix, CloudShadowmapFarDepthKm, CloudShadowmapTexture, CloudShadowmapSampler, OutOpticalDepth), CloudShadowmapStrength);
				ShadowMaskRay.bShadowFactorComplete = !IsVisibleToLight(ShadowMaskRay.ShadowFactor);
			}
			#endif

			if (!ShadowMaskRay.bShadowFactorComplete && LightData.IESAtlasIndex >= 0)
			{
				ShadowMaskRay.ShadowFactor *= ComputeLightProfileMultiplier(TranslatedWorldPosition, LightData.TranslatedWorldPosition, -LightData.Direction, LightData.Tangent, LightData.IESAtlasIndex);
				ShadowMaskRay.bShadowFactorComplete = !IsVisibleToLight(ShadowMaskRay.ShadowFactor);
			}

			const bool bCastShadows = LightData.ShadowedBits >= 3;
			if (!bCastShadows)
			{
				ShadowMaskRay.bShadowFactorComplete = true;
			}
		}
		else
		{
			ShadowMaskRay.ShadowFactor = 0.0f;
			ShadowMaskRay.bShadowFactorComplete = true;
		}
	}

	GroupThreadId.xy = Unflatten2D(LinearGroupThreadId, uint2(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y));

	WriteSharedShadowMaskRay(ShadowMaskRay, GroupThreadId.xy, false);

	#if COMPACT_SHADOW_TRACES
	{
		bool bWriteTrace = !ShadowMaskRay.bShadowFactorComplete;
		uint DownSampleLevel = 0;

		if (bAdaptiveShadowTracing != 0)
		{
			DownSampleLevel = ComputeTileTraceDownsampleLevel(GroupThreadId.xy, bWriteTrace, bTileHasUniformlyVisibility);

			if (DownSampleLevel > 0)
			{
				const uint2 SubQuadId = GroupThreadId.xy % 2;
				const uint SubQuadIndex = SubQuadId.y * 2 + SubQuadId.x;
				bWriteTrace = SubQuadIndex == CardPage.DirectLightingTemporalIndex % 4;
			}
		}

		if (bWriteTrace)
		{
			uint ShadowTraceIndex = 0;
			InterlockedAdd(SharedNumShadowTraces, 1, ShadowTraceIndex);

			FShadowTrace ShadowTrace;
			ShadowTrace.LightTileIndex = LightTileIndex;
			ShadowTrace.LightTileCoord = GroupThreadId.xy + uint2(0, ShadowMaskOffset ? 4 : 0);
			ShadowTrace.DownSampleLevel = DownSampleLevel;

			SharedShadowTraces[ShadowTraceIndex] = PackShadowTrace(ShadowTrace);
		}
	}
	#endif

	GroupMemoryBarrierWithGroupSync();

	#if COMPACT_SHADOW_TRACES
	{
		if (LinearGroupThreadId == 0)
		{
			InterlockedAdd(RWShadowTraceAllocator[0], SharedNumShadowTraces, SharedGlobalShadowTraceOffset);
		}
	}
	#endif

	GroupMemoryBarrierWithGroupSync();

	#if COMPACT_SHADOW_TRACES
	{
		if (LinearGroupThreadId < SharedNumShadowTraces)
		{
			RWShadowTraces[SharedGlobalShadowTraceOffset + LinearGroupThreadId] = SharedShadowTraces[LinearGroupThreadId];
		}
	}
	#endif

	if (LinearGroupThreadId < ShadowMaskSize)
	{
		RWShadowMaskTiles[SHADOW_MASK_CARD_TILE_DWORDS * LightTileIndex + ShadowMaskSize * ShadowMaskOffset + LinearGroupThreadId] = SharedShadowMask[LinearGroupThreadId];
	}
}