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;
        float4 gCameraRight;
        float4 gCameraUp;
    };

    cbuffer MaterialConstants
    {
        float4 gSplatParams;
    };

    struct GaussianSplatOtherData
    {
        float4 rotation;
        float4 scaleReserved;
    };

    StructuredBuffer<float3> GaussianSplatPositions;
    StructuredBuffer<GaussianSplatOtherData> GaussianSplatOther;
    StructuredBuffer<float4> GaussianSplatColor;

    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;
        const float2 corner = ResolveQuadCorner(vertexId);
        const float3 localCenter = GaussianSplatPositions[instanceId];
        const GaussianSplatOtherData otherData = GaussianSplatOther[instanceId];
        const float4 colorOpacity = GaussianSplatColor[instanceId];

        const float3 worldCenter = mul(gModelMatrix, float4(localCenter, 1.0)).xyz;
        const float maxAxisScale =
            max(max(otherData.scaleReserved.x, otherData.scaleReserved.y), otherData.scaleReserved.z);
        const float radius = max(maxAxisScale * gSplatParams.x, 0.0001);
        const float3 worldPosition =
            worldCenter +
            gCameraRight.xyz * (corner.x * radius) +
            gCameraUp.xyz * (corner.y * radius);

        output.position = mul(gProjectionMatrix, mul(gViewMatrix, float4(worldPosition, 1.0)));
        output.localUv = corner;
        output.colorOpacity = colorOpacity;
        return output;
    }

    float4 MainPS(VSOutput input) : SV_TARGET
    {
        const float radiusSq = dot(input.localUv, input.localUv);
        if (radiusSq > 1.0)
        {
            discard;
        }

        const float gaussianFalloff = exp(-radiusSq * 2.5);
        const float alpha = saturate(input.colorOpacity.a * gSplatParams.y * gaussianFalloff);
        if (alpha <= 0.001)
        {
            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
        }
    }
}