/*
 * 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 <Atom/Features/PostProcessing/FullscreenPixelInfo.azsli>
#include <Atom/Features/PostProcessing/FullscreenVertex.azsli>
#include "DepthOfField.azsli"

#include <viewsrg.srgi>

ShaderResourceGroup PassSrg : SRG_PerPass_WithFallback
{
    Texture2D<float4> m_InputColor;
    Texture2D<float4> m_InputDepth;
    Texture2D<float4> m_FrontResult_division2;
    Texture2D<float4> m_FrontResult_division4;
    Texture2D<float4> m_FrontResult_division8;
    Texture2D<float4> m_BackResult_division2;
    Texture2D<float4> m_BackResult_division4;
    Texture2D<float4> m_BackResult_division8;

    // scale / offset to convert DofFactor to blend ratio for back buffer.
    float2 m_backBlendFactorDivision2;
    float2 m_backBlendFactorDivision4;
    float2 m_backBlendFactorDivision8;
    // scale / offset to convert DofFactor to blend ratio for front buffer.
    float2 m_frontBlendFactorDivision2;
    float2 m_frontBlendFactorDivision4;
    float2 m_frontBlendFactorDivision8;

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

option enum class DepthOfFieldDebugColoring
{
    Disabled,
    Enabled
}
o_dofMode = DepthOfFieldDebugColoring::Disabled;

inline float3 ConvertDebugColor(float3 color, float3 debugColor)
{
    switch(o_dofMode)
    {
        case DepthOfFieldDebugColoring::Disabled:
        {
            return color;
        }
        case DepthOfFieldDebugColoring::Enabled:
        {
            return color * debugColor;
        }
    }
}

PSOutput MainPS(VSOutput IN)
{
    PSOutput OUT;

    float4 readColor = (float4)0;
    float3 colorSum = (float3)0;

    // Combine from the back buffer to the front.

    ////////////////////////////////////////////////////////////////////////
    // Background

    // Calculate the blend ratio.
    // On the background side, calculate the blend ratio from the depth value on the depth buffer.
    // The blend ratio can also be calculated from DofFactor on the blur buffer, but the depth buffer has a higher resolution.

    float depth = InvertDepth(PassSrg::m_InputDepth.Sample(PassSrg::LinearSampler, IN.m_texCoord).r);
    float far = ViewSrg::m_dof.m_cameraParameters.x;
    float near = ViewSrg::m_dof.m_cameraParameters.y;
    float focusDistance = ViewSrg::m_dof.m_cameraParameters.z;
    float dofFactor = ConvertDofFactor(depth, far, near, focusDistance);

    float blendBackDivision8 = ConvertBlend(dofFactor, PassSrg::m_backBlendFactorDivision8);
    float blendBackDivision4 = ConvertBlend(dofFactor, PassSrg::m_backBlendFactorDivision4);
    float blendBackDivision2 = ConvertBlend(dofFactor, PassSrg::m_backBlendFactorDivision2);

    // Buffer composition is done from back to front.
    // On the other hand, the blend ratio is a value from the center toward the back.
    // Therefore, when overlapping buffers, it is necessary to invert the blend ratio.
    float blendBackFromDiv8ToDiv4     = saturate(1.0f - blendBackDivision8);
    float blendBackFromDiv4ToDiv2     = saturate(1.0f - blendBackDivision4);
    float blendBackFromDiv2ToOriginal = saturate(1.0f - blendBackDivision2);

    // back Division8
    readColor = PassSrg::m_BackResult_division8.Sample(PassSrg::LinearSampler, IN.m_texCoord);
    colorSum = ConvertDebugColor(readColor.rgb,float3(1, 0, 0));
    // back Division4
    readColor = PassSrg::m_BackResult_division4.Sample(PassSrg::LinearSampler, IN.m_texCoord);
    readColor.rgb = ConvertDebugColor(readColor.rgb, float3(0.9, 0.5, 0));
    colorSum = lerp(colorSum, readColor.rgb, blendBackFromDiv8ToDiv4);
    // back Division2
    readColor = PassSrg::m_BackResult_division2.Sample(PassSrg::LinearSampler, IN.m_texCoord);
    readColor.rgb = ConvertDebugColor(readColor.rgb, float3(0.7, 0.7, 0));
    colorSum = lerp(colorSum, readColor.rgb, blendBackFromDiv4ToDiv2);

    // original
    readColor = PassSrg::m_InputColor.Sample(PassSrg::LinearSampler, IN.m_texCoord);
    readColor.rgb = ConvertDebugColor(readColor.rgb, float3(0, 1, 0));
    colorSum = lerp(colorSum, readColor.rgb, blendBackFromDiv2ToOriginal);
    // output original alpha
    OUT.m_color.a = readColor.a;

    ////////////////////////////////////////////////////////////////////////
    // Foreground

    // Calculate the blend ratio.
    // On the foreground side, calculate the blend ratio from the alpha value (DofFactor) on the blurred buffer.
    // This alpha value is blurred, so can blur the front and back border.

    float4 readColorFrontDivision2 = PassSrg::m_FrontResult_division2.Sample(PassSrg::LinearSampler, IN.m_texCoord);
    float4 readColorFrontDivision4 = PassSrg::m_FrontResult_division4.Sample(PassSrg::LinearSampler, IN.m_texCoord);
    float4 readColorFrontDivision8 = PassSrg::m_FrontResult_division8.Sample(PassSrg::LinearSampler, IN.m_texCoord);
    float blendFrontDivision2 = ConvertBlend(readColorFrontDivision2.a, PassSrg::m_frontBlendFactorDivision2);
    float blendFrontDivision4 = ConvertBlend(readColorFrontDivision4.a, PassSrg::m_frontBlendFactorDivision4);
    float blendFrontDivision8 = ConvertBlend(readColorFrontDivision8.a, PassSrg::m_frontBlendFactorDivision8);
    blendFrontDivision2 = max(0.0f, blendFrontDivision2);
    blendFrontDivision4 = max(0.0f, blendFrontDivision4);
    blendFrontDivision8 = max(0.0f, blendFrontDivision8);

    // The blend ratio used is subtracted from the previous buffer's blend ratio.
    // Because it blends the most blurred buffers.
    // The blend ratio of the frontmost buffer is used as is.
    blendFrontDivision2 = saturate(blendFrontDivision2 - blendFrontDivision4);
    blendFrontDivision4 = saturate(blendFrontDivision4 - blendFrontDivision8);
    blendFrontDivision8 = saturate(blendFrontDivision8);

    readColorFrontDivision2.rgb = ConvertDebugColor(readColorFrontDivision2.rgb, float3(0, 0.7, 0.7));
    readColorFrontDivision4.rgb = ConvertDebugColor(readColorFrontDivision4.rgb, float3(0, 0.5, 0.9));
    readColorFrontDivision8.rgb = ConvertDebugColor(readColorFrontDivision8.rgb, float3(0, 0, 1));
    colorSum = lerp(colorSum, readColorFrontDivision2.rgb, blendFrontDivision2);
    colorSum = lerp(colorSum, readColorFrontDivision4.rgb, blendFrontDivision4);
    colorSum = lerp(colorSum, readColorFrontDivision8.rgb, blendFrontDivision8);

    OUT.m_color.rgb = colorSum;

    return OUT;
}