#version 430 core
////////////////////////////////////////////////////////////////////////////////
// <copyright>
// Copyright (c) 2015 Hue AS. All rights reserved.
//
// All rights are reserved. Reproduction or transmission in whole or in part, in
// any form or by any means, electronic, mechanical or otherwise, is prohibited
// without the prior written permission of the copyright owner.
// </copyright>
////////////////////////////////////////////////////////////////////////////////

#define W $COMPUTE_WIDTH$
#define HALF_H $HALF_COMPUTE_HEIGHT$
#define N (W * 2 * HALF_H)

uniform usampler2D counts;
layout(binding=0) writeonly uniform uimage2D offsets;
layout(binding=1) writeonly uniform uimage2D sums;
shared uint acc[N];

#define NUM_BANKS 16
uint shuffle(uint index)
{
  return (index & ~(NUM_BANKS - 1)) | ((index + index / NUM_BANKS) & (NUM_BANKS - 1));
}

layout(local_size_x = W, local_size_y = HALF_H) in;

void main()
{
    ivec2 imgPos1 = ivec2(gl_GlobalInvocationID.x, gl_WorkGroupID.y * 2 * HALF_H + gl_LocalInvocationID.y);
    ivec2 imgPos2 = ivec2(imgPos1.x, imgPos1.y + HALF_H);
    uint accIndex1 = gl_LocalInvocationIndex;
    uint accIndex2 = gl_LocalInvocationIndex + W * HALF_H;
    acc[shuffle(accIndex1)] = texelFetch(counts, imgPos1, 0).x;
    acc[shuffle(accIndex2)] = texelFetch(counts, imgPos2, 0).x;
    for (uint d = N / 2, offset = 1; d > 0u; d >>= 1u, offset <<= 1u)
    {
        barrier();
        if (gl_LocalInvocationIndex < d)
        {
            accIndex1 = offset * (2 * gl_LocalInvocationIndex + 1) - 1;
            accIndex2 = accIndex1 + offset;
            acc[shuffle(accIndex2)] += acc[shuffle(accIndex1)];
        }
    }
    barrier();
    if (gl_LocalInvocationIndex == 0)
    {
        accIndex1 = shuffle(N - 1);
        imageStore(sums, ivec2(gl_WorkGroupID.xy), uvec4(acc[accIndex1]));
        acc[accIndex1] = 0;
    }
    for (uint d = 1, offset = N / 2; d < uint(N); d <<= 1u, offset >>= 1u)
    {
        barrier();
        if (gl_LocalInvocationIndex < d)
        {
            accIndex1 = offset * (2 * gl_LocalInvocationIndex + 1) - 1;
            accIndex2 = accIndex1 + offset;
            accIndex1 = shuffle(accIndex1);
            accIndex2 = shuffle(accIndex2);
            uint temp = acc[accIndex1];
            acc[accIndex1] = acc[accIndex2];
            acc[accIndex2] += temp;
        }
    }
    barrier();
    accIndex1 = gl_LocalInvocationIndex;
    accIndex2 = gl_LocalInvocationIndex + W * HALF_H;
    imageStore(offsets, imgPos1, uvec4(acc[shuffle(accIndex1)]));
    imageStore(offsets, imgPos2, uvec4(acc[shuffle(accIndex2)]));
}
