layout(set=0, binding=0) writeonly buffer destBuffer{
    FLOAT4 data[];
}uOutput;

layout(set=0, binding=1) readonly buffer sourceBuffer0{
    FLOAT4 data[];
} uInput;

layout(set=0, binding=2) readonly buffer sourceBuffer1{
    FLOAT data[];
} uGrid;

layout(set=0, binding=3) uniform gridSampleBuffer{
	ivec4 inShape;  // inW, inH
	ivec4 outShape; // outW, outH
	bool alignCorners;
}uGridSampleParam;

layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;

FLOAT4 LoadSample(int positionX, int positionY, int c, int n) {
	FLOAT4 value;
	int width = uGridSampleParam.inShape.x;
	int height = uGridSampleParam.inShape.y;
#ifdef PAD_MODE_ZEROS
	if (positionX < 0 || positionX >= width || positionY < 0 || positionY >= height) {
		value = FLOAT4(0.0);
	} else {
		value = uInput.data[0
			+ positionX
			+ positionY * width
			+ n * width * height
			+ c * width * height * uGridSampleParam.inShape.w
		];
	}
#else
	positionX = clamp(positionX, 0, width - 1);
	positionY = clamp(positionY, 0, height - 1);
	value = uInput.data[0
		+ positionX
		+ positionY * width
		+ n * width * height
		+ c * width * height * uGridSampleParam.inShape.w
	];
#endif
	return value;
}

void main()
{
	int pos = int(gl_GlobalInvocationID.x);
	// input output grid layout is NC4HW4

	ivec4 inputShape = uGridSampleParam.inShape;
	ivec4 outputShape = uGridSampleParam.outShape;
	int total = outputShape.x * outputShape.y * outputShape.z * outputShape.w;

	if(pos < total)
	{
		// get nchw num of output 
		int x = pos % outputShape.x;
		int tmp = pos / outputShape.x;
		int y = tmp % outputShape.y;
		tmp = tmp / outputShape.y;
		int z = tmp % outputShape.z;
		int n = tmp / outputShape.z;

		// get position in grid
		int gridPosition = n * outputShape.x * outputShape.y + y * outputShape.x + x;
		FLOAT gridX = uGrid.data[2 * gridPosition + 0];
		FLOAT gridY = uGrid.data[2 * gridPosition + 1];

		// compute position of input
		FLOAT a = FLOAT(uGridSampleParam.alignCorners);
		FLOAT b = FLOAT(1.0f) - a;
		int positionX = int(floor(((FLOAT(1.0) + gridX) * (FLOAT(inputShape.x) - a) - b) * FLOAT(0.5f) + FLOAT(0.5f)));
		int positionY = int(floor(((FLOAT(1.0) + gridY) * (FLOAT(inputShape.y) - a) - b) * FLOAT(0.5f) + FLOAT(0.5f)));		FLOAT4 value = LoadSample(positionX, positionY, z, n);

		uOutput.data[0
			+ x
			+ y * outputShape.x
			+ z * outputShape.x * outputShape.y * outputShape.w
			+ n * outputShape.x * outputShape.y
		] = value;
	}
}