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

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

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

layout(set = 0, binding = 3) uniform constBuffer
{
	ivec4 inputImgSize;
	ivec4 outputImgSize;
	float spatialScale;
}uRoiPoolingParam;

layout(local_size_x = 8, local_size_y = 8) in;

void main()
{
	ivec3 pos = ivec3(gl_GlobalInvocationID);
	ivec3 outputImgSize = uRoiPoolingParam.outputImgSize.xyz;

	if(pos.x < outputImgSize.x && pos.y < outputImgSize.y)
	{
		ivec3 inputImgSize = uRoiPoolingParam.inputImgSize.xyz;
		int roiBatchIndex = pos.z % uRoiPoolingParam.outputImgSize.w;
		int inputZIndex = pos.z / uRoiPoolingParam.outputImgSize.w;
		// 0, xmin, ymin, xmax, ymax
		FLOAT roiData0 = uRoI.data[roiBatchIndex * 5 + 0];
		FLOAT roiData1 = uRoI.data[roiBatchIndex * 5 + 1];
		FLOAT roiData2 = uRoI.data[roiBatchIndex * 5 + 2];
		FLOAT roiData3 = uRoI.data[roiBatchIndex * 5 + 3];
		FLOAT roiData4 = uRoI.data[roiBatchIndex * 5 + 4];
		int inputBatch = int(roiData0);
		int x1 = int(round(FLOAT(roiData1) * FLOAT(uRoiPoolingParam.spatialScale)));
		int y1 = int(round(FLOAT(roiData2) * FLOAT(uRoiPoolingParam.spatialScale)));
		int x2 = int(round(FLOAT(roiData3) * FLOAT(uRoiPoolingParam.spatialScale)));
		int y2 = int(round(FLOAT(roiData4) * FLOAT(uRoiPoolingParam.spatialScale)));

		int roiW = max(x2 - x1 + 1, 1);
		int roiH = max(y2 - y1 + 1, 1);
		FLOAT binSizeW = FLOAT(roiW) / FLOAT(outputImgSize.x);
		FLOAT binSizeH = FLOAT(roiH) / FLOAT(outputImgSize.y);

		int wStart = clamp(x1 + int(floor(FLOAT(pos.x) * binSizeW)), 0, inputImgSize.x);
		int wEnd = clamp(x1 + int(ceil(FLOAT(pos.x + 1) * binSizeW)), 0, inputImgSize.x);
		int hStart = clamp(y1 + int(floor(FLOAT(pos.y) * binSizeH)), 0, inputImgSize.y);
		int hEnd = clamp(y1 + int(ceil(FLOAT(pos.y + 1) * binSizeH)), 0, inputImgSize.y);

		bool isEmpty = (hEnd <= hStart) || (wEnd <= wStart);

		int inputBasicOffset = inputBatch * uRoiPoolingParam.inputImgSize.x * uRoiPoolingParam.inputImgSize.y
			+ inputZIndex * uRoiPoolingParam.inputImgSize.x * uRoiPoolingParam.inputImgSize.y * uRoiPoolingParam.inputImgSize.w
			;
		int inputStartOffset = wStart
					+ hStart * uRoiPoolingParam.inputImgSize.x
					+ inputBasicOffset;
		FLOAT4 res = isEmpty ? FLOAT4(0.0) : uInput.data[inputStartOffset];
		for(int i = hStart; i < hEnd; ++i)
		{
			for(int j = wStart; j < wEnd; ++j)
			{
				int inputOffset = j
					+ i * uRoiPoolingParam.inputImgSize.x
					+ inputBasicOffset;
				res = max(res, uInput.data[inputOffset]);
			}
		}
		uOutput.data[0
			+ pos.x
			+ pos.y * uRoiPoolingParam.outputImgSize.x
			+ pos.z * uRoiPoolingParam.outputImgSize.x * uRoiPoolingParam.outputImgSize.y
		] = res;
	}
}