#version 430 core

// Parameters
// Size of work group (must produce at least 256 invocations)
// Must match pass/workgroup-size
#define LOCAL_SIZEX     256
#define LOCAL_SIZEY     2
// Number of rows of image to read per invocation
// Must match inverse of pass/global-size/y-screen-scale) in compositor pass
#define INVOC_ROWS      16
// Number of parallel histograms per workgroup
#define LOCAL_HIST_ROWS 8

layout(local_size_x = LOCAL_SIZEX, local_size_y = LOCAL_SIZEY) in;

#define LOCAL_BINS (256*LOCAL_HIST_ROWS)
#define INVOCATIONS (LOCAL_SIZEX*LOCAL_SIZEY)
#define BINS_PER_INVOC (LOCAL_BINS / INVOCATIONS)

// We share a number of fast local histograms in the workgroup, to reduce
// atomics contention of the global histogram.
shared uint local_histogram[LOCAL_BINS];

uniform usampler2D bins_tex;
layout(r32ui) uniform uimage2D histogram_tex;

void main()
{
    // Initialise local histograms
#if LOCAL_BINS > INVOCATIONS
    // Each invocation has multiple bins to initialise
    for (int i = 0; i < BINS_PER_INVOC; ++i)
        local_histogram[gl_LocalInvocationIndex * BINS_PER_INVOC + i] = 0;
#elif LOCAL_BINS == INVOCATIONS
    // Each invocation has exactly 1 bin to initialise
    local_histogram[gl_LocalInvocationIndex] = 0;
#else
    // Only some invocations need to initialise bins
    if (gl_LocalInvocationIndex < LOCAL_BINS)
        local_histogram[gl_LocalInvocationIndex] = 0;
#endif
    barrier();

    // Check the invocation is actually inside the bin image
    ivec2 tex_size = ivec2(textureSize(bins_tex, 0));
    if (gl_GlobalInvocationID.x < tex_size.x) {
        // Iterate over multiple rows
        ivec2 texcoord = ivec2(gl_GlobalInvocationID);
        texcoord.y *= INVOC_ROWS;
        uint local_hist = gl_LocalInvocationIndex & (LOCAL_HIST_ROWS - 1);
        for (int i = 0; i < INVOC_ROWS; ++i) {
            if (texcoord.y >= tex_size.y)
                break;

            // Extract the bin number from the bin image
            uint bin = texelFetch(bins_tex, texcoord, 0).r;

            // Add 1 to the appropriate bin in the local histogram
            atomicAdd(local_histogram[bin * LOCAL_HIST_ROWS + local_hist], 1);

            ++texcoord.y;
        }
    }

    barrier();
    if (gl_LocalInvocationIndex < 256) {
        // Sum each row of the local histograms
        uint sum = 0u;
        for (int i = 0; i < LOCAL_HIST_ROWS; ++i)
            sum += local_histogram[gl_LocalInvocationIndex * LOCAL_HIST_ROWS + i];
        // Add to the global histogram
        if (sum > 0)
            imageAtomicAdd(histogram_tex, ivec2(gl_LocalInvocationIndex, 0), sum);
    }
}