/*
 * 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/PostProcessUtil.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 = Normal (Encoded), A2 = Flags
    Texture2DMS<float4> m_specularF0;   // RGB8 = SpecularF0, A8 = Roughness
    Texture2DMS<float4> m_reflection;
    Texture2D<float4>   m_previousFrame;

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

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

struct VSInput
{
    uint m_vertexID : SV_VertexID;
};

struct VSOutput
{
    float4 m_position : SV_Position;
    float2 m_texCoord : TEXCOORD0;
};

// 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
struct PSOutput
{    
    float4 m_color  : SV_Target0;
    float m_depth   : SV_Depth;
};

PSOutput MainPS(VSOutput IN)
{
    // compute screen coords based on a half-res render target
    float2 screenCoords = IN.m_position.xy * 2.0f;
       
    uint2 dimensions;
    uint samples;
    PassSrg::m_depth.GetDimensions(dimensions.x, dimensions.y, samples);
    float2 UV = saturate(screenCoords / dimensions.xy);

    float depth = PassSrg::m_depth.Load(screenCoords, 0).r;
    if (depth == 0.0f)
    {
        // skip tracing rays at max scene depth
        discard;
    }
        
    // compute view space surface position
    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;

    // normalize the viewspace position to get the camera-to-position vector
    float3 cameraToPositionVS = normalize(positionVS);  

    // retrieve surface normal
    float4 encodedNormal = PassSrg::m_normal.Load(screenCoords, 0);
    float3 normalWS = DecodeNormalSignedOctahedron(encodedNormal.rgb);
    float3 normalVS = normalize(mul(ViewSrg::m_viewMatrix, float4(normalWS, 0.0f)).xyz);

    // reflect view ray around surface normal
    float3 reflectDirVS = normalize(reflect(cameraToPositionVS, normalVS));
 
    // check to see if the reflected direction is approaching the camera
    float rdotv = dot(reflectDirVS, -cameraToPositionVS);
    bool fallbackEdge = false;
    if (rdotv >= -0.05f)
    {
        if (rdotv >= 0.0f)
        {
            // ray points back to camera, fallback to cubemaps
            discard;            
        }

        // ray is approaching the camera direction, but not there yet - trace the reflection and set this
        // as a non-reflected pixel, which will prevent artifacts at the boundary
        fallbackEdge = true;
    }

    // trace screenspace rays against the depth buffer to find the screenspace intersection coordinates    
    float4 result = float4(0.0f, 0.0f, 0.0f, 0.0f);
    float2 hitCoords = float2(0.0f, 0.0f);
    if (TraceRayScreenSpace(positionVS, reflectDirVS, dimensions, hitCoords))
    {
        // reconstruct the world space position of the trace coordinates
        float2 traceUV = saturate(hitCoords / dimensions);
        float traceDepth = PassSrg::m_depth.Load(hitCoords, 0).r;
        float2 traceNDC = float2(traceUV.x, 1.0f - traceUV.y) * 2.0f - 1.0f;
        float4 traceProjectedPos = float4(traceNDC, traceDepth, 1.0f);
        float4 tracePositionVS = mul(ViewSrg::m_projectionMatrixInverse, traceProjectedPos);
        tracePositionVS /= tracePositionVS.w;
        float4 tracePositionWS = mul(ViewSrg::m_viewMatrixInverse, tracePositionVS);
        
        // reproject to the previous frame image coordinates
        float4 tracePrevNDC = mul(ViewSrg::m_viewProjectionPrevMatrix, tracePositionWS);
        tracePrevNDC /= tracePrevNDC.w;
        float2 tracePrevUV = float2(tracePrevNDC.x, -1.0f * tracePrevNDC.y) * 0.5f + 0.5f;

        // sample the previous frame image
        result.rgb = PassSrg::m_previousFrame.SampleLevel(PassSrg::LinearSampler, tracePrevUV, 0).rgb;

        // apply surface specular
        float3 specularF0 = PassSrg::m_specularF0.Load(screenCoords, 0).rgb;
        result.rgb *= specularF0;

        // fade rays close to screen edge
        const float ScreenFadeDistance = 0.95f;
        float2 fadeAmount = max(max(0.0f, traceUV - ScreenFadeDistance), max(0.0f, 1.0f - traceUV - ScreenFadeDistance));
        fadeAmount /= (1.0f - ScreenFadeDistance);
        result.a = fallbackEdge ? 0.0f : 1.0f - max(fadeAmount.x, fadeAmount.y);
    }
    else
    {
        // ray miss, add in the IBL/probe reflections from the specular pass
        float4 positionWS = mul(ViewSrg::m_viewMatrixInverse, positionVS);
        float3 cameraToPositionWS = normalize(positionWS - ViewSrg::m_worldPosition);  
        float3 reflectDirWS = normalize(reflect(cameraToPositionWS, normalWS));
        
        result.rgb += PassSrg::m_reflection.Load(screenCoords, 0).rgb;
        result.a = fallbackEdge ? 0.0f : 1.0f;
    }

    // downsample depth
    float downsampledDepth = 0.0f;
    for (int y = -2; y < 2; ++y)
    {
        for (int x = -2; x < 2; ++x)
        {
            float depth = PassSrg::m_depth.Load(screenCoords + int2(x, y), 0).r;

            // take the closest depth sample (larger depth value due to reverse depth)
            if (depth > downsampledDepth)
            {
                downsampledDepth = depth;
            }
        }
    }

    PSOutput OUT;
    OUT.m_color = result;
    OUT.m_depth = downsampledDepth;
    return OUT;
}