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

/*=============================================================================
	LumenRadiosityProbeGather.ush
=============================================================================*/

#pragma once

#include "../../MonteCarlo.ush"
#include "../../BlueNoise.ush"

// Must match LumenRadiosityProbeGather.cpp
#define RADIOSITY_ATLAS_DOWNSAMPLE_FACTOR 1

StructuredBuffer<uint> CardTileAllocator;
StructuredBuffer<uint> CardTileData;

StructuredBuffer<uint> RadiosityTileAllocator;
StructuredBuffer<uint> RadiosityTileData;

// Probe trace radiance in hemisphere layout
Texture2D<float3> TraceRadianceAtlas;

// Probe trace hit distance
Texture2D<float> TraceHitDistanceAtlas;

// Number of texels between Radiosity probes, always a power of two
uint ProbeSpacingInCardTexels;

// Shift to apply to divide by ProbeSpacingInCardTexels
uint ProbeSpacingInCardTexelsDivideShift;

// CARD_TILE_SIZE / ProbeSpacingInCardTexels
uint CardTileSizeInProbes;

// CardTileSizeInProbes <= 2 ? 1 : CardTileSizeInProbes
uint RadiosityTileSizeInProbes;

// Resolution of the hemisphere trace layout in one dimension
uint HemisphereProbeResolution;

// HemisphereProbeResolution * HemisphereProbeResolution
uint NumTracesPerProbe;

float ProbeOcclusionStrength;

// When >= 0, specifies a fixed temporal index for debugging
int FixedJitterIndex;

// Size of the Radiosity physical texture
uint2 RadiosityAtlasSize;

// Length of the temporal jitter pattern, and controls the weight of the temporal accumulation history
uint MaxFramesAccumulated;

// Number of views in the renderer
uint NumViews;

// Index of the current view for the dispatch
uint ViewIndex;

// Maximum number of radiosity tiles for one view
uint MaxRadiosityTiles;

typedef FCardTileData FRadiosityTileData;

uint PackRadiosityTileData(FRadiosityTileData TileData)
{
	uint Packed = TileData.CardPageIndex;
	Packed |= (TileData.TileCoord.x << 22u);
	Packed |= (TileData.TileCoord.y << 27u);
	return Packed;
}

FRadiosityTileData UnpackRadiosityTileData(uint Packed)
{
	FRadiosityTileData TileData;
	TileData.CardPageIndex = Packed & 0x003FFFFF;
	TileData.TileCoord.x = BitFieldExtractU32(Packed, 5, 22);
	TileData.TileCoord.y = BitFieldExtractU32(Packed, 5, 27);
	return TileData;
}

uint GetRadiosityTileCount(uint ViewIndex)
{
	return RadiosityTileAllocator[ViewIndex];
}

// Returns the jitter offset in the range [0, ProbeSpacingInRadiosityTexels - 1]
uint2 GetProbeJitter(uint IndirectLightingTemporalIndex)
{
	uint TemporalIndex = (FixedJitterIndex < 0 ? IndirectLightingTemporalIndex : FixedJitterIndex);
	return Hammersley16(TemporalIndex % MaxFramesAccumulated, MaxFramesAccumulated, uint2(0x4ae4, 0x9bdb)) * ProbeSpacingInCardTexels;
}

FRadiosityTileData GetRadiosityTile(uint TileIndex)
{
	return UnpackRadiosityTileData(RadiosityTileData[ViewIndex * MaxRadiosityTiles + TileIndex]);
}

void UnswizzleRadiosityTileIndex(
	uint RadiosityProbeIndex,
	inout uint RadiosityTileIndex,
	inout uint2 CoordInRadiosityTile)
{
	uint NumProbesPerTile = RadiosityTileSizeInProbes * RadiosityTileSizeInProbes;
	RadiosityTileIndex = RadiosityProbeIndex / NumProbesPerTile;
	uint LinearIndexInRadiosityTile = RadiosityProbeIndex - RadiosityTileIndex * NumProbesPerTile;
	uint2 ProbeCoord = uint2(LinearIndexInRadiosityTile % RadiosityTileSizeInProbes, LinearIndexInRadiosityTile / RadiosityTileSizeInProbes);

	FRadiosityTileData RadiosityTile = GetRadiosityTile(RadiosityTileIndex);
	FLumenCardPageData CardPage = GetLumenCardPageData(RadiosityTile.CardPageIndex);
	CoordInRadiosityTile = ProbeCoord * ProbeSpacingInCardTexels + GetProbeJitter(CardPage.IndirectLightingTemporalIndex);
}

void UnswizzleTexelTraceCoords(
	uint DispatchThreadId,
	inout uint RadiosityTileIndex,
	inout uint2 CoordInRadiosityTile,
	inout uint2 TraceTexelCoord)
{
	uint RadiosityProbeIndex = DispatchThreadId / NumTracesPerProbe;

	UnswizzleRadiosityTileIndex(RadiosityProbeIndex, RadiosityTileIndex, CoordInRadiosityTile);

	uint LinearTexelIndex = DispatchThreadId - RadiosityProbeIndex * NumTracesPerProbe;
	TraceTexelCoord = uint2(LinearTexelIndex % HemisphereProbeResolution, LinearTexelIndex / HemisphereProbeResolution);
}

struct FRadiosityTexel
{
	bool bInsideAtlas;
	bool bHeightfield;
	bool bValid;

	float3 WorldPosition;
	float3 WorldNormal;
	float3x3 WorldToLocalRotation;

	uint2 AtlasCoord;
	uint2 CardCoord;
	uint IndirectLightingTemporalIndex;
};

FRadiosityTexel GetRadiosityTexel(FLumenCardPageData CardPage, uint2 CoordInCardPage)
{
	FRadiosityTexel RadiosityTexel = (FRadiosityTexel)0;
	RadiosityTexel.bInsideAtlas = false;
	RadiosityTexel.bHeightfield = false;
	RadiosityTexel.bValid = false;
	RadiosityTexel.WorldPosition = float3(0.0f, 0.0f, 0.0f);
	RadiosityTexel.WorldNormal = float3(0.0f, 0.0f, 0.0f);

	FLumenCardData Card = GetLumenCardData(CardPage.CardIndex);
	float2 AtlasUV = CardPage.PhysicalAtlasUVRect.xy + CardPage.PhysicalAtlasUVTexelScale * (CoordInCardPage + 0.5f * RADIOSITY_ATLAS_DOWNSAMPLE_FACTOR);
	float2 CardUV = CardPage.CardUVRect.xy + CardPage.CardUVTexelScale * (CoordInCardPage + 0.5f * RADIOSITY_ATLAS_DOWNSAMPLE_FACTOR);

	if (all(CoordInCardPage < (uint2)CardPage.SizeInTexels))
	{
		RadiosityTexel.bInsideAtlas = true;
		RadiosityTexel.bHeightfield = Card.bHeightfield;
		RadiosityTexel.AtlasCoord = AtlasUV * RadiosityAtlasSize;
		RadiosityTexel.CardCoord = (CardPage.CardUVRect.xy * CardPage.ResLevelSizeInTiles) * CARD_TILE_SIZE + CoordInCardPage;
		RadiosityTexel.IndirectLightingTemporalIndex = CardPage.IndirectLightingTemporalIndex;
		RadiosityTexel.WorldToLocalRotation = Card.WorldToLocalRotation;

		#if RADIOSITY_ATLAS_DOWNSAMPLE_FACTOR == 2
		{
			float4 TexelDepths = LumenCardScene.DepthAtlas.Gather(GlobalPointClampedSampler, AtlasUV, 0.0f);
			float4 NormalX4 = LumenCardScene.NormalAtlas.GatherRed(GlobalPointClampedSampler, AtlasUV);
			float4 NormalY4 = LumenCardScene.NormalAtlas.GatherGreen(GlobalPointClampedSampler, AtlasUV);

			float NumValidTexels = 0.0f;
			float DepthSum = 0.0f;
			float2 EncodedNormalSum = float2(0.0f, 0.0f);
			for (uint TexelIndex = 0; TexelIndex < 4; ++TexelIndex)
			{
				if (IsSurfaceCacheDepthValid(TexelDepths[TexelIndex]))
				{
					NumValidTexels += 1.0f;
					DepthSum += DepthData[TexelIndex].Depth;
					EncodedNormalSum.x += NormalX4[TexelIndex];
					EncodedNormalSum.y += NormalY4[TexelIndex];
				}
			}

			if (NumValidTexels > 0.0f)
			{
				float AvgDepth = DepthSum / NumValidTexels;
				float2 AvgEncodedNormal = EncodedNormalSum / NumValidTexels;

				RadiosityTexel.bValid = true;
				RadiosityTexel.WorldPosition = GetCardWorldPosition(Card, CardUV, AvgDepth);
				RadiosityTexel.WorldNormal = DecodeSurfaceCacheNormal(Card, AvgEncodedNormal);
			}
		}
		#else
		{
			FLumenSurfaceCacheData SurfaceCacheData = GetSurfaceCacheData(Card, CardUV, AtlasUV);
			RadiosityTexel.bValid = SurfaceCacheData.bValid;
			RadiosityTexel.WorldPosition = SurfaceCacheData.WorldPosition;
			RadiosityTexel.WorldNormal = SurfaceCacheData.WorldNormal;
		}
		#endif
	}

	return RadiosityTexel;
}

FRadiosityTexel GetRadiosityTexelFromRadiosityTile(uint RadiosityTileIndex, uint2 CoordInRadiosityTile)
{
	FRadiosityTexel RadiosityTexel = (FRadiosityTexel)0;
	RadiosityTexel.bInsideAtlas = false;

	if (RadiosityTileIndex < GetRadiosityTileCount(ViewIndex))
	{
		FRadiosityTileData RadiosityTile = GetRadiosityTile(RadiosityTileIndex);
		uint RadiosityTileSize = RadiosityTileSizeInProbes * ProbeSpacingInCardTexels;
		uint2 CoordInCardPage = RadiosityTile.TileCoord * RadiosityTileSize + CoordInRadiosityTile;

		FLumenCardPageData CardPage = GetLumenCardPageData(RadiosityTile.CardPageIndex);
		RadiosityTexel = GetRadiosityTexel(CardPage, CoordInCardPage);
	}

	return RadiosityTexel;
}

// Coord in persistent radiosity probe atlas
uint2 GetRadiosityProbeAtlasCoord(FLumenCardPageData CardPage, FRadiosityTileData RadiosityTile, uint2 CoordInRadiosityTile)
{
	uint RadiosityTileSize = RadiosityTileSizeInProbes * ProbeSpacingInCardTexels;
	uint2 AtlasCoord = CardPage.PhysicalAtlasCoord + RadiosityTile.TileCoord * RadiosityTileSize + CoordInRadiosityTile;
	return AtlasCoord >> ProbeSpacingInCardTexelsDivideShift;
}

float2 GetProbeTexelCenter(uint IndirectLightingTemporalIndex, uint2 ProbeTileCoord)
{
#define JITTER_RAY_DIRECTION 1
#if JITTER_RAY_DIRECTION
	uint TemporalIndex = (FixedJitterIndex < 0 ? IndirectLightingTemporalIndex : FixedJitterIndex);
	#define BLUE_NOISE 1
	#if BLUE_NOISE
		return BlueNoiseVec2(ProbeTileCoord, TemporalIndex); 
	#else
		uint2 RandomSeed = Rand3DPCG16(int3(ProbeTileCoord, 0)).xy;
		return Hammersley16(TemporalIndex % MaxFramesAccumulated, MaxFramesAccumulated, RandomSeed);
	#endif
#else
	return float2(0.5, 0.5);
#endif
}

float2 CosineSampleHemisphereInverseFast(float3 Vector)
{
	float CosTheta = Vector.z;
	float SinTheta = sqrt(1 - CosTheta * CosTheta);
	float Phi = atan2Fast(-Vector.y, -Vector.x) + PI;

	float2 E;
	E.x = Phi / (2 * PI);
	E.y = Vector.z * Vector.z;
	return E;
}

float2 UniformSampleHemisphereInverseFast(float3 Vector)
{
	float CosTheta = Vector.z;
	float SinTheta = sqrt(1 - CosTheta * CosTheta);
	float Phi = atan2Fast(-Vector.y, -Vector.x) + PI;

	float2 E;
	E.x = Phi / (2 * PI);
	E.y = Vector.z;
	return E;
}

#define PROBE_HEMISPHERE_HEMI_OCTAHEDRON	0
#define PROBE_HEMISPHERE_UNIFORM			1
#define PROBE_HEMISPHERE_COSINE				2

#define RADIOSITY_PROBE_MAPPING				PROBE_HEMISPHERE_UNIFORM

void GetRadiosityRay(FRadiosityTexel RadiosityTexel, uint2 ProbeCoord, uint2 TracingTexelCoord, out float3 WorldRayDirection, out float ConeHalfAngle, out float PDF)
{
	float2 ProbeTexelCenter = GetProbeTexelCenter(RadiosityTexel.IndirectLightingTemporalIndex, ProbeCoord);
	float2 ProbeUV = (TracingTexelCoord + ProbeTexelCenter) / float(HemisphereProbeResolution);

	float3 LocalRayDirection;

	uint RadiosityProbeHemisphereMapping = RADIOSITY_PROBE_MAPPING;
	// Sample generation must match probe occlusion
	if (RadiosityProbeHemisphereMapping == PROBE_HEMISPHERE_HEMI_OCTAHEDRON)
	{
		LocalRayDirection = HemiOctahedronToUnitVector(ProbeUV * 2 - 1);
		//@todo - hemi octahedron solid angle
		PDF = 1.0 / (2 * PI);
	}
	else if (RadiosityProbeHemisphereMapping == PROBE_HEMISPHERE_UNIFORM)
	{
		float4 Sample = UniformSampleHemisphere(ProbeUV);
		LocalRayDirection = Sample.xyz;
		PDF = Sample.w;
	}
	else
	{
		float4 Sample = CosineSampleHemisphere(ProbeUV);
		LocalRayDirection = Sample.xyz;
		PDF = Sample.w;
	}

	float3x3 TangentBasis = GetTangentBasisFrisvad(RadiosityTexel.WorldNormal);
	WorldRayDirection = mul(LocalRayDirection, TangentBasis);

	ConeHalfAngle = acosFast(1.0f - 1.0f / (float)(NumTracesPerProbe));
}

float CalculateProbeVisibility(float3 WorldPositionBeingTested, FRadiosityTexel ProbeTexel, uint2 ProbeAtlasCoord)
{
	float VisibilityWeight = 1.0f;
	float3 ProbeToTexel = WorldPositionBeingTested - ProbeTexel.WorldPosition;

	if (dot(ProbeToTexel, ProbeToTexel) > .01f)
	{
		float3x3 TangentBasis = GetTangentBasisFrisvad(ProbeTexel.WorldNormal);
		float3 LocalProbeToTexel = mul(ProbeToTexel, transpose(TangentBasis));

		uint RadiosityProbeHemisphereMapping = RADIOSITY_PROBE_MAPPING;
		float2 TexelUV;

		if (RadiosityProbeHemisphereMapping == PROBE_HEMISPHERE_HEMI_OCTAHEDRON)
		{
			TexelUV = UnitVectorToHemiOctahedron(LocalProbeToTexel) * .5f + .5f;
		}
		else if (RadiosityProbeHemisphereMapping == PROBE_HEMISPHERE_UNIFORM)
		{
			TexelUV = UniformSampleHemisphereInverseFast(LocalProbeToTexel);
		}
		else
		{
			TexelUV = CosineSampleHemisphereInverseFast(LocalProbeToTexel);
		}

		TexelUV = clamp(TexelUV, 0.0f, .99f);

		uint2 RadiosityProbeTracingAtlasCoord = (ProbeAtlasCoord + TexelUV) * HemisphereProbeResolution;
		float TraceHitDistance = TraceHitDistanceAtlas[RadiosityProbeTracingAtlasCoord];

		float ProbeToTexelDistance = sqrt(dot(ProbeToTexel, ProbeToTexel));
		float TransitionScale = 1.0f / max((1.0f - ProbeOcclusionStrength) * ProbeToTexelDistance, .0001f);
		VisibilityWeight = saturate((TraceHitDistance - ProbeToTexelDistance) * TransitionScale + 1);

		// Binary test for debugging
		//VisibilityWeight = TraceHitDistance * TraceHitDistance < dot(ProbeToTexel, ProbeToTexel) ? 0 : 1;
	}

	return VisibilityWeight;
}