#include "bindless_structs.hlsli"

struct Block
{
    uint HdrEnvHandle;
    uint OutputTextureHandle;
    int CubeSize;
};

[[vk::push_constant]]
Block pcb;

float2 SampleSphericalMap(float3 v)
{
    const float2 invAtan = float2(0.1591f, 0.3183f);
    float2 uv = float2(atan2(v.x, v.z), asin(v.y));
    uv *= invAtan;
    uv += 0.5;
    return uv;
}

/*
| Axis | Layer | Up |
|:----:|:-----:|:--:|
|  +x  |   0   | -y |
|  -x  |   1   | -y |
|  +y  |   2   | +z |
|  -y  |   3   | -z |
|  -z  |   4   | -y |
|  +z  |   5   | -y |
*/

[numthreads(16, 16, 1)]
void main(uint3 GlobalInvocationID : SV_DispatchThreadID)
{
    float3 LocalDir;
    float HalfSide = pcb.CubeSize / 2;
    float AxisSign = 2.0f * (0.5f - (GlobalInvocationID.z % 2));
    
    if (GlobalInvocationID.z < 2)
    {
        LocalDir = float3(AxisSign * HalfSide, GlobalInvocationID.y - HalfSide, (GlobalInvocationID.x - HalfSide) * AxisSign);
    }
    else if (GlobalInvocationID.z < 4)
    {
        LocalDir = float3(GlobalInvocationID.x - HalfSide, -AxisSign * HalfSide, -AxisSign * (GlobalInvocationID.y - HalfSide));
    }
    else
    {
        LocalDir = float3((GlobalInvocationID.x - HalfSide) * AxisSign, GlobalInvocationID.y - HalfSide, -AxisSign * HalfSide);
    }
    
    float2 UV = SampleSphericalMap(normalize(LocalDir));
    StorageTextureArrays[pcb.OutputTextureHandle][GlobalInvocationID.xyz] = Textures[pcb.HdrEnvHandle].SampleLevel(ImmutableSamplers[pcb.HdrEnvHandle], UV, 0);
}