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

#ifndef TILE_TYPE
	#define TILE_TYPE TILE_MODE_EMPTY
#endif

#ifndef NUM_SAMPLES_PER_PIXEL_1D
	#define NUM_SAMPLES_PER_PIXEL_1D 1
	#define NUM_SAMPLES_PER_PIXEL_2D_X 1
	#define NUM_SAMPLES_PER_PIXEL_2D_Y 1
#endif

// 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

#define USE_IES_PROFILE 1
#define USE_LIGHT_FUNCTION_ATLAS 1

#include "../Common.ush"
#include "/Engine/Shared/MegaLightsDefinitions.h"

// Substrate tile for faster shading
#if SUBSTRATE_ENABLED 
#if TILE_TYPE == TILE_MODE_SIMPLE_SHADING || TILE_TYPE == TILE_MODE_SIMPLE_SHADING_RECT || TILE_TYPE == TILE_MODE_SIMPLE_SHADING_RECT_TEXTURED || TILE_TYPE == TILE_MODE_EMPTY
	#define SUBSTRATE_FASTPATH 1
	#define SUBSTRATE_SINGLEPATH 0
	#define SUBSTRATE_COMPLEXSPECIALPATH 0
#elif TILE_TYPE == TILE_MODE_SINGLE_SHADING || TILE_TYPE == TILE_MODE_SINGLE_SHADING_RECT || TILE_TYPE == TILE_MODE_SINGLE_SHADING_RECT_TEXTURED
	#define SUBSTRATE_FASTPATH 0 
	#define SUBSTRATE_SINGLEPATH 1
	#define SUBSTRATE_COMPLEXSPECIALPATH 0
#elif TILE_TYPE == TILE_MODE_COMPLEX_SHADING || TILE_TYPE == TILE_MODE_COMPLEX_SHADING_RECT || TILE_TYPE == TILE_MODE_COMPLEX_SHADING_RECT_TEXTURED
	#define SUBSTRATE_FASTPATH 0 
	#define SUBSTRATE_SINGLEPATH 0
	#define SUBSTRATE_COMPLEXSPECIALPATH 0
#else // Special TILE_MODE_COMPLEX_SPECIAL_SHADING || TILE_MODE_SPECIAL_COMPLEX_SHADING_RECT || TILE_MODE_SPECIAL_COMPLEX_SHADING_RECT_TEXTURED
	#define SUBSTRATE_FASTPATH 0
	#define SUBSTRATE_SINGLEPATH 0
	#define SUBSTRATE_COMPLEXSPECIALPATH 1
#endif
#endif

#if SUBSTRATE_FASTPATH || SUBSTRATE_SINGLEPATH
#define SUBSTRATE_LOAD_FROM_MATERIALCONTAINER 0
#endif

#include "../BlueNoise.ush"
#include "MegaLightsShading.ush"
#include "../Lumen/LumenReflectionDenoiserCommon.ush"
#include "../StochasticLighting/StochasticLightingCommon.ush"

// For now, allow register spilling when Substrate is enabled, until we reduce register usage
#if SUBSTRATE_ENABLED
#pragma warning(disable:7203)
#endif

ADAPTIVE_LICM

#define WAVE_LIGHT_HASH_SIZE 4

RWTexture2D<float3> RWResolvedDiffuseLighting;
RWTexture2D<float3> RWResolvedSpecularLighting;
RWTexture2D<UNORM float> RWShadingConfidence;
RWTexture2D<float3> RWOutputColor;

int2 SampleViewMin;
int2 SampleViewSize;
uint2 DownsampledViewMin;
uint2 DownsampledViewSize;
Texture2D<float> DownsampledSceneDepth;
Texture2D<UNORM float3> DownsampledSceneWorldNormal;
float2 DownsampledBufferInvSize;

StructuredBuffer<uint> TileAllocator;
StructuredBuffer<uint> TileData;
uint TileDataStride;

Texture2D<uint> LightSamples;
Texture2D<uint> HairTransmittanceMaskTexture;
Texture2D<uint> PackedPixelDataTexture;

uint bSubPixelShading;
uint ShadingSampleIndex;
float MaxShadingWeight;
uint UseShadingConfidence;
uint ShadingPassIndex;
int UseIESProfiles;
int UseLightFunctionAtlas;
uint DebugLightId;

float GetLightSampleWeightRatio(FLightSample LightSample)
{
	return LightSample.bGuidedAsVisible ? 1.0f / LightSample.Weight : 0.0f;
}

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

void AccumulateLightSample(uint PackedLightSamples[NUM_SAMPLES_PER_PIXEL_1D], uint LocalLightIndex, inout uint NextLocalLightIndex, inout float SampleWeightSum, inout float WeightRatioSum, inout uint ValidSampleMask, inout FShaderPrintContext DebugContext)
{
	uint NumMergedSamples = 1;

	for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_PIXEL_1D; ++SampleIndex)
	{
		FLightSample LightSample = UnpackLightSample(PackedLightSamples[SampleIndex]);

		if (LightSample.LocalLightIndex == LocalLightIndex)
		{
			// Identical rays were skipped by marking such samples as weight=0
			if (LightSample.Weight == 0.0f)
			{
				++NumMergedSamples;
			}
			else
			{
				if (LightSample.bVisible)
				{
					SampleWeightSum += LightSample.Weight * NumMergedSamples;
					WeightRatioSum += GetLightSampleWeightRatio(LightSample) * NumMergedSamples;
					ValidSampleMask |= 1u << SampleIndex;
					NumMergedSamples = 1;
				}
				else
				{
					// Decrease confidence on ray miss
					WeightRatioSum -= GetLightSampleWeightRatio(LightSample) * NumMergedSamples;
				}
			}
		}

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

bool IsValidDownsampledCoord(uint2 DownsampledScreenCoord)
{
	return all(DownsampledScreenCoord.xy < DownsampledViewMin + DownsampledViewSize);
}

float GetNormalWeight(float3 SceneWorldNormal, uint2 DownsampledScreenCoord)
{
	float3 SampleWorldNormal = normalize(DecodeNormal(DownsampledSceneWorldNormal[DownsampledScreenCoord]));

	float AngleBetweenNormals = acosFast(saturate(dot(SampleWorldNormal, SceneWorldNormal)));
	float NormalWeight = 1.0f - saturate(AngleBetweenNormals);

	return Pow2(NormalWeight);
}

/**
 * Upsample light samples and apply all lights per pixel to affected tiles
 */
[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void ShadeLightSamplesCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	uint TileIndex = GroupId.x;
	if (TileIndex < TileAllocator[TILE_TYPE])
	{
		uint2 TileCoord = UnpackTile(TileData[TileIndex + TILE_TYPE * TileDataStride]);
		uint2 ScreenCoord = TileCoord * TILE_SIZE + GroupThreadId.xy + View.ViewRectMinAndSize.xy;
		uint LocalLightHiMask = 0;

		if (all(ScreenCoord < View.ViewRectMinAndSize.xy + View.ViewRectMinAndSize.zw))
		{
			FShaderPrintContext DebugContext = InitDebugContext(ScreenCoord, /*bDownsampled*/ false, float2(0.65, 0.05));
			DebugContext.bIsActive = DebugContext.bIsActive;

			const float2 ScreenUV = (ScreenCoord + 0.5f) * View.BufferSizeAndInvSize.zw;
			const bool bForceSimpleShading = IsSimpleShadingTileType(TILE_TYPE);
			FMegaLightsMaterial Material = LoadMaterial(ScreenUV, ScreenCoord, bForceSimpleShading, ShadingSampleIndex);

			#if INPUT_TYPE == INPUT_TYPE_GBUFFER
			Material.LightingChannelMask = GetSceneLightingChannel(ScreenCoord);
			#endif

			// #ml_todo: handle this path when shadows are disabled
			float ScreenSpaceAO = 1.0f;//Texture2DSampleLevel(SceneTexturesStruct.ScreenSpaceAOTexture, SceneTexturesStruct_ScreenSpaceAOTextureSampler, ScreenUV, 0).x;

			float3 DiffuseLighting = INVALID_LIGHTING;
			float3 SpecularLighting = INVALID_LIGHTING;
			float WeightRatioSum = 0.0f;

			if (Material.Depth > 0.0f && Material.IsValid())
			{
				float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, Material.Depth);
				float3 CameraVector = normalize(TranslatedWorldPosition - View.TranslatedWorldCameraOrigin);
				bool bValidPixel = true;

				float4 DownsampledScreenWeights = 0.0f;
				uint2 DownsampledScreenCoord = ScreenCoord;

				#if DOWNSAMPLE_FACTOR_X == 2 && DOWNSAMPLE_FACTOR_Y == 1
				{
					int2 SampleOffsets[4];
					SampleOffsets[0] = int2(0, 0);
					SampleOffsets[1] = int2(select(ScreenCoord.x % DOWNSAMPLE_FACTOR_X == 0, -1, 1), 0);
					SampleOffsets[2] = int2(0, -1);
					SampleOffsets[3] = int2(0, +1);

					FMegaLightsPackedPixelData PackedPixelData = (FMegaLightsPackedPixelData)PackedPixelDataTexture[ScreenCoord];
					int2 StochasticBilinearOffset = PackedPixelData.GetStochasticSampleOffset(SampleOffsets);
					int2 DownsampledScreenCoord00 = ScreenCoord / uint2(DOWNSAMPLE_FACTOR_X, DOWNSAMPLE_FACTOR_Y);

					bValidPixel = PackedPixelData.HasValidStochasticSample();
					DownsampledScreenCoord = clamp(DownsampledScreenCoord00 + StochasticBilinearOffset, int2(DownsampledViewMin), int2(DownsampledViewMin) + int2(DownsampledViewSize) - 1);
				}
				#elif DOWNSAMPLE_FACTOR_X == 2
				{
					FMegaLightsPackedPixelData PackedPixelData = (FMegaLightsPackedPixelData)PackedPixelDataTexture[ScreenCoord];
					int2 StochasticBilinearOffset = PackedPixelData.GetStochasticSampleOffset();
					int2 DownsampledScreenCoord00 = floor(ScreenUV * View.BufferSizeAndInvSize.xy / DOWNSAMPLE_FACTOR_X - 0.5f);

					bValidPixel = PackedPixelData.HasValidStochasticSample();
					DownsampledScreenCoord = clamp(DownsampledScreenCoord00 + StochasticBilinearOffset, int2(DownsampledViewMin), int2(DownsampledViewMin) + int2(DownsampledViewSize) - 1);
				}
				#endif

				if (DebugContext.bIsActive)
				{
					Print(DebugContext, TEXT("ShadeSamples"), FontTitle);
					Newline(DebugContext);
					Print(DebugContext, TEXT("ScreenCoord         : "));
					Print(DebugContext, ScreenCoord, FontValue);
				#if DOWNSAMPLE_FACTOR_X != 1
					Newline(DebugContext);
					Print(DebugContext, TEXT("DownsampledWeights  : "));
					Print(DebugContext, DownsampledScreenWeights, FontValue);
				#endif
					Newline(DebugContext);
					Print(DebugContext, TEXT("LightId    | Weight  | Diffuse                 | Specular"), FontSilver); 
				}

				if (bValidPixel)
				{
					DiffuseLighting = 0.0f;
					SpecularLighting = 0.0f;

					// Load samples into registers
					uint PackedLightSamples[NUM_SAMPLES_PER_PIXEL_1D];
					uint PackedTransmittanceMasks[NUM_SAMPLES_PER_PIXEL_1D];
					for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_PIXEL_1D; ++SampleIndex)
					{
						uint2 LightSampleCoord = DownsampledScreenCoord * uint2(NUM_SAMPLES_PER_PIXEL_2D_X, NUM_SAMPLES_PER_PIXEL_2D_Y) + uint2(SampleIndex % NUM_SAMPLES_PER_PIXEL_2D_X, SampleIndex / NUM_SAMPLES_PER_PIXEL_2D_X);				
						PackedLightSamples[SampleIndex] = LightSamples[LightSampleCoord];
						#if INPUT_TYPE == INPUT_TYPE_HAIRSTRANDS
						PackedTransmittanceMasks[SampleIndex] = HairTransmittanceMaskTexture[LightSampleCoord];
						#endif
					}

					uint NextLocalLightIndex = MAX_LOCAL_LIGHT_INDEX;
					FindNextLocalLightIndex(PackedLightSamples, NextLocalLightIndex);

					while (NextLocalLightIndex < MAX_LOCAL_LIGHT_INDEX)
					{
						uint LocalLightIndex = NextLocalLightIndex;
						NextLocalLightIndex = MAX_LOCAL_LIGHT_INDEX;

						float SampleWeight = 0.0f;
						uint ValidSampleMask = 0;
						AccumulateLightSample(PackedLightSamples, LocalLightIndex, NextLocalLightIndex, SampleWeight, WeightRatioSum, ValidSampleMask, DebugContext);

						// Normalize weight by a number of samples taken
						SampleWeight /= float(NUM_SAMPLES_PER_PIXEL_1D);

						SampleWeight = min(SampleWeight, MaxShadingWeight);

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

						if ((Material.LightingChannelMask & UnpackLightingChannelMask(ForwardLightData)) != 0)
						{
							if (UseLightFunctionAtlas == 0)
							{
								LightData.LightFunctionAtlasLightIndex = 0;
							}

							if (UseIESProfiles == 0)
							{
								LightData.IESAtlasIndex = -1;
							}
						
							if (!IsRectLightTileType(TILE_TYPE))
							{
								LightData.bRectLight = false;
							}

							if (!IsTexturedLightTileType(TILE_TYPE))
							{
								LightData.RectLightData.AtlasData.AtlasMaxLevel = MAX_RECT_ATLAS_MIP;
							}
						
							if (SampleWeight > 0.01f && LightData.IESAtlasIndex >= 0)
							{
								SampleWeight *= ComputeLightProfileMultiplier(TranslatedWorldPosition, LightData.TranslatedWorldPosition, -LightData.Direction, LightData.Tangent, LightData.IESAtlasIndex);
							}

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

							if (SampleWeight > 0.0f)
							{
								float4 LightAttenuation = 1.0f;
								float Dither = 0.5f;
								float SurfaceShadow = 1;
								float AmbientOcclusion = ScreenSpaceAO;
								LightData.ShadowedBits = 0;

								#if INPUT_TYPE == INPUT_TYPE_HAIRSTRANDS
								for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_PIXEL_1D; ++SampleIndex)
								{
									if ((ValidSampleMask & (1u << SampleIndex)) != 0)
									{
										const FHairTransmittanceMask HairTransmittanceMask = UnpackTransmittanceMask(PackedTransmittanceMasks[SampleIndex]); // TODO_HAIR_ML: handle multi light per pixel 
										LightData.HairTransmittance = GetTransmittanceDataFromTransmitttanceMask(LightData, Material, HairTransmittanceMask, TranslatedWorldPosition, CameraVector);
										break;
									}
								}
								#elif USE_HAIR_COMPLEX_TRANSMITTANCE
								if (Material.bNeedsComplexTransmittance)
								{
									LightData.HairTransmittance = EvaluateDualScattering(Material.DiffuseColor, Material.WorldNormal, Material.Roughness, CameraVector, -LightData.Direction);
								}
								#endif

								FDeferredLightingSplit SplitLighting = GetMegaLightsSplitLighting(
									TranslatedWorldPosition, CameraVector, Material, AmbientOcclusion, 
									LightData, LightAttenuation, Dither, ScreenCoord, 
									SurfaceShadow);

								DiffuseLighting += SplitLighting.DiffuseLighting.xyz * SampleWeight;
								SpecularLighting += SplitLighting.SpecularLighting.xyz * SampleWeight;

								if (DebugContext.bIsActive)
								{
									Print(DebugContext, SplitLighting.DiffuseLighting.xyz, FontValue, 8, 3);
									Print(DebugContext, SplitLighting.SpecularLighting.xyz, FontValue, 8, 3);
								}
							}				
						}
					}

					// Apply pre-exposure
					DiffuseLighting *= View.PreExposure;
					SpecularLighting *= View.PreExposure;

					#if INPUT_TYPE != INPUT_TYPE_HAIRSTRANDS
					{
						FDenoisingModulateFactors Factors = GetDenoisingModulateFactors(Material, TranslatedWorldPosition);
						DiffuseLighting /= Factors.Diffuse;
						SpecularLighting /= Factors.Specular;
					}
					#endif
				}
			}

			// Normalize weight by a number of samples taken
			WeightRatioSum /= float(NUM_SAMPLES_PER_PIXEL_1D);

			#if INPUT_TYPE == INPUT_TYPE_HAIRSTRANDS
			if (ShadingSampleIndex == 0)
			#endif
			{
				#if REFERENCE_MODE
				//This is not the first shading pass, hence we try to accumulate from previous one
				if (ShadingPassIndex > 0)
				{
					float3 PrevDiffuseLighting = RWResolvedDiffuseLighting[ScreenCoord];
					float3 PrevSpecularLighting = RWResolvedSpecularLighting[ScreenCoord];
					float PrevWeightRatioSum = RWShadingConfidence[ScreenCoord];
					const bool PrevLightingValid = all(PrevDiffuseLighting != INVALID_LIGHTING) && all(PrevSpecularLighting != INVALID_LIGHTING);
					const bool LightingValid = all(DiffuseLighting != INVALID_LIGHTING) && all(SpecularLighting != INVALID_LIGHTING);

					float BlendFactor = 1.0f / float(ShadingPassIndex + 1.0f); //True as all passes have the same number of samples
					BlendFactor = PrevLightingValid ? BlendFactor : 1.0f;
					BlendFactor = LightingValid ? BlendFactor : 0.0f;

					DiffuseLighting = PrevDiffuseLighting * (1.0f - BlendFactor) + DiffuseLighting * BlendFactor;
					SpecularLighting = PrevSpecularLighting * (1.0f - BlendFactor) + SpecularLighting * BlendFactor;
					WeightRatioSum = PrevWeightRatioSum * (1.0f - BlendFactor) + WeightRatioSum * BlendFactor;

					if (DebugContext.bIsActive)
					{
						Newline(DebugContext);
						Print(DebugContext, TEXT("DiffusePrev         : "));
						Print(DebugContext, PrevDiffuseLighting, FontValue);
						Newline(DebugContext);
						Print(DebugContext, TEXT("DiffuseAcc          : "));
						Print(DebugContext, DiffuseLighting, FontValue);
						Newline(DebugContext);
						Print(DebugContext, TEXT("SpecularPrev        : "));
						Print(DebugContext, PrevSpecularLighting, FontValue);
						Newline(DebugContext);
						Print(DebugContext, TEXT("SpecularAcc         : "));
						Print(DebugContext, SpecularLighting, FontValue);
						Newline(DebugContext);
						Print(DebugContext, TEXT("Blend: "));
						Print(DebugContext, BlendFactor, FontValue);
					}
				}
				#endif

				RWResolvedDiffuseLighting[ScreenCoord] = DiffuseLighting;
				RWResolvedSpecularLighting[ScreenCoord] = SpecularLighting;
				RWShadingConfidence[ScreenCoord] = UseShadingConfidence != 0 ? WeightRatioSum : 0.0f;
			}

			if (DebugContext.bIsActive)
			{
				Newline(DebugContext);
				Print(DebugContext, TEXT("Diffuse             : "));
				Print(DebugContext, DiffuseLighting, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("Specular            : "));
				Print(DebugContext, SpecularLighting, FontValue);
				Newline(DebugContext);
				Print(DebugContext, TEXT("WeightRatioSum      : "));
				Print(DebugContext, WeightRatioSum, FontValue);
				AddTextBackground(DebugContext, FontBackground);
			}

			#if INPUT_TYPE == INPUT_TYPE_HAIRSTRANDS
			if (Material.IsValid() && bSubPixelShading > 0)
			{
				RWOutputColor[Material.GlobalIndex2D] = RWOutputColor[Material.GlobalIndex2D] + DiffuseLighting * Material.Coverage;
			}
			#endif
		}
	}
}

/**
 * Clear some data for empty tiles, which won't be processed by ShadeLightSamplesCS
 */
[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void ClearResolvedLightingCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	uint TileIndex = GroupId.x;
	if (TileIndex < TileAllocator[TILE_MODE_EMPTY])
	{
		uint2 TileCoord = UnpackTile(TileData[TileIndex + TILE_MODE_EMPTY * TileDataStride]);
		uint2 ScreenCoord = TileCoord * TILE_SIZE + GroupThreadId.xy + View.ViewRectMinAndSize.xy;

		if (all(ScreenCoord < View.ViewRectMinAndSize.xy + View.ViewRectMinAndSize.zw))
		{
			RWResolvedDiffuseLighting[ScreenCoord] = INVALID_LIGHTING;
			RWResolvedSpecularLighting[ScreenCoord] = INVALID_LIGHTING;
		}
	}
}