// Ref: A Scalable and Production Ready Sky and Atmosphere Rendering Technique - Hillaire, ESGR 2020
// https://sebh.github.io/publications/egsr2020.pdf

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch
//#pragma enable_d3d11_debug_symbols

#pragma kernel MultiScatteringLUT   OUTPUT_MULTISCATTERING
#pragma kernel SkyViewLUT
#pragma kernel AtmosphericScatteringLUTCamera AtmosphericScatteringLUT=AtmosphericScatteringLUTCamera CAMERA_SPACE
#pragma kernel AtmosphericScatteringLUTWorld  AtmosphericScatteringLUT=AtmosphericScatteringLUTWorld  DISABLE_ATMOS_EVALUATE_ARTIST_OVERRIDE

#define DIRECTIONAL_SHADOW_ULTRA_LOW  // Different options are too expensive.

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Sampling/Hammersley.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/LightDefinition.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Sky/PhysicallyBasedSky/PhysicallyBasedSkyEvaluation.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Sky/SkyUtils.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/LightLoop/VolumetricCloudsShadowSampling.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/AtmosphericScattering/AtmosphericScattering.hlsl"

// This is the main function that integrates atmosphere along a ray
// It is baked in various LUTs by all the kernels below

// O is position in planet space, V is view dir in world space
void EvaluateAtmosphericColor(float3 O, float3 V, float tExit,
#ifdef OUTPUT_MULTISCATTERING
    float3 L, out float3 multiScattering,
#endif
    out float3 skyColor, out float3 skyTransmittance)
{
    skyColor = 0.0f;
    skyTransmittance = 1.0f;

#ifdef OUTPUT_MULTISCATTERING
    multiScattering = 0.0f;
#endif

    const uint sampleCount = 16;

    for (uint s = 0; s < sampleCount; s++)
    {
        float t, dt;
        GetSample(s, sampleCount, tExit, t, dt);

        const float3 P = O + t * V;
        const float  r = max(length(P), _PlanetaryRadius);
        const float3 N = P * rcp(r);
        const float  height = r - _PlanetaryRadius;

        const float3 sigmaE       = AtmosphereExtinction(height);
        const float3 scatteringMS = AirScatter(height) + AerosolScatter(height);
        const float3 transmittanceOverSegment = TransmittanceFromOpticalDepth(sigmaE * dt);

#ifdef OUTPUT_MULTISCATTERING
        multiScattering += IntegrateOverSegment(scatteringMS, transmittanceOverSegment, skyTransmittance, sigmaE);

        const float3 phaseScatter = scatteringMS * IsotropicPhaseFunction();
        const float3 S = EvaluateSunColorAttenuation(dot(N, L), r) * phaseScatter;
        skyColor += IntegrateOverSegment(S, transmittanceOverSegment, skyTransmittance, sigmaE);
#else
        for (uint i = 0; i < _CelestialLightCount; i++)
        {
            CelestialBodyData light = _CelestialBodyDatas[i];
            float3 L = -light.forward.xyz;

            const float3 sunTransmittance = EvaluateSunColorAttenuation(dot(N, L), r);
            const float3 phaseScatter = AirScatter(height) * AirPhase(-dot(L, V)) + AerosolScatter(height) * AerosolPhase(-dot(L, V));
            const float3 multiScatteredLuminance = EvaluateMultipleScattering(dot(N, L), height);

            float3 S = sunTransmittance * phaseScatter + multiScatteredLuminance * scatteringMS;
            skyColor += IntegrateOverSegment(light.color * S, transmittanceOverSegment, skyTransmittance, sigmaE);
        }
#endif

        skyTransmittance *= transmittanceOverSegment;
    }
}

// Multiple-Scattering LUT

#ifdef OUTPUT_MULTISCATTERING

#define SAMPLE_COUNT 64

RW_TEXTURE2D(float3, _MultiScatteringLUT_RW);

groupshared float3 gs_radianceMS[SAMPLE_COUNT];
groupshared float3 gs_radiance[SAMPLE_COUNT];

float3 RenderPlanet(float3 P, float3 L)
{
    float3 N = normalize(P);

    float3 albedo = _GroundAlbedo.xyz;
    float3 gBrdf = INV_PI * albedo;

    float cosHoriz = ComputeCosineOfHorizonAngle(_PlanetaryRadius);
    float cosTheta = dot(N, L);

    float3 intensity = 0.0f;
    if (cosTheta >= cosHoriz)
    {
        float3 opticalDepth = ComputeAtmosphericOpticalDepth(_PlanetaryRadius, cosTheta, true);
        intensity = TransmittanceFromOpticalDepth(opticalDepth);
    }

    return gBrdf * (saturate(dot(N, L)) * intensity);
}

void ParallelSum(uint threadIdx, inout float3 radiance, inout float3 radianceMS)
{
#ifdef PLATFORM_SUPPORTS_WAVE_INTRINSICS
    radiance   = float3(WaveActiveSum(radiance.x),   WaveActiveSum(radiance.y),   WaveActiveSum(radiance.z));
    radianceMS = float3(WaveActiveSum(radianceMS.x), WaveActiveSum(radianceMS.y), WaveActiveSum(radianceMS.z));
#else
    gs_radiance[threadIdx]   = radiance;
    gs_radianceMS[threadIdx] = radianceMS;
    GroupMemoryBarrierWithGroupSync();

    UNITY_UNROLL
    for (uint s = SAMPLE_COUNT / 2u; s > 0u; s >>= 1u)
    {
        if (threadIdx < s)
        {
            gs_radiance[threadIdx]   += gs_radiance[threadIdx + s];
            gs_radianceMS[threadIdx] += gs_radianceMS[threadIdx + s];
        }

        GroupMemoryBarrierWithGroupSync();
    }

    radiance   = gs_radiance[0];
    radianceMS = gs_radianceMS[0];
#endif
}

[numthreads(1, 1, SAMPLE_COUNT)]
void MultiScatteringLUT(uint3 coord : SV_DispatchThreadID)
{
    const uint threadIdx = coord.z;

    /// Map thread id to position in planet space + light direction

    float sunZenithCosAngle, radialDistance;
    UnmapMultipleScattering(coord.xy, sunZenithCosAngle, radialDistance);

    float3 L = float3(0.0, sunZenithCosAngle, SinFromCos(sunZenithCosAngle));
    float3 O = float3(0.0f, radialDistance, 0.0f);

    float2 U = Hammersley2d(threadIdx, SAMPLE_COUNT);
    float3 V = SampleSphereUniform(U.x, U.y);

    /// Compute single scattering light in direction V

    float3 N; float r; // These params correspond to the entry point
    float tEntry = IntersectAtmosphere(O, -V, N, r).x;
    float tExit  = IntersectAtmosphere(O, -V, N, r).y;

    float cosChi = dot(N, V);
    float cosHor = ComputeCosineOfHorizonAngle(r);

    bool rayIntersectsAtmosphere = (tEntry >= 0);
    bool lookAboveHorizon        = (cosChi >= cosHor);
    bool seeGround               = rayIntersectsAtmosphere && !lookAboveHorizon;

    if (seeGround)
        tExit = tEntry + IntersectSphere(_PlanetaryRadius, cosChi, r).x;

    float3 multiScattering = 0.0f, skyColor = 0.0f, skyTransmittance = 1.0f;
    if (tExit > 0.0f)
        EvaluateAtmosphericColor(O, V, tExit, L, multiScattering, skyColor, skyTransmittance);

    if (seeGround)
        skyColor += RenderPlanet(O + tExit * V, L) * skyTransmittance;

    const float dS = FOUR_PI * IsotropicPhaseFunction() / SAMPLE_COUNT;
    float3 radiance = skyColor * dS;
    float3 radianceMS = multiScattering * dS;

    /// Accumulate light from all directions using LDS

    ParallelSum(threadIdx, radiance, radianceMS);
    if (threadIdx > 0)
        return;

    /// Approximate infinite multiple scattering

    const float3 F_ms = 1.0f * rcp(1.0 - radianceMS); // Equation 9
    const float3 MS   = radiance * F_ms;           // Equation 10

    _MultiScatteringLUT_RW[coord.xy] = MS;
}

#else

// Sky View LUT

RW_TEXTURE2D(float3, _SkyViewLUT_RW);

[numthreads(8, 8, 1)]
void SkyViewLUT(uint2 coord : SV_DispatchThreadID)
{
    const float3 N = float3(0, 1, 0);
    const float r = _PlanetaryRadius;
    const float3 O = r * N;

    float3 V;
    UnmapSkyView(coord, V);

    float tExit = IntersectSphere(_AtmosphericRadius, dot(N, V), r).y;

    float3 skyColor, skyTransmittance;
    EvaluateAtmosphericColor(O, V, tExit, skyColor, skyTransmittance);

    _SkyViewLUT_RW[coord] = skyColor / _CelestialLightExposure;
}

// Atmospheric Scattering LUT

RW_TEXTURE3D(float3, _AtmosphericScatteringLUT_RW);

groupshared float3 gs_data[PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_DEPTH];

float3 ParallelPrefixProduct(uint threadIdx, float3 transmittance)
{
    // For some reason WavePrefixProduct doesn't compile on gamecore
#if defined(PLATFORM_SUPPORTS_WAVE_INTRINSICS) && !defined(SHADER_API_GAMECORE)
    return float3(WavePrefixProduct(transmittance.x), WavePrefixProduct(transmittance.y), WavePrefixProduct(transmittance.z));
#else
    if (threadIdx == PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_DEPTH-1) gs_data[0] = 1;
    else gs_data[threadIdx+1] = transmittance;
    GroupMemoryBarrierWithGroupSync();

    [unroll]
    for (uint s = 1u; s < PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_DEPTH; s <<= 1u)
    {
        uint k = s << 1;
        if (threadIdx % k >= s)
            gs_data[threadIdx] *= gs_data[(threadIdx & ~(k - 1)) + s - 1];

        GroupMemoryBarrierWithGroupSync();
    }
    return gs_data[threadIdx];
#endif
}

float3 ParallelPostfixSum(uint threadIdx, float3 radiance)
{
#ifdef PLATFORM_SUPPORTS_WAVE_INTRINSICS
    // for some reason, the sum has to be done per component
    return float3(WavePrefixSum(radiance.x), WavePrefixSum(radiance.y), WavePrefixSum(radiance.z)) + radiance;
#else
    gs_data[threadIdx] = radiance;
    GroupMemoryBarrierWithGroupSync();

    [unroll]
    for (uint s = 1u; s < PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_DEPTH; s <<= 1u)
    {
        uint k = s << 1;
        if (threadIdx % k >= s)
            gs_data[threadIdx] += gs_data[(threadIdx & ~(k - 1)) + s - 1];

        GroupMemoryBarrierWithGroupSync();
    }
    return gs_data[threadIdx];
#endif
}

[numthreads(1, 1, PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_DEPTH)]
void AtmosphericScatteringLUT(uint2 coord : SV_GroupID, uint s : SV_GroupIndex)
{
    const float2 res = float2(PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_WIDTH, PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_HEIGHT);
    const float2 uv = (coord + 0.5) / res;

    float3 V = -GetSkyViewDirWS(uv * _ScreenSize.xy);

    float3 O;
    float t, dt;
    UnmapAtmosphericScattering(s, V, O, t, dt);

    float3 skyColor = 0.0f;
    float3 skyTransmittance = 1.0f;

    // Following is the loop from EvaluateAtmosphericColor, with each iteration evaluated on a thread

    float3 P = O + t * V;
#ifndef CAMERA_SPACE
    // When ray starts to intersect the planet, don't stop but move the point to the surface
    // This is important because we bilinear sample the LUT and don't want garbage values anywhere
    if (length(P) < _PlanetaryRadius)
    {
        P = normalize(P) * _PlanetaryRadius;
        V = normalize(P - O);
    }
#endif

    const float  r = max(length(P), _PlanetaryRadius + 1);
    const float3 N = P * rcp(r);
    const float  height = r - _PlanetaryRadius;

    const float3 sigmaE         = AtmosphereExtinction(height);
    const float3 scatteringMS   = AirScatter(height) + AerosolScatter(height);
    const float3 transmittanceOverSegment = TransmittanceFromOpticalDepth(sigmaE * dt);

    skyTransmittance = ParallelPrefixProduct(s, transmittanceOverSegment);

    for (uint i = 0; i < _CelestialLightCount; i++)
    {
        CelestialBodyData light = _CelestialBodyDatas[i];
        float3 L = -light.forward.xyz;

        float shadow = 1.0f;

        /*
        // Disabled because the LUT is too low res to get interesting results
        if (light.shadowIndex >= 0)
        {
            HDShadowContext shadowContext = InitShadowContext();
            shadow *= GetDirectionalShadowAttenuation(shadowContext,
                    coord, P + _PlanetCenterPosition, -V,
                    light.shadowIndex, L);
        }
        */

        if (_VolumetricCloudsShadowOriginToggle.w == 1.0)
        {
            DirectionalLightData dirLight;
            dirLight.forward = light.forward;
            dirLight.right = light.right;
            dirLight.up = light.up;
            shadow *= EvaluateVolumetricCloudsShadows(dirLight, P + _PlanetCenterPosition);
        }

        const float3 sunTransmittance = shadow * EvaluateSunColorAttenuation(dot(N, L), r);
        const float3 phaseScatter = AirScatter(height) * AirPhase(-dot(L, V)) + AerosolScatter(height) * AerosolPhase(-dot(L, V));
        const float3 multiScatteredLuminance = EvaluateMultipleScattering(dot(N, L), height);

        // Compute color
        float3 S = sunTransmittance * phaseScatter + multiScatteredLuminance * scatteringMS;
        skyColor += IntegrateOverSegment(light.color * S, transmittanceOverSegment, skyTransmittance, sigmaE);
    }

    skyColor = ParallelPostfixSum(s, skyColor);

    // For debug: no LDS
    //GetSample(s, PBRSKYCONFIG_ATMOSPHERIC_SCATTERING_LUT_DEPTH, ATMOSPHERIC_SCATTERING_MAX_DISTANCE, t, dt);
    //EvaluateAtmosphericColor(O, V, t, skyColor, skyTransmittance);

    // Make sure first slice is all black. Looks better for bilinear at close range
    if (s == 0) skyColor = 0.0f;

    skyColor = Desaturate(skyColor, _ColorSaturation);
    _AtmosphericScatteringLUT_RW[uint3(coord, s)] = skyColor * _IntensityMultiplier * GetCurrentExposureMultiplier();
}

#endif