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

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Texture2D<float4> m_inputTexture;
    RWTexture2D<float4> m_outputTexture;
}

[numthreads(8,8,1)]
void MainCS(uint3 dispatch_id: SV_DispatchThreadID)
{
    // Get input and output texture dimensions
    uint2 inputDimensions;
    uint2 outputDimenstions;
    
    PassSrg::m_inputTexture.GetDimensions(inputDimensions.x, inputDimensions.y);
    PassSrg::m_outputTexture.GetDimensions(outputDimenstions.x, outputDimenstions.y);

    uint2 outPixel = dispatch_id.xy;

    // Early out if thread is outside of the target image
    if(outPixel.x >= outputDimenstions.x || outPixel.y >= outputDimenstions.y)
    {
        return;
    }
    
    // Calculate input pixel position in higher detail mip. If position
    // is out of bounds of input texture, output black and early out.
    uint2 inPixel = outPixel * 2;
    if(inPixel.x >= inputDimensions.x || inPixel.y >= inputDimensions.y)
    {
        PassSrg::m_outputTexture[outPixel] = float4(0.0f, 0.0f, 0.0f, 0.0f);
        return;
    }

    // Default sample is 4 pixels in a 2x2 quad
    uint xLoopCount = 2;
    uint yLoopCount = 2;
    
    // If we're at the output border pixel, adjust the amount we sample to extend to the input border
    if(outPixel.x == (outputDimenstions.x - 1))
    {
        xLoopCount = inputDimensions.x - inPixel.x;
    }
    if(outPixel.y == (outputDimenstions.y - 1))
    {
        yLoopCount = inputDimensions.y - inPixel.y;
    }

    float minValue = FLOAT_32_MAX;
    float maxValue = 0;
    float avgValue = 0;
    float sampleCount = 0;
    
    // Gather the input samples
    for(uint i = 0; i < xLoopCount; ++i)
    {
        for(uint j = 0; j < yLoopCount; ++j)
        {
            uint2 samplePoint = inPixel + uint2(i, j);
            float3 value = PassSrg::m_inputTexture[samplePoint].rgb;
            minValue = min(minValue, value.x);
            maxValue = max(maxValue, value.z);
            avgValue += value.y;
            ++sampleCount;
        }
    }
    
    avgValue = avgValue / max(1.0f, sampleCount);

    float4 output = float4(minValue, avgValue, maxValue, 0.0f);

    // Output the color
    PassSrg::m_outputTexture[outPixel] = output;
}