#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/PostProcessing/Shaders/DepthOfFieldCommon.hlsl"

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#pragma kernel KParametricBlurKernel          MAIN=KParametricBlurKernel        GROUP_SIZE=64

// Unused, keeping this just in case
#pragma kernel KParametricFloodfillKernel     MAIN=KParametricFloodfillKernel   GROUP_SIZE=16

// Each uint holds two fp16 (x,y) normalized coordinates
RWStructuredBuffer<uint> _BokehKernel;

CBUFFER_START(cb0)
    float4 _Params1;
    float4 _Params2;
CBUFFER_END

#define SampleCount    _Params1.x
#define NGonFactor     _Params1.y
#define BladeCount     _Params1.z
#define Rotation       _Params1.w
#define Anamorphism    _Params2.x

// Input is coordinates on a normalized square; returns sample coordinates in range [-1,1]
float2 GetSample(float2 coords)
{
    // TODO: Reference implementation, reduce all of this
    // TODO: Fix uneven distribution of samples when moving from circle to NGon

    // Returns a sample point for a given coordinates in a normalized square
    // "A Low Distortion Map Between Disk and Square" [Shirley97] [Chiu97]
    float phi, r;
    float a = coords.x * 2.0 - 1.0;
    float b = coords.y * 2.0 - 1.0;

    if (a > -b)
    {
        if (a > b)
        {
            r = a;
            phi = (PI / 4.0) * (b / a);
        }
        else
        {
            r = b;
            phi = (PI / 4.0) * (2.0 - (a / b));
        }
    }
    else
    {
        if (a < b)
        {
            r = -a;
            phi = (PI / 4.0) * (4.0 + (b / a));
        }
        else
        {
            r = -b;

            if (b != 0.0)
            {
                phi = (PI / 4.0) * (6.0 - (a / b));
            }
            else
            {
                phi = 0.0;
            }
        }
    }

    // Transform to rotated ngon
    // "CryEngine 3 Graphics Gems" [Sousa13]
    float n = BladeCount;
    float nt = cos(PI / n);
    float dt = cos(phi - (TWO_PI / n) * floor((n * phi + PI) / TWO_PI));
    r = r * PositivePow(nt / dt, NGonFactor);
    float u = r * cos(phi - Rotation);
    float v = r * sin(phi - Rotation);

    v *= 1.0 + Anamorphism;
    u *= 1.0 - Anamorphism;

    return float2(u, v);
}

[numthreads(GROUP_SIZE, 1, 1)]
void MAIN(uint dispatchThreadId : SV_DispatchThreadID)
{
    if (dispatchThreadId >= uint(SampleCount * SampleCount))
        return;

    // 1D -> 2D for regular samples
    uint2 kernelCoords = uint2(
        dispatchThreadId % uint(SampleCount),
        dispatchThreadId / uint(SampleCount)
    );

    // Sample coordinates on a normalized square
    float2 pos = kernelCoords / (SampleCount - 1.0).xx;

    // Grab the sample (range [-1,1])
    float2 smp = GetSample(pos);

    // Pack two fp16 in a single uint, we don't need full fp32 precision
    _BokehKernel[dispatchThreadId] = PackKernelCoord(smp);
}