#version 450

#include "random.glsl"

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

layout(set = 0, binding = 0) uniform Locals {
    mat4 View;
    mat4 Proj;
    mat4 matrix_2d;
    uvec4 light_count;
    vec4 cam_pos;
};

layout(set = 1, binding = 0, r16f) uniform writeonly image2D SSAOTexture;
layout(set = 1, binding = 1) uniform sampler Sampler;
layout(set = 1, binding = 2) uniform texture2D ScreenSpace;
layout(set = 1, binding = 3) uniform texture2D Normal;

#define SAMPLES 4
#define BIAS 0.05
#define RADIUS 0.5

vec3 randomSample(inout uint seed) {
    const float r1 = RandomFloat(seed);
    const float r2 = RandomFloat(seed);

    const float r = sqrt(1.0 - r1 * r1);
    const float phi = 2.0 * 3.14159265358979323846f * r2;
    return normalize(vec3(cos(phi) * r, sin(phi) * r, r1));
}

void main()
{
    const ivec2 image_size = imageSize(SSAOTexture).xy;
    const ivec2 pixel = ivec2(gl_GlobalInvocationID.x % image_size.x, gl_GlobalInvocationID.x / image_size.x);
    if (pixel.x >= image_size.x || pixel.y >= image_size.y) {
        return;
    }

    uint seed = WangHash((pixel.x + pixel.y * image_size.x) * 16789 * 1791);
    const vec2 uv = vec2(pixel) / textureSize(sampler2D(ScreenSpace, Sampler), 0).xy;

    vec4 norz = texture(sampler2D(ScreenSpace, Sampler), uv);
    float depth = norz.z / norz.w;
    float scale = RADIUS / depth;

    float ao = 0.0;
    for(int i = 0; i < SAMPLES; i++)
    {
        vec2 randUv = (vec2(pixel) + 23.71 * float(i)) / image_size.xy;
        vec3 randNor = randomSample(seed) * 2.0 - 1.0;
        if(dot(norz.xyz, randNor) < 0.0) {
            randNor *= -1.0;
        }

        vec2 off = randNor.xy * scale;
        vec4 sampleNorz = texture(sampler2D(ScreenSpace, Sampler), uv + off);
        float depthDelta = depth - sampleNorz.z / sampleNorz.w;

        vec3 sampleDir = vec3(randNor.xy * RADIUS, depthDelta);
        float occ = max(0.0, dot(normalize(norz.xyz), normalize(sampleDir)) - BIAS) / (length(sampleDir) + 1.0);
        ao += occ;
    }

    ao = 1.0 - (ao / float(SAMPLES));
    imageStore(SSAOTexture, pixel, vec4(vec3(ao * ao), 1.0));
}