//--------------------------------------------------------------------------------------------------
// Definitions
//--------------------------------------------------------------------------------------------------

// #pragma enable_d3d11_debug_symbols
#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#pragma kernel VolumeVoxelization

#pragma multi_compile _ SUPPORT_WATER_ABSORPTION

#define GROUP_SIZE_1D     8

//--------------------------------------------------------------------------------------------------
// Included headers
//--------------------------------------------------------------------------------------------------

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/GeometricTools.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/VolumeRendering.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Core/Utilities/GeometryUtils.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/VolumetricLighting/HDRenderPipeline.VolumetricLighting.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/LightLoop/LightLoopDef.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Water/Shaders/UnderWaterUtilities.hlsl"

//--------------------------------------------------------------------------------------------------
// Inputs & outputs
//--------------------------------------------------------------------------------------------------

StructuredBuffer<OrientedBBox>            _VolumeBounds;
StructuredBuffer<LocalVolumetricFogEngineData> _VolumeData;

RW_TEXTURE3D(float4, _VBufferDensity); // RGB = sqrt(scattering), A = sqrt(extinction)

//--------------------------------------------------------------------------------------------------
// Implementation
//--------------------------------------------------------------------------------------------------

// Jittered ray with screen-space derivatives.
struct JitteredRay
{
    float3 originWS;
    float3 centerDirWS;
    float3 jitterDirWS;
    float3 xDirDerivWS;
    float3 yDirDerivWS;
};

void FillVolumetricDensityBuffer(uint2 voxelCoord2D, JitteredRay ray)
{
    float t0 = DecodeLogarithmicDepthGeneralized(0, _VBufferDistanceDecodingParams);
    float de = _VBufferRcpSliceCount; // Log-encoded distance between slices

    #ifdef SUPPORT_WATER_ABSORPTION
    float waterDistance = FLT_MIN;
    float2 positionSS = voxelCoord2D * _VBufferViewportSize.zw * _ScreenSize.xy;
    bool underWater = ray.originWS.y < _UnderWaterUpHeight.w || IsUnderWater(positionSS);
    if (underWater && (GetStencilValue(LOAD_TEXTURE2D_X(_StencilTexture, positionSS)) & STENCILUSAGE_WATER_SURFACE) != 0)
    {
        waterDistance = LinearEyeDepth(LOAD_TEXTURE2D_X(_RefractiveDepthBuffer, positionSS).r, _ZBufferParams);
        waterDistance = waterDistance * rcp(dot(ray.centerDirWS, GetViewForwardDir()));
    }
    else
        waterDistance = underWater ? FLT_MAX : 0.0f;
    #endif


    for (uint slice = 0; slice < _VBufferSliceCount; slice++)
    {
        uint3 voxelCoord = uint3(voxelCoord2D, slice + _VBufferSliceCount * unity_StereoEyeIndex);

        float e1 = slice * de + de; // (slice + 1) / sliceCount
        float t1 = DecodeLogarithmicDepthGeneralized(e1, _VBufferDistanceDecodingParams);
        float dt = t1 - t0;
        float t  = t0 + 0.5 * dt;

        float3 voxelCenterWS = ray.originWS + t * ray.centerDirWS;

        // TODO: the fog value at the center is likely different from the average value across the voxel.
        // Compute the average value.
        float fragmentHeight   = voxelCenterWS.y;
        float heightMultiplier = ComputeHeightFogMultiplier(fragmentHeight, _HeightFogBaseHeight, _HeightFogExponents);

        // Start by sampling the height fog.
        float3 voxelScattering = _HeightFogBaseScattering.xyz * heightMultiplier;
        float  voxelExtinction = _HeightFogBaseExtinction * heightMultiplier;

        #if SUPPORT_WATER_ABSORPTION
        if (t < waterDistance)
        {
            voxelScattering = _UnderWaterScatteringExtinction.xyz;
            voxelExtinction = _UnderWaterScatteringExtinction.w * (1 - ray.centerDirWS.y);
        }
        #endif

        _VBufferDensity[voxelCoord] = float4(voxelScattering, voxelExtinction);

        t0 = t1;
    }
}

[numthreads(GROUP_SIZE_1D, GROUP_SIZE_1D, 1)]
void VolumeVoxelization(uint3 dispatchThreadId : SV_DispatchThreadID,
                        uint2 groupId          : SV_GroupID,
                        uint2 groupThreadId    : SV_GroupThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    // Reminder: our voxels are sphere-capped right frustums (truncated right pyramids).
    // The curvature of the front and back faces is quite gentle, so we can use
    // the right frustum approximation (thus the front and the back faces are squares).
    // Note, that since we still rely on the perspective camera model, pixels at the center
    // of the screen correspond to larger solid angles than those at the edges.
    // Basically, sizes of front and back faces depend on the XY coordinate.
    // https://www.desmos.com/calculator/i3rkesvidk

    float3 F = GetViewForwardDir();
    float3 U = GetViewUpDir();

    uint2 voxelCoord = dispatchThreadId.xy;
    float2 centerCoord = voxelCoord + float2(0.5, 0.5);

    // Compute a ray direction s.t. ViewSpace(rayDirWS).z = 1.
    float3 rayDirWS       = mul(-float4(centerCoord, 1, 1), _VBufferCoordToViewDirWS[unity_StereoEyeIndex]).xyz;
    float3 rightDirWS     = cross(rayDirWS, U);
    float  rcpLenRayDir   = rsqrt(dot(rayDirWS, rayDirWS));
    float  rcpLenRightDir = rsqrt(dot(rightDirWS, rightDirWS));

    JitteredRay ray;
    ray.originWS    = GetCurrentViewPosition();
    ray.centerDirWS = rayDirWS * rcpLenRayDir; // Normalize

    float FdotD = dot(F, ray.centerDirWS);
    float unitDistFaceSize = _VBufferUnitDepthTexelSpacing * FdotD * rcpLenRayDir;

    ray.xDirDerivWS = rightDirWS * (rcpLenRightDir * unitDistFaceSize); // Normalize & rescale
    ray.yDirDerivWS = cross(ray.xDirDerivWS, ray.centerDirWS); // Will have the length of 'unitDistFaceSize' by construction
    ray.jitterDirWS = ray.centerDirWS; // TODO

    ApplyCameraRelativeXR(ray.originWS);

    FillVolumetricDensityBuffer(voxelCoord, ray);
}