#pragma kernel ProbeVolumeAtlasOctahedralDepthConvolveKernel PROBE_VOLUME_ATLAS_OCTAHEDRAL_DEPTH_CONVOLVE_KERNEL=ProbeVolumeAtlasOctahedralDepthConvolveKernel

#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"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Packing.hlsl"

RWTexture2D<float2> _ProbeVolumeAtlasOctahedralDepthRWTexture;
float4 _ProbeVolumeAtlasOctahedralDepthScaleBiasTexels;
int _FilterSampleCount;
float _FilterSharpness;

// http://research.microsoft.com/en-us/um/people/johnsny/papers/sg.pdf
float SphericalGaussianEvaluateFromDirection(float3 sgMean, float sgSharpness, float3 evaluationDirection)
{
    return exp(dot(sgMean, evaluationDirection) * sgSharpness - sgSharpness);
}

float SphericalGaussianIntegralFromSharpness(float sharpness)
{
    float b = 2.0f * PI / sharpness;
    return exp(-2.0f * sharpness) * -b + b;
}

// Warning this needs to match with kBatchSize in ProbeVolumeLighting.cs
// Warning: This kernel numthreads must be an integer multiple of PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION
// We read + write from the same texture, so any partial work would pollute / cause feedback into neighboring chunks.
#define PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION (8)
[numthreads(PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION, PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION, 1)]
void PROBE_VOLUME_ATLAS_OCTAHEDRAL_DEPTH_CONVOLVE_KERNEL(uint2 groupThreadId : SV_GroupThreadID, uint2 groupId : SV_GroupID)
{
    uint2 octahedralDepthTexel = groupThreadId.xy;
    float2 octahedralDepthUV = (float2)octahedralDepthTexel * (1.0f / (float)PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION) + (0.5f / (float)PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION);

    uint2 probeIndex2D = groupId;
    uint2 writeIndexSW = probeIndex2D * PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION + _ProbeVolumeAtlasOctahedralDepthScaleBiasTexels.zw;
    uint2 writeIndex = octahedralDepthTexel + writeIndexSW;// + octahedralDepthTexel;

    // In order to generate a {mean, variance} distribution, we will convolve a weighted neighborhood of samples
    // via fibonacci spiral spherical gaussian importance sampling.
    // Really we are just importance sampling a spherical cap, whos solid angle is approximately equal to the
    // solid angle of our spherical gaussian filter.
    float3 filterCenterWS = UnpackNormalOctRectEncode(octahedralDepthUV * 2.0f - 1.0f);

    // Generate full basis for our spherical gaussian filter to sample within.
    // TODO: Dither rotate basis per pixel around filterCenterWS to reduce bias at the cost of noise.
    float3 basisNormal = filterCenterWS;
    float3 basisUp = (basisNormal.y < 0.999f) ? float3(0.0f, 1.0f, 0.0f) : float3(1.0f, 0.0f, 0.0f);
    float3 basisTangent = normalize(cross(basisUp, basisNormal));
    float3 basisBitangent = cross(basisNormal, basisTangent);
    float3x3 basis = float3x3(basisTangent, basisBitangent, basisNormal);

    float sampleCountInverse = 1.0f / (float)_FilterSampleCount;
    float GOLDEN_ANGLE = PI * (3.0f - sqrt(5.0f));
    float GOLDEN_ANGLE_HALF = GOLDEN_ANGLE * 0.5f;

    float filterSolidAngleSG = SphericalGaussianIntegralFromSharpness(_FilterSharpness);
    const float SOLID_ANGLE_SPHERE = 4.0f * PI;
    float cosThetaMaxRatio = filterSolidAngleSG / SOLID_ANGLE_SPHERE;

    float mean = 0.0f;
    float mean2 = 0.0f;
    float weightTotal = 0.0f;

    for (int i = 0; i < _FilterSampleCount; ++i)
    {
        float offset = float(i) + 0.5f;
        float theta = offset * GOLDEN_ANGLE + GOLDEN_ANGLE_HALF;
        float z = -(offset * sampleCountInverse * cosThetaMaxRatio * 2.0f - 1.0f);
        float r = sqrt(1.0f - z * z);

        float3 sampleCurrentDirectionOS = float3(
            r * cos(theta),
            r * sin(theta),
            z
        );

        float3 sampleCurrentDirectionWS = mul(sampleCurrentDirectionOS, basis);
        float sampleCurrentWeight = SphericalGaussianEvaluateFromDirection(filterCenterWS, _FilterSharpness, sampleCurrentDirectionWS);

        float2 sampleCurrentOctahedralDepthUV = PackNormalOctRectEncode(sampleCurrentDirectionWS) * 0.5f + 0.5f;
        uint2 sampleCurrentOctahedralDepthTexel = (uint2)(sampleCurrentOctahedralDepthUV * (float)PROBE_VOLUME_OCTAHEDRAL_DEPTH_RESOLUTION);

        uint2 sampleCurrentReadIndex2D = sampleCurrentOctahedralDepthTexel + writeIndexSW;

        float sampleCurrentDepth = _ProbeVolumeAtlasOctahedralDepthRWTexture[sampleCurrentReadIndex2D].x;

        mean += sampleCurrentDepth * sampleCurrentWeight;
        mean2 += sampleCurrentDepth * sampleCurrentDepth * sampleCurrentWeight;
        weightTotal += sampleCurrentWeight;
    }
    float filterNormalization = rcp(weightTotal);
    mean *= filterNormalization;
    mean2 *= filterNormalization;
    float variance = max(0.0f, mean2 - mean * mean); // avoid small negative values due to precision.

    // Sync group before writing to ensure we dont create race conditions for neighboring pixels that might still be resolving.
    GroupMemoryBarrierWithGroupSync();

    _ProbeVolumeAtlasOctahedralDepthRWTexture[writeIndex] = float2(mean, variance);
}