/*
 * 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/RPI/Math.azsli>
#include <Atom/Features/PostProcessing/FullscreenVertex.azsli>
#include <Atom/Features/PostProcessing/FullscreenPixelInfo.azsli>
#include <Atom/Features/PostProcessing/PostProcessUtil.azsli>
#include <Atom/Features/MatrixUtility.azsli>
#include <Atom/Features/PBR/LightingUtils.azsli>
#include <Atom/Features/PBR/Microfacet/Fresnel.azsli>

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Texture2D<float4>   m_reflection;
    Texture2D<float>    m_downsampledDepth;
    Texture2DMS<float4> m_normal;       // RGB10 = Normal (Encoded), A2 = Flags
    Texture2DMS<float4> m_specularF0;   // RGB8 = SpecularF0, A8 = Roughness
    Texture2DMS<float>  m_depth;
    Texture2D<float4>   m_previousFrame;

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

    // the max roughness mip level for sampling the previous frame image
    uint m_maxMipLevel;
}

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

float3 SampleReflection(float2 reflectionUV, float mip, float depth, float3 normal, uint2 invDimensions)
{  
    const float DepthTolerance = 0.001f;

    // attempt to trivially accept the downsampled reflection texel
    float downsampledDepth = PassSrg::m_downsampledDepth.SampleLevel(PassSrg::LinearSampler, reflectionUV, floor(mip)).r;
    if (abs(depth - downsampledDepth) <= DepthTolerance)
    {
        // use this reflection sample
        float3 reflection = PassSrg::m_reflection.SampleLevel(PassSrg::LinearSampler, reflectionUV, mip).rgb;
        return reflection;
    }
    
    // neighborhood search surrounding the downsampled texel, searching for the closest matching depth
    float closestDepthDelta = 1.0f;
    int2 closestOffsetUV = float2(0.0f, 0.0f);
    for (int y = -4; y <= 4; ++y)
    {
        for (int x = -4; x <= 4; ++x)
        {
            float2 offsetUV = float2(x * invDimensions.x, y * invDimensions.y);
            float downsampledDepth = PassSrg::m_downsampledDepth.SampleLevel(PassSrg::LinearSampler, reflectionUV + offsetUV, floor(mip)).r;
            float depthDelta = abs(depth - downsampledDepth);

            if (depthDelta <= DepthTolerance)
            {
                // depth is within tolerance, use this texel
                float3 reflection = PassSrg::m_reflection.SampleLevel(PassSrg::LinearSampler, reflectionUV + offsetUV, mip).rgb;
                return reflection;
            }

            if (closestDepthDelta > depthDelta)
            {
                closestDepthDelta = depthDelta;
                closestOffsetUV = offsetUV;
            }
        }
    }
    
    float3 reflection = PassSrg::m_reflection.SampleLevel(PassSrg::LinearSampler, reflectionUV + closestOffsetUV, mip).rgb;
    return reflection;
}

// Pixel Shader
PSOutput MainPS(VSOutput IN, in uint sampleIndex : SV_SampleIndex)
{
    float2 screenCoords = IN.m_position.xy;

    uint2 dimensions;
    uint samples;
    PassSrg::m_depth.GetDimensions(dimensions.x, dimensions.y, samples);

    // compute trace image coordinates for the half-res image
    float2 traceCoords = screenCoords * 0.5f;

    // check reflection data mip0 to see if there was a hit
    float4 reflectionData = PassSrg::m_reflection.Load(uint3(traceCoords, 0));
    if (reflectionData.w <= 0.0f)
    {
        // fallback to the cubemap reflections currently in the reflection buffer
        discard;
    }

    // load specular and roughness
    float4 specularF0 = PassSrg::m_specularF0.Load(screenCoords, sampleIndex);
    float roughness = specularF0.a;
    const float MaxRoughness = 0.5f;
    if (roughness > MaxRoughness)
    {
        // fallback to the cubemap reflections currently in the reflection buffer
        discard;
    }

    // reconstruct world space position of current pixel
    float2 UV = IN.m_texCoord;
    float depth = PassSrg::m_depth.Load(screenCoords, sampleIndex).r;
    float2 ndcPos = float2(UV.x, 1.0f - UV.y) * 2.0f - 1.0f;
    float4 projectedPos = float4(ndcPos, depth, 1.0f);
    float4 positionVS = mul(ViewSrg::m_projectionMatrixInverse, projectedPos);
    positionVS /= positionVS.w;
    float3 positionWS = mul(ViewSrg::m_viewMatrixInverse, positionVS).xyz;  

    // compute ray from camera to surface position
    float3 cameraToPositionWS = normalize(positionWS.xyz - ViewSrg::m_worldPosition);

    // retrieve surface normal
    float4 encodedNormal = PassSrg::m_normal.Load(screenCoords, sampleIndex);
    float3 normalWS = DecodeNormalSignedOctahedron(encodedNormal.rgb);
    float NdotV = dot(normalWS, -cameraToPositionWS);

    // compute the roughness mip to use in the reflection image
    // remap the roughness mip into a lower range to more closely match the material roughness values
    float mip = saturate(roughness / MaxRoughness) * PassSrg::m_maxMipLevel;
    
    // sample reflection color from the mip chain
    float3 reflectionColor = SampleReflection(IN.m_texCoord, mip, depth, normalWS, 1.0f / dimensions);

    PSOutput OUT;
    OUT.m_color = float4(reflectionColor, reflectionData.w);
    return OUT;
}