/*
 * 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/ColorManagement/TransformColor.azsli>
#include <Atom/Features/PostProcessing/FullscreenVertexUtil.azsli>
#include <Atom/Features/MatrixUtility.azsli>

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

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Sampler LinearSampler
    {
        MinFilter = Linear;
        MagFilter = Linear;
        MipFilter = Linear;
        AddressU = Clamp;
        AddressV = Clamp;
        AddressW = Clamp;
    };

    float m_sunIntensityMultiplier;
}

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)
{
    if(!SceneSrg::m_enable)
    {
        discard;
    }

    float3 color = float3(0.5,0.5,0.5);
    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
    {
        color = SceneSrg::m_skyboxCubemap.Sample(PassSrg::LinearSampler, GetCubemapCoords(input.m_cubemapCoord)).rgb;

        // apply the exposure for HDRi texture
        float exposureFactor = pow(2.0, 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.000001f, color);

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