import "includes/math";

[[vk::binding(0, 0)]] SamplerState _EquirectangularSampler;
[[vk::binding(1, 0)]] Texture2D _EquirectangularTexture;
[[vk::binding(2, 0)]] WTexture2DArray _CubemapImage;

static const float3x3 MATRIX_2D_TO_CUBE[6] = {
    // Positive X
    {
         0,  0, 1,
         0, -1, 0,
        -1,  0, 0,
    },
    // Negative X
    {
        0,  0, -1,
        0, -1,  0,
        1,  0,  0,
    },
    // Positive Y
    {
        1, 0, 0,
        0, 0, 1,
        0, 1, 0,
    },
    // Negative Y
    {
        1,  0,  0,
        0,  0, -1,
        0, -1,  0,
    },
    // Positive Z
    {
        1,  0, 0,
        0, -1, 0,
        0,  0, 1,
    },
    // Negative Z
    {
        -1,  0,  0,
         0, -1,  0,
         0,  0, -1,
    }
};

float2 cubemapSize() {
    uint width, height, elements;
    _CubemapImage.GetDimensions(width, height, elements);
    return float2(width, height);
}

[shader("compute")]
[numthreads(8, 8, 1)]
void main(uint3 threadID: SV_DispatchThreadID)
{
    let size = cubemapSize();
    var texCoord = (float2(threadID.xy) + 0.5) / size;
    let layerIndex = threadID.z;

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

    let cubeCoord = mul(MATRIX_2D_TO_CUBE[layerIndex], float3(texCoord, 1.0));
    let cubeDir = normalize(cubeCoord);

    let theta = atan2(cubeDir.y, cubeDir.x);
    let phi = asin(cubeDir.z);

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

    let irradiance = _EquirectangularTexture.SampleLevel(_EquirectangularSampler, equirectCoord, 0.0);
    _CubemapImage.Store(threadID, irradiance);
}