#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#pragma kernel ComputeVolumetricMaterialRenderingParameters

// #pragma enable_d3d11_debug_symbols

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/Builtin/BuiltinData.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/VolumetricLighting/HDRenderPipeline.VolumetricLighting.cs.hlsl"

uint _ViewCount;

// Readonly buffers
StructuredBuffer<OrientedBBox> _VolumeBounds;
StructuredBuffer<LocalVolumetricFogEngineData> _VolumeData;
ByteAddressBuffer _VolumetricVisibleGlobalIndicesBuffer;

RWBuffer<uint> _VolumetricGlobalIndirectArgsBuffer : register( u0 );   // Indirect arguments have to be in a _buffer_, not a structured buffer
RWByteAddressBuffer _VolumetricGlobalIndirectionBuffer : register( u1 );
RWStructuredBuffer<VolumetricMaterialRenderingData> _VolumetricMaterialData;

int DistanceToSlice(float distance)
{
    float de = _VBufferRcpSliceCount; // Log-encoded distance between slices

    float e1 = EncodeLogarithmicDepthGeneralized(distance, _VBufferDistanceEncodingParams);
    e1 -= de;
    e1 /= de;

    return int(e1 - 0.5);
}

float DepthDistance(float3 position)
{
    float cameraDistance = length(position);
    float3 viewDirWS = normalize(position);

    float3 cameraForward = -UNITY_MATRIX_V[2].xyz;
    float depth = cameraDistance * dot(viewDirWS, cameraForward);

    return depth;
}

// https://iquilezles.org/www/articles/distfunctions/distfunctions.htm
float DistanceToAABB(float3 p, float3 b)
{
    float3 q = abs(p) - b;
    return length(max(q, 0.0)) + min(max(q.x, max(q.y, q.z)), 0.0);
}

// Optimized version of https://www.sciencedirect.com/topics/computer-science/oriented-bounding-box
float DistanceDistanceToOBB(float3 p, OrientedBBox obb)
{
    float3 offset = p - obb.center;
    float3 boxForward = normalize(cross(obb.right, obb.up));
    float3 axisAlignedPoint = float3(dot(offset, normalize(obb.right)), dot(offset, normalize(obb.up)), dot(offset, boxForward));

    return DistanceToAABB(axisAlignedPoint, float3(obb.extentX, obb.extentY, obb.extentZ));
}

float3 ComputeCubeVertexPositionRWS(OrientedBBox obb, float3 minPositionRWS, float3 vertexMask)
{
    float3x3 obbFrame   = float3x3(obb.right, obb.up, cross(obb.right, obb.up));
    float3   obbExtents = float3(obb.extentX, obb.extentY, obb.extentZ);
    return mul((vertexMask * 2 - 1) * obbExtents, (obbFrame)) + obb.center;
}

// vertexMask contains the vertices of the cube in this order
//    6 .+------+ 7
//    .' |    .'|
// 2 +---+--+'3 |
//   |   |  |   |
//   | 4.+--+---+ 5
//   |.'    | .'
// 0 +------+' 1

static const float3 vertexMask[8] = {
    float3(0, 0, 0),
    float3(1, 0, 0),
    float3(0, 1, 0),
    float3(1, 1, 0),
    float3(0, 0, 1),
    float3(1, 0, 1),
    float3(0, 1, 1),
    float3(1, 1, 1),
};

// Sort cube indices to respect the winding order described in 'A COMPARISON OF GPU BOX-PLANE INTERSECTION ALGORITHMSFOR DIRECT VOLUME RENDERING' section 3
// We do this in a compute so we don't need to sample 3 index tables in the vertex shader
static const uint vertexIndicesMap[8 * 8] = {
    0, 1, 4, 2, 3, 5, 6, 7, // front vertex = 1
    1, 5, 3, 0, 4, 7, 2, 6, // front vertex = 1
    2, 3, 0, 6, 7, 1, 4, 5, // front vertex = 2
    3, 7, 1, 2, 6, 5, 0, 4, // front vertex = 3
    4, 0, 5, 6, 2, 1, 7, 3, // front vertex = 4
    5, 1, 7, 4, 0, 3, 6, 2, // front vertex = 5
    6, 2, 4, 7, 3, 0, 5, 1, // front vertex = 6
    7, 6, 5, 3, 2, 4, 1, 0, // front vertex = 7
};

void ComputeCubeVerticesOrder(uint volumeIndex)
{
    OrientedBBox obb = _VolumeBounds[volumeIndex]; // the OBB is in world space
    float3 obbExtents = float3(obb.extentX, obb.extentY, obb.extentZ);

    // Find the min and max distance value from vertices
    float3 minPositionRWS = obb.center - obbExtents;
    float minVertexDistance = -1;
    int frontVertexIndex = 0; // Compute the index of the vertex in front of the camera, needed for the slicing algorithm
    int i;
    for (i = 0; i < 8; i++)
    {
        // Select the correct starting vertex to sort the cube vertices
        float3 vertexNormal = normalize(vertexMask[i] * 2 - 1);
        float3 cameraForward = -UNITY_MATRIX_V[2].xyz;
        float d = dot(vertexNormal, cameraForward);
        if (d > minVertexDistance)
        {
            minVertexDistance = d;
            frontVertexIndex = i;
        }
    }

    // Write the cube vertices in a certain order respecting the winding described in vertexIndicesMap.
    uint vertexIndexMapOffset = frontVertexIndex * 8;
    for (i = 0; i < 8; i++)
    {
        float3 vertex = vertexMask[vertexIndicesMap[vertexIndexMapOffset + i]];
        _VolumetricMaterialData[volumeIndex].obbVertexPositionWS[i].xyz = ComputeCubeVertexPositionRWS(obb, minPositionRWS, vertex);
    }
}

// Generate the compute buffer to dispatch the indirect draw
[numthreads(32, 1, 1)]
void ComputeVolumetricMaterialRenderingParameters(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    if (dispatchThreadID.x >= _VolumeCount * _ViewCount)
        return;

    uint volumeIndex = dispatchThreadID.x % _VolumeCount;
    uint globalBufferWriteIndex = _VolumetricVisibleGlobalIndicesBuffer.Load(volumeIndex << 2);
    uint viewIndex = dispatchThreadID.x / _VolumeCount;
    uint materialDataIndex = volumeIndex + viewIndex * _VolumeCount;

#if defined(UNITY_STEREO_INSTANCING_ENABLED)
    unity_StereoEyeIndex = viewIndex;
#endif

    // Sort cube vertices for slicing in vertex
#if USE_VERTEX_CUBE_SLICING
    ComputeCubeVerticesOrder(volumeIndex);
#endif

    OrientedBBox obb = _VolumeBounds[volumeIndex]; // the OBB is in world space
    float3 obbExtents = float3(obb.extentX, obb.extentY, obb.extentZ);

    float2 minPositionVS = 1;
    float2 maxPositionVS = -1;

    float cameraDistanceToOBB = DistanceDistanceToOBB(GetCameraPositionWS(), obb);

    if (cameraDistanceToOBB <= 0)
    {
        minPositionVS = -1;
        maxPositionVS = 1;
    }

    // Find the min and max distance value from vertices
    float3 minPositionRWS = obb.center - obbExtents;
    float maxVertexDepth = 0;
    int i;
    for (i = 0; i < 8; i++)
    {
        float3 position = ComputeCubeVertexPositionRWS(obb, minPositionRWS, vertexMask[i]);
        float distance = length(position);
        maxVertexDepth = max(maxVertexDepth, distance);

        if (cameraDistanceToOBB > 0)
        {
            float4 positionCS = TransformWorldToHClip(position);

            // clamp positionCS inside the view in case the point is behind the camera
            if (positionCS.w < 0)
            {
                minPositionVS = -1;
                maxPositionVS = 1;
            }
            else
            {
                positionCS.xy /= positionCS.w;
                minPositionVS = min(positionCS.xy, minPositionVS);
                maxPositionVS = max(positionCS.xy, maxPositionVS);
            }
        }
    }

    minPositionVS = clamp(minPositionVS, -1, 1);
    maxPositionVS = clamp(maxPositionVS, -1, 1);

    float vBufferNearPlane = DecodeLogarithmicDepthGeneralized(0, _VBufferDistanceDecodingParams);
    float minBoxDistance = max(vBufferNearPlane, cameraDistanceToOBB);
    int startSliceIndex = clamp(DistanceToSlice(minBoxDistance), 0, int(_MaxSliceCount));
    int stopSliceIndex = DistanceToSlice(maxVertexDepth);
    uint sliceCount = clamp(stopSliceIndex - startSliceIndex, 0, int(_MaxSliceCount) - startSliceIndex);

    _VolumetricGlobalIndirectArgsBuffer[globalBufferWriteIndex * 5 + 0] = 6; // IndexCountPerInstance
    _VolumetricGlobalIndirectArgsBuffer[globalBufferWriteIndex * 5 + 1] = sliceCount * _ViewCount; // InstanceCount
    _VolumetricGlobalIndirectArgsBuffer[globalBufferWriteIndex * 5 + 2] = 0; // StartIndexLocation
    _VolumetricGlobalIndirectArgsBuffer[globalBufferWriteIndex * 5 + 3] = 0; // BaseVertexLocation
    _VolumetricGlobalIndirectArgsBuffer[globalBufferWriteIndex * 5 + 4] = 0; // StartInstanceLocation

    // Provide smaller buffer index in the global indirection buffer to sample those buffers in the vertex during the voxelization.
    _VolumetricGlobalIndirectionBuffer.Store(globalBufferWriteIndex << 2, volumeIndex);

    _VolumetricMaterialData[materialDataIndex].sliceCount = sliceCount;
    _VolumetricMaterialData[materialDataIndex].startSliceIndex = startSliceIndex;
    _VolumetricMaterialData[materialDataIndex].viewSpaceBounds = float4(minPositionVS, maxPositionVS - minPositionVS);
}