#version 440 core
layout(std140) buffer;

layout(set=0, binding=0) writeonly uniform image2D uOutput;
layout(set=0, binding=1) uniform sampler2D uInput;

layout(set=0, binding=2) readonly uniform constBuffer {
    ivec2 pad;
    ivec2 kernelSize;
    ivec2 stride;
    ivec2 dilate;
    ivec4 inputSize;
    ivec4 outputSize;
    ivec4 offset;//batchOffset, hOffset, outputHeight, other
} uConstant;

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

#define UP_DIV(x, y) (((x)+(y)-1)/(y))

void main()
{
    ivec3 posTmp = ivec3(gl_GlobalInvocationID);
    ivec3 pos;
    pos.x = posTmp.x % uConstant.outputSize.x;
    int subIndex = posTmp.x / uConstant.outputSize.x;
    pos.y = subIndex % uConstant.outputSize.y;
    pos.z = subIndex / uConstant.outputSize.y;
    ivec3 inputSize = uConstant.inputSize.xyz;
    ivec3 outputSize = uConstant.outputSize.xyz;
    int batchOffset = uConstant.offset.x * uConstant.inputSize.z;
    int fz = pos.z % uConstant.inputSize.z;
    int sb = pos.z / uConstant.inputSize.z;

    if (pos.x < outputSize.x && pos.y < outputSize.y && sb < uConstant.outputSize.w)
    {
        ivec2 s0 = (pos.xy + ivec2(0, uConstant.offset.y))*uConstant.stride-uConstant.pad;
        ivec2 sfxy = max(ivec2(0), (UP_DIV(-s0, uConstant.dilate)));
        ivec2 efxy = min(uConstant.kernelSize, UP_DIV(inputSize.xy-s0, uConstant.dilate));

        int destYOrigin = pos.x + (pos.y)*uConstant.outputSize.x + sb*uConstant.outputSize.x*uConstant.outputSize.y;
        int destY = destYOrigin / 4;
        int destXOffset = destYOrigin % 4;
        for (int fy=0; fy<uConstant.kernelSize.y; ++fy)
        {
            int sy = fy*uConstant.dilate.y + s0.y;
            for (int fx=0; fx<uConstant.kernelSize.x; ++fx)
            {
                ivec2 fxy = ivec2(fx, fy);
                int destX = 0
                    + fx
                    + fy*uConstant.kernelSize.x
                    + fz*uConstant.kernelSize.x * uConstant.kernelSize.y
                    ;
                if (all(lessThan(fxy, efxy)) && all(greaterThanEqual(fxy, sfxy)))
                {
                    int sx = fx*uConstant.dilate.x + s0.x;
                    int szb_offset = (pos.z + batchOffset);
                    int sz_offset = szb_offset % uConstant.inputSize.z;
                    int sb_offset = szb_offset / uConstant.inputSize.z;
                    vec4 color = texelFetch(uInput, ivec2(sx+sz_offset*uConstant.inputSize.x, sy+sb_offset*uConstant.inputSize.y), 0);
                    imageStore(uOutput, ivec2(4*destX+destXOffset, destY), color);
                }
                else
                {
                    imageStore(uOutput, ivec2(4*destX+destXOffset, destY), vec4(0));
                }
            }
        }
    }
}