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

#pragma once

#include "LumenSceneLighting.ush"
#include "../EyeAdaptationCommon.ush"
#include "/Engine/Shared/LightDefinitions.h"

#define NUM_BATCHABLE_LIGHT_TYPES 3

RWStructuredBuffer<uint> RWShadowMaskTiles;
StructuredBuffer<uint> ShadowMaskTiles;

struct FShadowTrace
{
	uint LightTileIndex;
	uint2 LightTileCoord;
	uint DownSampleLevel;
};

FShadowTrace UnpackShadowTrace(uint PackedValue)
{
	FShadowTrace ShadowTrace;
	ShadowTrace.LightTileIndex = PackedValue & 0xFFFFFF;
	ShadowTrace.LightTileCoord.x = BitFieldExtractU32(PackedValue, 3, 24);
	ShadowTrace.LightTileCoord.y = BitFieldExtractU32(PackedValue, 3, 27);
	ShadowTrace.DownSampleLevel = BitFieldExtractU32(PackedValue, 2, 30);
	return ShadowTrace;
}

uint PackShadowTrace(FShadowTrace ShadowTrace)
{
	uint PackedValue = 0;
	PackedValue = ShadowTrace.LightTileIndex & 0xFFFFFF;
	PackedValue |= (ShadowTrace.LightTileCoord.x & 0x7) << 24;
	PackedValue |= (ShadowTrace.LightTileCoord.y & 0x7) << 27;
	PackedValue |= (ShadowTrace.DownSampleLevel & 0x3) << 30;
	return PackedValue;
}

#define LUMEN_COLORED_LIGHT_FUNCTION_ATLAS (COLORED_LIGHT_FUNCTION_ATLAS)

#if LUMEN_COLORED_LIGHT_FUNCTION_ATLAS

// When colored light function atlas is enabled, we pack each sample as 16bits instead of 8: 5bit per color component and 1 bit for complete.
// 5 bits looks still acceptable but we can pivote to a wider representation if needed later. 
// => If this is changed, ShadowMaskTilesSizeFactor must also be updated when allocating Lumen.DirectLighting.ShadowMaskTiles.
#define SHADOW_MASK_RAY_BITS 16
#define SHADOW_MASK_RAY_BITS_MASK ((1u << SHADOW_MASK_RAY_BITS) - 1)
#define SHADOW_MASK_CARD_TILE_DWORDS (SHADOW_MASK_RAY_BITS * CARD_TILE_SIZE * CARD_TILE_SIZE / 32)
#define SHADOW_FACTOR_BITS 15
#define SHADOW_FACTOR_BITS_MASK ((1u << SHADOW_FACTOR_BITS) - 1)
#define SHADOW_FACTOR_COMPLETE_BITS 1
#define SHADOW_FACTOR_COMPLETE_BITS_MASK (((1u << SHADOW_FACTOR_COMPLETE_BITS) - 1) << SHADOW_FACTOR_BITS)

#else

#define SHADOW_MASK_RAY_BITS 8
#define SHADOW_MASK_RAY_BITS_MASK ((1u << SHADOW_MASK_RAY_BITS) - 1)
#define SHADOW_MASK_CARD_TILE_DWORDS (SHADOW_MASK_RAY_BITS * CARD_TILE_SIZE * CARD_TILE_SIZE / 32)
#define SHADOW_FACTOR_BITS 7
#define SHADOW_FACTOR_BITS_MASK ((1u << SHADOW_FACTOR_BITS) - 1)
#define SHADOW_FACTOR_COMPLETE_BITS 1
#define SHADOW_FACTOR_COMPLETE_BITS_MASK (((1u << SHADOW_FACTOR_COMPLETE_BITS) - 1) << SHADOW_FACTOR_BITS)

#endif

struct FShadowMaskRay
{
#if LUMEN_COLORED_LIGHT_FUNCTION_ATLAS
	float3 ShadowFactor;
#else
	float ShadowFactor;
#endif
	bool bShadowFactorComplete;
};

#if LUMEN_COLORED_LIGHT_FUNCTION_ATLAS
uint GetMaskFromColoredShadowFactor(float3 ShadowFactor)
{
	uint3 RGB = ShadowFactor * 31.0f;
	return (RGB.r & 0x1F) | ((RGB.g & 0x1F) << 5) | ((RGB.b & 0x1F) << 10);
}
float3 GetColoredShadowFactorFromMask(uint Mask)
{
	float3 RGB = float3(BitFieldExtractU32(Mask, 5, 0), BitFieldExtractU32(Mask, 5, 5), BitFieldExtractU32(Mask, 5, 10));
	return RGB * (1.0f / 31.0f);
}

bool IsVisibleToLight(float3 ShadowFactor)
{
	return any(ShadowFactor > 0.01f);
}
#else
bool IsVisibleToLight(float ShadowFactor)
{
	return ShadowFactor > 0.01f;
}
#endif

void ReadShadowMaskRayRW(uint CardTileIndex, uint2 CoordInCardTile, inout FShadowMaskRay ShadowMaskRay)
{
	uint BitOffset = SHADOW_MASK_RAY_BITS * (CoordInCardTile.x + CoordInCardTile.y * CARD_TILE_SIZE);

	uint ShadowMask = RWShadowMaskTiles[SHADOW_MASK_CARD_TILE_DWORDS * CardTileIndex + BitOffset / 32];
	ShadowMask = ShadowMask >> (BitOffset % 32);

#if LUMEN_COLORED_LIGHT_FUNCTION_ATLAS
	ShadowMaskRay.ShadowFactor = GetColoredShadowFactorFromMask(ShadowMask);
#else
	ShadowMaskRay.ShadowFactor = float(ShadowMask & SHADOW_FACTOR_BITS_MASK) / SHADOW_FACTOR_BITS_MASK;
#endif
	ShadowMaskRay.bShadowFactorComplete = (ShadowMask & SHADOW_FACTOR_COMPLETE_BITS_MASK) != 0;
}

void ReadShadowMaskRay(uint CardTileIndex, uint2 CoordInCardTile, inout FShadowMaskRay ShadowMaskRay)
{
	uint BitOffset = SHADOW_MASK_RAY_BITS * (CoordInCardTile.x + CoordInCardTile.y * CARD_TILE_SIZE);

	uint ShadowMask = ShadowMaskTiles[SHADOW_MASK_CARD_TILE_DWORDS * CardTileIndex + BitOffset / 32];
	ShadowMask = ShadowMask >> (BitOffset % 32);

#if LUMEN_COLORED_LIGHT_FUNCTION_ATLAS
	ShadowMaskRay.ShadowFactor = GetColoredShadowFactorFromMask(ShadowMask);
#else
	ShadowMaskRay.ShadowFactor = float(ShadowMask & SHADOW_FACTOR_BITS_MASK) / SHADOW_FACTOR_BITS_MASK;
#endif
	ShadowMaskRay.bShadowFactorComplete = (ShadowMask & SHADOW_FACTOR_COMPLETE_BITS_MASK) != 0;
}

void WriteShadowMaskRay(FShadowMaskRay Ray, uint CardTileIndex, uint2 CoordInCardTile, const bool bClearExistingMask)
{
#if LUMEN_COLORED_LIGHT_FUNCTION_ATLAS
	uint Mask = GetMaskFromColoredShadowFactor(Ray.ShadowFactor);
#else
	uint Mask = uint(Ray.ShadowFactor * SHADOW_FACTOR_BITS_MASK);
#endif

	if (Ray.bShadowFactorComplete)
	{
		Mask |= SHADOW_FACTOR_COMPLETE_BITS_MASK;
	}

	uint BitOffset = SHADOW_MASK_RAY_BITS * (CoordInCardTile.x + CoordInCardTile.y * CARD_TILE_SIZE);

	if (bClearExistingMask)
	{
		InterlockedAnd(RWShadowMaskTiles[SHADOW_MASK_CARD_TILE_DWORDS * CardTileIndex + BitOffset / 32], ~(SHADOW_MASK_RAY_BITS_MASK << (BitOffset % 32)));
	}

	if (Mask != 0)
	{
		InterlockedOr(RWShadowMaskTiles[SHADOW_MASK_CARD_TILE_DWORDS * CardTileIndex + BitOffset / 32], Mask << (BitOffset % 32));
	}
}

// Must match FLumenPackedLight in LumenSceneDirectLighting.cpp
struct FLumenPackedLight
{
	float3 WorldPositionHigh;
	uint LightingChannelMask;

	float3 WorldPositionLow;
	float InvRadius;

	float3 Color;
	float FalloffExponent;

	float3 Direction;
	uint DiffuseAndSpecularScale;

	float3 Tangent;
	float SourceRadius;

	float2 SpotAngles;
	float SoftSourceRadius;
	float SourceLength;

	float RectLightBarnCosAngle;
	float RectLightBarnLength;
	float RectLightAtlasMaxLevel;
	uint LightType;
	
	float2 SinCosConeAngleOrRectLightAtlasUVScale;
	float2 RectLightAtlasUVOffset;

	uint LightFunctionAtlasIndex_bHasShadowMask_bIsStandalone_bCastDynamicShadows;
	float IESAtlasIndex;
	float InverseExposureBlend;
	float Padding0;
};

struct FLumenLight
{
	FDeferredLightData DeferredLightData;

	float4 InfluenceSphere;
	float3 ProxyPosition;
	float3 ProxyDirection;
	float ProxyRadius;
	float CosConeAngle;
	float SinConeAngle;

	uint VirtualShadowMapId;
	uint LightingChannelMask;
	uint bHasShadowMask;
	uint bIsStandaloneLight;
	uint Type;
};

StructuredBuffer<FLumenPackedLight> LumenPackedLights;

// Storing influence spheres in a separate buffer to accelerate light culling.
// Most lights can be culled by testing only their influence spheres so no point loading other data for culled lights.
StructuredBuffer<float4> LumenLightInfluenceSpheres;

float4 LoadLumenLightInfluenceSphere(uint LightIndex)
{
	return LumenLightInfluenceSpheres[LightIndex];
}

FLumenLight LoadLumenLight(uint LightIndex, float3 PreViewTranslation, float Exposure)
{
	FLumenPackedLight PackedLight = LumenPackedLights[LightIndex];

	FDFVector3 WorldPosition = MakeDFVector3(PackedLight.WorldPositionHigh, PackedLight.WorldPositionLow);

	FDeferredLightData DeferredLightData = (FDeferredLightData)0;
	DeferredLightData.TranslatedWorldPosition = DFHackToFloat(WorldPosition) + PreViewTranslation; // LUMEN_LWC_TODO
	DeferredLightData.InvRadius = PackedLight.InvRadius;
	DeferredLightData.Color = PackedLight.Color * InverseExposureLerp(Exposure, PackedLight.InverseExposureBlend);
	DeferredLightData.FalloffExponent = PackedLight.FalloffExponent;
	DeferredLightData.Direction = PackedLight.Direction;
	DeferredLightData.Tangent = PackedLight.Tangent;
	DeferredLightData.SpotAngles = PackedLight.SpotAngles;
	DeferredLightData.SourceRadius = PackedLight.SourceRadius;
	DeferredLightData.SourceLength = PackedLight.SourceLength;
	DeferredLightData.SoftSourceRadius = PackedLight.SoftSourceRadius;	
	DeferredLightData.DiffuseScale = UnpackFloat2FromUInt(PackedLight.DiffuseAndSpecularScale).x;
	DeferredLightData.SpecularScale = UnpackFloat2FromUInt(PackedLight.DiffuseAndSpecularScale).y;
	DeferredLightData.ContactShadowLength = 0.0f;
	DeferredLightData.ContactShadowLengthInWS = false;
	DeferredLightData.DistanceFadeMAD = 0.0f;
	DeferredLightData.ShadowMapChannelMask = 0.0f;
	DeferredLightData.ShadowedBits = BitFieldExtractU32(PackedLight.LightFunctionAtlasIndex_bHasShadowMask_bIsStandalone_bCastDynamicShadows, 1, 29) != 0 ? 3 : 0;
	DeferredLightData.RectLightData.BarnCosAngle = PackedLight.RectLightBarnCosAngle;
	DeferredLightData.RectLightData.BarnLength = PackedLight.RectLightBarnLength;
	DeferredLightData.RectLightData.AtlasData.AtlasUVScale = PackedLight.SinCosConeAngleOrRectLightAtlasUVScale;
	DeferredLightData.RectLightData.AtlasData.AtlasUVOffset = PackedLight.RectLightAtlasUVOffset;
	DeferredLightData.RectLightData.AtlasData.AtlasMaxLevel = PackedLight.RectLightAtlasMaxLevel;
	DeferredLightData.bInverseSquared = PackedLight.FalloffExponent == 0.0f;
	DeferredLightData.bRadialLight = PackedLight.LightType != LIGHT_TYPE_DIRECTIONAL;
	DeferredLightData.bSpotLight = PackedLight.LightType == LIGHT_TYPE_SPOT;
	DeferredLightData.bRectLight = PackedLight.LightType == LIGHT_TYPE_RECT;
	DeferredLightData.IESAtlasIndex = PackedLight.IESAtlasIndex;
	DeferredLightData.LightFunctionAtlasLightIndex = PackedLight.LightFunctionAtlasIndex_bHasShadowMask_bIsStandalone_bCastDynamicShadows & 0x1FFFFFFF;

	FLumenLight LumenLight = (FLumenLight)0;
	LumenLight.DeferredLightData = DeferredLightData;
	LumenLight.InfluenceSphere = LoadLumenLightInfluenceSphere(LightIndex); // LUMEN_LWC_TODO: rebase to translated world
	LumenLight.ProxyPosition = DFHackToFloat(WorldPosition); // LUMEN_LWC_TODO: rebase to translated world
	LumenLight.ProxyRadius = rcp(PackedLight.InvRadius);
	LumenLight.ProxyDirection = -PackedLight.Direction;
	LumenLight.CosConeAngle = PackedLight.SinCosConeAngleOrRectLightAtlasUVScale.y;
	LumenLight.SinConeAngle = PackedLight.SinCosConeAngleOrRectLightAtlasUVScale.x;
	LumenLight.VirtualShadowMapId = 0;
	LumenLight.bHasShadowMask = PackedLight.LightFunctionAtlasIndex_bHasShadowMask_bIsStandalone_bCastDynamicShadows >> 31;
	LumenLight.bIsStandaloneLight  = BitFieldExtractU32(PackedLight.LightFunctionAtlasIndex_bHasShadowMask_bIsStandalone_bCastDynamicShadows, 1, 30);
	LumenLight.LightingChannelMask = PackedLight.LightingChannelMask;
	LumenLight.Type = PackedLight.LightType;
	return LumenLight;
}

#if LIGHT_FUNCTION
/** Fade distance in x, disabled brightness in y, output for preview shadows mask in z. */
float3 LightFunctionParameters2;

float GetLumenLightFunction(float3 TranslatedWorldPosition) 
{
	float4 LightVector = mul(float4(TranslatedWorldPosition, 1), LightFunctionTranslatedWorldToLight);
	float3 LightFunction = GetLightFunctionColor(LightVector.xyz / LightVector.w, TranslatedWorldPosition);

	float GreyScale = dot(LightFunction, .3333f).x;

	// Calculate radial view distance for stable fading
	float ViewDistance = GetDistanceToCameraFromViewVector(PrimaryView.TranslatedWorldCameraOrigin - TranslatedWorldPosition);

	float DistanceFadeAlpha = saturate((LightFunctionParameters2.x - ViewDistance) / (LightFunctionParameters2.x * .2f));
	// Fade to disabled based on LightFunctionFadeDistance
	GreyScale = lerp(LightFunctionParameters2.y, GreyScale, DistanceFadeAlpha);

	// Fade to disabled based on ShadowFadeFraction
	GreyScale = lerp(LightFunctionParameters2.y, GreyScale, LightFunctionParameters.y);

	// We want to get the same light function contribution as if it was written to a unorm8 render target so we saturate. This will also avoid amplifying other shadow factors.
	return saturate(GreyScale);
}
#endif

struct FLumenSampleCoord
{
	uint2 Coord;
	uint LayerIndex;
};

uint PackLumenSampleCoord(uint2 Coord, uint LayerIndex)
{
	return PackTileCoord12bits(Coord) | (LayerIndex << 24u);
}

uint PackLumenSampleCoord(FLumenSampleCoord In)
{
	return PackLumenSampleCoord(In.Coord, In.LayerIndex);
}

FLumenSampleCoord UnpackLumenSampleCoord(uint In)
{
	FLumenSampleCoord Out;
	Out.Coord = UnpackTileCoord12bits(In);
	Out.LayerIndex = (In >> 24u);
	return Out;
}

struct FLumenSampleSceneData
{
	float3 TranslatedWorldPosition;
	float3 WorldNormal;
	uint ViewIndex;
	bool bHeightfield;
};

float4 PackLumenSampleSceneData(float3 InPosition, float3 InNormal, uint InViewIndex, bool bHeightfield)
{
	uint Packed = 0;
	{
		const float2 Result = UnitVectorToOctahedron(InNormal);
		const uint2 PackedXY = uint2(clamp(Result * 2047.0f + 2048.0f, 0.0f, 4095.0f));
		Packed = PackedXY.x | (PackedXY.y << 12);
	}
	Packed |= InViewIndex << 24;
	Packed |= bHeightfield ? (1u<<26) : 0u;
	return float4(InPosition, asfloat(Packed));
}

FLumenSampleSceneData UnpackLumenSampleSceneData(float4 In)
{
	const uint Packed = asuint(In.w);
	const int2 XY12Bits = int2(0xFFF & Packed, 0xFFF & (Packed >> 12));
	const float2 xy = float2(XY12Bits - 2048) / 2047.0f;

	FLumenSampleSceneData Out;
	Out.TranslatedWorldPosition = In.xyz;
	Out.WorldNormal = OctahedronToUnitVector(xy);
	Out.ViewIndex = BitFieldExtractU32(Packed, 2, 24);
	Out.bHeightfield = BitFieldExtractU32(Packed, 1, 26);

	return Out;
}