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

#pragma once

#include "../MonteCarlo.ush"
#include "../BlueNoise.ush"
#include "LumenPosition.ush"
#include "LumenMaterial.ush"
#include "LumenReflectionsCombine.ush"

// Change this to force recompilation of all Lumen reflection shaders
#pragma message("UESHADERMETADATA_VERSION A2EBD51-AD20-4353-B91A-4EE5FF02A797")

#define REFLECTION_THREADGROUP_SIZE_2D 8 
#define REFLECTION_THREADGROUP_SIZE_1D 64

#define FLOAT_32_MAX float(+3.402823466e+38f)

uint2 ReflectionTracingBufferSize;

// Downsample factor from full res to tracing res
uint2 ReflectionDownsampleFactorXY;
float MaxRayIntensity;
float ReflectionSmoothBias;

// 0: Opaque, 1: SingleLayerWater, 2: FrontLayerTranslucency
uint ReflectionPass;
uint UseJitter;
uint UseHighResSurface;

// Note: negative if invalid ray
Texture2DArray DownsampledDepth;
Texture2DArray<float4> RayBuffer;
// The distance to trace each ray before interpolating from the Radiance Cache
Texture2DArray<uint> RayTraceDistance;

Texture2DArray<float3> TraceRadiance;
Texture2DArray<float> TraceHit;

uint ReflectionsStateFrameIndex;
uint ReflectionsStateFrameIndexMod8;
uint ReflectionsRayDirectionFrameIndex;

/** 
 * Returns sample jitter offset in the range [0, ReflectionDownsampleFactorXY - 1]
 */
float2 GetScreenTileJitter(uint2 DownsampledScreenCoord)
{
	if (ReflectionDownsampleFactorXY.x > 1)
	{
		if (ReflectionDownsampleFactorXY.y > 1)
		{
			uint2 CellIndex = DownsampledScreenCoord % 2;
			uint LinearIndex = CellIndex.x + CellIndex.y * 2;
			LinearIndex = (LinearIndex + ReflectionsStateFrameIndex) % 4;

			// 4 rooks sampling pattern
			uint2 Jitter;
			Jitter.x = LinearIndex & 0x02 ? 1 : 0;
			Jitter.y = LinearIndex & 0x01 ? 0 : 1;

			return Jitter;
		}
		else
		{
			return float2((DownsampledScreenCoord.y + ReflectionsStateFrameIndex) % 2, 0);
		}
	}
	else
	{
		return 0;	
	}
}

float2 GetScreenUVFromReflectionTracingCoord(uint2 ReflectionTracingCoord)
{
	float2 ScreenCoord = ReflectionTracingCoord * ReflectionDownsampleFactorXY + GetScreenTileJitter(ReflectionTracingCoord) + 0.5f;
	// ScreenUV can be outside of valid viewport, since viewport is downsampled with DivideAndRoundUp
	ScreenCoord = min(ScreenCoord, View.ViewRectMin.xy + View.ViewSizeAndInvSize.xy - 1.0f);
	return ScreenCoord * View.BufferSizeAndInvSize.zw;
}

float GetMaxHitDistance()
{
	return MaxHalfFloat;
}

float EncodeRayDistance(float HitDistance, bool bHit)
{
	HitDistance = max(HitDistance, 0.0f);
	return HitDistance * (bHit ? -1.0f : 1.0f);
}

float DecodeRayDistance(float Encoded, out bool bHit)
{
	bHit = asuint(Encoded) & 0x80000000;
	return abs(Encoded);
}

float DecodeProbeRayDistance(float Encoded)
{
	return abs(Encoded);
}

struct FRayData
{
	float3 Direction;
	float PDF;
	float ConeHalfAngle;

	// Radiance cache max trace distance. Don't need to load this data if radiance cache isn't used.
	float RadianceCacheMaxTraceDistance;
	bool bUseRadianceCache;
	bool bIsFirstPersonPixel;
};

// Must be larger than FLOAT16_MIN
const static float MinReflectionConeAngle = 0.00001f;

uint PackRayTraceDistance(float RayTraceDistance, bool bUseRadianceCache)
{
	uint PackedData =
		(f32tof16(min(RayTraceDistance, GetMaxHitDistance())) & 0x7FFF) |
		(bUseRadianceCache << 15);
	return PackedData;
}

void UnpackRayTraceDistance(uint PackedData, inout float RadianceCacheMaxTraceDistance, inout bool bUseRadianceCache)
{
	RadianceCacheMaxTraceDistance = f16tof32(PackedData & 0x7FFF);
	bUseRadianceCache = (PackedData >> 15) != 0;
}

FRayData GetRayData(uint3 TracingCoord)
{
	float4 PackedRayData = RayBuffer[TracingCoord];

	FRayData RayData;
	RayData.Direction = PackedRayData.xyz;
	RayData.ConeHalfAngle = abs(PackedRayData.w);
	RayData.PDF = 1.0f / abs(PackedRayData.w);
	RayData.bIsFirstPersonPixel = (PackedRayData.w < 0.0f);

	UnpackRayTraceDistance(RayTraceDistance[TracingCoord], RayData.RadianceCacheMaxTraceDistance, RayData.bUseRadianceCache);

	return RayData;
}

Buffer<uint> ReflectionClearTileData;
Buffer<uint> ReflectionResolveTileData;

struct FReflectionTileData
{
	uint2 Coord;
	uint ClosureIndex;
};

uint PackTileData(FReflectionTileData In)
{
	return 
	#if SUBTRATE_GBUFFER_FORMAT==1
		((In.ClosureIndex & 0x7) << 24) |
	#endif
		PackTileCoord12bits(In.Coord);
}

FReflectionTileData UnpackTileData(uint In)
{
	FReflectionTileData Out;
	Out.Coord 	  = UnpackTileCoord12bits(In);
	Out.ClosureIndex = SUBTRATE_GBUFFER_FORMAT == 1 ? ((In>>24) & 0x7) : 0;
	return Out;
}

uint2 GetReflectionClearScreenCoord(uint TileIndex, uint2 ThreadIndex, inout FReflectionTileData OutTileData)
{
	OutTileData = UnpackTileData(ReflectionClearTileData[TileIndex]);
	uint2 TileOffset = OutTileData.Coord;
	return TileOffset * REFLECTION_THREADGROUP_SIZE_2D + ThreadIndex + View.ViewRectMin.xy;
}

uint2 GetReflectionResolveScreenCoord(uint TileIndex, uint2 ThreadIndex, inout FReflectionTileData OutTileData)
{
	OutTileData = UnpackTileData(ReflectionResolveTileData[TileIndex]);
	uint2 TileOffset = OutTileData.Coord;
	return TileOffset * REFLECTION_THREADGROUP_SIZE_2D + ThreadIndex + View.ViewRectMin.xy;
}

Buffer<uint> ReflectionTracingTileData;

uint2 ReflectionTracingViewMin;
uint2 ReflectionTracingViewSize;

uint3 GetReflectionTracingScreenCoord(uint TileIndex, uint ThreadIndex, inout FReflectionTileData OutTileData)
{
	OutTileData = UnpackTileData(ReflectionTracingTileData[TileIndex]);
	uint2 TileOffset = OutTileData.Coord;
	return uint3(TileOffset * REFLECTION_THREADGROUP_SIZE_2D + uint2(ThreadIndex % REFLECTION_THREADGROUP_SIZE_2D, ThreadIndex / REFLECTION_THREADGROUP_SIZE_2D) + ReflectionTracingViewMin, OutTileData.ClosureIndex);
}

uint2 GetReflectionTracingScreenCoordZOrder(uint TileIndex, uint ThreadIndex, inout FReflectionTileData OutTileData)
{
	OutTileData = UnpackTileData(ReflectionTracingTileData[TileIndex]);
	uint2 TileOffset = OutTileData.Coord;
	return TileOffset * REFLECTION_THREADGROUP_SIZE_2D + ZOrder2D(ThreadIndex, log2(REFLECTION_THREADGROUP_SIZE_2D)) + ReflectionTracingViewMin;
}

Buffer<uint> CompactedTraceTexelAllocator;
Buffer<uint> CompactedTraceTexelData;

struct FTraceMaterialId
{
	bool bNeedsHitLightingPass;
	uint MaterialId;
};

uint PackTraceMaterialId(bool bNeedsHitLightingPass, uint MaterialShaderIndex)
{
	uint PackedMaterialId = MaterialShaderIndex & 0x7FFF;
	PackedMaterialId |= bNeedsHitLightingPass ? 0x8000 : false;
	return PackedMaterialId;
}

FTraceMaterialId UnpackTraceMaterialId(uint PackedMaterialId)
{
	FTraceMaterialId TraceMaterialId;
	TraceMaterialId.bNeedsHitLightingPass = PackedMaterialId & 0x8000 ? true : false;
	TraceMaterialId.MaterialId = PackedMaterialId & 0x7FFF;
	return TraceMaterialId;
}

float ApplySmoothBias(float Roughness)
{
	float NewRoughness = Roughness;

	if (ReflectionSmoothBias > 0)
	{
		// SmoothStep-like function up to SmoothBias, original value above
		float X = saturate(Roughness / ReflectionSmoothBias);
		NewRoughness = Roughness * X * X * (3.0 - 2.0 * X);
	}

	// Only opaque and translucent reflections get the denoiser, otherwise force to mirror
	return ReflectionPass == 0 || ReflectionPass == 2 ? NewRoughness : 0.0f;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Abstract coord for reflection tracing

struct FReflectionTracingCoord
{
	uint2 Coord;
	uint3 CoordFlatten;
	uint  ClosureIndex;
};

FReflectionTracingCoord GetReflectionTracingCoord(uint2 InReflectionTracingCoord, uint InClosureIndex)
{
	FReflectionTracingCoord Out;
	Out.Coord = InReflectionTracingCoord;
	Out.ClosureIndex = SUBTRATE_GBUFFER_FORMAT == 1 && SUBSTRATE_MATERIAL_CLOSURE_COUNT > 1? InClosureIndex : 0u;
	Out.CoordFlatten = uint3(Out.Coord, Out.ClosureIndex);
	return Out;
}

uint EncodeTraceTexel(uint2 ReflectionTracingCoord, uint InClosureIndex)
{
	// SUBSTRATE_CLOSURE_OFFSET_BIT_COUNT (4) bits is needed for ClosureIndex
	return PackTileCoord14bits(ReflectionTracingCoord) 
		#if SUBTRATE_GBUFFER_FORMAT == 1 && SUBSTRATE_MATERIAL_CLOSURE_COUNT > 1
		| ((InClosureIndex & 0xF) << 28)
		#endif
		;
}

FReflectionTracingCoord DecodeTraceTexel(uint TraceTexelData)
{
	const uint2 TraceCoord = UnpackTileCoord14bits(TraceTexelData);
	const uint ClosureIndex = BitFieldExtractU32(TraceTexelData, 4, 28);
	return GetReflectionTracingCoord(TraceCoord, ClosureIndex);
}

FLumenMaterialCoord InternalGetLumenMaterialCoord_Reflection(in uint2 SvPosition, FReflectionTileData InTileData, bool bForceClosureIndex, uint InForcedClosureIndex)
{
	FLumenMaterialCoord Out = (FLumenMaterialCoord)0;
	Out.ClosureIndex = SUBTRATE_GBUFFER_FORMAT == 1 && SUBSTRATE_MATERIAL_CLOSURE_COUNT > 1? InTileData.ClosureIndex : 0;
	Out.SvPosition = SvPosition;
	Out.SvPositionFlatten = uint3(Out.SvPosition, Out.ClosureIndex);

	#if SUBSTRATE_STOCHASTIC_LIGHTING_ALLOWED
	if (bForceClosureIndex)
	{
		Out.ClosureIndex = InForcedClosureIndex;
		Out.SvPositionFlatten.z = InForcedClosureIndex;
	}
	#endif
	return Out;
}

FLumenMaterialCoord GetLumenMaterialCoord_Reflection(in uint2 SvPosition, FReflectionTileData InTileData, uint InForcedClosureIndex)
{
	return InternalGetLumenMaterialCoord_Reflection(SvPosition, InTileData, Substrate.bStochasticLighting, InForcedClosureIndex);
}

FLumenMaterialCoord GetLumenMaterialCoord_Reflection(in uint2 SvPosition, FReflectionTileData InTileData)
{
	return InternalGetLumenMaterialCoord_Reflection(SvPosition, InTileData, false /*bAllowStochastic*/, 0);
}

FLumenMaterialData ApplySmoothBias(FLumenMaterialData In, bool bTopLayerRoughness)
{
	FLumenMaterialData Out = In;
	if (bTopLayerRoughness)
	{
		Out.TopLayerRoughness = ApplySmoothBias(In.TopLayerRoughness);
	}
	else
	{
		Out.Roughness = ApplySmoothBias(In.Roughness);
	}
	return Out;
}

bool NeedRayTracedReflections(float InRoughness, const FLumenMaterialData In)
{
	bool bNeedReflections = false;

	if (ReflectionPass == 0)
	{
		bNeedReflections = IsValid(In) && NeedsLumenReflections(InRoughness, In.TopLayerRoughness, HasBackfaceDiffuse(In), IsClearCoat(In));
	}
	else if (ReflectionPass == 1)
	{
		bNeedReflections = IsSingleLayerWater(In);
	}
	else
	{
		bNeedReflections = IsFrontLayerTranslucency(In);
	}

	return bNeedReflections;
}