#version 450

layout(location = 0) in vec2 inTextCoord;
layout(location = 0) out float outAO;

layout(set = 0, binding = 0) uniform sampler2D posSampler;
layout(set = 0, binding = 1) uniform sampler2D normalSampler;
layout(set = 0, binding = 2) uniform sampler2D noiseSampler;

layout(set = 1, binding = 0) uniform  SSAOKernel {
    vec4 samples[64];
} kernel;

layout(set = 2, binding = 0) uniform SSAOInfo {
    mat4 projection; //64
    mat4 view; //128
    vec2 screenSize; // 136
    float radius; //140
    float bias; //144
    int kernelSize; //148
} ssao;

void main() {
    int KERNEL_SIZE = clamp(ssao.kernelSize, 1, 64);
    vec3 fragPos = texture(posSampler, inTextCoord).xyz;
    vec3 worldNorm = texture(normalSampler, inTextCoord).xyz;

    if (dot(worldNorm, worldNorm) < 0.001 || fragPos.z >= 0.0) {
        outAO = 1.0;
        return;
    }

    vec3 normal = normalize(mat3(ssao.view) *  worldNorm);

    vec3 randomVec = texelFetch(noiseSampler, ivec2(gl_FragCoord.xy) % 4, 0).xyz;
    vec3 tangent = randomVec - normal * dot(randomVec, normal);
    if (dot(tangent, tangent) < 0.0001) {
        tangent = cross(normal, abs(normal.z) < 0.9 ? vec3(0, 0, 1) : vec3(0, 1, 0));
    }
    tangent = normalize(tangent);
    vec3 bitangent = cross(normal, tangent);
    mat3 TBN = mat3(tangent, bitangent, normal);
    float fragDepth = -fragPos.z;

    float occlusion = 0.0;

    float Passes = 0.0f;

    for (int i = 0; i < KERNEL_SIZE; i++) {
        // Cover the complete radius distribution, even with a smaller sample budget.
        int sampleIndex = i * 64 / KERNEL_SIZE;
        vec3 samplePos = fragPos + (TBN * kernel.samples[sampleIndex].xyz) * ssao.radius;
        vec4 offset = ssao.projection * vec4(samplePos, 1.0);
        if (offset.w <= 0.0) continue;
        offset.xyz /= offset.w;
        offset.xyz = offset.xyz * 0.5 + 0.5;
        vec2 sampleUV = vec2(offset.x, 1.0 - offset.y);
        if (any(lessThan(sampleUV, vec2(0))) || any(greaterThan(sampleUV, vec2(1)))) continue;
        vec3 sampleViewPos = texture(posSampler, sampleUV).xyz;
        if (sampleViewPos.z >= 0.0) continue;
        float sampleDepth = -sampleViewPos.z;
        float depthDelta = abs(fragDepth - sampleDepth);
        if (depthDelta > ssao.radius * 2.0) {
            continue;
        }
        float rangeCheck = smoothstep(0.0, 1.0, ssao.radius / max(depthDelta, 0.0001));
        float isOccluded = step(samplePos.z + ssao.bias, sampleViewPos.z);
        occlusion += isOccluded * rangeCheck;
        Passes += 1.0;
    }

    float Subtraction = 0.0;
    if(Passes > 0){
        Subtraction =  (occlusion / Passes);
    }
    occlusion = 1.0 - Subtraction;
    outAO = clamp(pow(occlusion, 3), 0.0, 1.0);
}