o3de/Gems/AtomTressFX/Assets/Shaders/HairLighting.azsli

/*
 * Modifications 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
 *
 */
 
//------------------------------------------------------------------------------
// Shader code related to lighting and shadowing for TressFX
//------------------------------------------------------------------------------
//
// Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

#pragma once




#include <HairUtilities.azsli>
//------------------------------------------------------------------------------
struct HairShadeParams
{
    float3 m_color;
    float m_hairShadowAlpha;
    float m_fiberRadius;
    float m_fiberSpacing;

    // Original TressFX Kajiya lighting model parameters
    float m_Ka;
    float m_Kd;
    float m_Ks1;
    float m_Ex1;
    float m_Ks2;
    float m_Ex2;

    // Marschner lighting model parameters
    float m_cuticleTilt;
    float m_roughness;
};

//! Original TressFX enhanced Kajiya-Kay lighting model code
//!  
//! Returns a float3 which is the scale for diffuse, spec term, and colored spec term.
//! 
//! The diffuse term is from Kajiya.
//! 
//! The spec term is what Marschner refers to as "R", reflecting directly off the surface 
//! of the hair, taking the color of the light like a dielectric specular term.  This 
//! highlight is shifted towards the root of the hair.
//! 
//! The colored spec term is caused by light passing through the hair, bouncing off the 
//! back, and coming back out. It therefore picks up the color of the light.  
//! Marschner refers to this term as the "TRT" term.  This highlight is shifted towards the 
//! tip of the hair.
//! 
//! vEyeDir, vLightDir and vTangentDir are all pointing out.
//! coneAngleRadians explained below.
//! 
//! hair has a tiled-conical shape along its lenght.  Sort of like the following.
//! 
//! \    /
//!  \  /
//! \    /
//!  \  /  
//! 
//! The angle of the cone is the last argument, in radians.  
//! It's typically in the range of 5 to 10 degrees
float3 ComputeDiffuseSpecFactors(
    float3 vEyeDir, float3 vLightDir, float3 vTangentDir, HairShadeParams params, 
    float coneAngleRadians = 10 * AMD_PI / 180)
{
    // In Kajiya's model: diffuse component: sin(t, l)
    float cosTL = (dot(vTangentDir, vLightDir));
    float sinTL = sqrt(1 - cosTL * cosTL);
    float diffuse = sinTL;

    float cosTRL = -cosTL;
    float sinTRL = sinTL;
    float cosTE = (dot(vTangentDir, vEyeDir));
    float sinTE = sqrt(1 - cosTE * cosTE);

    // Primary highlight: reflected direction shift towards root (2 * coneAngleRadians)
    float cosTRL_root = cosTRL * cos(2 * coneAngleRadians) - sinTRL * sin(2 * coneAngleRadians);
    float sinTRL_root = sqrt(1 - cosTRL_root * cosTRL_root);
    float specular_root = max(0, cosTRL_root * cosTE + sinTRL_root * sinTE);

    // Secondary highlight: reflected direction shifted toward tip (3*coneAngleRadians)
    float cosTRL_tip = cosTRL * cos(-3 * coneAngleRadians) - sinTRL * sin(-3 * coneAngleRadians);
    float sinTRL_tip = sqrt(1 - cosTRL_tip * cosTRL_tip);
    float specular_tip = max(0, cosTRL_tip * cosTE + sinTRL_tip * sinTE);

    return float3( 
        params.m_Kd * diffuse, 
        params.m_Ks1 * pow(specular_root, params.m_Ex1), 
        params.m_Ks2 * pow(specular_tip, params.m_Ex2));
}

float LinearizeDepth(float depthNDC, float fNear, float fFar)
{
    return fNear * fFar / (fFar - depthNDC * (fFar - fNear));
}

//! The following code is for reference only and should be removed once the 
//! Kajiya-Kay original TressFX lighting model is connected to Atom as was 
//! done for the Marshcner lighting model.
#define DEMO_NUMBER_OF_LIGHTS 3
#define DEMO_NUMBER_OF_LIGHTS 3
float3 SimplifiedHairLighting(float3 vTangent, float3 vPositionWS, float3 vViewDirWS, in HairShadeParams params, float3 vNDC)
{
    // Initialize information needed for all lights
    float3 V = normalize(vViewDirWS);
    float3 T = normalize(vTangent);

    float3 accumulatedHairColor = float3(0.0, 0.0, 0.0);

    float4 lightPosWSVec[DEMO_NUMBER_OF_LIGHTS] = {
        float4(3, 0, 3, 1.5f),  // Sun
        float4(-.5, 0, 0.5, .5f),
        float4(.5, 0, 0.5, .5f),
    };

    float3 lightColorVec[DEMO_NUMBER_OF_LIGHTS] = {
        float3(1,1,.95f),        // Sun
        float3(1,1,1),
        float3(1,1,1)
    };

    // Static lights loop for reference - not connected to Atom
    // [To Do] - connect to Atom lighting ala HairLightTypes loop
    for (int l = 0; l < DEMO_NUMBER_OF_LIGHTS ; l++)
    {
        float3 lightPosWS = lightPosWSVec[l].xyz;
        float3 LightVector = normalize( vPositionWS - lightPosWS );
        float lightIntensity = lightPosWSVec[l].w;
        float3 LightColor = lightColorVec[l]; 
        float3 L = LightVector;

        // Reference usage of shadow
//        float shadowTerm = ComputeLightShadow(l, vPositionWS, params);
//        if (shadowTerm <= 0.f)
//              continue;

        float3 lightSurfaceCoeffs = ComputeDiffuseSpecFactors(V, L, T, params);

        // The diffuse coefficient here is a rough approximation as per the Kajiya model 
        float3 diffuseCompoenent = lightSurfaceCoeffs.x * params.m_color;

        // This is the approximation to Marschner R but azimuthal only
        float3 specularAtPos = lightSurfaceCoeffs.y;    

        // This is the approximation to Marschner TRT but azimuthal only
        // Notice the base color gather due to the trsmittance within the hair
        float3 specularAtBase = lightSurfaceCoeffs.z * params.m_color;

        // Final result
        float3 lightContribution = (diffuseCompoenent + specularAtPos + specularAtBase) * lightIntensity * LightColor; // * shadowTerm;
 
        accumulatedHairColor += max(float3(0, 0, 0), lightContribution );
    }
    return accumulatedHairColor;
}

//==============================================================================
//                             Atom Lighting
//==============================================================================
#include <Atom/Features/Shadow/DirectionalLightShadow.azsli>
#include <HairLightTypes.azsli>
//------------------------------------------------------------------------------

float3 CalculateLighting( 
    float4 screenCoords,    // XY - screen coords 0..max pix res, Z - depth 0..1
    float3 vPositionWS, float3 vViewDirWS, float3 vTangent, 
    float thickness, in HairShadeParams material )
{    
    //-------- Surface init --------
    Surface surface;

    const float specularF0Factor = 0.04f; // set this to 0.04?!

    surface.position = vPositionWS;
    surface.tangent = vTangent;         // Redundant - will be calculated per light
    surface.normal = float3(0, 0, 0);   // Will fail lights that did not initialize properly. 
    surface.vertexNormal = float3(0,0,0); // [To Do] - vertex normals are not handled yet in the hair shader.
    surface.roughnessLinear = material.m_roughness;
    surface.cuticleTilt = material.m_cuticleTilt;
    surface.thickness = thickness;
    surface.CalculateRoughnessA();
    surface.SetAlbedoAndSpecularF0( material.m_color, specularF0Factor);

    // The trasmission / back lighting does not seem to work!
    surface.transmission.InitializeToZero();    // Assuming thin layer
    surface.transmission.tint = material.m_color;
    surface.transmission.thickness = 0.001;     // 1 mm settings
    surface.transmission.transmissionParams = float4(1.0, 1.0, 1.0, 32.0);  // for thin surface XYZ are partials * W that is the exponent mult

    //------- LightingData init -------
    LightingData lightingData;

    float4 screenPositionForLighting = mul(ViewSrg::m_viewProjectionMatrix, float4(vPositionWS, 1.0));
    uint2 dimensions;
    PassSrg::m_linearDepth.GetDimensions(dimensions.x, dimensions.y);
    screenPositionForLighting.y = 1.0 - screenPositionForLighting.y;
    screenPositionForLighting.xy = (screenPositionForLighting.xy * 0.5 + 0.5) * dimensions;

    // Light iterator - required for the init but the culling is not used yet.
    lightingData.tileIterator.Init(screenCoords, PassSrg::m_lightListRemapped, PassSrg::m_tileLightData);

    // The normal assignment will be overriden afterwards per light
    lightingData.Init(surface.position, surface.GetSpecularNormal(), surface.roughnessLinear, ViewSrg::m_worldPosition.xyz);

    ApplyLighting(surface, lightingData, screenCoords);

    return lightingData.diffuseLighting + lightingData.specularLighting;
}

float3 TressFXShading(float2 pixelCoord, float depth, float3 tangent, float3 baseColor, float thickness, int shaderParamIndex)
{
    float3 vNDC;    // normalized device / screen coordinates: [-1..1, -1..1, 0..1]
    float3 vPositionWS = ScreenPosToWorldPos(PassSrg::m_linearDepth, pixelCoord, depth, vNDC);

    // [To Do] the follwing two lines are a hack to make the tile lighting work for now
    #define _BIG_HACK_FOR_TESTING_ 1//32
    float4 screenCoords = float4( _BIG_HACK_FOR_TESTING_ * pixelCoord, depth, depth);   // screen space position - XY in pixels - ZW are depth 0..1

    float3 vViewDirWS = g_vEye - vPositionWS;

    //---- TressFX original lighting params setting ----
    HairShadeParams params;
    params.m_color = baseColor;
    params.m_hairShadowAlpha = HairParams[shaderParamIndex].m_shadowAlpha;
    params.m_fiberRadius = HairParams[shaderParamIndex].m_fiberRadius;
    params.m_fiberSpacing = HairParams[shaderParamIndex].m_fiberSpacing;

    params.m_Ka = HairParams[shaderParamIndex].m_matKValue.x;
    params.m_Kd = HairParams[shaderParamIndex].m_matKValue.y;
    params.m_Ks1 = HairParams[shaderParamIndex].m_matKValue.z;
    params.m_Ex1 = HairParams[shaderParamIndex].m_matKValue.w;
    params.m_Ks2 = HairParams[shaderParamIndex].m_hairKs2;
    params.m_Ex2 = HairParams[shaderParamIndex].m_hairEx2;

    params.m_cuticleTilt = HairParams[shaderParamIndex].m_cuticleTilt;
    params.m_roughness = HairParams[shaderParamIndex].m_roughness;
    //---------------------------------------------------

    float3 accumulatedLight = float3(0, 0, 1);
    if (o_hairLightingModel == HairLightingModel::Kajiya)
    {   // This option should be removed and the Kajiya-Kay model should be operated from within 
        // the Atom lighting loop.
        accumulatedLight = SimplifiedHairLighting(tangent, vPositionWS, vViewDirWS, params, vNDC);    
    }
    else
    {
        accumulatedLight = CalculateLighting(screenCoords, vPositionWS, vViewDirWS, tangent, thickness, params);
    }
    return accumulatedLight;
}

float3 TressFXShadingFullScreen(float2 pixelCoord, float depth, float3 compressedTangent, float3 baseColor, float thickness, int shaderParamIndex)
{
    // The tangent that was compressed to store in the PPLL structure
    float3 tangent = normalize(compressedTangent.xyz * 2.f - 1.f);

    return TressFXShading(pixelCoord, depth, tangent, baseColor, thickness, shaderParamIndex);
}