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

#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/PBR/LightingUtils.azsli>
#include <Atom/RPI/Math.azsli>

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Texture2DMS<float4> m_downsampledProbeIrradiance;
    Texture2DMS<float>  m_downsampledDepth;
    Texture2DMS<float4> m_downsampledNormal;
    Texture2DMS<float4> m_albedo;       // RGB8 = Albedo with pre-multiplied factors, A = Unused here
    Texture2DMS<float4> m_normal;       // RGB10 = Normal (Encoded), A2 = Flags
    Texture2DMS<float>  m_depth;
    
    Sampler LinearSampler
    {
        MinFilter = Linear;
        MagFilter = Linear;
        MipFilter = Linear;
        AddressU = Clamp;
        AddressV = Clamp;
        AddressW = Clamp;
    };

    // scale multiplier of the downsampled size to the fullscreen size (e.g., 4)
    uint m_imageScale;
}

#include <Atom/RPI/ShaderResourceGroups/DefaultDrawSrg.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;
}

// retrieve probe irradiance from the downsampled irradiance image generated by the DiffuseProbeGridRender pass
float3 SampleProbeIrradiance(uint sampleIndex, uint2 probeIrradianceCoords, float depth, float3 normal, float3 albedo, uint imageScale)
{  
    const float NormalMatchTolerance = 0.99f;
    const float DepthTolerance = 0.001f;

    // attempt to trivially accept the downsampled irradiance texel
    float3 downsampledNormal = PassSrg::m_downsampledNormal.Load(probeIrradianceCoords, sampleIndex).rgb;
    downsampledNormal = downsampledNormal * 2.0f - 1.0f;
    if (dot(downsampledNormal, normal) > NormalMatchTolerance)
    {
        // normal is almost identical, test depth to see if it's within tolerance
        float downsampledDepth = PassSrg::m_downsampledDepth.Load(probeIrradianceCoords, sampleIndex).r;
        if (abs(depth - downsampledDepth) <= DepthTolerance)
        {
            // use this irradiance sample
            return PassSrg::m_downsampledProbeIrradiance.Load(probeIrradianceCoords, sampleIndex).rgb;
        }
    }

    // perform a normal-aware upsample:
    // - search for the downsampled normal sample that most closely matches the high resolution normal
    // - secondary check downsampled depth against high resolution depth
    // - if both checks pass, use the downsampled irradiance from that offset in the downsampled irradiance texture
    float closestDot = -1.0f;
    int2 closestOffset = int2(0, 0);
    
    // neighborhood search surrounding the downsampled texel, searching for the closest matching normal
    int extent = floor(imageScale * 0.5f);
    for (int y = -extent; y <= extent; ++y)
    {
        for (int x = -extent; x <= extent; ++x)
        {
            float3 downsampledNormal = PassSrg::m_downsampledNormal.Load(probeIrradianceCoords + int2(x, y), sampleIndex).rgb;
            downsampledNormal = downsampledNormal * 2.0f - 1.0f;
            float normalDot = dot(downsampledNormal, normal);

            if (normalDot >= closestDot)
            {              
                if (normalDot > NormalMatchTolerance)
                {
                    // the normals are almost identical, if the depth is within the tolerance we can optimize by just taking this sample
                    float downsampledDepth = PassSrg::m_downsampledDepth.Load(probeIrradianceCoords + int2(x, y), sampleIndex).r;
                    if (abs(depth - downsampledDepth) <= DepthTolerance)
                    {
                        return PassSrg::m_downsampledProbeIrradiance.Load(probeIrradianceCoords + int2(x, y), sampleIndex).rgb;
                    }

                    closestDot = normalDot;
                    closestOffset = int2(x, y);
                }
            }
        }
    }
    
    return PassSrg::m_downsampledProbeIrradiance.Load(probeIrradianceCoords + closestOffset, sampleIndex).rgb;
}

// retrieve irradiance from the global IBL diffuse cubemap
float3 SampleGlobalIBL(uint sampleIndex, uint2 screenCoords, float depth, float3 normal)
{
    uint2 dimensions;
    uint samples;
    PassSrg::m_depth.GetDimensions(dimensions.x, dimensions.y, samples);

    // reconstruct world space position
    float2 UV = saturate((float2)screenCoords / dimensions.xy);

    float x = UV.x * 2.0f - 1.0f;
    float y = (1.0f - UV.y) * 2.0f - 1.0f;
    float4 projectedPos = float4(x, y, depth, 1.0f);
    float4 positionVS = mul(ViewSrg::m_projectionMatrixInverse, projectedPos);
    positionVS /= positionVS.w;
    float3 positionWS = mul(ViewSrg::m_viewMatrixInverse, positionVS).xyz;  

    // apply global IBL
    float3 dirToCamera = normalize(ViewSrg::m_worldPosition.xyz - positionWS);
    float NdotV = dot(normal, dirToCamera);
    NdotV = max(NdotV, 0.01f);    // [GFX TODO][ATOM-4466] This is a current band-aid for specular noise at grazing angles.

    float3 irradianceDir = MultiplyVectorQuaternion(normal, SceneSrg::m_iblOrientation);
    float3 irradiance = SceneSrg::m_diffuseEnvMap.Sample(SceneSrg::m_samplerEnv, GetCubemapCoords(irradianceDir)).rgb;
    return irradiance;
}

// Pixel Shader
PSOutput MainPS(VSOutput IN, in uint sampleIndex : SV_SampleIndex)
{
    uint2 screenCoords = IN.m_position.xy;
    float imageScaleInverse = 1.0f / PassSrg::m_imageScale;

    // compute image coords for the downsampled probe irradiance image
    uint2 probeIrradianceCoords = screenCoords * imageScaleInverse;

    float depth = PassSrg::m_depth.Load(screenCoords, sampleIndex).r;
    float4 encodedNormal = PassSrg::m_normal.Load(screenCoords, sampleIndex);
    float3 normal = DecodeNormalSignedOctahedron(encodedNormal.rgb);
    float4 albedo = PassSrg::m_albedo.Load(screenCoords, sampleIndex);
    float probeIrradianceBlendWeight = saturate(PassSrg::m_downsampledProbeIrradiance.Load(probeIrradianceCoords, sampleIndex).a);
  
    float3 diffuse = float3(0.0f, 0.0f, 0.0f);
    if (probeIrradianceBlendWeight > 0.0f)
    {
        float3 probeIrradiance = SampleProbeIrradiance(sampleIndex, probeIrradianceCoords, depth, normal, albedo, PassSrg::m_imageScale);
        diffuse = (albedo.rgb / PI) * probeIrradiance * probeIrradianceBlendWeight;
    }

    if (probeIrradianceBlendWeight < 1.0f)
    {
        float3 globalIrradiance = SampleGlobalIBL(sampleIndex, screenCoords, depth, normal);

        // adjust IBL lighting by exposure
        float3 globalDiffuse = (albedo * globalIrradiance) * pow(2.0, SceneSrg::m_iblExposure);

        diffuse += globalDiffuse * (1.0f - probeIrradianceBlendWeight);
    }

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