o3de/Gems/Atom/Feature/Common/Assets/Shaders/Materials/StandardMultilayerPBR/StandardMultilayerPBR_Common.azsli

/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#pragma once 

#include <Atom/RPI/Math.azsli>

// ------ Shader Options ----------------------------------------

option bool o_layer2_enabled;
option bool o_layer3_enabled;

COMMON_OPTIONS_PARALLAX(o_layer1_)
COMMON_OPTIONS_PARALLAX(o_layer2_)
COMMON_OPTIONS_PARALLAX(o_layer3_)

enum class DebugDrawMode { None, BlendMask, Displacement, FinalBlendWeights };
option DebugDrawMode o_debugDrawMode;

// If you modify this enum, you must update the BlendSourceUsesDisplacement function in StandardMultilayerPBR_Displacement.lua
enum class LayerBlendSource { BlendMaskTexture, BlendMaskVertexColors, Displacement, Displacement_With_BlendMaskTexture, Displacement_With_BlendMaskVertexColors, Fallback };
option LayerBlendSource o_layerBlendSource;

// ------ Blend Utilities ----------------------------------------
// This is mainly used to pass extra data to the GetDepth callback function during the parallax depth search.
// But since we have it, we use it in some other functions as well rather than passing it around.
static real3 s_blendMaskFromVertexStream;

// TODO: Consider storing the result of GetFinalLayerBlendSource() in a static similar to s_blendMaskFromVertexStream. That might give better performance when variants aren't used.

//! Returns the LayerBlendSource that will actually be used when rendering (not necessarily the same LayerBlendSource specified by the user)
LayerBlendSource GetFinalLayerBlendSource()
{    
    if(o_layerBlendSource == LayerBlendSource::BlendMaskTexture)
    {
        return o_layerBlendSource;
    }
    else if(o_layerBlendSource == LayerBlendSource::BlendMaskVertexColors)
    {
        if(o_blendMask_isBound)
        {
            return o_layerBlendSource;
        }
        else
        {
            return LayerBlendSource::Fallback;
        }
    }
    else if(o_layerBlendSource == LayerBlendSource::Displacement)
    {
        return o_layerBlendSource;
    }
    else if(o_layerBlendSource == LayerBlendSource::Displacement_With_BlendMaskTexture)
    {
        return o_layerBlendSource;
    }
    else if(o_layerBlendSource == LayerBlendSource::Displacement_With_BlendMaskVertexColors)
    {
        if(o_blendMask_isBound)
        {
            return o_layerBlendSource;
        }
        else
        {
            return LayerBlendSource::Displacement;
        }
    }
    else
    {
        return LayerBlendSource::Fallback;
    }
}

//! Return the applicable blend mask values from the blend mask texture or vertex colors, and filters out any that don't apply.
//!     layer1 is an implicit base layer
//!     layer2 mask is in the r channel
//!     layer3 mask is in the g channel
//!     b is reserved for perhaps a dedicated puddle layer
//! @param blendSource indicates where to get the blend mask from
//! @param blendMaskUv for sampling a blend mask texture, if that's the blend source
//! @param blendMaskVertexColors the vertex color values to use for the blend mask, if that's the blend source
//! @param uvDxDy for sampling the textures if derivatives are customized
//! @param customDerivatives whether or not using customized derivatives
//! @return the blend mask values, or 0 if there is no blend mask
real3 GetApplicableBlendMaskValues(const MaterialParameters params, LayerBlendSource blendSource, float2 blendMaskUv, float3 blendMaskVertexColors, float4 uvDxDy = float4(0.0, 0.0, 0.0, 0.0), bool customDerivatives = false)
{
    real3 blendSourceValues = real3(0,0,0);

    if(o_layer2_enabled || o_layer3_enabled)
    {
        switch(blendSource)
        {
            case LayerBlendSource::Displacement:
                // In this case the blend mask has no effect, returning (1,1,1) disables any impact of the mask.
                blendSourceValues = real3(1,1,1);
                break;
            case LayerBlendSource::BlendMaskTexture:
            case LayerBlendSource::Displacement_With_BlendMaskTexture:
                if (customDerivatives)
                {
                    blendSourceValues = real3(GetMaterialTexture(params.m_blendMaskTexture).SampleGrad(GetMaterialTextureSampler(), blendMaskUv, uvDxDy.xy, uvDxDy.zw).rgb);
                }
                else
                {
#ifdef CS_SAMPLERS
                    blendSourceValues = real3(GetMaterialTexture(params.m_blendMaskTexture).SampleLevel(GetMaterialTextureSampler(), blendMaskUv, 0).rgb);
#else
                    blendSourceValues = real3(GetMaterialTexture(params.m_blendMaskTexture).Sample(GetMaterialTextureSampler(), blendMaskUv).rgb);
#endif /* CS_SAMPLERS */
                }
                break;
            case LayerBlendSource::BlendMaskVertexColors:
            case LayerBlendSource::Displacement_With_BlendMaskVertexColors:
                blendSourceValues = blendMaskVertexColors;
                break;
            case LayerBlendSource::Fallback:
                break;
        }

        if(!o_layer2_enabled)
        {
            blendSourceValues.r = 0.0;
        }

        if(!o_layer3_enabled)
        {
            blendSourceValues.g = 0.0;
        }
    }

    return blendSourceValues;
}

//! When dealing with masks for displacement-based blending, we sometimes need to push the value below the min displacement to make it disappear.
real GetSubMinDisplacement(const MaterialParameters params)
{
    return real(params.m_displacementMin) - 0.001;
    //return params.m_displacementMin - max(params.m_displacementBlendDistance, 0.001);
}

real3 ApplyBlendMaskToDepthValues(const MaterialParameters params, real3 blendMaskValues, real3 layerDepthValues, real zeroMaskDisplacement);

//! Return the final blend weights to be used for rendering, based on the available data and configuration.
//! @param blendSource - indicates where to get the blend mask from
//! @param blendMaskValues - blend mask values as returned by GetApplicableBlendMaskValues()
//! @param layerDepthValues - the depth values for each layer, used if the blend source includes displacement. See GetLayerDepthValues().
//!                           Note the blendMaskValues will not be applied here, those should have already been applied to layerDepthValues.
//! @param layerDepthBlendDistance - controls how smoothly to blend layers 2 and 3 with the base layer, when the blend source includes displacement.
//!                                  When layers are close together their weights will be blended together, otherwise the highest layer will have the full weight.
//! @return The blend weights for each layer.
//!         Even though layer1 not explicitly specified in the blend mask data, it is explicitly included with the returned values.
//!         layer1 = r
//!         layer2 = g
//!         layer3 = b
real3 GetBlendWeights(LayerBlendSource blendSource, real3 blendMaskValues, real3 layerDepthValues, real layerDepthBlendDistance)
{
    real3 blendWeights;
    
    if(o_layer2_enabled || o_layer3_enabled)
    {
        if(LayerBlendSource::Displacement == blendSource || 
           LayerBlendSource::Displacement_With_BlendMaskTexture == blendSource || 
           LayerBlendSource::Displacement_With_BlendMaskVertexColors == blendSource)
        {
            // Calculate the blend weights based on displacement values...

            // The inputs are depth values, but we change them to height values to make the code a bit more intuitive.
            real3 layerHeightValues = -layerDepthValues;

            real highestPoint = layerHeightValues.x;
            if(o_layer2_enabled)
            {
                highestPoint = max(highestPoint, layerHeightValues.y);
            }
            if(o_layer3_enabled)
            {
                highestPoint = max(highestPoint, layerHeightValues.z);
            }
            
            if(layerDepthBlendDistance > 0.0001)
            {
                real lowestVisiblePoint = highestPoint - layerDepthBlendDistance;
                blendWeights = smoothstep(lowestVisiblePoint, highestPoint, layerHeightValues);
            
                if(!o_layer2_enabled)
                {
                    blendWeights.y = 0.0;
                }
            
                if(!o_layer3_enabled)
                {
                    blendWeights.z = 0.0;
                }
            }
            else
            {
                blendWeights = real3(layerHeightValues.x >= highestPoint ? 1.0 : 0.0,
                                        layerHeightValues.y >= highestPoint && o_layer2_enabled ? 1.0 : 0.0,
                                        layerHeightValues.z >= highestPoint && o_layer3_enabled ? 1.0 : 0.0);
            }
            
            // Calculate blend weights such that multiplying and adding them with layer data is equivalent
            // to lerping between each layer.
            //     final = lerp(final, layer1, blendWeights.r)
            //     final = lerp(final, layer2, blendWeights.g)
            //     final = lerp(final, layer3, blendWeights.b)

            blendWeights.y = (1 - blendWeights.z) * blendWeights.y;
            blendWeights.x = 1 - blendWeights.y - blendWeights.z;
        }
        else
        {
            // Calculate blend weights such that multiplying and adding them with layer data is equivalent
            // to lerping between each layer.
            //     final = lerp(final, layer1, blendWeights.r)
            //     final = lerp(final, layer2, blendWeights.g)
            //     final = lerp(final, layer3, blendWeights.b)

            blendWeights.b = blendMaskValues.g;
            blendWeights.g = (1.0 - blendMaskValues.g) * blendMaskValues.r;
            blendWeights.r = (1.0 - blendMaskValues.g) * (1.0 - blendMaskValues.r);
        }

    }
    else
    {
        blendWeights = real3(1,0,0);
    }

    return blendWeights;
}

float3 GetLayerDepthValues(const MaterialParameters params, float2 uv, float2 uv_ddx, float2 uv_ddy);

//! Return the final blend weights to be used for rendering, based on the available data and configuration.
//! Note this will sample the displacement maps in the case of LayerBlendSource::Displacement. If you have already
//! the layer depth values, use the GetBlendWeights() overload that takes layerDepthValues instead.
real3 GetBlendWeights(const MaterialParameters params, LayerBlendSource blendSource, float2 uv, float3 blendMaskVertexColors, float4 uvDxDy, bool customDerivatives)
{
    real3 layerDepthValues = real3(0,0,0);
    
    real3 blendMaskValues = GetApplicableBlendMaskValues(params, blendSource, uv, blendMaskVertexColors, uvDxDy, customDerivatives);

    if(blendSource == LayerBlendSource::Displacement ||
       blendSource == LayerBlendSource::Displacement_With_BlendMaskTexture ||
       blendSource == LayerBlendSource::Displacement_With_BlendMaskVertexColors)
    {
        bool useBlendMask = 
           LayerBlendSource::Displacement_With_BlendMaskTexture == blendSource || 
           LayerBlendSource::Displacement_With_BlendMaskVertexColors == blendSource;

#ifdef CS_SAMPLERS
        layerDepthValues = real3(GetLayerDepthValues(params, uv, float2(0, 0), float2(0, 0)));
#else
        layerDepthValues = real3(GetLayerDepthValues(params, uv, ddx_fine(uv), ddy_fine(uv)));
#endif /* CS_SAMPLERS */

        if(useBlendMask)
        {
            // Unlike the GetDepth() callback used for parallax, we don't just shift the values toward GetSubMinDisplacement(),
            // we shift extra to ensure that completely masked-out layers are not blended onto upper layers.
            layerDepthValues = ApplyBlendMaskToDepthValues(params, blendMaskValues, layerDepthValues, GetSubMinDisplacement(params) - real(params.m_displacementBlendDistance));
        }
    }
    
    return GetBlendWeights(blendSource, blendMaskValues, layerDepthValues, real(params.m_displacementBlendDistance));
}

real BlendLayers(real layer1, real layer2, real layer3, real3 blendWeights)
{
    return dot(real3(layer1, layer2, layer3), blendWeights);
}
real2 BlendLayers(real2 layer1, real2 layer2, real2 layer3, real3 blendWeights)
{
    return layer1 * blendWeights.r + layer2 * blendWeights.g + layer3 * blendWeights.b;
}
real3 BlendLayers(real3 layer1, real3 layer2, real3 layer3, real3 blendWeights)
{
    return layer1 * blendWeights.r + layer2 * blendWeights.g + layer3 * blendWeights.b;
}

// ------ Parallax Utilities ----------------------------------------

//! Returns the depth values for each layer.
float3 GetLayerDepthValues(const MaterialParameters params, float2 uv, float2 uv_ddx, float2 uv_ddy)
{
    float3 layerDepthValues = float3(0,0,0);
    
    // layer1
    {
        if(o_layer1_o_useHeightmap)
        {
            float2 layerUv = uv;
            if(params.m_parallaxUvIndex == 0)
            {
                layerUv = mul(params.m_layer1_m_uvMatrix, float3(uv, 1.0)).xy;
            }

            layerDepthValues.r = SampleDepthFromHeightmap(GetMaterialTexture(params.m_layer1_m_heightmap), GetMaterialTextureSampler(), layerUv, uv_ddx, uv_ddy).m_depth;
            layerDepthValues.r *= params.m_layer1_m_heightmapScale;
        }

        layerDepthValues.r -= params.m_layer1_m_heightmapOffset;
    }
    
    if(o_layer2_enabled)
    {
        if(o_layer2_o_useHeightmap)
        {
            float2 layerUv = uv;
            if(params.m_parallaxUvIndex == 0)
            {
                layerUv = mul(params.m_layer2_m_uvMatrix, float3(uv, 1.0)).xy;
            }

            layerDepthValues.g = SampleDepthFromHeightmap(GetMaterialTexture(params.m_layer2_m_heightmap), GetMaterialTextureSampler(), layerUv, uv_ddx, uv_ddy).m_depth;
            layerDepthValues.g *= params.m_layer2_m_heightmapScale;
        }

        layerDepthValues.g -= params.m_layer2_m_heightmapOffset;

    }
    
    if(o_layer3_enabled)
    {
        if(o_layer3_o_useHeightmap)
        {
            float2 layerUv = uv;
            if(params.m_parallaxUvIndex == 0)
            {
                layerUv = mul(params.m_layer3_m_uvMatrix, float3(uv, 1.0)).xy;
            }

            layerDepthValues.b = SampleDepthFromHeightmap(GetMaterialTexture(params.m_layer3_m_heightmap), GetMaterialTextureSampler(), layerUv, uv_ddx, uv_ddy).m_depth;
            layerDepthValues.b *= params.m_layer3_m_heightmapScale;
        }

        layerDepthValues.b -= params.m_layer3_m_heightmapOffset;

    }
    
    return layerDepthValues;
}


//! Uses a layer blend mask to further displace each layer's surface so that it disappears beyond the other surfaces.
//! Note the blend mask does not apply to the first layer, it is the implicit base layer. Layers 2 and 3 are masked by the r and g channels of the mask.
//! @param blendMaskValues layer mask values as returned by GetApplicableBlendMaskValues()
//! @param layerDepthValues layer depth values as returned by GetLayerDepthValues()
//! @param zeroMaskDisplacement the target displacement value that corresponds to a mask value of 0
//! @return new layer depth values that have been adjusted according to the layerMaskValues
real3 ApplyBlendMaskToDepthValues(const MaterialParameters params, real3 blendMaskValues, real3 layerDepthValues, real zeroMaskDisplacement)
{
    if(o_layer2_enabled || o_layer3_enabled)
    {        
        // We add to the depth value rather than lerp toward m_displacementMin to avoid squashing the topology, but instead lower it out of sight.

        if(o_layer2_enabled)
        {
            real dropoffRange = real(params.m_layer2_m_heightmapOffset) - zeroMaskDisplacement;
            layerDepthValues.g += dropoffRange * (1-blendMaskValues.r);
        }
        
        if(o_layer3_enabled)
        {
            real dropoffRange = real(params.m_layer3_m_heightmapOffset) - zeroMaskDisplacement;
            layerDepthValues.b += dropoffRange * (1-blendMaskValues.g);
        }
    }

    return layerDepthValues;
}

//! Callback function for ParallaxMapping.azsli
DepthResult GetDepth(const MaterialParameters params, Texture2D heightmap, sampler mapSampler, float2 uv, float2 uv_ddx, float2 uv_ddy)
{
    LayerBlendSource blendSource = GetFinalLayerBlendSource();

    // heightmap and mapSampler ignored here, since each layer can have it's own heightmap
    // we sample and blend the height from our configured layers
    real3 layerDepthValues = real3(GetLayerDepthValues(params, uv, uv_ddx, uv_ddy));
    
    // Note, when the blend source uses the blend mask from the vertex colors, parallax will not be able to blend correctly between layers. It will end up using the same blend mask values
    // for every UV position when searching for the intersection. This leads to smearing artifacts at the transition point, but these won't be as noticeable if
    // you have a small depth factor relative to the size of the blend transition.
    real3 blendMaskValues = GetApplicableBlendMaskValues(params, blendSource, uv, s_blendMaskFromVertexStream);
        
    bool useBlendMask = 
       LayerBlendSource::Displacement_With_BlendMaskTexture == blendSource || 
       LayerBlendSource::Displacement_With_BlendMaskVertexColors == blendSource;

    if(useBlendMask)
    {
        // Regarding GetSubMinDisplacement(), when a mask of 0 pushes the surface all the way to the bottom, we want that 
        // to go a little below the min so it will disappear if there is something else right at the min.
        
        layerDepthValues = ApplyBlendMaskToDepthValues(params, blendMaskValues, layerDepthValues, GetSubMinDisplacement(params));
    }
    
    // When blending the depth together, we don't use params.m_displacementBlendDistance. The intention is that m_displacementBlendDistance
    // is for transitioning the appearance of the surface itself, but we still want a distinct change in the heightmap. If someday we want to 
    // support smoothly blending the depth as well, there is a bit more work to do to get it to play nice with the blend mask code in GetLayerDepthValues().
    real layerDepthBlendDistance = 0.0;
    real3 blendWeightValues = GetBlendWeights(blendSource, blendMaskValues, layerDepthValues, layerDepthBlendDistance);
    
    real depth = BlendLayers(layerDepthValues.r, layerDepthValues.g, layerDepthValues.b, blendWeightValues);

    return DepthResultAbsolute(float(depth));
}