#version 450 core

layout(std430) uniform;
layout(set=0, binding=0) uniform lowp usampler3D uInput;

layout(std430, set=0, binding=1) writeonly buffer destBuffer{
    uint data[];
} uOutBuffer;

layout(set=0, binding=2) uniform constBuffer{
    int width;
    int height;
    int channel;
    int batch;
} uConstant;

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

void main()
{
    ivec3 pos = ivec3(gl_GlobalInvocationID);
    int channelC4 = (uConstant.channel + 3) / 4;
    int batchIndex = pos.z / channelC4;
    int zDiv4 = pos.z % channelC4;

    int lastZ = uConstant.channel / 4;
    int cIndex = uConstant.channel % 4;

    if (pos.x < uConstant.width && pos.y < uConstant.height)
    {
        uvec4 color = uvec4(texelFetch(uInput, pos, 0));
        int z = zDiv4*4;
        int basicOffset = 0
            + z
            + uConstant.width*uConstant.channel*pos.y
            + uConstant.height*batchIndex*uConstant.channel*uConstant.width
            + uConstant.channel*pos.x;
        
        if(zDiv4 == lastZ)
        {
            if(cIndex == 1)
            {
                uOutBuffer.data[basicOffset+0] = color.r;
            }
            else if(cIndex == 2)
            {
                uOutBuffer.data[basicOffset+0] = color.r;
                uOutBuffer.data[basicOffset+1] = color.g;
            }
            else
            {
                uOutBuffer.data[basicOffset+0] = color.r;
                uOutBuffer.data[basicOffset+1] = color.g;
                uOutBuffer.data[basicOffset+2] = color.b;
            }
        }
        else
        {
            uOutBuffer.data[basicOffset+0] = color.r;
            uOutBuffer.data[basicOffset+1] = color.g;
            uOutBuffer.data[basicOffset+2] = color.b;
            uOutBuffer.data[basicOffset+3] = color.a;
        }

    }
}