Shader "Builtin Gaussian Splat"
{
    Properties
    {
        _PointScale ("Point Scale", Float) = 1.0
        _OpacityScale ("Opacity Scale", Float) = 1.0
    }
    HLSLINCLUDE
    cbuffer PerObjectConstants
    {
        float4x4 gProjectionMatrix;
        float4x4 gViewMatrix;
        float4x4 gModelMatrix;
        float4x4 gWorldToObjectMatrix;
        float4 gCameraRight;
        float4 gCameraUp;
        float4 gCameraWorldPos;
        float4 gScreenParams;
        float4 gSplatParams;
    };

    cbuffer MaterialConstants
    {
        float4 gPointScaleParams;
        float4 gOpacityScaleParams;
    };

    struct GaussianSplatViewData
    {
        float4 clipCenter;
        float4 ellipseAxisU;
        float4 ellipseAxisV;
        float4 colorOpacity;
    };

    StructuredBuffer<uint> GaussianSplatOrderBuffer;
    StructuredBuffer<GaussianSplatViewData> GaussianSplatViewDataBuffer;

    struct VSOutput
    {
        float4 position : SV_POSITION;
        float2 localUv : TEXCOORD0;
        float4 colorOpacity : TEXCOORD1;
    };

    float2 ResolveQuadCorner(uint vertexId)
    {
        switch (vertexId)
        {
        case 0u: return float2(-1.0, -1.0);
        case 1u: return float2( 1.0, -1.0);
        case 2u: return float2( 1.0,  1.0);
        case 3u: return float2(-1.0, -1.0);
        case 4u: return float2( 1.0,  1.0);
        default: return float2(-1.0,  1.0);
        }
    }

    VSOutput MainVS(uint vertexId : SV_VertexID, uint instanceId : SV_InstanceID)
    {
        VSOutput output = (VSOutput)0;
        const uint splatIndex = GaussianSplatOrderBuffer[instanceId];
        const GaussianSplatViewData viewData = GaussianSplatViewDataBuffer[splatIndex];

        if (viewData.clipCenter.w <= 0.0)
        {
            const float nanValue = asfloat(0x7fc00000u);
            output.position = float4(nanValue, nanValue, nanValue, nanValue);
            return output;
        }

        const float2 quadPos = ResolveQuadCorner(vertexId) * 2.0;
        const float2 ellipseOffsetPixels =
            (quadPos.x * viewData.ellipseAxisU.xy + quadPos.y * viewData.ellipseAxisV.xy) * gPointScaleParams.x;
        const float2 clipOffset =
            ellipseOffsetPixels * (2.0 / max(gScreenParams.xy, float2(1.0, 1.0))) * viewData.clipCenter.w;

        output.position = viewData.clipCenter;
        output.position.xy += clipOffset;
        output.localUv = quadPos;
        output.colorOpacity = viewData.colorOpacity;
        return output;
    }

    float4 MainPS(VSOutput input) : SV_TARGET
    {
        const float alpha =
            saturate(exp(-dot(input.localUv, input.localUv)) * input.colorOpacity.a * gOpacityScaleParams.x);
        if (alpha <= (1.0 / 255.0))
        {
            discard;
        }

        return float4(input.colorOpacity.rgb, alpha);
    }
    ENDHLSL

    SubShader
    {
        Tags { "Queue" = "Transparent" }
        Pass
        {
            Name "GaussianSplat"
            Tags { "LightMode" = "GaussianSplat" }
            Cull Off
            ZWrite Off
            ZTest LEqual
            Blend SrcAlpha OneMinusSrcAlpha
            HLSLPROGRAM
            #pragma target 4.5
            #pragma vertex MainVS
            #pragma fragment MainPS
            ENDHLSL
        }
    }
}