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

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

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;

float MaxShadingWeight;
uint UseShadingConfidence;
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, inout float MissWeightRatioSum)
{
	for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_PIXEL_1D; ++SampleIndex)
	{
		FLightSample LightSample = UnpackLightSample(PackedLightSamples[SampleIndex]);
		if (LightSample.bVisible)
		{	
			NextLocalLightIndex = min(NextLocalLightIndex, LightSample.LocalLightIndex);
		}
		else if (LightSample.LocalLightIndex != MAX_LOCAL_LIGHT_INDEX)
		{
			MissWeightRatioSum += GetLightSampleWeightRatio(LightSample);
		}
	}
}

void AccumulateLightSample(uint PackedLightSamples[NUM_SAMPLES_PER_PIXEL_1D], uint LocalLightIndex, inout uint NextLocalLightIndex, inout float SampleWeightSum, inout float WeightRatioSum, inout uint ValidSampleMask)
{
	for (uint SampleIndex = 0; SampleIndex < NUM_SAMPLES_PER_PIXEL_1D; ++SampleIndex)
	{
		FLightSample LightSample = UnpackLightSample(PackedLightSamples[SampleIndex]);
		if (LightSample.bVisible)
		{
			if (LightSample.LocalLightIndex == LocalLightIndex)
			{
				SampleWeightSum += LightSample.Weight;
				WeightRatioSum += GetLightSampleWeightRatio(LightSample);
				ValidSampleMask |= 1u << SampleIndex;
			}

			if (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;
		uint LocalLightHiMask = 0;

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

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

			#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())
			{
				DiffuseLighting = 0.0f;
				SpecularLighting = 0.0f;

				float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, Material.Depth);
				float3 CameraVector = normalize(TranslatedWorldPosition - View.TranslatedWorldCameraOrigin);

#if	DOWNSAMPLE_FACTOR == 2
				// Prepare weights for upsampling
				float4 DownsampledScreenWeights = 0.0f;
				{
					int2 DownsampledScreenCoord00 = floor(ScreenUV * View.BufferSizeAndInvSize.xy / DOWNSAMPLE_FACTOR - 0.5f);
					float2 DownsampledGatherUV = (DownsampledScreenCoord00 + 1.0f) * DownsampledBufferInvSize;

					float4 CornerDepths = DownsampledSceneDepth.GatherRed(GlobalPointClampedSampler, DownsampledGatherUV).wzxy;

					int2 ScreenCoordOffset = ScreenCoord - DownsampledScreenCoord00 * 2;
					int2 SampleOffset00 = GetSampleScreenCoordJitter(DownsampledScreenCoord00 + uint2(0, 0)) + uint2(0, 0) * 2 - ScreenCoordOffset;
					int2 SampleOffset10 = GetSampleScreenCoordJitter(DownsampledScreenCoord00 + uint2(1, 0)) + uint2(1, 0) * 2 - ScreenCoordOffset;
					int2 SampleOffset01 = GetSampleScreenCoordJitter(DownsampledScreenCoord00 + uint2(0, 1)) + uint2(0, 1) * 2 - ScreenCoordOffset;
					int2 SampleOffset11 = GetSampleScreenCoordJitter(DownsampledScreenCoord00 + uint2(1, 1)) + uint2(1, 1) * 2 - ScreenCoordOffset;

					// Triangle filter weights between pixel and 4 samples
					float4 InterpolationWeights;
					InterpolationWeights.x = (2.0f - abs(SampleOffset00.x)) * (2.0f - abs(SampleOffset00.y));
					InterpolationWeights.y = (2.0f - abs(SampleOffset10.x)) * (2.0f - abs(SampleOffset10.y));
					InterpolationWeights.z = (2.0f - abs(SampleOffset01.x)) * (2.0f - abs(SampleOffset01.y));
					InterpolationWeights.w = (2.0f - abs(SampleOffset11.x)) * (2.0f - abs(SampleOffset11.y));

					float4 DepthWeights = 0.0f;	
					{
						float4 ScenePlane = float4(Material.WorldNormalForPositionBias, dot(TranslatedWorldPosition, Material.WorldNormalForPositionBias));

						float3 Position00 = GetTranslatedWorldPositionFromScreenUV(DownsampledScreenCoordToScreenUV(DownsampledScreenCoord00 + uint2(0, 0)), CornerDepths.x);
						float3 Position10 = GetTranslatedWorldPositionFromScreenUV(DownsampledScreenCoordToScreenUV(DownsampledScreenCoord00 + uint2(1, 0)), CornerDepths.y);
						float3 Position01 = GetTranslatedWorldPositionFromScreenUV(DownsampledScreenCoordToScreenUV(DownsampledScreenCoord00 + uint2(0, 1)), CornerDepths.z);
						float3 Position11 = GetTranslatedWorldPositionFromScreenUV(DownsampledScreenCoordToScreenUV(DownsampledScreenCoord00 + uint2(1, 1)), CornerDepths.w);

						float4 PlaneDistances;
						PlaneDistances.x = abs(dot(float4(Position00, -1), ScenePlane));
						PlaneDistances.y = abs(dot(float4(Position10, -1), ScenePlane));
						PlaneDistances.z = abs(dot(float4(Position01, -1), ScenePlane));
						PlaneDistances.w = abs(dot(float4(Position11, -1), ScenePlane));

						float4 RelativeDepthDifference = PlaneDistances / Material.Depth;

						DepthWeights = select(CornerDepths > 0.0f, exp2(-10000.0f * (RelativeDepthDifference * RelativeDepthDifference)), 0.0f);
					}
					InterpolationWeights *= DepthWeights;

					float4 NormalWeights = 1.0f;
					{
						NormalWeights.x = GetNormalWeight(Material.WorldNormal, DownsampledScreenCoord00 + uint2(0, 0));
						NormalWeights.y = GetNormalWeight(Material.WorldNormal, DownsampledScreenCoord00 + uint2(1, 0));
						NormalWeights.z = GetNormalWeight(Material.WorldNormal, DownsampledScreenCoord00 + uint2(0, 1));
						NormalWeights.w = GetNormalWeight(Material.WorldNormal, DownsampledScreenCoord00 + uint2(1, 1));
					}
					InterpolationWeights *= NormalWeights;

					// Skip out of view samples
					InterpolationWeights.x = IsValidDownsampledCoord(DownsampledScreenCoord00.xy + uint2(0, 0)) ? InterpolationWeights.x : 0.0f;
					InterpolationWeights.y = IsValidDownsampledCoord(DownsampledScreenCoord00.xy + uint2(1, 0)) ? InterpolationWeights.y : 0.0f;
					InterpolationWeights.z = IsValidDownsampledCoord(DownsampledScreenCoord00.xy + uint2(0, 1)) ? InterpolationWeights.z : 0.0f;
					InterpolationWeights.w = IsValidDownsampledCoord(DownsampledScreenCoord00.xy + uint2(1, 1)) ? InterpolationWeights.w : 0.0f;

					DownsampledScreenWeights = InterpolationWeights;
				}

				const float RandomScalar = BlueNoiseScalar(ScreenCoord, MegaLightsStateFrameIndex);
				const uint2 StochasticBilinearOffset = GetStochasticBilinearOffset(RandomScalar, DownsampledScreenWeights);
#endif

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

#if DOWNSAMPLE_FACTOR == 2
				int2 DownsampledScreenCoord = (int2(ScreenCoord) - 1) / DOWNSAMPLE_FACTOR + StochasticBilinearOffset;
				DownsampledScreenCoord = clamp(DownsampledScreenCoord, int2(DownsampledViewMin), int2(DownsampledViewMin) + int2(DownsampledViewSize) - 1);
#else
				uint2 DownsampledScreenCoord = ScreenCoord;
#endif

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

				float MissWeightRatioSum = 0.0f;
				uint NextLocalLightIndex = MAX_LOCAL_LIGHT_INDEX;
				FindNextLocalLightIndex(PackedLightSamples, NextLocalLightIndex, MissWeightRatioSum);
				WeightRatioSum -= MissWeightRatioSum;

				while (NextLocalLightIndex < MAX_LOCAL_LIGHT_INDEX)
				{
					const uint LocalLightIndex = WaveActiveMin(NextLocalLightIndex);
					if (LocalLightIndex == NextLocalLightIndex)
					{
						NextLocalLightIndex = MAX_LOCAL_LIGHT_INDEX;
						float SampleWeight = 0.0f;
						uint ValidSampleMask = 0;
						AccumulateLightSample(PackedLightSamples, LocalLightIndex, NextLocalLightIndex, SampleWeight, WeightRatioSum, ValidSampleMask);
						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;
									}
								}
								#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);
									Print(DebugContext, SplitLighting.SpecularLighting.xyz, FontValue);
								}
							}				
						}
					}
				}

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

				#if INPUT_TYPE != INPUT_TYPE_HAIRSTRANDS
				DemodulateLighting(Material, TranslatedWorldPosition, DiffuseLighting, SpecularLighting);
				#endif
			}

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

			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);
			}

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

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

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