#version 330 core

in vec2 v_uv;
out vec4 frag_color;

uniform mat4 inv_view;
uniform mat4 inv_projection;
uniform vec3 camera_pos;
uniform vec3 light_dir;
uniform vec3 light_color;
uniform float draw_planet;

const float PI = 3.14159265359;
const float PLANET_RADIUS = 6371000.0;
const float ATMOSPHERE_HEIGHT = 100000.0;
const float RAYLEIGH_HEIGHT = (ATMOSPHERE_HEIGHT * 0.08);
const float MIE_HEIGHT = (ATMOSPHERE_HEIGHT * 0.012);
const float EXPOSURE = 20.0;
const vec3  C_RAYLEIGH = vec3(5.802, 13.558, 33.100) * 1e-6;
const vec3  C_MIE      = vec3(3.996,  3.996,  3.996) * 1e-6;
const vec3  C_OZONE    = vec3(0.650,  1.881,  0.085) * 1e-6;

float AtmosphereHeight(vec3 positionWS) {
    return length(positionWS - vec3(0.0, -PLANET_RADIUS, 0.0)) - PLANET_RADIUS;
}
float DensityRayleigh(float h) { return exp(-max(0.0, h / RAYLEIGH_HEIGHT)); }
float DensityMie     (float h) { return exp(-max(0.0, h / MIE_HEIGHT)); }
float DensityOzone   (float h) { return max(0.0, 1.0 - abs(h - 25000.0) / 15000.0); }
vec3  AtmosphereDensity(float h) { return vec3(DensityRayleigh(h), DensityMie(h), DensityOzone(h)); }

vec2 SphereIntersection(vec3 rayStart, vec3 rayDir, vec3 sphereCenter, float sphereRadius) {
    vec3 o = rayStart - sphereCenter;
    float a = dot(rayDir, rayDir);
    float b = 2.0 * dot(o, rayDir);
    float c = dot(o, o) - (sphereRadius * sphereRadius);
    float d = b*b - 4.0*a*c;
    if (d < 0.0) return vec2(-1.0);
    d = sqrt(d);
    return vec2(-b - d, -b + d) / (2.0 * a);
}
vec2 PlanetIntersection(vec3 rayStart, vec3 rayDir) {
    return SphereIntersection(rayStart, rayDir, vec3(0.0, -PLANET_RADIUS, 0.0), PLANET_RADIUS);
}
vec2 AtmosphereIntersection(vec3 rayStart, vec3 rayDir) {
    return SphereIntersection(rayStart, rayDir, vec3(0.0, -PLANET_RADIUS, 0.0), PLANET_RADIUS + ATMOSPHERE_HEIGHT);
}

float PhaseRayleigh(float costh) {
    return 3.0 * (1.0 + costh*costh) / (16.0 * PI);
}
float PhaseMie(float costh, float g) {
    g = min(g, 0.9381);
    float k = 1.55*g - 0.55*g*g*g;
    float kcosth = k*costh;
    return (1.0 - k*k) / ((4.0*PI) * (1.0 - kcosth) * (1.0 - kcosth));
}

vec3 IntegrateOpticalDepth(vec3 rayStart, vec3 rayDir) {
    vec2 intersection = AtmosphereIntersection(rayStart, rayDir);
    float rayLength = intersection.y;
    int sampleCount = 8;
    float stepSize = rayLength / float(sampleCount);
    vec3 opticalDepth = vec3(0.0);
    for (int i = 0; i < sampleCount; ++i) {
        vec3 localPosition = rayStart + rayDir * (float(i) + 0.5) * stepSize;
        float localHeight  = AtmosphereHeight(localPosition);
        vec3 localDensity  = AtmosphereDensity(localHeight);
        opticalDepth += localDensity * stepSize;
    }
    return opticalDepth;
}

vec3 Absorb(vec3 opticalDepth) {
    return exp(-(opticalDepth.x * C_RAYLEIGH + opticalDepth.y * C_MIE * 1.1 + opticalDepth.z * C_OZONE));
}

vec3 IntegrateScattering(vec3 rayStart, vec3 rayDir, float rayLength, vec3 lightDir, vec3 lightColor, out vec3 transmittance) {
    float rayHeight = AtmosphereHeight(rayStart);
    float sampleDistributionExponent = 1.0 + clamp(1.0 - rayHeight / ATMOSPHERE_HEIGHT, 0.0, 1.0) * 8.0;

    vec2 intersection = AtmosphereIntersection(rayStart, rayDir);
    rayLength = min(rayLength, intersection.y);
    if (intersection.x > 0.0) {
        rayStart += rayDir * intersection.x;
        rayLength -= intersection.x;
    }

    float costh = dot(rayDir, lightDir);
    float phaseR = PhaseRayleigh(costh);
    float phaseM = PhaseMie(costh, 0.85);

    int sampleCount = 64;
    vec3 opticalDepth = vec3(0.0);
    vec3 rayleigh = vec3(0.0);
    vec3 mie = vec3(0.0);
    float prevRayTime = 0.0;
    for (int i = 0; i < sampleCount; ++i) {
        float t = pow(float(i) / float(sampleCount), sampleDistributionExponent) * rayLength;
        float stepSize = (t - prevRayTime);
        vec3 localPosition = rayStart + rayDir * t;
        float localHeight = AtmosphereHeight(localPosition);
        vec3 localDensity = AtmosphereDensity(localHeight);
        opticalDepth += localDensity * stepSize;
        vec3 viewTransmittance = Absorb(opticalDepth);
        vec3 opticalDepthLight = IntegrateOpticalDepth(localPosition, lightDir);
        vec3 lightTransmittance = Absorb(opticalDepthLight);
        rayleigh += viewTransmittance * lightTransmittance * phaseR * localDensity.x * stepSize;
        mie      += viewTransmittance * lightTransmittance * phaseM * localDensity.y * stepSize;
        prevRayTime = t;
    }
    transmittance = Absorb(opticalDepth);
    return (rayleigh * C_RAYLEIGH + mie * C_MIE) * lightColor * EXPOSURE;
}

void main() {
    // Reconstruct view ray from NDC
    vec2 ndc = v_uv * 2.0 - 1.0;
    vec4 clip = vec4(ndc, 1.0, 1.0);
    vec4 view = inv_projection * clip;
    view = vec4(view.xy, -1.0, 0.0);
    vec3 worldDir = normalize((inv_view * view).xyz);

    float rayLength = 1.0/0.0; // infinity
    if (draw_planet == 1.0) {
        vec2 isect = PlanetIntersection(camera_pos, worldDir);
        if (isect.x > 0.0) {
            rayLength = min(rayLength, isect.x);
        }
    }

    vec3 transmittance;
    vec3 color = IntegrateScattering(camera_pos, worldDir, rayLength, normalize(light_dir), light_color, transmittance);
    frag_color = vec4(color, 1.0);
}