// Copyright 2020 Phyronnaz

#include "/Engine/Public/Platform.ush"
#include "/Engine/Generated/GeneratedUniformBuffers.ush"

#if __RSCPP_VERSION
struct FConstants
{
    uint SizeX;
    uint SizeY;
    uint SizeZ;
	uint Step;
};
FConstants VoxelDistanceFieldParameters;
#endif

RWBuffer<float> RWSrc;
RWBuffer<float> RWDst;

float3 GetCoord(uint3 Position)
{
	const int Index = Position.x + VoxelDistanceFieldParameters.SizeX * Position.y + VoxelDistanceFieldParameters.SizeX * VoxelDistanceFieldParameters.SizeY * Position.z;
	return float3(
		RWSrc[3 * Index + 0],
		RWSrc[3 * Index + 1],
		RWSrc[3 * Index + 2]);
}
void SetCoord(uint3 Position, float3 Coord)
{
	const int Index = Position.x + VoxelDistanceFieldParameters.SizeX * Position.y + VoxelDistanceFieldParameters.SizeX * VoxelDistanceFieldParameters.SizeY * Position.z;
	RWDst[3 * Index + 0] = Coord.x;
	RWDst[3 * Index + 1] = Coord.y;
	RWDst[3 * Index + 2] = Coord.z;
}

bool IsCoordValid(float3 Coord)
{
	return Coord.x <= 1e9;
}

[numthreads(NUM_THREADS_CS, NUM_THREADS_CS, NUM_THREADS_CS)]
void ExpandDistanceField(uint3 ThreadId : SV_DispatchThreadID)
{
    const uint3 Position = ThreadId.xyz;

    if (Position.x >= VoxelDistanceFieldParameters.SizeX || 
        Position.y >= VoxelDistanceFieldParameters.SizeY || 
        Position.z >= VoxelDistanceFieldParameters.SizeZ)
    {
        return;
	}

	float BestDistance = 1e9;
	float3 BestCoord = float3(1e9, 1e9, 1e9);

	for (int x = -1; x <= 1; ++x)
	{
		for (int y = -1; y <= 1; ++y)
		{
			for (int z = -1; z <= 1; ++z)
			{
				const int3 NeighborPosition = Position + int3(x, y, z) * VoxelDistanceFieldParameters.Step;
				
				if (NeighborPosition.x < 0 ||
					NeighborPosition.y < 0 ||
					NeighborPosition.z < 0 ||
					NeighborPosition.x >= int(VoxelDistanceFieldParameters.SizeX) ||
					NeighborPosition.y >= int(VoxelDistanceFieldParameters.SizeY) ||
					NeighborPosition.z >= int(VoxelDistanceFieldParameters.SizeZ))
				{
					// We could clamp the index below instead of jumping, but it probably has the same cost anyways
					// and produce less nice patterns in the first passes
					continue;
				}
				
				const float3 NeighborCoord = GetCoord(NeighborPosition);

				if (IsCoordValid(NeighborCoord))
				{
					const float Dist = length(NeighborCoord - float3(Position));

					if (Dist < BestDistance)
					{
						BestDistance = Dist;
						BestCoord = NeighborCoord;
					}
				}
			}
		}
	}

	SetCoord(Position, BestCoord);
}