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

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

uint2 InputDimensions;
float2 OneOverInputDimensions;
float2 UVRatioAdjustment;
uint4 StepSize; // per-channel
uint4 ChannelMask; // whether to calculate for this step
Buffer<float4> InputBuffer;
RWBuffer<float4> OutputBuffer;

[numthreads(TILE_SIZE, TILE_SIZE, 1)]
void MainCS(
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 GroupId : SV_GroupID)
{
	uint2 PixelCoord = DispatchThreadId.xy;
	if (PixelCoord.x > InputDimensions.x - 1 || PixelCoord.y > InputDimensions.y - 1)
	{
		return;
	}

	const uint RowLength = InputDimensions.x * NUM_CHANNELS;
	uint PixelCoord1D = (PixelCoord.y * RowLength) + (PixelCoord.x * NUM_CHANNELS);

	static uint2 NeighborOffsetIndices[] =
	{
		{-1,  1},  {0,  1},  {1,  1},
		{-1,  0},  {0,  0},  {1,  0},
		{-1, -1},  {0, -1},  {1, -1}
	};

#if NUM_CHANNELS == 1
	float MinDistance[NUM_CHANNELS] = { 100000000 };
	float4 MinSample[NUM_CHANNELS] = { float4(0,0,0,0) };
#elif NUM_CHANNELS == 2
	float MinDistance[NUM_CHANNELS] = { 100000000, 100000000 };
	float4 MinSample[NUM_CHANNELS] = { float4(0,0,0,0), float4(0,0,0,0) };
#elif NUM_CHANNELS == 3
	float MinDistance[NUM_CHANNELS] = { 100000000, 100000000, 100000000 };
	float4 MinSample[NUM_CHANNELS] = { float4(0,0,0,0), float4(0,0,0,0), float4(0,0,0,0) };
#elif NUM_CHANNELS == 4
	float MinDistance[NUM_CHANNELS] = { 100000000, 100000000, 100000000, 100000000 };
	float4 MinSample[NUM_CHANNELS] = { float4(0,0,0,0), float4(0,0,0,0), float4(0,0,0,0), float4(0,0,0,0) };
#endif

	float X = (float)PixelCoord.x + 0.5;
	float Y = (float)PixelCoord.y + 0.5;
	const float2 UV = float2(X, Y) * OneOverInputDimensions * UVRatioAdjustment;

	// Perform on each of R,G,B,A
	[unroll]
	for (int ChannelIdx = 0; ChannelIdx < NUM_CHANNELS; ++ChannelIdx)
	{
		float4 Input = InputBuffer[PixelCoord1D + ChannelIdx];
		const uint2 StepSizeWithChannel = StepSize[ChannelIdx] * uint2(1, 1);

		// Check if we should calculate this channel (for this step)
		if (ChannelMask[ChannelIdx] == 0)
		{
			MinSample[ChannelIdx] = Input[ChannelIdx];
			continue;
		}

		MinDistance[ChannelIdx] = 1000000000;
		MinSample[ChannelIdx] = float4(0,0,0,0);

		[unroll]
		for (int HorizontalIdx = 0; HorizontalIdx < 3; ++HorizontalIdx)
		{
			[unroll]
			for (int VerticalIdx = 0; VerticalIdx < 3; ++VerticalIdx)
			{
				uint NeighborIdx = (VerticalIdx * KERNEL_SIZE) + HorizontalIdx;
				uint2 NeighborPixelCoord = PixelCoord + (NeighborOffsetIndices[NeighborIdx] * StepSizeWithChannel);

				// Prevents border banding
				NeighborPixelCoord = min(InputDimensions - 1, max(uint2(1,1), NeighborPixelCoord));

				uint NeighborPixelCoord1D = (NeighborPixelCoord.y * RowLength) + (NeighborPixelCoord.x * NUM_CHANNELS);

				float4 NeighborSample = InputBuffer[NeighborPixelCoord1D + ChannelIdx];
				if (NeighborSample.x + NeighborSample.y == 0.0)
				{
					continue;
				}

				const float Distance = length(NeighborSample.xy - UV);
				if (Distance < MinDistance[ChannelIdx])
				{
					MinDistance[ChannelIdx] = Distance;
					MinSample[ChannelIdx] = NeighborSample;
				}
			}
		}

		MinSample[ChannelIdx].a = Input.a;
		OutputBuffer[PixelCoord1D + ChannelIdx] = MinSample[ChannelIdx];
	}
}