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

#pragma kernel main

#define _PlanetaryRadius _GroundAlbedo_PlanetRadius.w

#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.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"

#define TABLE_SIZE uint3(PBRSKYCONFIG_IN_SCATTERED_RADIANCE_TABLE_SIZE_X, \
                         PBRSKYCONFIG_IN_SCATTERED_RADIANCE_TABLE_SIZE_Y, \
                         PBRSKYCONFIG_IN_SCATTERED_RADIANCE_TABLE_SIZE_Z)

// Emulate a 4D texture with a "deep" 3D texture.
RW_TEXTURE3D(float3, _AirSingleScatteringTable);
RW_TEXTURE3D(float3, _AerosolSingleScatteringTable);
RW_TEXTURE3D(float3, _MultipleScatteringTable);

[numthreads(4, 4, 4)]
void main(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    const float A = _AtmosphericRadius;
    const float R = _PlanetaryRadius;

    const uint zTexSize = PBRSKYCONFIG_IN_SCATTERED_RADIANCE_TABLE_SIZE_Z;
    const uint zTexCnt  = PBRSKYCONFIG_IN_SCATTERED_RADIANCE_TABLE_SIZE_W;

    // We don't care about the extremal points for XY, but need the full range of Z values.
    const float3 scale = rcp(float3(TABLE_SIZE.xy, 1));
    const float3 bias  = float3(0.5 * scale.xy, 0);

    // Let the hardware and the driver handle the ordering of the computation.
    uint3 tableCoord = dispatchThreadId;
    uint  texId      = tableCoord.z / zTexSize;       // [0, zTexCnt  - 1]
    uint  texCoord   = tableCoord.z & (zTexSize - 1); // [0, zTexSize - 1]

    float3 uvw = tableCoord * scale + bias;

    // Convention:
    // V points towards the camera.
    // The normal vector N points upwards (local Z).
    // The view vector V spans the local XZ plane.
    // The light vector is represented as {phiL, cosThataL} w.r.t. the XZ plane.
    float cosChi = UnmapAerialPerspective(uvw.xy).x;                             // [-1, 1]
    float height = UnmapAerialPerspective(uvw.xy).y;                             // [0, _AtmosphericDepth]
    float phiL   = PI * saturate(texCoord * rcp(zTexSize - 1));                  // [-Pi, Pi]
    float NdotL  = UnmapCosineOfZenithAngle(saturate(texId * rcp(zTexCnt - 1))); // [-0.5, 1]

    float NdotV  = -cosChi;
    float r      = height + R;
    float cosHor = ComputeCosineOfHorizonAngle(r);

    bool lookAboveHorizon = (cosChi >= cosHor);
    bool viewAboveHorizon = (NdotV >= cosHor);

    float3 N = float3(0, 0, 1);
    float3 V = SphericalToCartesian(0, NdotV);
    float3 L = SphericalToCartesian(phiL, NdotL);

    float LdotV = dot(L, V);
    // float LdotV = SphericalDot(NdotL, phiL, NdotV, 0);

    // Set up the ray...
    float  h = height;
    float3 O = r * N;

    // Determine the region of integration.
    float tMax;

    if (lookAboveHorizon)
    {
        tMax = IntersectSphere(A, cosChi, r).y; // Max root
    }
    else
    {
        tMax = IntersectSphere(R, cosChi, r).x; // Min root
    }


    // Integrate in-scattered radiance along -V.
    // Note that we have to evaluate the transmittance integral along -V as well.
    // The transmittance integral is pretty smooth (I plotted it in Mathematica).
    // However, using a non-linear distribution of samples is still a good idea both
    // when looking up (due to the exponential falloff of the coefficients)
    // and for horizontal rays (due to the exponential transmittance term).
    // It's easy enough to use a simple quadratic remap.

    float3 airTableEntry     = 0;
    float3 aerosolTableEntry = 0;
    float3 msTableEntry = 0;
    float3 transmittance = 1.0f;

    // Eye-balled number of samples.
    const int numSamples = 16;

    for (int i = 0; i < numSamples; i++)
    {
        float t, dt;
        GetSample(i, numSamples, tMax, t, dt);

        float3 P = O + t * -V;

        // Update these for the step along the ray...
        r      = max(length(P), R);
        height = r - R;
        NdotV  = dot(normalize(P), V);
        NdotL  = dot(normalize(P), L);

        const float3 sigmaE       = AtmosphereExtinction(height);
        const float3 transmittanceOverSegment = TransmittanceFromOpticalDepth(sigmaE * dt);

        // Apply the phase function at runtime.
        float3 sunTransmittance = EvaluateSunColorAttenuation(NdotL, r);
        float3 airTerm      = sunTransmittance * AirScatter(height);
        float3 aerosolTerm  = sunTransmittance * AerosolScatter(height);
        float3 scatteringMS = AirScatter(height) + AerosolScatter(height);
        float3 msTerm       = EvaluateMultipleScattering(NdotL, height) * scatteringMS;

        airTableEntry     += IntegrateOverSegment(airTerm, transmittanceOverSegment, transmittance, sigmaE);
        aerosolTableEntry += IntegrateOverSegment(aerosolTerm, transmittanceOverSegment, transmittance, sigmaE);
        msTableEntry      += IntegrateOverSegment(msTerm, transmittanceOverSegment, transmittance, sigmaE);

        transmittance *= transmittanceOverSegment;
    }

    // TODO: deep compositing.
    // Note: ground reflection is computed at runtime.
    _AirSingleScatteringTable[tableCoord]     = airTableEntry;                    // One order
    _AerosolSingleScatteringTable[tableCoord] = aerosolTableEntry;                // One order
    _MultipleScatteringTable[tableCoord]      = msTableEntry * MS_EXPOSURE;
}