#pragma kernel ProbeVolumeAtlasOctahedralDepthBlitKernel PROBE_VOLUME_ATLAS_OCTAHEDRAL_DEPTH_BLIT_KERNEL=ProbeVolumeAtlasOctahedralDepthBlitKernel

#ifdef SHADER_API_PSSL
#   pragma argument( scheduler=minpressure ) // instruct the shader compiler to prefer minimizing vgpr usage
#endif

#pragma only_renderers d3d11 playstation xboxone vulkan metal switch

#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/ProbeVolume/ProbeVolumeLighting.cs.hlsl"

StructuredBuffer<float> _ProbeVolumeAtlasOctahedralDepthReadBuffer;
RWTexture2D<float2> _ProbeVolumeAtlasOctahedralDepthWriteTexture;

uint _ProbeVolumeAtlasOctahedralDepthReadBufferCount;
float3 _ProbeVolumeResolution;
float3 _ProbeVolumeResolutionInverse;
float4 _ProbeVolumeAtlasOctahedralDepthScaleBias;
float4 _ProbeVolumeAtlasOctahedralDepthResolutionAndInverse;

uint2 ComputeWriteIndexFromReadIndex(uint readIndex, float3 resolution, float3 resolutionInverse, float4 scaleBias, float4 atlasResolutionAndInverse)
{
    // return uint2(readIndex % 1024u, readIndex >> 10);

    // _ProbeVolumeAtlasOctahedralDepthReadBuffer[z * resolutionY * resolutionX + y * resolutionX + x]
    // TODO: Could implement as floating point operations, which is likely faster.
    // Would need to verify precision.
    uint x = readIndex % (uint)resolution.x;
    uint y = (readIndex / (uint)resolution.x) % (uint)resolution.y;
    uint z = readIndex / ((uint)resolution.y * (uint)resolution.x);

    // Atlas layout:
    // int width = resolution.x * resolution.z;
    // int height = resolution.y;
    // scaleBias: x,y is width and height (scale) z,w offset into atlas (bias)
    // Handle Z "vertical" flip, as in source, index z == 0 is "top" and index z == resolution.z - 1 is "bottom"
    // but in destination atlas, z coordinate min is "bottom", and z coordinate max is "top".
    uint atlasX = x + (resolution.z - 1 - z) * (uint)resolution.x;
    uint atlasY = y;

    float atlasU = (float)atlasX + scaleBias.z * atlasResolutionAndInverse.x;
    float atlasV = (float)atlasY + scaleBias.w * atlasResolutionAndInverse.y;

    return uint2((uint)atlasU, (uint)atlasV);
}

// Warning this needs to match with kBatchSize in ProbeVolumeLighting.cs
#define BATCH_SIZE 256
[numthreads(BATCH_SIZE, 1, 1)]
void PROBE_VOLUME_ATLAS_OCTAHEDRAL_DEPTH_BLIT_KERNEL(uint groupThreadId : SV_GroupThreadID, uint groupId : SV_GroupID, uint3 dispatchThreadId : SV_DispatchThreadID)
{
    uint readIndex = groupId * BATCH_SIZE + groupThreadId;
    if (readIndex >= _ProbeVolumeAtlasOctahedralDepthReadBufferCount) { return; }

    uint readIndexProbe = readIndex / (8 * 8);
    uint readIndexOctahedral = readIndex - (readIndexProbe * (8 * 8));

    uint2 writeIndexProbe = ComputeWriteIndexFromReadIndex(
        readIndexProbe,
        _ProbeVolumeResolution,
        _ProbeVolumeResolutionInverse,
        _ProbeVolumeAtlasOctahedralDepthScaleBias,
        _ProbeVolumeAtlasOctahedralDepthResolutionAndInverse * (1.0f / 8.0f) // TODO: Cleanup: since we are computing the probe, rather than the pixel we are in, we need to divide scale bias by octahedral map resolution.
    );

    uint2 writeIndexOctahedral = uint2(readIndexOctahedral % 8, readIndexOctahedral / 8);
    uint2 writeIndex = writeIndexProbe * 8 + writeIndexOctahedral;

    // Depth sample is placed in X component, in order to be filtered into a mean, variance distribution and stored into XY in a following pass.
    _ProbeVolumeAtlasOctahedralDepthWriteTexture[uint2(writeIndex.x, writeIndex.y)] = float2(_ProbeVolumeAtlasOctahedralDepthReadBuffer[readIndex], 0.0f);
}