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

#include "../Common.ush"
#include "../BlueNoise.ush"
#include "../SceneTexturesCommon.ush"
#include "../LightData.ush"
#include "../VirtualShadowMaps/VirtualShadowMapProjectionCommon.ush"
#include "../VirtualShadowMaps/VirtualShadowMapProjectionSpot.ush"
#include "../VirtualShadowMaps/VirtualShadowMapProjectionDirectional.ush"
#include "../VirtualShadowMaps/VirtualShadowMapScreenRayTrace.ush"
#include "MegaLights.ush"
#include "MegaLightsRayTracing.ush"

RWTexture2D<uint> RWLightSamples;
RWTexture2D<uint> RWLightSampleRays;

Buffer<uint> CompactedTraceTexelAllocator;
Buffer<uint> CompactedTraceTexelData;
Texture2D<float> DownsampledSceneDepth;
Texture2D<UNORM float3> DownsampledSceneWorldNormal;

SCHEDULER_MIN_PRESSURE
// Avoid for now due to issues with failing compilation when occupancy not reached
//MAX_OCCUPANCY

[numthreads(THREADGROUP_SIZE, 1, 1)]
void VirtualShadowMapTraceLightSamplesCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	uint TraceTexelIndex = DispatchThreadId.x;

	if (TraceTexelIndex < CompactedTraceTexelAllocator[0])
	{
		const uint2 SampleCoord = UnpackTraceTexel(CompactedTraceTexelData[TraceTexelIndex]);
		FLightSample LightSample = UnpackLightSample(RWLightSamples[SampleCoord]);

		// Get VSM from LocalLightIndex
		const FForwardLightData ForwardLightData = GetForwardLightData(LightSample.LocalLightIndex, 0);
		const int VirtualShadowMapId = ForwardLightData.VirtualShadowMapId;

		// Should we select for this separately in the tile classification instead?
		if (VirtualShadowMapId != INDEX_NONE)
		{
			const FDeferredLightData LightData = ConvertToDeferredLight(ForwardLightData);

			uint2 ScreenCoord = SampleCoordToScreenCoord(SampleCoord);
			uint2 DownsampledScreenCoord = SampleCoordToDownsampledScreenCoord(SampleCoord);

			float2 ScreenUV = (ScreenCoord + 0.5f) * View.BufferSizeAndInvSize.zw;
			float SceneDepth = DownsampledSceneDepth[DownsampledScreenCoord];
			float3 SceneWorldNormal = normalize(DecodeNormal(DownsampledSceneWorldNormal[DownsampledScreenCoord]));
			float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);

			bool bDirectionalLight = !LightData.bRadialLight;

			const float3 ToLightDirection =
				bDirectionalLight ?
				LightData.Direction :
				normalize(LightData.TranslatedWorldPosition - TranslatedWorldPosition);

			// Currently using the embedded SMRT ray sampling; could change this to use this one easily if needed
			//FLightSampleTrace LightSampleTrace = GetLightSampleTrace(TranslatedWorldPosition, LightSample.LocalLightIndex, SampleCoord);

			float StepOffset = BlueNoiseScalar(DownsampledScreenCoord, MegaLightsStateFrameIndex);
			const float ScreenRayLengthWorld = GetScreenRayLengthMultiplierForProjectionType(VirtualShadowMap.ScreenRayLength * SceneDepth).y;

			// TODO: Avoid doing this normal bias for subsurface pixels
			// Will need to consider both substrate and downsampling
			//if (!ShadingInfo.bIsSubsurface)
			{
				TranslatedWorldPosition += SceneWorldNormal * VirtualShadowMapGetNormalBiasLength(TranslatedWorldPosition);
			}

			// VSM screen ray trace
			float RayStartOffset = ScreenRayLengthWorld;
			if (ScreenRayLengthWorld > 0.0f)
			{
				RayStartOffset = VirtualShadowMapScreenRayCast(
					TranslatedWorldPosition,
					ToLightDirection,
					ScreenRayLengthWorld,
					StepOffset);
			}

			// TODO: Back-facing tests? (or BRDF test directly?)
			// In MegaLights this is already factored into the sample selection though, so less important
#if 0
			FVirtualShadowMapSampleResult VSMResult = SampleVirtualShadowMapLocal(
				VirtualShadowMapId,
				TranslatedWorldPosition,
				RayStartOffset,
				SceneWorldNormal
			);

			if (!VSMResult.bValid || VSMResult.ShadowFactor < 0.5f)
			{
				LightSample.bVisible = false;
			}

#else
			FLightShaderParameters LightParameters = ConvertFromForward(ForwardLightData);

			FVirtualShadowMapSampleResult Result;

			// TODO: Possibly worth a permutation for register pressure reasons
			if (bDirectionalLight)
			{
				FSMRTTraceSettings Settings = GetSMRTTraceSettingsDirectional();
				// Certain settings may be less desirable with the MegaLights path, but since there
				// is typically only one directional light, it makes sense to respect the cvars
				// which can be set up as desired for this path.
				//Settings.RayCount = 1;
				//Settings.AdaptiveRayCount = 0;

				Result = TraceDirectional(
					VirtualShadowMapId,
					LightParameters,
					SampleCoord,
					SceneDepth,
					TranslatedWorldPosition,
					RayStartOffset,
					StepOffset,
					SceneWorldNormal,
					Settings);
			}
			else
			{
				FSMRTTraceSettings Settings = GetSMRTTraceSettingsLocal();
				Settings.RayCount = 1;
				Settings.AdaptiveRayCount = 0;

				Result = TraceLocalLight(
					VirtualShadowMapId,
					LightParameters,
					SampleCoord,
					SceneDepth,
					TranslatedWorldPosition,
					RayStartOffset,
					StepOffset,
					SceneWorldNormal,
					Settings
				);
			}

			if (!Result.bValid || Result.ShadowFactor < 0.5f)
			{
				LightSample.bVisible = false;
			}
#endif

			// Mark sample as complete to skip continuation
			FLightSampleRay LightSampleRay = UnpackLightSampleRay(RWLightSampleRays[SampleCoord]);
			LightSampleRay.bCompleted = true;

			RWLightSampleRays[SampleCoord] = PackLightSampleRay(LightSampleRay);
			RWLightSamples[SampleCoord] = PackLightSample(LightSample);
		}
	}
}