#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#pragma multi_compile_local _ USE_SECOND_CLOUD_LAYER

#pragma kernel BakeCloudTexture KERNEL_NAME=BakeCloudTexture
//#pragma enable_d3d11_debug_symbols

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/VolumeRendering.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Sky/SkyUtils.hlsl"

#ifndef USE_SECOND_CLOUD_LAYER
    #define NUM_LAYERS  1
    #define LAYER       0
#else
    #define NUM_LAYERS  2
    #define LAYER       dispatchThreadId.z
#endif

#define NUM_MULTI_SCATTERING_OCTAVES 2
#define FORWARD_ECCENTRICITY    0.7
#define BACKWARD_ECCENTRICITY   0.3
#define MULTI_SCATTERING        0.75
#define _HighestAltitude        (_LowestAltitude + 800)

TEXTURE2D(_CloudMapA);
SAMPLER(sampler_CloudMapA);

#ifdef USE_SECOND_CLOUD_LAYER
TEXTURE2D(_CloudMapB);
SAMPLER(sampler_CloudMapB);
#endif

RW_TEXTURE2D_ARRAY(float2, _CloudTexture);

float4 _Params;
float4 _Params1[NUM_LAYERS];
float4 _Params2[NUM_LAYERS];
float _Resolution;

#define _SunDirection       _Params.xyz
#define _UpperHemisphere    (_Params.w > 0)
#define _PlanetaryRadius    abs(_Params.w)
#define _Opacities          _Params1[data.index]
#define _Rotation(index)    _Params2[index].x
#define _LightingSteps      _Params2[data.index].y
#define _SigmaT             _Params2[data.index].z
#define _LowestAltitude     (_Params2[data.index].w + _PlanetaryRadius)

#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Sky/CloudUtils.hlsl"

struct LayerSamplingData
{
    int index;
    float4 weights;
    float3 sunDirection;

    TEXTURE2D(cloudMap);
    SAMPLER(sampler_cloudMap);
};

void GetCloudVolumeIntersection(LayerSamplingData data, float3 dir, out float rangeStart, out float range)
{
    float tEntry, tExit;
    IntersectCloudVolume(float3(0, _PlanetaryRadius, 0), dir, _LowestAltitude, _HighestAltitude, tEntry, tExit);

    rangeStart = tEntry;
    range = tExit - tEntry;
}

float3 InverseLatLong(float2 uv, float upperHemisphereOnly)
{
    const float2 invAtan = float2(0.1591, 0.3183);

    uv.y = upperHemisphereOnly ? uv.y * 0.5 + 0.5 : uv.y;
    uv = (uv - 0.5) / invAtan;

    float y = sin(uv.y);
    float scale = sqrt(1.0 - y*y);

    return float3(sin(uv.x) * scale, y, cos(uv.x) * scale);
}

float SampleCloudMap(LayerSamplingData data, float2 uv)
{
    float4 cloudLayerColor = SAMPLE_TEXTURE2D_LOD(data.cloudMap, data.sampler_cloudMap, uv, 0);

    float4 clouds = cloudLayerColor * data.weights;
    return clouds.r + clouds.g + clouds.b + clouds.a;
}

float GetDensity(LayerSamplingData data, float3 positionWS, float3 dir, float thickness)
{
    float rangeStart, range;
    GetCloudVolumeIntersection(data, dir, rangeStart, range);

    float distToCenter = length(dir * (rangeStart + 0.5 * range) - positionWS);
    return (distToCenter > range * thickness) ? 0 : 0.1 * thickness;
}

float GetDensity(LayerSamplingData data, float3 positionWS)
{
    float3 dir = normalize(positionWS);
    float thickness = SampleCloudMap(data, GetLatLongCoords(dir, _UpperHemisphere));
    return GetDensity(data, positionWS, dir, thickness);
}

float EvaluateSunLuminance(LayerSamplingData data, float3 positionWS, float phaseFunction[NUM_MULTI_SCATTERING_OCTAVES])
{
    const int _NumLightSteps = (int)_LightingSteps;

    bool hit;
    float luminance = 0.0, extinction = 0.0, totalLightDistance;
    ExitCloudVolume(positionWS + float3(0, _PlanetaryRadius, 0), data.sunDirection, _HighestAltitude, totalLightDistance);
    totalLightDistance = clamp(totalLightDistance, 0, 500);

    // Compute the size of the current step
    float stepSize = totalLightDistance / (float)_NumLightSteps;

    // Collect total density along light ray.
    for (int j = 1; j <= _NumLightSteps; j++)
    {
        float3 currentSamplePointWS = positionWS + data.sunDirection * stepSize * j;
        extinction += GetDensity(data, currentSamplePointWS);
    }

    // Combine the luminance for each octave
    for (int o = 0; o < NUM_MULTI_SCATTERING_OCTAVES; ++o)
        luminance += exp(-stepSize * extinction * _SigmaT * PositivePow(MULTI_SCATTERING, o))
                     * phaseFunction[o]
                     * PositivePow(MULTI_SCATTERING, o);

    return luminance;
}

float2 ComputeCloudLighting(TEXTURE2D(cloudMap), SAMPLER(sampler_cloudMap), float2 uv, uint index)
{
    LayerSamplingData data;

    data.index = index;
    data.cloudMap = cloudMap;
    data.sampler_cloudMap = sampler_cloudMap;
    data.weights = _Opacities / max(_Opacities.r + _Opacities.g + _Opacities.b + _Opacities.a, 1.0);
    data.sunDirection = RotationUp(_SunDirection, float2(cos(_Rotation(index) * TWO_PI), sin(_Rotation(index) * TWO_PI)));

    const float opacity = SampleCloudMap(data, uv);
    float2 finalColor = opacity;

    if ((int)_LightingSteps != 0)
    {
        const float3 dir = InverseLatLong(uv, _UpperHemisphere);
        const float cosAngle = dot(dir, data.sunDirection);

        float phaseFunction[NUM_MULTI_SCATTERING_OCTAVES];
        for (int o = 0; o < NUM_MULTI_SCATTERING_OCTAVES; ++o)
        {
            const float forwardP = HenyeyGreensteinPhaseFunction(FORWARD_ECCENTRICITY * PositivePow(MULTI_SCATTERING, o), cosAngle);
            const float backwardsP = HenyeyGreensteinPhaseFunction(-BACKWARD_ECCENTRICITY * PositivePow(MULTI_SCATTERING, o), cosAngle);
            phaseFunction[o] = backwardsP + forwardP;
        }

        float rangeStart, range;
        GetCloudVolumeIntersection(data, dir, rangeStart, range);

        float3 currentPositionWS = dir * (rangeStart + 0.5 * (1 - opacity) * range);

        finalColor = 0;
        if (opacity != 0.0f)
        {
            const float sigmaE = 0.1 * opacity * _SigmaT;
            const float currentStepExtinction = exp(-sigmaE * range * opacity);
            const float luminance = EvaluateSunLuminance(data, currentPositionWS, phaseFunction);

            // Output (scattering, 1 - transmittance)
            finalColor.x = luminance - luminance * currentStepExtinction;
            finalColor.y = 1 - currentStepExtinction;
        }
    }

    return finalColor;
}

[numthreads(8, 8, NUM_LAYERS)]
void KERNEL_NAME(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    float2 uv = float2(dispatchThreadId.x * _Resolution, dispatchThreadId.y * _Resolution);

    uv.x -= _Rotation(LAYER);
    if (_UpperHemisphere)
        uv.y *= 2.0f;

#ifndef USE_SECOND_CLOUD_LAYER
    float2 cloudLayerColor = ComputeCloudLighting(_CloudMapA, sampler_CloudMapA, uv, LAYER);
#else
    float2 cloudLayerColor = LAYER == 0 ?
        ComputeCloudLighting(_CloudMapA, sampler_CloudMapA, uv, LAYER) :
        ComputeCloudLighting(_CloudMapB, sampler_CloudMapB, uv, LAYER);
#endif

    _CloudTexture[dispatchThreadId] = cloudLayerColor;
}