#include "bindless_structs.hlsli"

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

[[vk::push_constant]]
Block pcb;

float2 SampleSphericalMap(float3 v)
{
    const float2 InvTan = float2(0.1591f, 0.3183f);     // (1/2PI, 1/PI)
    float2 UV = float2(atan2(-v.x, v.z), asin(-v.y));   // (-PI, -PI/2) to (PI, PI/2)
    UV *= InvTan;                                       // (-1/2, -1/2) to (1/2, 1/2)
    UV += 0.5f.xx;                                      // (0, 0) to (1, 1)
    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, HalfSide - GlobalInvocationID.y, (HalfSide - GlobalInvocationID.x) * AxisSign);
    }
    else if (GlobalInvocationID.z < 4)
    {
        LocalDir = float3(GlobalInvocationID.x - HalfSide, AxisSign * HalfSide, AxisSign * (GlobalInvocationID.y - HalfSide));
    }
    else
    {
        LocalDir = float3((GlobalInvocationID.x - HalfSide) * AxisSign, HalfSide - GlobalInvocationID.y, HalfSide * AxisSign);
    }
    
    float2 UV = SampleSphericalMap(normalize(LocalDir));
    StorageTextureArrays[pcb.OutputTextureHandle][GlobalInvocationID.xyz] = Textures[pcb.HdrEnvHandle].SampleLevel(ImmutableSamplers[pcb.HdrEnvHandle], UV, 0);
}