#extension GL_EXT_buffer_reference : require
#extension GL_EXT_nonuniform_qualifier : require

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

layout(set = 0, binding = 0) uniform sampler samplers[];
layout(set = 0, binding = 1) uniform texture2D textures[];
layout(set = 0, binding = 2, rgba32f) uniform image2D images[];

layout(push_constant) uniform buffers {
    uint output_image_index;
    uvec2 tex_size;
    uint padding;
} params;

vec2 sphere_intersection ( in vec3 pos, in vec3 dir, float radius )
{
    float b = dot(pos, vec3(radius));
    float c = dot(pos, pos) - radius * radius;
    float h = b * b - c;
    if (h < 0.f) return vec2(-1.f); // no intersection
    h = sqrt( h );
    return vec2( -b - h, -b + h );
}

vec3 rayleigh_scattering = vec3(5.802f, 13.558f, 33.1f);
vec3 rayleigh_absorption = vec3(0.f);

vec3 mie_scattering = vec3(3.996f);
vec3 mie_absorption = vec3(4.40f);

void main() {
    ivec2 coords = ivec2(gl_GlobalInvocationID.xy);
    vec3 out_color = vec3(0);
    if (coords.x % 2 == 0 ) {
        out_color = vec3(1);
    }
    if (coords.y % 2 == 1 ) {
        out_color = vec3(1);
    }
    imageStore(images[nonuniformEXT(params.output_image_index)], coords, vec4(out_color, 1.0));

    const uint raymarching_steps = 10U;

    const float ground_radius = 6360.f;
    const float top_radius = 6420.f;
    const float atmosphere_radius = top_radius - ground_radius;

    vec2 uvs = coords / params.tex_size;

    float mu = uvs.y * 2.f - 1.f;
    float theta = acos(mu);
    float height = atmosphere_radius *  uvs.x;

    vec2 pos = vec2(0.f, ground_radius + height);
    vec2 dir = vec2(mu, sin(theta));
    vec3 transmittance = vec3(0.f);

    // find ground hit
    float t = 0.f;
    vec2 ground_intersection = sphere_intersection(vec3(pos, 0.f), vec3(dir, 0.f), ground_radius);
    t = ground_intersection.x;

    // find atmosphere hit
    if (ground_intersection.y < 0.f) {
        vec2 top_intersection = sphere_intersection(vec3(pos, 0.f), vec3(dir, 0.f), top_radius);
        t = top_intersection.y;
    }

    vec2 hit = pos + dir * t;
    vec3 extinction = rayleigh_scattering + rayleigh_absorption + mie_scattering + mie_absorption;
    vec3 total_extinction = vec3(0.f);

    // raymarching
    float step = t / raymarching_steps;
    for (float dist = 0.f; dist < t; dist += step) {
        vec2 cur = pos + dir * dist;
        dist += step;

        total_extinction += extinction * step;
    }

    vec4 color = vec4(exp(-total_extinction), 1.f);

    imageStore(images[nonuniformEXT(params.output_image_index)], coords, color);
}