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

// When loading SSS checkerboard pixel, do not adjust DiffuseColor/SpecularColor to preserve specular and diffuse lighting values for each pixel
#define ALLOW_SSS_MATERIAL_OVERRIDE 0

#include "../Common.ush"
#include "../BlueNoise.ush"
#include "../HeightFogCommon.ush"
#include "MegaLights.ush"
#include "MegaLightsVolume.ush"
#include "MegaLightsRayTracing.ush"
#include "MegaLightsSampling.ush"

uint VolumeDebugMode;
float3 VolumeFrameJitterOffset;
uint3 NumSamplesPerVoxel;
float VolumeMinSampleWeight;
uint3 DownsampledVolumeViewSize;
float VolumeInverseSquaredLightDistanceBiasScale;
uint DebugLightId;
float VolumeGuideByHistoryHiddenRatio;

uint3 VolumeVisibleLightHashTileSize;
uint3 HistoryVolumeVisibleLightHashViewSizeInTiles;
StructuredBuffer<uint> VolumeVisibleLightHashHistory;

RWTexture3D<uint> RWVolumeLightSamples;

uint3 GetSampleCoord(uint3 DownsampledVolumeCoord, uint LightSampleIndex)
{
	return DownsampledVolumeCoord * NumSamplesPerVoxel + uint3(LightSampleIndex % NUM_SAMPLES_PER_VOXEL_3D_X, LightSampleIndex / NUM_SAMPLES_PER_VOXEL_3D_X, 0);
}

void SampleLight(
	float3 TranslatedWorldPosition,
	float3 CameraVector,
	float DistanceBiasSqr,
	float LightVolumetricSoftFadeDistance,
	uint VisibleLightHash[VISIBLE_LIGHT_HASH_SIZE],
	bool bHasValidHistory, 
	uint ForwardLightIndex,
	uint LightSceneId,
	FDeferredLightData LightData,
	float VolumetricScatteringIntensity,
	uint PrevForwardLightIndex,
	inout FLightSampler VisibleLightSampler,
	inout FLightSampler HiddenLightSampler,
    inout FShaderPrintContext DebugContext)
{
	float3 L;
	float3 LightScattering = GetMegaLightsVolumeLighting(TranslatedWorldPosition, CameraVector, DistanceBiasSqr, LightVolumetricSoftFadeDistance, LightData, VolumetricScatteringIntensity, L) * View.PreExposure;
	float SampleWeight = log2(Luminance(LightScattering) + 1.0f);

	bool bWasVisibleInLastFrame = true;

	#if GUIDE_BY_HISTORY
	if (SampleWeight > VolumeMinSampleWeight && PrevForwardLightIndex >= 0)
	{
		bWasVisibleInLastFrame = GetLightVisibility(VisibleLightHash, PrevForwardLightIndex);
	}
	#endif
				
	if (DebugContext.bIsActive)
	{
		Newline(DebugContext);
		Print(DebugContext, LightSceneId, Select(LightSceneId == DebugLightId, FontSelected, FontValue));
		Print(DebugContext, ForwardLightIndex, Select(LightSceneId == DebugLightId, FontSelected, FontValue));
		Print(DebugContext, SampleWeight, Select(SampleWeight > VolumeMinSampleWeight, FontWhite, FontGrey));
		Print(DebugContext, bWasVisibleInLastFrame ? 1u : 0u, Select(bWasVisibleInLastFrame, FontWhite, FontGrey));
	}

	if (SampleWeight > VolumeMinSampleWeight)
	{
		if (bWasVisibleInLastFrame)
		{
			AddLightSample(VisibleLightSampler, SampleWeight, ForwardLightIndex, bWasVisibleInLastFrame);
		}
		else
		{
			AddLightSample(HiddenLightSampler, SampleWeight, ForwardLightIndex, bWasVisibleInLastFrame);
		}
	}
}

void SampleLight(
	float3 TranslatedWorldPosition,
	float3 CameraVector,
	float DistanceBiasSqr,
	float LightVolumetricSoftFadeDistance,
	uint VisibleLightHash[VISIBLE_LIGHT_HASH_SIZE],
	bool bHasValidHistory, 
	uint LocalLightIndex,
	inout FLightSampler VisibleLightSampler,
	inout FLightSampler HiddenLightSampler,
	inout FShaderPrintContext DebugContext)
{
	const FLocalLightData LocalLightData = GetLocalLightData(LocalLightIndex, 0);	
	checkSlow(UnpackIsHandledByMegaLights(LocalLightData));

	const FDeferredLightData LightData = ConvertToDeferredLight(LocalLightData);

	const float VolumetricScatteringIntensity = UnpackVolumetricScatteringIntensity(LocalLightData.Internal);

	SampleLight(
		TranslatedWorldPosition, CameraVector, DistanceBiasSqr, LightVolumetricSoftFadeDistance, 
		VisibleLightHash, bHasValidHistory, LocalLightIndex, LocalLightData.Internal.LightSceneId, LightData, VolumetricScatteringIntensity, LocalLightData.Internal.PrevLocalLightIndex,
		VisibleLightSampler,
		HiddenLightSampler,
		DebugContext);
}

void SampleDirectionalLight(
	uint DirectionalLightIndex,
	float3 TranslatedWorldPosition,
	float3 CameraVector,
	float DistanceBiasSqr,
	float LightVolumetricSoftFadeDistance,
	uint VisibleLightHash[VISIBLE_LIGHT_HASH_SIZE],
	bool bHasValidHistory, 
	inout FLightSampler VisibleLightSampler,
	inout FLightSampler HiddenLightSampler,
	inout FShaderPrintContext DebugContext)
{
	uint LightIndex = ForwardLightStruct.DirectionalLightIndices[DirectionalLightIndex];
	FForwardLightData ForwardLightData = GetForwardLightData(LightIndex, 0);
	checkSlow(UnpackIsHandledByMegaLights(ForwardLightData));

	FDeferredLightData LightData = ConvertToDeferredLight(ForwardLightData);
	LightData.bRadialLight = false;
	LightData.bInverseSquared = false;

	const float VolumetricScatteringIntensity = UnpackVolumetricScatteringIntensity(ForwardLightData);

	SampleLight(
		TranslatedWorldPosition, CameraVector, DistanceBiasSqr, LightVolumetricSoftFadeDistance, 
		VisibleLightHash, bHasValidHistory, LightIndex, ForwardLightData.LightSceneId, LightData, VolumetricScatteringIntensity, ForwardLightData.PrevLocalLightIndex,
		VisibleLightSampler,
		HiddenLightSampler,
		DebugContext);
}

/**
 * Run one thread per sample and generate new light samples for tracing
 */
[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, THREADGROUP_SIZE)]
void VolumeGenerateLightSamplesCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	const uint3 DownsampledVolumeCoord = DispatchThreadId.xyz;
	if (all(DownsampledVolumeCoord < DownsampledVolumeViewSize))
	{
		FShaderPrintContext DebugContext = InitVolumeDebugContext(DownsampledVolumeCoord, /*bDownsampled*/ true, float2(0.05, 0.05));

		const uint3 VolumeCoord = DownsampledVolumeCoordToVolumeCoord(DownsampledVolumeCoord);
#if !TRANSLUCENCY_LIGHTING_VOLUME
		// #ml_todo: FroxelFootprintMargin should be calculated based on DownsampleFactor
		const float FroxelFootprintMargin = 3.0f; // Sampling is done at half res, so we need a 3 froxel margin to cover all neighbors during the shading pass
		if (IsFroxelVisible(VolumeCoord, FroxelFootprintMargin))
#endif // !TRANSLUCENCY_LIGHTING_VOLUME
		{
			float SceneDepth;
			const float3 TranslatedWorldPosition = ComputeCellTranslatedWorldPosition(VolumeCoord, VolumeFrameJitterOffset, SceneDepth);
			const float3 CameraVector = normalize(TranslatedWorldPosition - View.TranslatedWorldCameraOrigin);

#if TRANSLUCENCY_LIGHTING_VOLUME
			float DistanceBiasSqr = 0.0f;
			float LightVolumetricSoftFadeDistance = 0.0f;
#else
			float SceneDepth2 = 0.0f;
			float SceneDepth3 = 0.0f;
			float CellRadius = length(TranslatedWorldPosition - ComputeCellTranslatedWorldPosition(VolumeCoord + uint3(1, 1, 1), VolumeFrameJitterOffset, SceneDepth2));
			float Cell2DRadius = length(TranslatedWorldPosition - ComputeCellTranslatedWorldPosition(VolumeCoord + uint3(1, 1, 0), VolumeFrameJitterOffset, SceneDepth3));
			float LightVolumetricSoftFadeDistance = LightSoftFading * Cell2DRadius;
			// Bias the inverse squared light falloff based on voxel size to prevent aliasing near the light source
			float DistanceBiasSqr = max(CellRadius * VolumeInverseSquaredLightDistanceBiasScale, 1);
			DistanceBiasSqr *= DistanceBiasSqr;
#endif

			bool bHasValidHistory = true;
			uint3 PrevDownsampledVolumeCoord = DownsampledVolumeCoord;

#define REPROJECT_HISTORY_FOR_GUIDING (!TRANSLUCENCY_LIGHTING_VOLUME)
#if GUIDE_BY_HISTORY && REPROJECT_HISTORY_FOR_GUIDING
				{
					bHasValidHistory = false;
					float3 HistoryUV = ComputeHistoryVolumeUVFromTranslatedPos(TranslatedWorldPosition, UnjitteredPrevTranslatedWorldToClip);

					bool bHistoryWasOnScreen = all(HistoryUV >= 0.0f) && all(HistoryUV < float3(View.VolumetricFogPrevUVMaxForTemporalBlend, 1.0f));

					HistoryUV = clamp(HistoryUV, 0.0f, float3(View.VolumetricFogPrevUVMaxForTemporalBlend, 1.0f));

					uint3 HistoryDownsampledVolumeCoord = floor(HistoryUV * DownsampledVolumeViewSize - 0.5f);

					if (bHistoryWasOnScreen)
					{
						bHasValidHistory = true;
					}

					PrevDownsampledVolumeCoord = HistoryDownsampledVolumeCoord;
				}
#endif // GUIDE_BY_HISTORY && REPROJECT_HISTORY_FOR_GUIDING

			if (DebugContext.bIsActive)
			{
				Print(DebugContext, TEXT("VolumeGenerateLightSamples"), FontTitle);
				Newline(DebugContext);
				Print(DebugContext, TEXT("Coord             : "));
				Print(DebugContext, DownsampledVolumeCoord, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("TWS               : "));
				Print(DebugContext, TranslatedWorldPosition, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("DebugLightId      : "));
				Print(DebugContext, DebugLightId, FontValue);

				Newline(DebugContext);
				Print(DebugContext, TEXT("NoiseCoord        : "));
				Print(DebugContext, VolumeCoordToNoiseCoord(DownsampledVolumeCoord), FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("ValidGuideHistory : "));
				Print(DebugContext, bHasValidHistory, FontValue);

				AddCrossTWS(DebugContext, TranslatedWorldPosition, 5.0f, float4(1, 0, 0, 1));

				Newline(DebugContext);
				Print(DebugContext, TEXT("LightId | LocalLightId | Weight  | History"), FontSilver);
			}

			uint VisibleLightHash[VISIBLE_LIGHT_HASH_SIZE];
			for (uint IndexInHash = 0; IndexInHash < VISIBLE_LIGHT_HASH_SIZE; ++IndexInHash)
			{
				VisibleLightHash[IndexInHash] = 0xFFFFFFFF;

				#if GUIDE_BY_HISTORY
				{
					const uint3 PrevVisibilityTileCoord = clamp(PrevDownsampledVolumeCoord / VolumeVisibleLightHashTileSize, 0u, HistoryVolumeVisibleLightHashViewSizeInTiles - 1);
					const uint HistoryBufferBase = VISIBLE_LIGHT_HASH_SIZE * ((PrevVisibilityTileCoord.z * HistoryVolumeVisibleLightHashViewSizeInTiles.y + PrevVisibilityTileCoord.y) * HistoryVolumeVisibleLightHashViewSizeInTiles.x + PrevVisibilityTileCoord.x);
					VisibleLightHash[IndexInHash] = VolumeVisibleLightHashHistory[HistoryBufferBase + IndexInHash];
				}
				#endif
			}

			uint3 CellIndex = DownsampledVolumeCoord % 2;
			uint LinearIndex = CellIndex.x + CellIndex.y * 2 + CellIndex.z * 4;
			LinearIndex = (LinearIndex + MegaLightsStateFrameIndex) % 8;

			const float RandomScalar = BlueNoiseScalar(VolumeCoordToNoiseCoord(DownsampledVolumeCoord), MegaLightsStateFrameIndex);
			FLightSampler LightSampler = InitLightSampler(RandomScalar);
			FLightSampler HiddenLightSampler = InitLightSampler(1.0f - RandomScalar);

			const uint EyeIndex = 0;
#if TRANSLUCENCY_LIGHTING_VOLUME
			const float4 FroxelClipSpacePosition = mul(float4(TranslatedWorldPosition, 1), View.TranslatedWorldToClip);
			const float2 SvPosition = (FroxelClipSpacePosition.xy / FroxelClipSpacePosition.w * float2(.5f, -.5f) + .5f) * View.ViewSizeAndInvSize.xy;

			const uint GridIndex = ComputeLightGridCellIndex((uint2)(SvPosition * View.LightProbeSizeRatioAndInvSizeRatio.zw - View.ViewRectMin.xy), FroxelClipSpacePosition.w, EyeIndex);
#else
			const uint GridIndex = ComputeLightGridCellIndex(VolumeCoord.xy * MegaLightsVolumePixelSize, SceneDepth, EyeIndex);
#endif
			const FCulledLightsGridHeader CulledLightsGridHeader = GetCulledLightsGridHeader(GridIndex);
			const uint NumLightsInGridCell = min(CulledLightsGridHeader.NumMegaLights, GetMaxLightsPerCell());

			const uint ScalarGridIndex = WaveReadLaneFirst(GridIndex);
			const bool bScalarGridCell = WaveActiveAllTrue(ScalarGridIndex == GridIndex);

			if (bScalarGridCell)
			{
				const FCulledLightsGridHeader CulledLightsGridHeader = GetCulledLightsGridHeader(ScalarGridIndex);
				const uint NumLightsInGridCell = min(CulledLightsGridHeader.NumMegaLights, GetMaxLightsPerCell());

				uint GridLightIndex = 0;
				while(GridLightIndex < NumLightsInGridCell)
				{
					uint LocalLightIndex = GetCulledLightDataGrid(CulledLightsGridHeader.MegaLightsDataStartIndex + GridLightIndex);
					if (LocalLightIndex >= MAX_LOCAL_LIGHT_INDEX)
					{
						break;
					}
						
					++GridLightIndex;
					SampleLight(TranslatedWorldPosition, CameraVector, DistanceBiasSqr, LightVolumetricSoftFadeDistance, VisibleLightHash, bHasValidHistory, LocalLightIndex, LightSampler, HiddenLightSampler, DebugContext);
				}
			}
			else
			{
				uint GridLightIndex = 0;
				while(GridLightIndex < NumLightsInGridCell)
				{
					const uint VectorLocalLightIndex = GetCulledLightDataGrid(CulledLightsGridHeader.MegaLightsDataStartIndex + GridLightIndex);
					if (VectorLocalLightIndex >= MAX_LOCAL_LIGHT_INDEX)
					{
						break;
					}

					uint LocalLightIndex = WaveActiveMin(VectorLocalLightIndex);
					if (LocalLightIndex == VectorLocalLightIndex)
					{
						++GridLightIndex;
						SampleLight(TranslatedWorldPosition, CameraVector, DistanceBiasSqr, LightVolumetricSoftFadeDistance, VisibleLightHash, bHasValidHistory, LocalLightIndex, LightSampler, HiddenLightSampler, DebugContext);
					}
				}
			}

			// sample directional lights
			for (uint Index = ForwardLightStruct.DirectionalMegaLightsSupportedStartIndex; Index < ForwardLightStruct.NumDirectionalLights; ++Index)
			{
				SampleDirectionalLight(Index, TranslatedWorldPosition, CameraVector, DistanceBiasSqr, LightVolumetricSoftFadeDistance, VisibleLightHash, bHasValidHistory, LightSampler, HiddenLightSampler, DebugContext);
			}

			if (DebugContext.bIsActive)
			{
				Newline(DebugContext);
				Print(DebugContext, TEXT("Visible Weight sum : "));
				Print(DebugContext, LightSampler.WeightSum, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("Hidden Weight sum  : "));
				Print(DebugContext, HiddenLightSampler.WeightSum, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("Selected	: "));
				Newline(DebugContext);
				Print(DebugContext, TEXT("LightId    | Weight"), FontSilver);

				for (uint LightSampleIndex = 0; LightSampleIndex < NUM_SAMPLES_PER_VOXEL_1D; ++LightSampleIndex)
				{
					FCandidateLightSample LightSample = UnpackCandidateLightSample(LightSampler.PackedSamples[LightSampleIndex]);

					if (VolumeDebugMode == DEBUG_MODE_VISUALIZE_SAMPLING)
					{
						const uint3 SampleCoord = GetSampleCoord(DownsampledVolumeCoord, LightSampleIndex);

						const FForwardLightData ForwardLightData = GetForwardLightData(LightSample.LocalLightIndex, 0);
						const FDeferredLightData LightData = ConvertToDeferredLight(ForwardLightData);

						const float2 LightSampleUV = BlueNoiseVec2(VolumeCoordToNoiseCoord(SampleCoord), MegaLightsStateFrameIndex);
						const FLightSampleTrace LightSampleTrace = GetLightSampleTrace(TranslatedWorldPosition, LightSample.LocalLightIndex, LightSampleUV);
						float4 RayColor = float4(LightData.Color.xyz / Luminance(LightData.Color.xyz), 1.0f);
						AddLineTWS(DebugContext, TranslatedWorldPosition, TranslatedWorldPosition + LightSampleTrace.Direction * LightSampleTrace.Distance, RayColor);
					}
				}
			}

			CombineLightSamplers(LightSampler, HiddenLightSampler, bHasValidHistory, VolumeGuideByHistoryHiddenRatio, DebugContext);

			// Finalize samples
			for (uint LightSampleIndex = 0; LightSampleIndex < NUM_SAMPLES_PER_VOXEL_1D; ++LightSampleIndex)
			{
				FCandidateLightSample CandidateLightSample = UnpackCandidateLightSample(LightSampler.PackedSamples[LightSampleIndex]);

				FLightSample LightSample = InitLightSample();
				LightSample.bVisible = true;
				LightSample.bGuidedAsVisible = true;
				LightSample.LocalLightIndex = CandidateLightSample.LocalLightIndex;
				LightSample.Weight = CandidateLightSample.Weight;

				if (LightSample.LocalLightIndex != MAX_LOCAL_LIGHT_INDEX)
				{
#if TRANSLUCENCY_LIGHTING_VOLUME
					// Temporarily reuse bGuidedAsVisible as bCompleted until we don't implement guiding for volumes
					LightSample.bGuidedAsVisible = false;
#else
					const FForwardLightData ForwardLightData = GetForwardLightData(LightSample.LocalLightIndex, 0);
					const bool bCastVolumetricShadow = UnpackCastVolumetricShadow(ForwardLightData);

					// Temporarily reuse bGuidedAsVisible as bCompleted since it's not currently used for volumes
					LightSample.bGuidedAsVisible = bCastVolumetricShadow ? false : true;
#endif
					LightSample.Weight = LightSampler.WeightSum / (NUM_SAMPLES_PER_VOXEL_1D * LightSample.Weight);
				}

				if (DebugContext.bIsActive)
				{
					const FForwardLightData ForwardLightData = GetForwardLightData(LightSample.LocalLightIndex, 0);

					Newline(DebugContext);
					Print(DebugContext, ForwardLightData.LightSceneId, Select(ForwardLightData.LightSceneId == DebugLightId, FontSelected, FontValue));
					Print(DebugContext, LightSample.Weight, FontValue);
				}

				RWVolumeLightSamples[GetSampleCoord(DownsampledVolumeCoord, LightSampleIndex)] = PackLightSample(LightSample);
			}
		}
#if !TRANSLUCENCY_LIGHTING_VOLUME
		else
		{
			for (uint LightSampleIndex = 0; LightSampleIndex < NUM_SAMPLES_PER_VOXEL_1D; ++LightSampleIndex)
			{
				FLightSample LightSample = InitLightSample();
				// Temporarily reuse bGuidedAsVisible as bCompleted since it's not currently used for volumes
				LightSample.bGuidedAsVisible = true;
				RWVolumeLightSamples[GetSampleCoord(DownsampledVolumeCoord, LightSampleIndex)] = PackLightSample(LightSample);
			}
		}
#endif
	}
}