/***********************************************/
/*           Copyright (c) 2025 Belmu          */
/*             All Rights Reserved             */
/***********************************************/

/*
    [Credits]:
        Jessie - help with atmospheric scattering and providing data for aerosol asymmetry and Ozone and Mie cross section approximations (https://github.com/Jessie-LC)
        
    [References]:
        Nishita, T. (1993). Display of the earth taking into account atmospheric scattering. http://nishitalab.org/user/nis/cdrom/sig93_nis.pdf
        ASTM. (2000). 2000 ASTM Standard Extraterrestrial Spectrum Reference E-490-00. https://www.nrel.gov/grid/solar-resource/spectra-astm-e490.html
        Elek, O. (2009). Rendering Parametrizable Planetary Atmospheres with Multiple Scattering in Real-Time. https://old.cescg.org/CESCG-2009/papers/PragueCUNI-Elek-Oskar09.pdf
        Novak, J. (2014). Residual Ratio Tracking for Estimating Attenuation in Participating Media. https://cs.dartmouth.edu/~wjarosz/publications/novak14residual.pdf
        Wilmouth, M., D., & Sayres, S., D. (2019). Rayleigh scattering cross sections of argon, carbon dioxide, sulfur hexafluoride, and methane in the UV-A region using Broadband Cavity Enhanced Spectroscopy. https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6800691/
*/

#include "/include/utility/phase.glsl"

#define MAX_ATMOSPHERE_SCATTERING_ORDERS 100

#define SMALL_PARTICLE   0
#define AEROSOL_PARTICLE 1
#define OZONE_PARTICLE   2

vec3 calculateAtmosphereDensities(float centerDist) {
    float altitudeKm = (centerDist - planetRadius) / kilo;
    vec2 rayleighMie = exp(altitudeKm / -(scaleHeights / kilo));

    float o1 = 25.0 *     exp(( 0.0 - altitudeKm) * rcp(  8.0)) * 0.5;
    float o2 = 30.0 * pow(exp((18.0 - altitudeKm) * rcp( 80.0)), altitudeKm - 18.0);
    float o3 = 75.0 * pow(exp((25.3 - altitudeKm) * rcp( 35.0)), altitudeKm - 25.3);
    float o4 = 50.0 * pow(exp((30.0 - altitudeKm) * rcp(150.0)), altitudeKm - 30.0);
    float ozone = (o1 + o2 + o3 + o4) * rcp(134.628);

    return vec3(rayleighMie, ozone);
}

vec3 calculateMaximumAtmosphereDensities() {
    float peakRMAltitudeKm = (atmosphereLowerRadius - planetRadius) / kilo;
    vec2  peakRayleighMie  = exp(peakRMAltitudeKm / -(scaleHeights / kilo));

    float peakOzoneAltitudeKm = (atmosphereLowerRadius + 25e3 - planetRadius) / kilo;

    float o1 = 25.0 *     exp(( 0.0 - peakOzoneAltitudeKm) * rcp(  8.0)) * 0.5;
    float o2 = 30.0 * pow(exp((18.0 - peakOzoneAltitudeKm) * rcp( 80.0)), peakOzoneAltitudeKm - 18.0);
    float o3 = 75.0 * pow(exp((25.3 - peakOzoneAltitudeKm) * rcp( 35.0)), peakOzoneAltitudeKm - 25.3);
    float o4 = 50.0 * pow(exp((30.0 - peakOzoneAltitudeKm) * rcp(150.0)), peakOzoneAltitudeKm - 30.0);
    float peakOzone = (o1 + o2 + o3 + o4) * rcp(134.628);

    return vec3(peakRayleighMie, peakOzone);
}

int deltaTrackingEstimator(inout vec3 position, vec3 direction, vec3 attenuationCoefficients, out int particle) {
    vec2 dists = intersectSphericalShell(position, direction, atmosphereLowerRadius, atmosphereUpperRadius);

    vec3  maxAttenuation      = attenuationCoefficients * calculateMaximumAtmosphereDensities();
    float attenuationMajorant = saturate(maxAttenuation.x + maxAttenuation.y + maxAttenuation.z);

    float t = 0.0;

    while (t < dists.y) {
        float stepSize = -log(randF()) / attenuationMajorant;

        t += stepSize;
        position += direction * stepSize;

        float centerDist = length(position);

        if (centerDist <= atmosphereLowerRadius) return -1;
        if (centerDist >  atmosphereUpperRadius) break;

        vec3 attenuation = attenuationCoefficients * calculateAtmosphereDensities(centerDist);

        float rng = randF();

        if (rng < (attenuation.x + attenuation.y + attenuation.z) / attenuationMajorant) {
            float fraction = 0.0;
            particle = 0;
            while (particle < 2) {
                fraction += attenuation[particle];
                if (rng * attenuationMajorant < fraction) break;
                particle++;
            }
            return 1;
        }
    }
    return 0;
}

float ratioTrackingEstimator(vec3 position, vec3 direction, const vec3 attenuationCoefficients) {
    vec2 dists = intersectSphericalShell(position, direction, atmosphereLowerRadius, atmosphereUpperRadius);
    if (dists.y < 0.0) return 0.0;

    vec3  maxAttenuation      = attenuationCoefficients * calculateMaximumAtmosphereDensities();
    float attenuationMajorant = saturate(maxAttenuation.x + maxAttenuation.y + maxAttenuation.z);

    float t = 0.0;
    float transmittance = 1.0;

    while (t < dists.y) {
        float stepSize = -log(randF()) / attenuationMajorant;

        t += stepSize;
        position += direction * stepSize;

        float centerDist = length(position);

        if (centerDist <= atmosphereLowerRadius) return 0.0;
        if (centerDist >  atmosphereUpperRadius) break;

        vec3 attenuation = attenuationCoefficients * calculateAtmosphereDensities(centerDist);
        transmittance   *= 1.0 - (attenuation.x + attenuation.y + attenuation.z) / attenuationMajorant;
    }
    return transmittance;
}

float estimateAtmosphericScattering(vec3 viewPosition, vec3 viewDirection, vec3 sunDirection, vec3 moonDirection, const vec3 attenuationCoefficients, 
    float sunRadianceSpectrum, float sunIrradianceSpectrum, float moonIrradianceSpectrum, float asymmetry, float wavelength) {
    vec2 dists = intersectSphericalShell(viewPosition, viewDirection, atmosphereLowerRadius, atmosphereUpperRadius);
    if (dists.y < 0.0) return 0.0;

    vec3 rayPosition  = viewPosition + viewDirection * dists.x;
    vec3 rayDirection = viewDirection;

    float throughput = 1.0;
    float estimate   = 0.0;

    int scatteringOrder = 0;
    int particle;
    
    while (scatteringOrder < MAX_ATMOSPHERE_SCATTERING_ORDERS) {
        int interactedWithParticle = deltaTrackingEstimator(rayPosition, rayDirection, attenuationCoefficients, particle);

        if (interactedWithParticle == 1) {
            bool canScatter = false;

            switch (particle) {
                case SMALL_PARTICLE:
                    canScatter = true;
                    break;
                case AEROSOL_PARTICLE:
                    canScatter = randF() < mieScatteringAlbedo;
                    break;
                case OZONE_PARTICLE:
                    canScatter = false;
                    break;
            }

            if (!canScatter) break;

            float phaseSun  = 0.0, phaseMoon = 0.0;
            switch (particle) {
                case SMALL_PARTICLE:
                    phaseSun     = rayleighPhase(dot(rayDirection, sunDirection ));
                    phaseMoon    = rayleighPhase(dot(rayDirection, moonDirection));
                    rayDirection = sampleRayleighPhase(); 
                    break;
                case AEROSOL_PARTICLE:
                    phaseSun     = kleinNishinaPhase(dot(rayDirection, sunDirection ), asymmetry);
                    phaseMoon    = kleinNishinaPhase(dot(rayDirection, moonDirection), asymmetry);
                    rayDirection = sampleKleinNishinaPhase(asymmetry); 
                    break;
            }

            estimate += throughput * phaseSun  * sunIrradianceSpectrum  * ratioTrackingEstimator(rayPosition, sunDirection , attenuationCoefficients);
            estimate += throughput * phaseMoon * moonIrradianceSpectrum * ratioTrackingEstimator(rayPosition, moonDirection, attenuationCoefficients);

            scatteringOrder++;
            
        } else {
            break;
        }
    }
    if (isnan(estimate) || isinf(estimate)) return 0.0;

    return estimate;
}

float rayleighCrossSectionApproximation(float lambda) {
    return 1.0 / (pow4(lambda) * 1e-6);
}

float rayleighCrossSection(float lambda, float nAir) {
    float lambdaMeter = lambda * 1e-9; // To meters
    float lambdaMicro = lambda * 1e-3; // To micrometers

    float rayleigh = (8.0 * PI * PI * PI * pow2(nAir * nAir - 1.0)) / (3.0 * pow4(lambdaMeter) * airNumberDensity);

    float waveNumber = 1.0 / lambdaMicro;

    float N2  = 1.034  + 3.17e-12 * waveNumber * waveNumber;
    float CO2 = 1.1364 + 2.53e-11 * waveNumber * waveNumber;

    float kingCorrectionFactor = (N2 + CO2) * 0.5;
    
    return rayleigh * kingCorrectionFactor;
}

float averageMieCrossSection() {
    const float junge = 4.0;

    float mie = 0.0;
    for (int i = 0; i < 441; i++) {
        float wl = (float(i + 1) / 441.0) * 441.0 + 390.0;
        float c  = (0.6544 * turbidity - 0.6510) * 5e-18;
        float K  = (0.773335 - 0.00386891 * wl) / (1.0 - 0.00546759 * wl);
        mie += (0.434 * c * PI * pow(TAU / (wl * 1e-9), junge - 2.0) * K) / 441.0;
    }
    return mie;
}

float mieCrossSection(float lambda) {
    const float junge = 4.0;

    float c = (0.6544 * turbidity - 0.6510) * 5e-18;
    float K = (0.773335 - 0.00386891 * lambda) / (1.0 - 0.00546759 * lambda);
    return (0.434 * c * PI * pow(TAU / (lambda * 1e-9), junge - 2.0) * K + averageMieCrossSection()) * 0.5;
}

uniform sampler2D atmosphericSpectralData;

float solarIrradianceAM0(float wavelength) {
    wavelength -= 390.0;
    return inVisibleRange(wavelength) ? texelFetch(atmosphericSpectralData, ivec2(wavelength, 0), 0).r : 0.0;
}

float aerosolAsymmetry(float wavelength) {
    wavelength -= 390.0;
    return inVisibleRange(wavelength) ? texelFetch(atmosphericSpectralData, ivec2(wavelength, 1), 0).r : 0.0;
}

float ozoneCrossSection(float wavelength) {
    wavelength -= 390.0;
    float crossSection = inVisibleRange(wavelength) ? texelFetch(atmosphericSpectralData, ivec2(wavelength, 2), 0).r : 0.0;
    return crossSection * 1e-21 * ozoneNumberDensity * ozoneUnitConversion; // 1e-21 is the encoding scale inverse factor
}
