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

// --- Static Options Available ---
// SKYBOX_TWO_OUTPUTS - Skybox renders to two rendertargets instead of one (SkyBox_TwoOutputs.pass writes to specular and reflection targets)

//#include <Atom/Features/Pipeline/Forward/ForwardPassSrg.azsli> // GALIB
#include <Atom/Features/ColorManagement/TransformColor.azsli>
#include <Atom/Features/PostProcessing/FullscreenVertexUtil.azsli>
#include <Atom/Features/PostProcessing/Tonemap.azsli>
#include <Atom/Features/MatrixUtility.azsli>
#include <Atom/Features/ShaderQualityOptions.azsli>

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

#ifndef ENABLE_PHYSICAL_SKY
#define ENABLE_PHYSICAL_SKY              1
#endif

#ifndef ENABLE_MERGE_FILMIC_TONEMAP
#define ENABLE_MERGE_FILMIC_TONEMAP              0
#endif

ShaderResourceGroup PassSrg : SRG_PerPass
{
    [[input_attachment_index(0)]]
    SubpassInput m_depthInput;

    // Required for sampling the skybox texture.
    Sampler LinearSampler
    {
        MinFilter = Linear;
        MagFilter = Linear;
        MipFilter = Linear;
        AddressU = Clamp;
        AddressV = Clamp;
        AddressW = Clamp;
    };
}

struct VSInput
{
    uint m_vertexID : SV_VertexID;
};

struct VSOutput
{
    float4 m_position : SV_Position;
    float3 m_cubemapCoord : TEXCOORD0;
};

// The Hosek-Wilkie version of Perez formula
// Parameters defination can be found in SkyBoxSystem.h
// https://cgg.mff.cuni.cz/projects/SkylightModelling/HosekWilkie_SkylightModel_SIGGRAPH2012_Preprint_lowres.pdf

float3 HosekWilkie(float cosGamma, float gamma, float cosTheta)
{
    float3 A = SceneSrg::m_physicalSkyData.m_physicalSkyParameterA.xyz;
    float3 B = SceneSrg::m_physicalSkyData.m_physicalSkyParameterB.xyz;
    float3 C = SceneSrg::m_physicalSkyData.m_physicalSkyParameterC.xyz;
    float3 D = SceneSrg::m_physicalSkyData.m_physicalSkyParameterD.xyz;
    float3 E = SceneSrg::m_physicalSkyData.m_physicalSkyParameterE.xyz;
    float3 F = SceneSrg::m_physicalSkyData.m_physicalSkyParameterF.xyz;
    float3 G = SceneSrg::m_physicalSkyData.m_physicalSkyParameterG.xyz;
    float3 H = SceneSrg::m_physicalSkyData.m_physicalSkyParameterH.xyz;
    float3 I = SceneSrg::m_physicalSkyData.m_physicalSkyParameterI.xyz;

    float3 chi = (1 + cosGamma * cosGamma) / pow(1.0 + H * H - 2.0 * cosGamma * H, float3(1.5, 1.5, 1.5));
    return (1.0 + A * exp(B / (cosTheta + 0.01))) * (C + D * exp(E * gamma) + F * (cosGamma * cosGamma) + G * chi + I * sqrt(cosTheta));
}

VSOutput MainVS(VSInput input)
{
    VSOutput OUT;

    float4 posTex = GetVertexPositionAndTexCoords(input.m_vertexID);
    OUT.m_position = float4(posTex.x, posTex.y, 0.0, 1.0);

    // camera transform matrix without translation
    float4x4 viewToWorldNoTranslation = ViewSrg::m_viewMatrixInverse;
    viewToWorldNoTranslation[0][3] = 0.0;
    viewToWorldNoTranslation[1][3] = 0.0;
    viewToWorldNoTranslation[2][3] = 0.0;
    viewToWorldNoTranslation[3][3] = 1.0;

    float4x4 skyboxMatrix = mul(viewToWorldNoTranslation, ViewSrg::m_projectionMatrixInverse);
    float usePhysicalSky = (float)SceneSrg::m_physicalSky;

    // Workaround for Qualcomm bug. This is the same as:
    // if(!SceneSrg::m_physicalSky)
    //    skyboxMatrix = mul(SceneSrg::m_cubemapRotationMatrix, skyboxMatrix);
    skyboxMatrix = (mul(SceneSrg::m_cubemapRotationMatrix, skyboxMatrix) * (1.0 - usePhysicalSky)) + (skyboxMatrix * usePhysicalSky);

    // calculate cubemap coordinate
    float4 clipPosition = float4(posTex.x, posTex.y, 1.0, 1.0);
    float4 worldPosition = mul(skyboxMatrix, clipPosition);
    OUT.m_cubemapCoord = worldPosition.xyz / worldPosition.w;

    return OUT;
}

float3 GetCubemapCoords(float3 original)
{
    // Environment cubemaps created by IBL Baker have to be rotated to be right-side-up
    return float3(-original.x, original.z, -original.y);
}

struct PSOutput
{
    float4 m_specular : SV_Target0;
#ifdef SKYBOX_TWO_OUTPUTS
    float4 m_reflection : SV_Target1;
#endif
};

PSOutput MainPS(VSOutput input)
{
    PSOutput OUT;

    if(!SceneSrg::m_enable)
    {
        discard;
    }

    int3 loadPos = int3(input.m_position.x, input.m_position.y, 0);
    const float zDepth = PassSrg::m_depthInput.SubpassLoad(loadPos).x;
    if (zDepth > 0.0)
    {
        //TODO: discard or clip(-1)?
        // In reverse depth, any value bigger than 0.0 means there's something already rendered here.
        // Because the skybox only renders in empty space, we skip this pixel.
        discard;
        return OUT;
    }

    real3 color = real3(0.5,0.5,0.5);
#if ENABLE_PHYSICAL_SKY
    if(SceneSrg::m_physicalSky)
    {
        if(input.m_cubemapCoord.z >= 0.0)
        {
            // Default sun direction should be in +Y direction, so sun position should be in -Y
            input.m_cubemapCoord.y = -input.m_cubemapCoord.y;
            // Atom is Z-up, but the value is calculated in Y-up
            float3 worldPosition = normalize(input.m_cubemapCoord.yzx);

            // Theta is the angle between viewing direction and the zenith
            float cosTheta = worldPosition.y;
            
            // Gamma is the angle between viewing direction and the sun
            float cosGamma = dot(worldPosition, SceneSrg::m_physicalSkyData.m_physicalSkySunDirection.xyz);

            float gamma = acos(cosGamma);
            float3 sunParameters = SceneSrg::m_physicalSkyData.m_physicalSkySunParameters.xyz;

            // if gamma is smaller then anular diameter of the sun
            if(cosGamma > sunParameters.z)
            {
                // Render Sun
                // m_physicalSkyAndSunIntensity.y is the sun intensity assigned by user
                color = SceneSrg::m_physicalSkyData.m_physicalSkySunRGB.xyz * SceneSrg::m_physicalSkyData.m_physicalSkyAndSunIntensity.y * PassSrg::m_sunIntensityMultiplier;
            }
            else
            {
                // Render Sky
                // m_physicalSkyParameterZ is the sky color value at zenith
                // m_physicalSkyAndSunIntensity.x is the sky intensity assigned by user
                float3 Z = SceneSrg::m_physicalSkyData.m_physicalSkyParameterZ.xyz;
                float3 srgbColor = Z * HosekWilkie(cosGamma, gamma, cosTheta) * SceneSrg::m_physicalSkyData.m_physicalSkyAndSunIntensity.x;
                color = TransformColor(srgbColor, ColorSpaceId::LinearSRGB, ColorSpaceId::ACEScg);
            }
        } 

        if (SceneSrg::m_fogEnable)
        {
            if (input.m_cubemapCoord.z >= 0.0 && input.m_cubemapCoord.z <= SceneSrg::m_fogTopHeight)
            {
                color = lerp(SceneSrg::m_fogColor.rgb, color, input.m_cubemapCoord.z > 0.0 ? input.m_cubemapCoord.z/SceneSrg::m_fogTopHeight : 0.0);
            }
            else if (input.m_cubemapCoord.z < 0.0 && input.m_cubemapCoord.z >= -SceneSrg::m_fogBottomHeight)
            {
                color = SceneSrg::m_fogColor.rgb;
            }
        }
    }
    else
#endif
    {
        color = real3(SceneSrg::m_skyboxCubemap.Sample(PassSrg::LinearSampler, GetCubemapCoords(input.m_cubemapCoord)).rgb);

        //No need to go to acescg if the tonemap will be applied right away.
#if !ENABLE_MERGE_FILMIC_TONEMAP
        color = TransformColor(color, ColorSpaceId::LinearSRGB, ColorSpaceId::ACEScg);
#endif

        // apply the exposure for HDRi texture
        real exposureFactor = pow(2.0, real(SceneSrg::m_cubemapExposure));
        color *= exposureFactor;
    }
    
    // Clamp lower bounds of sky to prevent zero or negative values, which can leads to NaNs in other shaders like SMAA
    color = max(0.000001, color);

#if ENABLE_MERGE_FILMIC_TONEMAP
    // Apply manual exposure compensation
    color = ApplyManualExposure(color, real(ViewSrg::GetExposureValueCompensation()));

    // We could add Aces support here as well if perf allows.
    color = TonemapFilmic(color);
#endif

    OUT.m_specular = float4(color, 1.0);
#ifdef SKYBOX_TWO_OUTPUTS
    OUT.m_reflection = float4(color, 1.0);
#endif
    return OUT;
}
