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

// This shader file will be removed when AZSLc can automatically generate nomsaa shader variants
// [GFX TODO][ATOM-13646] AZSLc support for generating non-MSAA shader variants

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

#include <Atom/Features/PostProcessing/FullscreenVertexUtil.azsli>
#include <Atom/Features/PostProcessing/FullscreenVertexInfo.azsli>
#include <Atom/RPI/Math.azsli>

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Texture2D<float4> m_normal;       // RGB10 = Normal (Encoded), A2 = Flags
    Texture2D<float>  m_depth;

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

#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;
}

struct PSOutput
{
    float m_depth : SV_Depth;
    float4 m_normal : SV_Target0;
};

// Pixel Shader
PSOutput MainPS(VSOutput IN)
{  
    // the downsample is 1/4 resolution
    // [GFX TODO][ATOM-6172] Add image scale PassSrg constant to the DiffuseProbeGrid downsample/upsample
    const uint ImageScale = 4;
    uint2 screenCoords = IN.m_position.xy * ImageScale;

    float downsampledDepth = 0;
    float4 downsampledEncodedNormal;
    for (uint y = 0; y < ImageScale; ++y)
    {
        for (uint x = 0; x < ImageScale; ++x)
        {
            float depth = PassSrg::m_depth.Load(int3(screenCoords + int2(x, y), 0)).r;
            float4 encodedNormal = PassSrg::m_normal.Load(int3(screenCoords + int2(x, y), 0));

            // take the closest depth sample to ensure we're getting the normal closest to the viewer 
            // (larger depth value due to reverse depth)
            if (depth > downsampledDepth)
            {
                downsampledDepth = depth;
                downsampledEncodedNormal = encodedNormal;
            }
        }
    } 

    float3 downsampledNormal = DecodeNormalSignedOctahedron(downsampledEncodedNormal.rgb);

    PSOutput OUT;
    OUT.m_depth = downsampledDepth;
    OUT.m_normal = float4(downsampledNormal * 0.5f + 0.5f, 0.0f);
    return OUT;
}