/*
 * 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/SrgSemantics.azsli>

#include <Atom/Features/PostProcessing/FullscreenPixelInfo.azsli>
#include <Atom/Features/PostProcessing/FullscreenVertex.azsli>
#include <Atom/Features/PostProcessing/PostProcessUtil.azsli>

static const int SHAPER_LINEAR = 0;
static const int SHAPER_LOG2 = 1;

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Texture2D<float4> m_colorTexture;
    Texture3D<float4> m_lut;
    Sampler LinearSampler
    {
        MinFilter = Linear;
        MagFilter = Linear;
        MipFilter = Linear;
        AddressU = Clamp;
        AddressV = Clamp;
        AddressW = Clamp;
    };
    
    int m_shaperType;
    float m_shaperBias;
    float m_shaperScale;
}

PSOutput MainPS(VSOutput IN)
{
    PSOutput OUT;

    float2 uvCoord = float2(IN.m_texCoord.x, IN.m_texCoord.y);
    float3 color = PassSrg::m_colorTexture.Sample(PassSrg::LinearSampler, uvCoord).rgb;

    // Convert from working color space to lut coordinates by applying the shaper function
    float3 lutCoordinate = color;
    if (PassSrg::m_shaperType == SHAPER_LINEAR)
    {
        lutCoordinate = color * PassSrg::m_shaperScale + PassSrg::m_shaperBias;
    }
    else if (PassSrg::m_shaperType == SHAPER_LOG2)
    {
        lutCoordinate = log2(color) * PassSrg::m_shaperScale + PassSrg::m_shaperBias;
    }

    // Adjust coordinate to the domain excluding the outer half texel in all directions
    uint3 outputDimensions;
    PassSrg::m_lut.GetDimensions(outputDimensions.x, outputDimensions.y, outputDimensions.z);
    float3 coordBias = 1.0/(2.0 * outputDimensions);
    float3 coordScale = (outputDimensions-1.0)/outputDimensions;
    lutCoordinate = (lutCoordinate * coordScale) + coordBias;

    float3 lutColor = PassSrg::m_lut.Sample(PassSrg::LinearSampler, lutCoordinate).rgb;

    // Apply the inverse of the shaper function to give the color in the working color space
    float3 finalColor = lutColor;
    if (PassSrg::m_shaperType == SHAPER_LINEAR)
    {
        finalColor = (lutColor - PassSrg::m_shaperBias)/PassSrg::m_shaperScale;
    }
    else if (PassSrg::m_shaperType == SHAPER_LOG2)
    {
        finalColor = pow(2.0, (lutColor - PassSrg::m_shaperBias)/PassSrg::m_shaperScale);
    }

    OUT.m_color.rgb = finalColor;
    OUT.m_color.a = 1.0;

    return OUT;
}