#version 450
// FSR Ray Regeneration backend, stage 3: the denoised diffuse and specular signals remodulated by
// their albedos, plus the residual the denoiser never saw (particles, what demodulation could not
// carry, sky). The result is render-resolution colour for the Super Sampling stage.
layout(local_size_x = 8, local_size_y = 8) in;

layout(binding = 0, rgba16f) uniform readonly image2D denoisedDiffuse;
layout(binding = 1, rgba16f) uniform readonly image2D denoisedSpecular;
layout(binding = 2, rgba16f) uniform readonly image2D diffuseAlbedo;
layout(binding = 3, rgba16f) uniform readonly image2D specularAlbedo;
layout(binding = 4, rgba16f) uniform readonly image2D residual;
layout(binding = 5, rgba16f) uniform writeonly image2D composed;
layout(binding = 6, rgba16f) uniform readonly image2D noisyDiffuse;
layout(binding = 7, rgba16f) uniform readonly image2D noisySpecular;

layout(push_constant) uniform Constants
{
    uvec2 renderSize;
    float haloLimit; // > 0: each denoised lobe at most this multiple of the brightest noisy input of that lobe around it
    uint unused;
} pc;

void main()
{
    const ivec2 px = ivec2(gl_GlobalInvocationID.xy);
    if (px.x >= int(pc.renderSize.x) || px.y >= int(pc.renderSize.y))
        return;
    const vec3 da = imageLoad(diffuseAlbedo, px).rgb, sa = imageLoad(specularAlbedo, px).rgb;
    vec3 diffuse = imageLoad(denoisedDiffuse, px).rgb * da;
    vec3 specular = imageLoad(denoisedSpecular, px).rgb * sa;
    if (pc.haloLimit > 0.0)
    {
        // The denoiser can carry a thin object's large demodulated signal (dark albedo) across its edge, where this
        // texel's albedo turns it into a bright fringe. A denoised value above everything around it in the noisy
        // input is such a leak: capped, by luminance (the hue is kept).
        const vec3 w = vec3(0.2126, 0.7152, 0.0722);
        float diffuseMax = 0.0, specularMax = 0.0;
        for (int dy = -1; dy <= 1; ++dy)
            for (int dx = -1; dx <= 1; ++dx)
            {
                const ivec2 q = clamp(px + ivec2(dx, dy), ivec2(0), ivec2(pc.renderSize) - 1);
                // what the lobe contributes there (the neighbour's own albedo): a leaked demodulated signal
                // remodulated by this texel's albedo exceeds all of these
                diffuseMax = max(diffuseMax, dot(imageLoad(noisyDiffuse, q).rgb * imageLoad(diffuseAlbedo, q).rgb, w));
                specularMax = max(specularMax, dot(imageLoad(noisySpecular, q).rgb * imageLoad(specularAlbedo, q).rgb, w));
            }
        const float dl = dot(diffuse, w), sl = dot(specular, w);
        if (dl > pc.haloLimit * diffuseMax)
            diffuse *= pc.haloLimit * diffuseMax / dl;
        if (sl > pc.haloLimit * specularMax)
            specular *= pc.haloLimit * specularMax / sl;
    }
    const vec3 color = diffuse + specular + imageLoad(residual, px).rgb;
    imageStore(composed, px, vec4(clamp(color, 0.0, 65000.0), 1.0));
}