/*
 * 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
 *
 */

// Specular IBL reflection pipeline:
// Stencil -> BlendWeight -> GlobalFullscreen -> RenderOuter -> RenderInner -> Composite
//                           ----------------
//
// This shader performs a fullscreen pass on stenciled pixels that require specular IBL.
// It outputs some amount of the Global IBL cubemap as determined by the blend weight texture
// from the previous pass.  If the probes are going to write a total of >= 100% of local IBL 
// to this pixel then no Global IBL will be written since it is fully covered by probe IBL. 
// Note that this shader only considers pixels stenciled to non-probe space or probe outer 
// volumes (stencil value == UseSpecularIBLPass).  Inner probe volumes never receive Global IBL.
//
// Additional notes:
// - MSAA Input: The shader calculates the specular from the sub-pixels of the incoming textures 
//   and box-filters the result.
// - MSAA Output: Since this is a fullscreen pass the shader will write identical values to all
//   sub-pixel samples that pass the sub-pixel stencil test. 
// - An output of (0.0f, 0.0f, 0.0f) will be written if no global IBL will blend at this location.  
//   This is necessary to clear the pixel since the probes additively write their specular values
//   in the following passes.

#include <scenesrg.srgi>
#include <viewsrg.srgi>

#include <Atom/Features/PostProcessing/FullscreenVertexUtil.azsli>
#include <Atom/Features/PostProcessing/FullscreenVertexInfo.azsli>
#include <Atom/Features/PostProcessing/FullscreenPixelInfo.azsli>
#include <Atom/Features/MatrixUtility.azsli>
#include <Atom/Features/PBR/LightingUtils.azsli>
#include <Atom/Features/PBR/Microfacet/Fresnel.azsli>

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Texture2DMS<float>  m_depth;
    Texture2DMS<float4> m_normal;       // RGB10 = EncodeNormalSignedOctahedron(worldNormal), A2 = multiScatterCompensationEnabled
    Texture2DMS<float4> m_specularF0;   // RGB8 = SpecularF0, A8 = Roughness
    Texture2DMS<float4> m_albedo;       // RGB8 = Albedo, A = specularOcclusion
    Texture2DMS<float>  m_blendWeight;
    Texture2D<float2>   m_brdfMap;

    Sampler LinearSampler
    {
        MinFilter = Linear;
        MagFilter = Linear;
        MipFilter = Linear;
        AddressU = Clamp;
        AddressV = Clamp;
        AddressW = Clamp;
    };
}

#include <Atom/RPI/ShaderResourceGroups/DefaultDrawSrg.azsli>
#include "ReflectionCommon.azsli"

// Vertex Shader
VSOutput MainVS(VSInput input)
{
    VSOutput OUT;

    float4 posTex = GetVertexPositionAndTexCoords(input.m_vertexID);
    OUT.m_texCoord = float2(posTex.z, posTex.w);
    OUT.m_position = float4(posTex.x, posTex.y, 0.0, 1.0);
    return OUT;
}

// Pixel Shader
PSOutput MainPS(VSOutput IN, in uint sampleIndex : SV_SampleIndex)
{
    float4 encodedNormal = PassSrg::m_normal.Load(IN.m_position.xy, sampleIndex);
    float3 normal = DecodeNormalSignedOctahedron(encodedNormal.rgb);
    bool multiScatterCompensationEnabled = (encodedNormal.a > 0.0f);
    
    float4 albedo = PassSrg::m_albedo.Load(IN.m_position.xy, sampleIndex);
    float specularOcclusion = albedo.a;

    // reconstruct world space position from the depth at this location in screenspace
    float3 positionWS = ReconstructWorldPositionFromDepth(IN.m_position.xy, sampleIndex).xyz;

    // sample blend weight at this location
    float blendWeight = PassSrg::m_blendWeight.Load(IN.m_position.xy, sampleIndex).r;

    // blend weight for this pixel is the remainder of total blends for local probes
    // E.g., 0.2f probes = 0.8f global, 0.0f probes = 1.0f global, >1.0f probes = 0.0f global
    blendWeight = 1.0f - min(blendWeight, 1.0f);

    float3 specularF0 = PassSrg::m_specularF0.Load(IN.m_position.xy, sampleIndex).rgb;
    float roughness = PassSrg::m_specularF0.Load(IN.m_position.xy, sampleIndex).a;

    float3 dirToCamera = normalize(ViewSrg::m_worldPosition.xyz - positionWS);
    float3 reflectDir = reflect(-dirToCamera, normal);
    reflectDir = MultiplyVectorQuaternion(reflectDir, SceneSrg::m_iblOrientation);
    float3 globalSpecular = SceneSrg::m_specularEnvMap.SampleLevel(SceneSrg::m_samplerEnv, GetCubemapCoords(reflectDir), GetRoughnessMip(roughness)).rgb;

    float NdotV = saturate(dot(normal, dirToCamera));
    NdotV = max(NdotV, 0.01f);    // [GFX TODO][ATOM-4466] This is a current band-aid for specular noise at grazing angles.
    float2 brdf = PassSrg::m_brdfMap.Sample(PassSrg::LinearSampler, GetBRDFTexCoords(roughness, NdotV)).rg;
    float3 multiScatterCompensation = GetMultiScatterCompensation(specularF0, brdf, multiScatterCompensationEnabled);
    float3 specular = blendWeight * globalSpecular * multiScatterCompensation * (specularF0 * brdf.x + brdf.y);

    // apply exposure setting
    specular *= pow(2.0, SceneSrg::m_iblExposure);

    // maybe attenuate the specular 
    specular *= specularOcclusion;

    PSOutput OUT;
    OUT.m_color = float4(specular, 1.0f);
    return OUT;
}