#version 460

layout(set = 0, binding = 0) uniform sampler _EquirectangularSampler;

layout(set = 0, binding = 1) uniform texture2D _EquirectangularTexture;
layout(set = 0, binding = 2, rgba16f) uniform writeonly restrict imageCube _CubemapImage;

const float INV_PI = 0.31830987;
const float HALF_INV_PI = 0.15915494;

const mat3x3 MATRIX_2D_TO_CUBE[6] = mat3x3[](
    // Positive X
    mat3x3(
        0,  0, -1,
        0, -1,  0,
        1,  0,  0
    ),
    // Negative X
    mat3x3(
         0,  0, 1,
         0, -1, 0,
        -1,  0, 0
    ),
    // Positive Y
    mat3x3(
        1, 0, 0,
        0, 0, 1,
        0, 1, 0
    ),
    // Negative Y
    mat3x3(
        1,  0,  0,
        0,  0, -1,
        0, -1,  0
    ),
    // Positive Z
    mat3x3(
        1,  0, 0,
        0, -1, 0,
        0, 0,  1
    ),
    // Negative Z
    mat3x3(
        -1,  0,  0,
         0, -1,  0,
         0,  0, -1
    )
);

layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
void main() {
    vec2 size = vec2(imageSize(_CubemapImage).xy);
    vec2 texCoord = (vec2(gl_GlobalInvocationID.xy) + vec2(0.5)) / size;
    uint layerIndex = gl_GlobalInvocationID.z;

    texCoord = texCoord * vec2(2.0) - vec2(1.0); // Map to range [-1, 1]

    vec3 cubeCoord = MATRIX_2D_TO_CUBE[layerIndex] * vec3(texCoord, 1.0);
    vec3 cubeDir = normalize(cubeCoord);

    float theta = atan(cubeDir.y, cubeDir.x);
    float phi = asin(cubeDir.z);

    vec2 equirectCoord = vec2(theta * HALF_INV_PI, phi * INV_PI) + vec2(0.5);

    vec4 irradiance = texture(sampler2D(_EquirectangularTexture, _EquirectangularSampler), equirectCoord);
    imageStore(_CubemapImage, ivec3(gl_GlobalInvocationID.xyz), irradiance);
}