// 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 "../SHCommon.ush"
#include "../BlueNoise.ush"
#include "MegaLights.ush"
#include "MegaLightsVolume.ush"
#include "../Lumen/LumenReflectionDenoiserCommon.ush"
#include "../StochasticLighting/StochasticLightingUpsample.ush"

float3 VolumeFrameJitterOffset;
float VolumeInverseSquaredLightDistanceBiasScale;
uint3 VolumeSampleViewSize;
uint3 DownsampledVolumeViewSize;
uint3 NumSamplesPerVoxel;
Texture3D<uint> VolumeLightSamples;

#if TRANSLUCENCY_LIGHTING_VOLUME
RWTexture3D<float4> RWTranslucencyVolumeResolvedLightingAmbient;
RWTexture3D<float4> RWTranslucencyVolumeResolvedLightingDirectional;
#else
RWTexture3D<float3> RWVolumeResolvedLighting;
#endif

float VolumeMaxShadingWeight;
uint DebugLightId;

void LoadPackedLightSamples(inout uint PackedLightSamples[NUM_SAMPLES_PER_VOXEL_1D], int3 DownsampledVolumeCoord)
{
	DownsampledVolumeCoord = clamp(DownsampledVolumeCoord, int3(0, 0, 0), int3(DownsampledVolumeViewSize - 1));

	for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_VOXEL_1D; ++SampleIndex)
	{
		uint3 LightSampleCoord = DownsampledVolumeCoord * uint3(NUM_SAMPLES_PER_VOXEL_3D_X, NUM_SAMPLES_PER_VOXEL_3D_Y, NUM_SAMPLES_PER_VOXEL_3D_Z) 
			+ uint3(SampleIndex % NUM_SAMPLES_PER_VOXEL_3D_X, SampleIndex / NUM_SAMPLES_PER_VOXEL_3D_X, 0);

		PackedLightSamples[SampleIndex] = VolumeLightSamples[LightSampleCoord];
	}
}

void AccumulateLightSample(uint PackedLightSamples[NUM_SAMPLES_PER_VOXEL_1D], inout uint NextLocalLightIndex)
{
	for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_VOXEL_1D; ++SampleIndex)
	{
		FLightSample LightSample = UnpackLightSample(PackedLightSamples[SampleIndex]);
		if (LightSample.bVisible)
		{
			NextLocalLightIndex = min(NextLocalLightIndex, LightSample.LocalLightIndex);
		}
	}
}

void AccumulateLightSample(uint PackedLightSamples[NUM_SAMPLES_PER_VOXEL_1D], uint LocalLightIndex, float UpsampleWeight, inout uint NextLocalLightIndex, inout float SampleWeightSum)
{
	for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_VOXEL_1D; ++SampleIndex)
	{
		FLightSample LightSample = UnpackLightSample(PackedLightSamples[SampleIndex]);
		if (LightSample.bVisible)
		{
			if (LightSample.LocalLightIndex == LocalLightIndex)
			{
				SampleWeightSum += LightSample.Weight * UpsampleWeight;
			}

			if (LightSample.LocalLightIndex > LocalLightIndex)
			{
				NextLocalLightIndex = min(NextLocalLightIndex, LightSample.LocalLightIndex);
			}
		}
	}
}

/**
 * Upsample light samples and apply all lights per voxel
 */
[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, THREADGROUP_SIZE)]
void VolumeShadeLightSamplesCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	uint3 VolumeCoord = DispatchThreadId.xyz;
	if (all(VolumeCoord < VolumeViewSize))
	{
		FShaderPrintContext DebugContext = InitVolumeDebugContext(VolumeCoord, /*bDownsampled*/ false, float2(0.5, 0.05));

		float SceneDepth = 0.0f;
		const float3 TranslatedWorldPosition = ComputeCellTranslatedWorldPosition(VolumeCoord, VolumeFrameJitterOffset, SceneDepth);
		const float3 CameraVector = normalize(TranslatedWorldPosition - View.TranslatedWorldCameraOrigin);

		if (DebugContext.bIsActive)
		{
			Print(DebugContext, TEXT("VolumeShadeLightSamples"), FontTitle);
			Newline(DebugContext);
			Print(DebugContext, TEXT("Coord    : "));
			Print(DebugContext, VolumeCoord, FontValue);
			Newline(DebugContext);
			Print(DebugContext, TEXT("TileCoord: "));
			Print(DebugContext, GroupId.xyz, FontValue);
			Newline(DebugContext);
			Print(DebugContext, TEXT("TWS      : "));
			Print(DebugContext, TranslatedWorldPosition, FontValue);
			AddCrossTWS(DebugContext, TranslatedWorldPosition, 5.0f, float4(0, 0, 1, 1));

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

#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

#if VOLUME_DOWNSAMPLE_FACTOR == 2
		int3 DownsampledVolumeCoord000 = int3(VolumeCoord - 1) / VOLUME_DOWNSAMPLE_FACTOR;
		int3 VolumeCoordOffset = VolumeCoord - DownsampledVolumeCoord000 * 2;
		int3 SampleOffset000 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(0, 0, 0)) + uint3(0, 0, 0) * 2 - VolumeCoordOffset;
		int3 SampleOffset100 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(1, 0, 0)) + uint3(1, 0, 0) * 2 - VolumeCoordOffset;
		int3 SampleOffset010 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(0, 1, 0)) + uint3(0, 1, 0) * 2 - VolumeCoordOffset;
		int3 SampleOffset110 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(1, 1, 0)) + uint3(1, 1, 0) * 2 - VolumeCoordOffset;
		int3 SampleOffset001 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(0, 0, 1)) + uint3(0, 0, 1) * 2 - VolumeCoordOffset;
		int3 SampleOffset101 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(1, 0, 1)) + uint3(1, 0, 1) * 2 - VolumeCoordOffset;
		int3 SampleOffset011 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(0, 1, 1)) + uint3(0, 1, 1) * 2 - VolumeCoordOffset;
		int3 SampleOffset111 = GetSampleVoxelCoordJitter(DownsampledVolumeCoord000 + uint3(1, 1, 1)) + uint3(1, 1, 1) * 2 - VolumeCoordOffset;

		// Triangle filter weight between the shaded voxel and 8 neighbors
		float4 InterpolationWeights0;
		InterpolationWeights0.x = (2.0f - abs(SampleOffset000.x)) * (2.0f - abs(SampleOffset000.y)) * (2.0f - abs(SampleOffset000.z));
		InterpolationWeights0.y = (2.0f - abs(SampleOffset100.x)) * (2.0f - abs(SampleOffset100.y)) * (2.0f - abs(SampleOffset100.z));
		InterpolationWeights0.z = (2.0f - abs(SampleOffset010.x)) * (2.0f - abs(SampleOffset010.y)) * (2.0f - abs(SampleOffset010.z));
		InterpolationWeights0.w = (2.0f - abs(SampleOffset110.x)) * (2.0f - abs(SampleOffset110.y)) * (2.0f - abs(SampleOffset110.z));
		float4 InterpolationWeights1;
		InterpolationWeights1.x = (2.0f - abs(SampleOffset001.x)) * (2.0f - abs(SampleOffset001.y)) * (2.0f - abs(SampleOffset001.z));
		InterpolationWeights1.y = (2.0f - abs(SampleOffset101.x)) * (2.0f - abs(SampleOffset101.y)) * (2.0f - abs(SampleOffset101.z));
		InterpolationWeights1.z = (2.0f - abs(SampleOffset011.x)) * (2.0f - abs(SampleOffset011.y)) * (2.0f - abs(SampleOffset011.z));
		InterpolationWeights1.w = (2.0f - abs(SampleOffset111.x)) * (2.0f - abs(SampleOffset111.y)) * (2.0f - abs(SampleOffset111.z));

		// Normalize weights
		InterpolationWeights0 /= 8.0f;
		InterpolationWeights1 /= 8.0f;
#endif

		float3 LightScattering = 0.0f;

		float4 LightAmbient = 0.0f;
		float4 LightDirectional = 0.0f;

#if !TRANSLUCENCY_LIGHTING_VOLUME
		const float FroxelFootprintMargin = 0.5f; // trilinear interpolation
		if (IsFroxelVisible(VolumeCoord, FroxelFootprintMargin))
#endif
		{
#if VOLUME_DOWNSAMPLE_FACTOR == 2
			// Stochastic sample interpolation. Need to use at least samples 2 for good quality.
			const float RandomScalar = BlueNoiseScalar(VolumeCoordToNoiseCoord(VolumeCoord), MegaLightsStateFrameIndex);
			uint3 StochasticTrilinearOffset0 = GetStochasticTrilinearOffset((RandomScalar + 0.0f) / 2.0f, InterpolationWeights0, InterpolationWeights1);
			uint3 StochasticTrilinearOffset1 = GetStochasticTrilinearOffset((RandomScalar + 1.0f) / 2.0f, InterpolationWeights0, InterpolationWeights1);

			if (DebugContext.bIsActive)
			{
				Newline(DebugContext);
				Print(DebugContext, StochasticTrilinearOffset0, FontValue);
				Print(DebugContext, StochasticTrilinearOffset1, FontValue);
			}
#endif

#if VOLUME_DOWNSAMPLE_FACTOR == 2
			uint PackedLightSamples0[NUM_SAMPLES_PER_VOXEL_1D];
			uint PackedLightSamples1[NUM_SAMPLES_PER_VOXEL_1D];
			LoadPackedLightSamples(PackedLightSamples0, DownsampledVolumeCoord000 + StochasticTrilinearOffset0);
			LoadPackedLightSamples(PackedLightSamples1, DownsampledVolumeCoord000 + StochasticTrilinearOffset1);
#else
			uint PackedLightSamples0[NUM_SAMPLES_PER_VOXEL_1D];
			LoadPackedLightSamples(PackedLightSamples0, VolumeCoord);
#endif

			uint NextLocalLightIndex = MAX_LOCAL_LIGHT_INDEX;
			AccumulateLightSample(PackedLightSamples0, NextLocalLightIndex);
#if VOLUME_DOWNSAMPLE_FACTOR == 2
			AccumulateLightSample(PackedLightSamples1, NextLocalLightIndex);
#endif

			while (NextLocalLightIndex < MAX_LOCAL_LIGHT_INDEX)
			{
				const uint LocalLightIndex = WaveActiveMin(NextLocalLightIndex);
				if (LocalLightIndex == NextLocalLightIndex)
				{
					NextLocalLightIndex = MAX_LOCAL_LIGHT_INDEX;
					float SampleWeight = 0.0f;

#if VOLUME_DOWNSAMPLE_FACTOR == 2
					AccumulateLightSample(PackedLightSamples0, LocalLightIndex, 0.5f, NextLocalLightIndex, SampleWeight);
					AccumulateLightSample(PackedLightSamples1, LocalLightIndex, 0.5f, NextLocalLightIndex, SampleWeight);
#else
					AccumulateLightSample(PackedLightSamples0, LocalLightIndex, 1.0f, NextLocalLightIndex, SampleWeight);
#endif
					SampleWeight = min(SampleWeight, VolumeMaxShadingWeight);

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

					if (DebugContext.bIsActive)
					{
						Newline(DebugContext);
						Print(DebugContext, ForwardLightData.LightSceneId, Select(ForwardLightData.LightSceneId == DebugLightId, FontSelected, FontValue));
						Print(DebugContext, SampleWeight, FontValue);
					}

					if (SampleWeight > 0.01f)
					{
						float3 L;
						const float VolumetricScatteringIntensity = UnpackVolumetricScatteringIntensity(ForwardLightData);
						float3 Lighting = GetMegaLightsVolumeLighting(TranslatedWorldPosition, CameraVector, DistanceBiasSqr, LightVolumetricSoftFadeDistance, LightData, VolumetricScatteringIntensity, L) * SampleWeight;
						LightScattering += Lighting;

						{
							FTwoBandSHVectorRGB SHLighting = MulSH(SHBasisFunction(L), Lighting);

							// Directional light contribution in w
							LightAmbient += float4(SHLighting.R.V.x, SHLighting.G.V.x, SHLighting.B.V.x, 0.0f);

							float3 LuminanceWeights = LuminanceFactors();
							float3 Coefficient0 = float3(SHLighting.R.V.y, SHLighting.G.V.y, SHLighting.B.V.y);
							float3 Coefficient1 = float3(SHLighting.R.V.z, SHLighting.G.V.z, SHLighting.B.V.z);
							float3 Coefficient2 = float3(SHLighting.R.V.w, SHLighting.G.V.w, SHLighting.B.V.w);
							LightDirectional += float4(dot(Coefficient0, LuminanceWeights), dot(Coefficient1, LuminanceWeights), dot(Coefficient2, LuminanceWeights), 0);
						}

						if (DebugContext.bIsActive)
						{
							Print(DebugContext, Lighting, FontValue);
						}
					}
				}
			}
		}

		if (DebugContext.bIsActive)
		{
			Newline(DebugContext);
			Print(DebugContext, TEXT("Lighting: "));
			Print(DebugContext, LightScattering, FontValue);
		}

#if TRANSLUCENCY_LIGHTING_VOLUME
		RWTranslucencyVolumeResolvedLightingAmbient[VolumeCoord] = LightAmbient;
		RWTranslucencyVolumeResolvedLightingDirectional[VolumeCoord] = LightDirectional;
#else
		RWVolumeResolvedLighting[VolumeCoord] = LightScattering * View.PreExposure;
#endif
	}
}