Shader "Builtin Volumetric"
{
    Properties
    {
        _Tint ("Tint", Color) = (1,1,1,1)
        _DensityScale ("Density Scale", Float) = 0.2
        _StepSize ("Step Size", Float) = 1.0
        _MaxSteps ("Max Steps", Float) = 2000.0
        _AmbientStrength ("Ambient Strength", Float) = 0.005
        _LightDirection ("Light Direction", Vector) = (0.5, 0.8, 0.3, 0.0)
        _LightSamples ("Light Samples", Float) = 8.0
    }
    HLSLINCLUDE
    #define PNANOVDB_HLSL
    #define PNANOVDB_ADDRESS_32
    #include "../../../third_party/nanovdb/shaders/PNanoVDB.hlsl"

    cbuffer PerObjectConstants
    {
        float4x4 gProjectionMatrix;
        float4x4 gViewMatrix;
        float4x4 gModelMatrix;
        float4x4 gInverseModelMatrix;
        float4 gCameraWorldPosition;
        float4 gVolumeBoundsMin;
        float4 gVolumeBoundsMax;
    };

    cbuffer LightingConstants
    {
        float4 gMainLightDirectionAndIntensity;
        float4 gMainLightColorAndFlags;
    };

    cbuffer MaterialConstants
    {
        float4 gVolumeTint;
        float4 gDensityScale;
        float4 gStepSize;
        float4 gMaxSteps;
        float4 gAmbientStrength;
        float4 gLightDirection;
        float4 gLightSamples;
    };

    StructuredBuffer<uint> VolumeData;

    struct VSInput
    {
        float3 position : POSITION;
    };

    struct PSInput
    {
        float4 position : SV_POSITION;
        float3 localPosition : TEXCOORD0;
    };

    float3 MapUnitCubeToVolumeBounds(float3 unitCubePosition)
    {
        const float3 uvw = saturate(unitCubePosition + float3(0.5, 0.5, 0.5));
        return lerp(gVolumeBoundsMin.xyz, gVolumeBoundsMax.xyz, uvw);
    }

    PSInput MainVS(VSInput input)
    {
        PSInput output;
        const float3 localPosition = MapUnitCubeToVolumeBounds(input.position);
        const float4 worldPosition = mul(gModelMatrix, float4(localPosition, 1.0f));
        output.position = mul(gProjectionMatrix, mul(gViewMatrix, worldPosition));
        output.localPosition = localPosition;
        return output;
    }

    bool IntersectAabb(
        float3 rayOrigin,
        float3 rayDirection,
        float3 boxMin,
        float3 boxMax,
        out float tEnter,
        out float tExit)
    {
        const float3 invDirection = 1.0f / max(abs(rayDirection), 1e-6f) * sign(rayDirection);
        const float3 t0 = (boxMin - rayOrigin) * invDirection;
        const float3 t1 = (boxMax - rayOrigin) * invDirection;
        const float3 tMin3 = min(t0, t1);
        const float3 tMax3 = max(t0, t1);

        tEnter = max(max(tMin3.x, tMin3.y), tMin3.z);
        tExit = min(min(tMax3.x, tMax3.y), tMax3.z);
        return tExit >= tEnter && tExit > 0.0f;
    }

    bool IsPointInsideAabb(float3 samplePosition, float3 boxMin, float3 boxMax)
    {
        return all(samplePosition >= boxMin) && all(samplePosition <= boxMax);
    }

    struct NanoVolume
    {
        pnanovdb_grid_handle_t grid;
        pnanovdb_readaccessor_t accessor;
    };

    void InitNanoVolume(out NanoVolume volume)
    {
        volume.grid.address.byte_offset = 0;
        const pnanovdb_tree_handle_t tree = pnanovdb_grid_get_tree(VolumeData, volume.grid);
        const pnanovdb_root_handle_t root = pnanovdb_tree_get_root(VolumeData, tree);
        pnanovdb_readaccessor_init(volume.accessor, root);
    }

    float GetValueCoord(inout pnanovdb_readaccessor_t accessor, float3 position)
    {
        const pnanovdb_vec3_t p = position;
        const pnanovdb_coord_t ijk = pnanovdb_hdda_pos_to_ijk(p);
        const pnanovdb_address_t address =
            pnanovdb_readaccessor_get_value_address(
                PNANOVDB_GRID_TYPE_FLOAT,
                VolumeData,
                accessor,
                ijk);
        return pnanovdb_read_float(VolumeData, address);
    }

    uint GetDimCoord(inout pnanovdb_readaccessor_t accessor, float3 position)
    {
        const pnanovdb_vec3_t p = position;
        const pnanovdb_coord_t ijk = pnanovdb_hdda_pos_to_ijk(p);
        return pnanovdb_readaccessor_get_dim(
            PNANOVDB_GRID_TYPE_FLOAT,
            VolumeData,
            accessor,
            ijk);
    }

    bool GetHddaHit(
        inout pnanovdb_readaccessor_t accessor,
        inout float tMin,
        float3 origin,
        float3 direction,
        float tMax,
        out float valueAtHit)
    {
        const pnanovdb_vec3_t pOrigin = origin;
        const pnanovdb_vec3_t pDirection = direction;
        float tHit = 0.0f;
        const bool hit = pnanovdb_hdda_tree_marcher(
            PNANOVDB_GRID_TYPE_FLOAT,
            VolumeData,
            accessor,
            pOrigin, tMin,
            pDirection, tMax,
            tHit,
            valueAtHit);
        tMin = tHit;
        return hit;
    }

    float ComputeVolumetricShadow(
        float3 localPosition,
        float densityScale,
        float3 localLightDirection,
        float lightSamples,
        inout pnanovdb_readaccessor_t accessor)
    {
        if (lightSamples < 1.0f) {
            return 1.0f;
        }

        float shadow = 1.0f;
        float stepSize = 1.0f;

        for (int stepIndex = 0; stepIndex < 10; ++stepIndex) {
            const float3 samplePosition = localPosition + stepSize * localLightDirection;
            const float sigmaS =
                max(GetValueCoord(accessor, samplePosition) * densityScale, 0.0f);
            const float sigmaE = max(0.000001f, sigmaS) * 0.3f;
            shadow *= exp(-sigmaE * stepSize);
            stepSize *= 2.0f;
        }

        return shadow;
    }

    float4 MainPS(PSInput input) : SV_TARGET
    {
        const float3 boxMin = gVolumeBoundsMin.xyz;
        const float3 boxMax = gVolumeBoundsMax.xyz;
        const float densityScale = max(gDensityScale.x, 0.0f);
        const float stepSize = max(gStepSize.x, 0.001f);
        const int maxSteps = max((int)gMaxSteps.x, 1);
        const float ambientStrength = max(gAmbientStrength.x, 0.0f);
        const float lightSamples = max(gLightSamples.x, 0.0f);

        const float3 cameraLocalPosition =
            mul(gInverseModelMatrix, float4(gCameraWorldPosition.xyz, 1.0f)).xyz;
        const float3 rayDirection = normalize(input.localPosition - cameraLocalPosition);

        float tEnter = 0.0f;
        float tExit = 0.0f;
        if (!IntersectAabb(cameraLocalPosition, rayDirection, boxMin, boxMax, tEnter, tExit)) {
            discard;
        }

        const bool cameraInside = IsPointInsideAabb(cameraLocalPosition, boxMin, boxMax);
        const float surfaceT = cameraInside ? tExit : max(tEnter, 0.0f);
        const float3 expectedSurfacePosition = cameraLocalPosition + rayDirection * surfaceT;
        if (distance(expectedSurfacePosition, input.localPosition) > stepSize * 1.5f + 0.01f) {
            discard;
        }

        float t = max(surfaceT, 0.01f);
        const float marchEnd = tExit;
        float hitValue = 0.0f;

        NanoVolume volume;
        InitNanoVolume(volume);
        if (!GetHddaHit(volume.accessor, t, cameraLocalPosition, rayDirection, marchEnd, hitValue)) {
            discard;
        }

        float skipDistance = 0.0f;
        float3 integratedLight = 0.0f.xxx;
        float transmittance = 1.0f;
        float accumulatedDensity = 0.0f;

        float3 resolvedLightDirectionWS = gLightDirection.xyz;
        float3 resolvedLightRadiance = 1.0f.xxx;
        if (gMainLightColorAndFlags.a >= 0.5f &&
            length(gMainLightDirectionAndIntensity.xyz) > 0.0001f) {
            resolvedLightDirectionWS = gMainLightDirectionAndIntensity.xyz;
            resolvedLightRadiance =
                max(gMainLightColorAndFlags.rgb, 0.0f.xxx) *
                max(gMainLightDirectionAndIntensity.w, 0.0f);
        }

        float3 localLightDirection =
            mul((float3x3)gInverseModelMatrix, resolvedLightDirectionWS);
        const float localLightLength = length(localLightDirection);
        if (localLightLength > 0.0001f) {
            localLightDirection /= localLightLength;
        } else {
            localLightDirection = normalize(float3(0.5f, 0.8f, 0.3f));
        }

        [loop]
        for (int stepIndex = 0; stepIndex < maxSteps; ++stepIndex) {
            if (t >= marchEnd) {
                break;
            }

            float3 localPosition = cameraLocalPosition + rayDirection * t;
            const uint dim = GetDimCoord(volume.accessor, localPosition);
            if (dim > 1u) {
                skipDistance = 15.0f;
                t += skipDistance;
                continue;
            }

            float density = GetValueCoord(volume.accessor, localPosition) * densityScale;
            if (density < 0.01f) {
                skipDistance = 5.0f;
                t += skipDistance;
                continue;
            }

            if (skipDistance > 0.0f) {
                t -= skipDistance * 0.8f;
                localPosition = cameraLocalPosition + rayDirection * t;
                skipDistance = 0.0f;
            }

            const float sigmaS = density;
            const float sigmaE = max(0.000001f, sigmaS);
            accumulatedDensity += sigmaS;

            const float shadow =
                ComputeVolumetricShadow(
                    localPosition,
                    densityScale,
                    localLightDirection,
                    lightSamples,
                    volume.accessor);
            const float3 S = sigmaS * shadow.xxx * resolvedLightRadiance;
            const float3 integratedSegment = (S - S * exp(-sigmaE * stepSize)) / sigmaE;
            integratedLight += transmittance * integratedSegment;
            transmittance *= exp(-sigmaE * stepSize);

            if (accumulatedDensity > 1.0f) {
                break;
            }

            if (transmittance < 0.05f) {
                transmittance = 0.0f;
                break;
            }

            t += stepSize;
        }

        const float3 ambientLight = ambientStrength.xxx;
        float3 finalColor = (integratedLight + ambientLight) * accumulatedDensity;
        finalColor *= gVolumeTint.rgb;
        finalColor = pow(max(finalColor, 0.0f.xxx), 1.0f / 2.2f);

        const float alpha = saturate(accumulatedDensity) * saturate(gVolumeTint.a);
        if (alpha <= 0.001f) {
            discard;
        }

        return float4(finalColor, alpha);
    }
    ENDHLSL

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