layout(set=0, binding=0) buffer destbuffer{
    int data[];
}uOutput;

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

layout(set = 0, binding = 2) uniform constBuffer {
    ivec4 size;//inside, axis, outside, reduceAxis
}uConst;

layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
shared FLOAT local_buffer[256];
shared int local_index[256];

void main()
{
    int index = int(gl_GlobalInvocationID.x) / uConst.size.w;
    int lidIndex = int(gl_LocalInvocationID.x);
    int lidx = lidIndex / uConst.size.w;
    int lid = lidIndex % uConst.size.w;

    // y: index in outside, x: index in inside
    int x = index % uConst.size.x;
    int y = index / uConst.size.x;
    int W = uConst.size.x;
    int H = uConst.size.y;
    int C = uConst.size.z;

    if(y < uConst.size.z && lid < uConst.size.y) {
        int offset = y * H * W + x;
        FLOAT maxValue = uInput.data[offset + lid * W];
        int maxIndex = lid;
        for(int i = lid+uConst.size.w; i < uConst.size.y; i+=uConst.size.w) {
            FLOAT value = uInput.data[offset + i * W];
#ifndef ARGMIN
            if (value > maxValue) {
                maxValue = value;
                maxIndex = i;
            }
#else
            if (value < maxValue) {
                maxValue = value;
                maxIndex = i;
            }
#endif
        }
        local_buffer[lid + lidx * uConst.size.w] = maxValue;
        local_index[lid + lidx * uConst.size.w] = maxIndex;
    }
    barrier();
    if(y < uConst.size.z && lid == 0) {
        FLOAT maxValue = local_buffer[lidx * uConst.size.w];
        int maxIndex = local_index[lidx * uConst.size.w];
        for (int t=1; t<uConst.size.w && t < uConst.size.y; ++t) {
            FLOAT next = local_buffer[t + lidx * uConst.size.w];
#ifndef ARGMIN
            if (next > maxValue) {
                maxValue = next;
                maxIndex = local_index[t + lidx * uConst.size.w];
            }
#else
            if (next < maxValue) {
                maxValue = next;
                maxIndex = local_index[t + lidx * uConst.size.w];
            }
#endif
        }
        uOutput.data[index] = maxIndex;
    }
}