/***********************************************/
/*           Copyright (c) 2025 Belmu          */
/*             All Rights Reserved             */
/***********************************************/

#include "/settings.glsl"

#if MANUAL_CAMERA == 0 && EXPOSURE == 2

    #include "/include/uniforms.glsl"
    #include "/include/utility/math.glsl"
    #include "/include/utility/color.glsl"

    layout (r32ui) uniform uimage1D histogram;

    #if defined HISTOGRAM_BUILD_PASS

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

        shared uint histogramLocal[HISTOGRAM_BINS];

        int getBinFromLuminance(float luminance) {
            return clamp(int(round((log(luminance) * rcpLogLuminanceRange - (minLogLuminance * rcpLogLuminanceRange)) * HISTOGRAM_BINS)), 0, HISTOGRAM_BINS - 1);
        }

        void main() {
            histogramLocal[gl_LocalInvocationIndex] = 0u;

            memoryBarrierShared();
            barrier();

            if (all(lessThan(gl_GlobalInvocationID.xy, textureSize(DOWNSAMPLED_BUFFER, 0)))) {
                float luminance = luminance(texelFetch(DOWNSAMPLED_BUFFER, ivec2(gl_GlobalInvocationID.xy * 0.5), 0).rgb);

                if (!isnan(luminance) && !isinf(luminance) && luminance > 0.0) {
                    atomicAdd(histogramLocal[getBinFromLuminance(luminance)], 1u);
                }
            }

            memoryBarrierShared();
            barrier();

            imageAtomicAdd(histogram, int(gl_LocalInvocationIndex), histogramLocal[gl_LocalInvocationIndex]);
        }

    #elif defined HISTOGRAM_MEDIAN_PASS

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

        const ivec3 workGroups = ivec3(1, 1, 1);

        uniform usampler1D histogramSampler;

        float getLuminanceFromBin(int bin) {
            return exp((bin * rcp(HISTOGRAM_BINS)) * logLuminanceRange + minLogLuminance);
        }

        int getClosestBinToMedian() {
            uint totalDensity = 0;
            for (int bin = 0; bin < HISTOGRAM_BINS; bin++) {
                totalDensity += texelFetch(histogramSampler, bin, 0).r;
            }

            uint median = totalDensity / 2;
            uint cumulativeDensity = 0;

            for (int bin = 0; bin < HISTOGRAM_BINS; bin++) {
                cumulativeDensity += texelFetch(histogramSampler, bin, 0).r;

                if (cumulativeDensity >= median) return bin;
            }

            return HISTOGRAM_BINS - 1;
        }

        float getGeometricMeanLuminance() {
            uint totalDensity = 0;

            for (int bin = 0; bin < HISTOGRAM_BINS; bin++) {
                totalDensity += texelFetch(histogramSampler, bin, 0).r;
            }

            float cumulativeDensity = 0.0;

            int lowerBound = 0;
            int upperBound = HISTOGRAM_BINS - 1;

            float minDensity = EXPOSURE_IGNORE_DARK           * totalDensity;
            float maxDensity = (1.0 - EXPOSURE_IGNORE_BRIGHT) * totalDensity;

            for (int bin = 0; bin < HISTOGRAM_BINS; bin++) {
                cumulativeDensity += texelFetch(histogramSampler, bin, 0).r;

                if (cumulativeDensity >= minDensity) { lowerBound = bin; break; }
            }

            cumulativeDensity = 0.0;

            for (int bin = 0; bin < HISTOGRAM_BINS; bin++) {
                cumulativeDensity += texelFetch(histogramSampler, bin, 0).r;

                if (cumulativeDensity >= maxDensity) { upperBound = bin; break; }
            }

            upperBound = max(upperBound, lowerBound);

            float logStep = logLuminanceRange / float(HISTOGRAM_BINS);

            float weightedSum = 0.0;
            float densitySum  = 0.0;

            for (int bin = lowerBound; bin <= upperBound; bin++) {
                float binDensity = texelFetch(histogramSampler, bin, 0).r;
                float logCenter  = minLogLuminance + (float(bin) + 0.5) * logStep;

                weightedSum += binDensity * logCenter;
                densitySum  += binDensity;
            }

            return exp(weightedSum / densitySum);
        }

        void main() {
            float averageLuminance = getGeometricMeanLuminance();

            imageStore(histogram, 0, uvec4(floatBitsToUint(averageLuminance), 0, 0, 0));
        }

    #endif

#else

    layout (local_size_x = 1) in;

    void main() {
        return;
    }

#endif
