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

#pragma kernel MAIN_1 main=MAIN_1 SINGLE_SCATTERING
#pragma kernel MAIN_S main=MAIN_S MULTIPLE_SCATTERING_GATHER     SRC_SS
#pragma kernel MAIN_M main=MAIN_M MULTIPLE_SCATTERING_GATHER     SRC_MS
#pragma kernel MAIN_A main=MAIN_A MULTIPLE_SCATTERING_ACCUMULATE

#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Sky/PhysicallyBasedSky/PhysicallyBasedSkyCommon.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(float4, _AirSingleScatteringTable);
RW_TEXTURE3D(float4, _AerosolSingleScatteringTable);
RW_TEXTURE3D(float4, _MultipleScatteringTable);
RW_TEXTURE3D(float4, _MultipleScatteringTableOrder);

[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
    }

#ifdef SINGLE_SCATTERING

    float4 airTableEntry     = 0;
    float4 aerosolTableEntry = 0;

    // For single scattering with the directional light assumption, planet's shadow forms a shadow volume
    // which is shaped like a cylinder with the planet's radius, aligned with the light direction.
    // We can use it to optimize the integration process, as the shadowed region has 0 contribution.

    float  tMin = 0;
    float2 tCyl = IntersectRayCylinder(L, R, r, -V);

    if (tCyl.y >= 0) // Will exit?
    {
        // We found an intersection with the shadow volume.
        // It could be either the entry or the exit.

        if (tCyl.x >= 0) // Will enter?
        {
            float zHit = dot(L, O + tCyl.x * -V);

            if (zHit < 0) // Below?
            {
                // Stop the ray at the entry.
                // We assume that the ray does not exit the shadow volume before exiting the atmosphere.
                // The error of that assumption should be reasonably low.
                tMax = min(tMax, tCyl.x);
            }
        }
        else // Already entered?
        {
            bool inShadow = (L.z < 0); // Inside the shadow volume?

            if (inShadow)
            {
                // Start the ray at the exit.
                tMin = max(tMin, tCyl.y);
            }
        }
    }
    else // Failed to intersect?
    {
        bool inShadow = (L.z < 0) && ((1 - Sq(L.z)) < (R * rcp(r)));

        // If we are inside the shadow volume, we should terminate the ray.
        if (inShadow)
        {
            tMax = -tMin;
        }
    }

    airTableEntry.a     = tMax - tMin;
    aerosolTableEntry.a = tMax - tMin;

    if (tMin >= tMax)
    {
        _AirSingleScatteringTable[tableCoord]     = 0; // One order
        _AerosolSingleScatteringTable[tableCoord] = 0; // One order
        _MultipleScatteringTable[tableCoord]      = 0; // MS orders
        return;
    }

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

    float  t0          = tMin;
    float3 viewODepthX = ComputeAtmosphericOpticalDepth(r, NdotV, viewAboveHorizon); // Back to start
    float3 viewODepth  = 0;                                         // Sample to start

    if (tMin > 0) // Start changed?
    {
        float3 P = O + tMin * -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);

        viewODepth = ComputeAtmosphericOpticalDepth(r, NdotV, viewAboveHorizon) - viewODepthX;
    }

    float3 lightODepth = ComputeAtmosphericOpticalDepth1(r, NdotL);

    // Apply the phase function at runtime.
    float3 transm0               = TransmittanceFromOpticalDepth(viewODepth + lightODepth);
    float3 airInScatterTerm0     = transm0 * AirScatter(height);
    float3 aerosolInScatterTerm0 = transm0 * AerosolScatter(height);

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

    const float startNdotV = NdotV;

    for (int i = 0; i < numSamples; i++)
    {
        // Use (n = 0.5) when looking down, (n = 1) when looking horizontally,
        // and (n = 2) when looking up.
        // TODO: use exponential media importance sampling...
        float n  = pow(2, -startNdotV);
        float x  = pow(abs((i + 1) * rcp(numSamples)), n); // TODO: Cranley-Patterson Rotation
        float t1 = tMin + (tMax - tMin) * x;
        float dt = t1 - t0;

        float3 P = O + t1 * -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);

        viewODepth  = ComputeAtmosphericOpticalDepth(r, NdotV, viewAboveHorizon) - viewODepthX;
        lightODepth = ComputeAtmosphericOpticalDepth1(r, NdotL);

        // Compute the amount of in-scattered radiance.
        // Evaluate the transmittance integral from 't0' to 't1' using the trapezoid rule.
        // Integral{a, b}{Transmittance(0, t) dt} ≈
        // (Transmittance(0, a) + Transmittance(0, b)) * (b - a) / 2 =
        // (Transmittance(0, a) + Transmittance(0, tMax) / Transmittance(b, tMax)) * (b - a) / 2.
        // We handle other terms in a similar way (by pre-multiplying transmittance).
        float3 transm1               = TransmittanceFromOpticalDepth(viewODepth + lightODepth);
        float3 airInScatterTerm1     = transm1 * AirScatter(height);
        float3 aerosolInScatterTerm1 = transm1 * AerosolScatter(height);

        // Apply the phase function at runtime.
        airTableEntry.rgb     += (airInScatterTerm0     + airInScatterTerm1    ) * (0.5 * dt);
        aerosolTableEntry.rgb += (aerosolInScatterTerm0 + aerosolInScatterTerm1) * (0.5 * dt);

        // Carry these over to the next iteration...
        t0                    = t1;
        airInScatterTerm0     = airInScatterTerm1;
        aerosolInScatterTerm0 = aerosolInScatterTerm1;
    }

    // TODO: deep compositing.
    // Note: ground reflection is computed at runtime.
    _AirSingleScatteringTable[tableCoord]     = airTableEntry;                    // One order
    _AerosolSingleScatteringTable[tableCoord] = aerosolTableEntry;                // One order
    _MultipleScatteringTable[tableCoord]      = float4(0, 0, 0, airTableEntry.a); // MS orders

#elif MULTIPLE_SCATTERING_GATHER

    float4 tableEntry;

    tableEntry.rgb = 0;
    tableEntry.a   = tMax;

    // Gather the ground contribution.
    // The ground is a Lambertian reflector. It is basically a textured sphere light.
    // TODO: 50% of the ground is black! And most of the time it's too far away!
    // Also, anisotropic Mie scattering will make sure the contribution is near-zero...
    {
        float cosAperture = -cosHor; // Cosine of the half-angle

        // Arbitrary number of samples...
        const int numGroundSamples = 89;

        for (int i = 0; i < numGroundSamples; i++)
        {
            float2 f = Fibonacci2d(i, numGroundSamples); // TODO: Cranley-Patterson Rotation

            // Construct a direction around the up vector.
            float3 dir; float rcpPdf;
            SampleCone(f, cosAperture, dir, rcpPdf);

            // Flip it upside-down to point towards the sphere light (planet).
            float3 gL = -dir;

            float  t  = IntersectSphere(R, gL.z, r).x;
            float3 gP = O + t * gL;
            float3 gN = normalize(gP);

            // Shade the ground.
            const float3 gBrdf = INV_PI * _GroundAlbedo.rgb;

            // TODO: we compute the same thing for many voxels,
            // the only real difference is the phase function...
            float3 phase  = AtmospherePhaseScatter(dot(gL, V), height);
            float3 oDepth = ComputeAtmosphericOpticalDepth1(r, gL.z);
            float3 transm = TransmittanceFromOpticalDepth(oDepth);

            float weight = rcpPdf * rcp(numGroundSamples);
            tableEntry.rgb += weight * phase * transm * gBrdf * SampleGroundIrradianceTexture(dot(gN, L));
        }
    }

    // Gather the volume contribution.
    {
        // Arbitrary number of samples...
#if defined(SHADER_API_XBOXONE)
        // This forces a different code path, the xbox compiler flusters with the path coming from 89 samples.
        // We don't care too much about perf as it is pre-computation, so we simply get more samples.
        const int numVolumeSamples = 144;
#else
        const int numVolumeSamples = 89;
#endif

        for (int i = 0; i < numVolumeSamples; i++)
        {
            float2 f = Fibonacci2d(i, numVolumeSamples); // TODO: Cranley-Patterson Rotation

            // TODO: anisotropy support without doubling the cost?
            float cosTheta = ImportanceSampleRayleighPhase(f.x);
            float phi      = TWO_PI * f.y;

            // Construct a direction around the V vector.
            float3x3 frame = GetLocalFrame(V); // Orientation around V is arbitrary, so maybe Cranley-Patterson Rotation is not needed

            // Note that this direction is likely no longer in the V-N plane.
            // Must add a fudge factor here, otherwise Unity's shader compiler does not work correctly...
            float3 V1 = SphericalToCartesian(phi, cosTheta + FLT_EPS);
                   V1 = mul(V1, frame);

            // TODO: solve in spherical coords?
            float NdotV1 = V1.z;
            float phiL1  = acos(clamp(dot(L.xy, V1.xy) * rsqrt(max(dot(L.xy, L.xy) * dot(V1.xy, V1.xy), FLT_EPS)), -1, 1));

            TexCoord4D tc = ConvertPositionAndOrientationToTexCoords(height, NdotV1, NdotL, phiL1);

            float3 radiance = 0;

        #ifdef SRC_SS
            // Single scattering does not contain the phase function.
            float LdotV1 = dot(L, V1);

            radiance += lerp(SAMPLE_TEXTURE3D_LOD(_AirSingleScatteringTexture,     s_linear_clamp_sampler, float3(tc.u, tc.v, tc.w0), 0).rgb,
                             SAMPLE_TEXTURE3D_LOD(_AirSingleScatteringTexture,     s_linear_clamp_sampler, float3(tc.u, tc.v, tc.w1), 0).rgb,
                             tc.a) * AirPhase(LdotV1);

            radiance += lerp(SAMPLE_TEXTURE3D_LOD(_AerosolSingleScatteringTexture, s_linear_clamp_sampler, float3(tc.u, tc.v, tc.w0), 0).rgb,
                             SAMPLE_TEXTURE3D_LOD(_AerosolSingleScatteringTexture, s_linear_clamp_sampler, float3(tc.u, tc.v, tc.w1), 0).rgb,
                             tc.a) * AerosolPhase(LdotV1);
        #else  // SRC_MS
            radiance += lerp(SAMPLE_TEXTURE3D_LOD(_MultipleScatteringTexture,      s_linear_clamp_sampler, float3(tc.u, tc.v, tc.w0), 0).rgb,
                             SAMPLE_TEXTURE3D_LOD(_MultipleScatteringTexture,      s_linear_clamp_sampler, float3(tc.u, tc.v, tc.w1), 0).rgb,
                             tc.a);
        #endif // SRC_MS

            // AtmospherPhaseScatter / AirPhase.
            // Assume that, after multiple bounces, the effect of anisotropy is lost.
            // The view direction is the opposite of the incoming direction.
            float3 phase  = AirScatter(height) + AerosolScatter(height);
            float  weight = rcp(numVolumeSamples);

            tableEntry.rgb += weight * phase * radiance;
        }
    }

    // TODO: deep compositing.
    _MultipleScatteringTable[tableCoord] = tableEntry;

#else // MULTIPLE_SCATTERING_ACCUMULATE

    float4 tableEntry;

    tableEntry.rgb = 0;
    tableEntry.a   = tMax;

    // Integrate in-scattered radiance along -V.

    float  t0          = 0;
    float3 viewODepthX = ComputeAtmosphericOpticalDepth(r, NdotV, viewAboveHorizon); // Back to start
    float3 viewODepth  = 0;                                         // Sample to start

    float3 inScatteredRadiance = LOAD_TEXTURE3D(_MultipleScatteringTexture, tableCoord).rgb;

    float3 inScatterTerm0 = TransmittanceFromOpticalDepth(viewODepth) * inScatteredRadiance;

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

    const float startNdotV = NdotV;

    for (int i = 0; i < numSamples; i++)
    {
        // Use (n = 0.5) when looking down, (n = 1) when looking horizontally,
        // and (n = 2) when looking up.
        // TODO: use exponential media importance sampling...
        float n  = pow(2, -startNdotV);
        float x  = pow(abs((i + 1) * rcp(numSamples)), n); // TODO: Cranley-Patterson Rotation
        float t1 = tMax * x;
        float dt = t1 - t0;

        float3 P = O + t1 * -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);

        viewODepth = ComputeAtmosphericOpticalDepth(r, NdotV, viewAboveHorizon) - viewODepthX;

        // 'phiL' is unchanged, so quadrilinear interpolation is actually unnecessary.
        TexCoord4D tc = ConvertPositionAndOrientationToTexCoords(height, NdotV, NdotL, phiL);

        inScatteredRadiance = SAMPLE_TEXTURE3D_LOD(_MultipleScatteringTexture, s_linear_clamp_sampler, float3(tc.u, tc.v, tc.w0), 0).rgb;

        float3 inScatterTerm1 = TransmittanceFromOpticalDepth(viewODepth) * inScatteredRadiance;

        // Apply the phase function at runtime.
        tableEntry.rgb += (inScatterTerm0 + inScatterTerm1) * (0.5 * dt);

        // Carry these over to the next iteration...
        t0             = t1;
        inScatterTerm0 = inScatterTerm1;
    }

    // TODO: deep compositing.
    _MultipleScatteringTableOrder[tableCoord] = tableEntry;                // One order
    _MultipleScatteringTable[tableCoord]     += float4(tableEntry.rgb, 0); // MS orders

#endif // MULTIPLE_SCATTERING_ACCUMULATE
}