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

// When loading Substrate material, needs to support all types of pixels (Simple/Single/Complex/Special) 
// as the temporal pass does not have tile types
#define SUBSTRATE_FASTPATH 0
#define SUBSTRATE_SINGLEPATH 0
#define SUBSTRATE_COMPLEXSPECIALPATH 1

#include "../Common.ush"
#include "MegaLights.ush"
#include "MegaLightsMaterial.ush"
#include "../Lumen/LumenReflectionDenoiserCommon.ush"

// Spatial denoise is the last pass writing to the SceneColor buffer. We allow SSS checkerboard pixel to adjust 
// DiffuseColor/SpecularColor to write specular & diffuse lighting in checkerboard pattern
#define ALLOW_SSS_MATERIAL_OVERRIDE 1

Texture2D<float3> DiffuseLightingTexture;
Texture2D<float3> SpecularLightingTexture;
Texture2D<float4> LightingMomentsTexture;
Texture2D<UNORM float> NumFramesAccumulatedTexture;
Texture2D<UNORM float> ShadingConfidenceTexture;
RWTexture2D<float4> RWSceneColor;

float SpatialFilterDepthWeightScale;
float SpatialFilterKernelRadius;
uint SpatialFilterNumSamples;
float SpatialFilterMaxDisocclusionFrames;
float TemporalMaxFramesAccumulated;
uint bSubPixelShading;

float3 TonemapLighting(float3 Lighting, float DisocclusionFactor)
{
	// Run heavy tonemapping for DisocclusionFactor in order to suppress fireflies in new areas
	return Lighting / (1.0f + DisocclusionFactor * Luminance(Lighting));
}

float3 InverseTonemapLighting(float3 TonemappedLighting, float DisocclusionFactor)
{
	return TonemappedLighting / (1.0f - DisocclusionFactor * Luminance(TonemappedLighting));
}

/**
 * Run a spatial filter in order to filter out noise based on the temporal variance.
 */
[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void DenoiserSpatialCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	uint2 ScreenCoord = DispatchThreadId.xy + View.ViewRectMinAndSize.xy;

	FShaderPrintContext DebugContext = InitDebugContext(ScreenCoord, /*bDownsampled*/ false, float2(0.65, 0.8));

	#if INPUT_TYPE == INPUT_TYPE_HAIRSTRANDS
	if (all(ScreenCoord < View.ViewRectMinAndSize.xy + View.ViewRectMinAndSize.zw))
	{
		float2 ScreenUV = (ScreenCoord + 0.5f) * View.BufferSizeAndInvSize.zw;

		float3 CenterDiffuseLighting = DiffuseLightingTexture[ScreenCoord];
		if (IsLightingValid(CenterDiffuseLighting))
		{
			// For hair, we don't apply proper spatial filtering at the moment.
			FMegaLightsMaterial Material = LoadMaterial(ScreenUV, ScreenCoord);
			if (Material.bIsValid)
			{
				float ShadingConfidence = ShadingConfidenceTexture[ScreenCoord];

				float4 CenterLightingMoments = LightingMomentsTexture[ScreenCoord];
				float CenterDiffuseStdDev = sqrt(max(CenterLightingMoments.y - Pow2(CenterLightingMoments.x), 0.0f));
	
				// Boost spatial filter until we accumulate enough frames to be able to rely on temporal variance
				float NumFramesAccumulated = UnpackNumFramesAccumulated(NumFramesAccumulatedTexture[ScreenCoord]);
				float DisocclusionFactor = 0.0f;
				if (TemporalMaxFramesAccumulated > 1 && ShadingConfidence <= 0.25f && SpatialFilterMaxDisocclusionFrames > 0.0f)
				{
					DisocclusionFactor = 1.0f - saturate((NumFramesAccumulated - 1.0f) / (SpatialFilterMaxDisocclusionFrames - 1.0f));
					CenterDiffuseStdDev = lerp(CenterDiffuseStdDev, 1.0f, DisocclusionFactor);
				}

				float3 DiffuseLightingSum = TonemapLighting(CenterDiffuseLighting, DisocclusionFactor);
				float DiffuseLightingWeightSum = 1.0f;

				uint2 RandomSeed = Rand3DPCG16(int3(ScreenCoord, MegaLightsStateFrameIndex)).xy;

				float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, Material.Depth);
				float4 ScenePlane = float4(Material.WorldNormalForPositionBias, dot(TranslatedWorldPosition, Material.WorldNormalForPositionBias));

				// Final pass outputs composites irradiance and outputs it to scene color
				float3 DiffuseLighting = InverseTonemapLighting(DiffuseLightingSum / DiffuseLightingWeightSum, DisocclusionFactor);
				float3 SpecularLighting = 0.f;

				// Do not modulate/demodulate hair lighting as GetDenoisingModulateFactors() does not work with hair shading (demodulation is done in ShadeLightSamplesCS())
				// FDenoisingModulateFactors Factors = GetDenoisingModulateFactors(Material, TranslatedWorldPosition);
				// DiffuseLighting *= Factors.Diffuse;
				// SpecularLighting *= Factors.Specular;

				// Composite diffuse and specular into scene color
				FLightAccumulator LightAccumulator = (FLightAccumulator)0;
				LightAccumulator_AddSplit(LightAccumulator, DiffuseLighting, SpecularLighting, /*ScatterableLight*/ DiffuseLighting, /*CommonMultiplier*/ 1.0f, Material.bNeedsSeparateSubsurfaceLightAccumulation);
				float4 AccumulatedSceneColor = LightAccumulator_GetResult(LightAccumulator);

				const uint TotalNodeCount = HairStrands.HairSampleCount[0];
				const uint2 Resolution = GetHairSampleResolution(TotalNodeCount);
				const FNodeDesc NodeDesc = DecodeNodeDesc(HairStrands.HairSampleOffset.Load(uint3(ScreenCoord, 0)));

				// If sub-pixel shading is:
				// * Disabled, splat the first-sample's lighting onto all other samples
				// * Enabled, lighting has already be written during the shading pass
				if (bSubPixelShading == 0)
				{
					LOOP
					for (uint SampleIt = 0; SampleIt < NodeDesc.Count; SampleIt++)
					{
						const uint LinearCoord = NodeDesc.Offset + SampleIt;
						const uint2 OutCoord = GetHairSampleCoord(LinearCoord, Resolution);

						const FHairSample Sample = UnpackHairSample(HairStrands.HairSampleData[LinearCoord]);
						const float SampleCoverage = saturate(From8bitCoverage(Sample.Coverage8bit));
						RWSceneColor[OutCoord] = RWSceneColor[OutCoord] + AccumulatedSceneColor * SampleCoverage;
					}
				}
			}
		}
	}
	#elif INPUT_TYPE == INPUT_TYPE_GBUFFER
	if (all(ScreenCoord < View.ViewRectMinAndSize.xy + View.ViewRectMinAndSize.zw))
	{
		float2 ScreenUV = (ScreenCoord + 0.5f) * View.BufferSizeAndInvSize.zw;

		float3 CenterDiffuseLighting = DiffuseLightingTexture[ScreenCoord];
		float3 CenterSpecularLighting = SpecularLightingTexture[ScreenCoord];
		if (IsLightingValid(CenterDiffuseLighting))
		{
			FMegaLightsMaterial Material = LoadMaterial(ScreenUV, ScreenCoord);

			float ShadingConfidence = ShadingConfidenceTexture[ScreenCoord];

			float4 CenterLightingMoments = LightingMomentsTexture[ScreenCoord];
			float CenterDiffuseStdDev = sqrt(max(CenterLightingMoments.y - Pow2(CenterLightingMoments.x), 0.0f));
			float CenterSpecularStdDev = sqrt(max(CenterLightingMoments.w - Pow2(CenterLightingMoments.z), 0.0f));
	
			// Boost spatial filter until we accumulate enough frames to be able to rely on temporal variance
			float NumFramesAccumulated = UnpackNumFramesAccumulated(NumFramesAccumulatedTexture[ScreenCoord]);
			float DisocclusionFactor = 0.0f;
			if (TemporalMaxFramesAccumulated > 1 && ShadingConfidence <= 0.25f && SpatialFilterMaxDisocclusionFrames > 0.0f)
			{
				DisocclusionFactor = 1.0f - saturate((NumFramesAccumulated - 1.0f) / (SpatialFilterMaxDisocclusionFrames - 1.0f));
				CenterDiffuseStdDev = lerp(CenterDiffuseStdDev, 1.0f, DisocclusionFactor);
			}

			float3 DiffuseLightingSum = TonemapLighting(CenterDiffuseLighting, DisocclusionFactor);
			float3 SpecularLightingSum = TonemapLighting(CenterSpecularLighting, DisocclusionFactor);
			float DiffuseLightingWeightSum = 1.0f;
			float SpecularLightingWeightSum = 1.0f;

			uint2 RandomSeed = Rand3DPCG16(int3(ScreenCoord, MegaLightsStateFrameIndex)).xy;

			float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, Material.Depth);
			float4 ScenePlane = float4(Material.WorldNormalForPositionBias, dot(TranslatedWorldPosition, Material.WorldNormalForPositionBias));

			bool bFilterDiffuse = false;
			bool bFilterSpecular = false;

			// #ml_todo: expose as CVars
			if (ShadingConfidence <= 0.2f)
			{
				// Don't run any filtering if relative error isn't big enough. TSR can clean it up without adding extra blur.
				bFilterDiffuse = CenterDiffuseStdDev / max(Luminance(CenterDiffuseLighting), 0.1f) > 0.5f;
				bFilterSpecular = CenterSpecularStdDev / max(Luminance(CenterSpecularLighting), 0.1f) > 0.5f;
			}

			if (DisocclusionFactor > 0.01f)
			{
				bFilterDiffuse = true;
			}

			const uint NumSamples = lerp(SpatialFilterNumSamples, 2 * SpatialFilterNumSamples, DisocclusionFactor);

			#if SPATIAL_FILTER
			if (Material.bAllowSpatialFilter && SpatialFilterKernelRadius > 0.0f && (bFilterDiffuse || bFilterSpecular))
			{
				for (uint NeighborIndex = 0; NeighborIndex < NumSamples; ++NeighborIndex)
				{
					float2 NeighborOffsetInRect = Hammersley16(NeighborIndex, NumSamples, RandomSeed);
					float2 NeighborOffset = UniformSampleDiskConcentric(NeighborOffsetInRect) * SpatialFilterKernelRadius;
					int2 NeighborCoord = (int2)(ScreenCoord + NeighborOffset);

					if (all(and(NeighborCoord >= View.ViewRectMinAndSize.xy, NeighborCoord < (View.ViewRectMinAndSize.xy + View.ViewRectMinAndSize.zw))))
					{
						// Depth weight
						float2 NeighborScreenUV = (NeighborCoord + 0.5f) * View.BufferSizeAndInvSize.zw;
						const FMegaLightsMaterial NeighborMaterial = LoadMaterial(NeighborScreenUV, NeighborCoord);
						const float3 NeighborTranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(NeighborScreenUV, NeighborMaterial.Depth);
						float PlaneDistance = abs(dot(float4(NeighborTranslatedWorldPosition, -1), ScenePlane));
						float RelativeDepthDifference = PlaneDistance / Material.Depth;
						float DepthWeight = exp2(-SpatialFilterDepthWeightScale * (RelativeDepthDifference * RelativeDepthDifference));

						// #sdl_todo: separate weight for specular accounting for roughness and expose as CVars
						// Normal weight
						float AngleBetweenNormals = acosFast(saturate(dot(ScenePlane.xyz, NeighborMaterial.WorldNormalForPositionBias)));
						float NormalWeight = 1.0f - saturate(AngleBetweenNormals);

						// Diffuse
						if (bFilterDiffuse)
						{
							float3 NeighborDiffuseLighting = DiffuseLightingTexture[NeighborCoord];

							if (IsLightingValid(NeighborDiffuseLighting))
							{
								float LuminanceDelta = abs(Luminance(CenterDiffuseLighting) - Luminance(NeighborDiffuseLighting));
								float LuminanceWeight = exp2(-(LuminanceDelta / max(CenterDiffuseStdDev, 0.001f)) * saturate(1.0f - DisocclusionFactor));

								const float DiffuseNeighborWeight = DepthWeight * NormalWeight * LuminanceWeight;
								DiffuseLightingSum += TonemapLighting(NeighborDiffuseLighting, DisocclusionFactor) * DiffuseNeighborWeight;
								DiffuseLightingWeightSum += DiffuseNeighborWeight;
							}
						}

						// Specular
						if (bFilterSpecular)
						{
							float3 NeighborSpecularLighting = SpecularLightingTexture[NeighborCoord];

							if (IsLightingValid(NeighborSpecularLighting))
							{
								float LuminanceDelta = abs(Luminance(CenterSpecularLighting) - Luminance(NeighborSpecularLighting));
								float LuminanceWeight = exp2(-(LuminanceDelta / max(CenterSpecularStdDev, 0.001f)) * saturate(1.0f - DisocclusionFactor));

								const float SpecularNeighborWeight = DepthWeight * NormalWeight * LuminanceWeight;
								SpecularLightingSum += TonemapLighting(NeighborSpecularLighting, DisocclusionFactor) * SpecularNeighborWeight;
								SpecularLightingWeightSum += SpecularNeighborWeight;
							}
						}
					}
				}
			}
			#endif

			// Final pass outputs composites irradiance and outputs it to scene color
			float3 DiffuseLighting = InverseTonemapLighting(DiffuseLightingSum / DiffuseLightingWeightSum, DisocclusionFactor);
			float3 SpecularLighting = InverseTonemapLighting(SpecularLightingSum / SpecularLightingWeightSum, DisocclusionFactor);

			FDenoisingModulateFactors Factors = GetDenoisingModulateFactors(Material, TranslatedWorldPosition);
			DiffuseLighting *= Factors.Diffuse;
			SpecularLighting *= Factors.Specular;

			// Composite diffuse and specular into scene color
			FLightAccumulator LightAccumulator = (FLightAccumulator)0;
			LightAccumulator_AddSplit(LightAccumulator, DiffuseLighting, SpecularLighting, /*ScatterableLight*/ DiffuseLighting, /*CommonMultiplier*/ 1.0f, Material.bNeedsSeparateSubsurfaceLightAccumulation);
			float4 AccumulatedSceneColor = LightAccumulator_GetResult(LightAccumulator);

			if (DebugContext.bIsActive)
			{
				Print(DebugContext, TEXT("Spatial pass"), FontTitle);
				Newline(DebugContext);
				Print(DebugContext, TEXT("AccumulatedSceneColor : "));
				Print(DebugContext, AccumulatedSceneColor, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("CenterDiffuseStdDev   : "));
				Print(DebugContext, CenterDiffuseStdDev, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("CenterSpecularStdDev  : "));
				Print(DebugContext, CenterSpecularStdDev, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("NumFramesAccumulated  : "));
				Print(DebugContext, NumFramesAccumulated, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("DisocclusionFactor    : "));
				Print(DebugContext, DisocclusionFactor, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("ShadingConfidence     : "));
				Print(DebugContext, ShadingConfidence, Select(ShadingConfidence > 0.75f, FontYellow, Select(ShadingConfidence > 0.25f, FontWhite, FontGrey)));
				Newline(DebugContext);
				Print(DebugContext, TEXT("Filter                : "));
				Print(DebugContext, bFilterDiffuse ? 1u : 0u, FontValue);
				Print(DebugContext, bFilterSpecular ? 1u : 0u, FontValue);
				AddTextBackground(DebugContext, FontBackground);
			}

			// Visualize where spatial filter gets applied
			if (false)
			{
				if (bFilterDiffuse || bFilterSpecular)
				{
					float4 DebugColor = 0;
					if (bFilterDiffuse)
						DebugColor.x = 1.0f;
					if (bFilterSpecular)
						DebugColor.y = 1.0f;

					AccumulatedSceneColor = DebugColor;
				}
			}

			RWSceneColor[ScreenCoord] = RWSceneColor[ScreenCoord] + AccumulatedSceneColor;
		}
	}
	#endif
}