// https://zhuanlan.zhihu.com/p/595576594

// https://zhuanlan.zhihu.com/p/419165090

// https://zhuanlan.zhihu.com/p/237502022


float getHeigth(vec3 p){
    return length(p) - earth_r;
}

float getRho(float h, float H){
    return min(fastExp(-(h / H)), 1.0);
}

vec3 RayleighCoefficient(float h){
    return RayleighSigma * getRho(h, H_R);
}

vec3 MieCoefficient(float h){
    return MieSigma * getRho(h, H_M);
}

vec3 MieAbsorption(float h){
    return MieAbsorptionSigma * getRho(h, H_M);
}

vec3 OzoneAbsorption(float h){
    float rho = max(0, 1.0 - abs(h - ozoneCenter) / ozoneWidth);
    return OzoneAbsorptionSigma * rho;
}

float RayleiPhase(float cos_theta){
    return (3.0 / (16.0 * PI)) * (1.0 + cos_theta * cos_theta);
}

float MiePhase(float cos_theta){
    float g = MIE_G;
    float g2 = g * g;

    // float a = 3.0 / (8.0 * PI);
    // float b = (1.0 - g2) / (2.0 + g2);
    // float c = 1.0 + cos_theta*cos_theta;
    // float d = pow(1.0 + g2 - 2*g*cos_theta, 1.5);
    // return a * b * (c / d);

    return (1 - g2) / (4.0 * PI * pow((1 + g2 - 2 * g * cos_theta), 1.5));
}

vec3 Scattering(float h, vec3 lightDir, vec3 worldDir){
    float cos_theta = dot(lightDir, worldDir);

    vec3 rayleigh = RayleighCoefficient(h) * RayleiPhase(cos_theta);
    vec3 mie = MieCoefficient(h) * MiePhase(cos_theta);

    return rayleigh + mie;
}

vec3 Transmittance(vec3 p1, vec3 p2, int N_SAMPLE){
    vec3 dir = normalize(p2 - p1);
    float distance = length(p2 - p1);
    float ds = distance / float(N_SAMPLE);
    vec3 sum = vec3(0.0);
    vec3 p = p1 + (dir * ds) * 0.5;

    for(int i=0; i<N_SAMPLE; i++){
        float h = getHeigth(p);

        vec3 scattering = RayleighCoefficient(h) + MieCoefficient(h);
        vec3 absorption = OzoneAbsorption(h) + MieAbsorption(h);
        vec3 extinction = scattering + absorption;

        sum += extinction * ds;
        p += dir * ds;
    }

    return fastExp(-sum);
}

void UvToTransmittanceLutParams(float bottomRadius, float topRadius, vec2 uv, out float mu, out float r){
    float x_mu = uv.x;
    float x_r = uv.y;

    float H = sqrt(max(0.0f, topRadius * topRadius - bottomRadius * bottomRadius));
    float rho = H * x_r;
    r = sqrt(max(0.0f, rho * rho + bottomRadius * bottomRadius));

    float d_min = topRadius - r;
    float d_max = rho + H;
    float d = d_min + x_mu * (d_max - d_min);
    mu = d == 0.0f ? 1.0f : (H * H - rho * rho - d * d) / (2.0f * r * d);

    mu = clamp(mu, -1.0f, 1.0f);
}

#define T1_T 0.0
vec3 DrawTransmittanceLut(){
    float bottomRadius = earth_r - T1_T;
    float topRadius = earth_r + atmosphere_h;

    float cos_theta = 0.0;
    float r = 0.0;
    UvToTransmittanceLutParams(bottomRadius, topRadius, texcoord, cos_theta, r);

    float sin_theta = sqrt(1.0 - cos_theta * cos_theta);
    vec3 worldDir = vec3(sin_theta, cos_theta, 0);
    vec3 eyePos = vec3(0, r, 0);

    float dis = RaySphereIntersection(eyePos, worldDir, vec3(0.0), topRadius).y;
    vec3 hitPoint = eyePos + worldDir * dis;

    return Transmittance(eyePos, hitPoint, 100);
}

vec2 GetTransmittanceLutUv(float bottomRadius, float topRadius, float mu, float r){
    float H = sqrt(max(0.0f, topRadius * topRadius - bottomRadius * bottomRadius));
    float rho = sqrt(max(0.0f, r * r - bottomRadius * bottomRadius));

    float discriminant = r * r * (mu * mu - 1.0f) + topRadius * topRadius;
	float d = max(0.0f, (-r * mu + sqrt(discriminant)));

    float d_min = topRadius - r;
    float d_max = rho + H;

    if(d > d_max) return vec2(-1.0, -1.0);

    float x_mu = (d - d_min) / (d_max - d_min);
    float x_r = rho / H;

    return vec2(x_mu, x_r);
}

vec3 TransmittanceToAtmosphere(vec3 p, vec3 dir){
    float bottomRadius = earth_r - T1_T;
    float topRadius = earth_r + atmosphere_h;

    vec3 upVector = normalize(p);
    float cos_theta = dot(upVector, dir);
    float r = length(p);

    vec2 uv = GetTransmittanceLutUv(bottomRadius, topRadius, cos_theta, r);
    if(outScreen(uv)) return BLACK;

    uv = vec2(remap(uv.x, 0.0, 1.0, T1_I.x - 1, T1_I.z + 1),
              remap(uv.y, 0.0, 1.0, T1_I.y - 1, T1_I.w + 1));
    uv /= 512.0;

    return texture(colortex7, uv).rgb;
}

vec3 IntegralMultiScattering(vec3 samplePoint, vec3 lightDir){
    vec3 RandomSphereSamples[64] = {
        vec3(0.99358869, 0.11260822, 0.00000000),
        vec3(0.91169626, 0.09264531, -0.40001094),
        vec3(0.67365669, 0.07622308, -0.73516846),
        vec3(0.32027131, 0.05980085, -0.94550312),
        vec3(-0.08461305, 0.04337863, -0.99548599),
        vec3(-0.47211121, 0.02695640, -0.88121127),
        vec3(-0.77287477, 0.01053417, -0.63441844),
        vec3(-0.93893703, -0.00588805, -0.34400967),
        vec3(0.95121443, 0.30510069, 0.04737789),
        vec3(0.87262333, 0.28513778, -0.39735537),
        vec3(0.64508557, 0.26517487, -0.71696598),
        vec3(0.30667517, 0.24521196, -0.91965048),
        vec3(-0.08097941, 0.22524904, -0.97095348),
        vec3(-0.45195152, 0.20528613, -0.86835698),
        vec3(-0.73988292, 0.18532322, -0.64610069),
        vec3(-0.89850573, 0.16536031, -0.40661972),
        vec3(0.81491274, 0.48256240, 0.32112631),
        vec3(0.74692797, 0.46259949, -0.47737804),
        vec3(0.55268120, 0.44263658, -0.70779533),
        vec3(0.26300195, 0.42267367, -0.86739253),
        vec3(-0.06941613, 0.40271076, -0.91272002),
        vec3(-0.38734470, 0.38274785, -0.83677632),
        vec3(-0.63384295, 0.36278493, -0.68292368),
        vec3(-0.77011555, 0.34282202, -0.53851903),
        vec3(0.64279879, 0.64279879, 0.41727299),
        vec3(0.58964950, 0.62283588, -0.51502062),
        vec3(0.43651702, 0.60287297, -0.66832678),
        vec3(0.20790049, 0.58291006, -0.78604042),
        vec3(-0.05489454, 0.56294714, -0.82452911),
        vec3(-0.30611887, 0.54298423, -0.78172057),
        vec3(-0.50124353, 0.52302132, -0.68862370),
        vec3(-0.60813201, 0.50305841, -0.61378117),
        vec3(0.48256240, 0.81491274, 0.32112631),
        vec3(0.44263658, 0.79494983, -0.41521868),
        vec3(0.32760242, 0.77498692, -0.54062333),
        vec3(0.15604192, 0.75502401, -0.63679218),
        vec3(-0.04121196, 0.73506110, -0.67671007),
        vec3(-0.22978144, 0.71509818, -0.65979621),
        vec3(-0.37596411, 0.69513527, -0.61310626),
        vec3(-0.45634954, 0.67517236, -0.57935863),
        vec3(0.30510069, 0.95121443, 0.04737789),
        vec3(0.28513778, 0.93125152, -0.22846527),
        vec3(0.21090146, 0.91128861, -0.35043767),
        vec3(0.10042936, 0.89132570, -0.44192607),
        vec3(-0.02645784, 0.87136279, -0.49013114),
        vec3(-0.14758408, 0.85139987, -0.50327897),
        vec3(-0.24189736, 0.83143696, -0.49920425),
        vec3(-0.29394532, 0.81147405, -0.50381371),
        vec3(0.11260822, 0.99358869, 0.00000000),
        vec3(0.09264531, 0.91169626, -0.40001094),
        vec3(0.07622308, 0.67365669, -0.73516846),
        vec3(0.05980085, 0.32027131, -0.94550312),
        vec3(0.04337863, -0.08461305, -0.99548599),
        vec3(0.02695640, -0.47211121, -0.88121127),
        vec3(0.01053417, -0.77287477, -0.63441844),
        vec3(-0.00588805, -0.93893703, -0.34400967),
        vec3(-0.11260822, 0.99358869, 0.00000000),
        vec3(-0.09264531, 0.91169626, -0.40001094),
        vec3(-0.07622308, 0.67365669, -0.73516846),
        vec3(-0.05980085, 0.32027131, -0.94550312),
        vec3(-0.04337863, -0.08461305, -0.99548599),
        vec3(-0.02695640, -0.47211121, -0.88121127),
        vec3(-0.01053417, -0.77287477, -0.63441844),
        vec3(0.00588805, -0.93893703, -0.34400967)
    };
    const int N_DIRECTION = 64;
    const int N_SAMPLE = 32;
    const float uniform_phase = 1.0 / (4.0 * PI);
    const float sphereSolidAngle = 4.0 * PI / float(N_DIRECTION);
    
    vec3 G_2 = vec3(0, 0, 0);
    vec3 f_ms = vec3(0, 0, 0);

    for(int i=0; i<N_DIRECTION; i++){
        vec3 worldDir = RandomSphereSamples[i];
        float dis = RaySphereIntersection(samplePoint, worldDir, vec3(0.0), earth_r + atmosphere_h).y;
        float d = RaySphereIntersection(samplePoint, worldDir, vec3(0.0), earth_r).x;
        if(d > 0) dis = min(dis, d);
        float ds = dis / float(N_SAMPLE);

        vec3 p = samplePoint + (worldDir * ds) * 0.5;
        vec3 opticalDepth = vec3(0, 0, 0);

        for(int j=0; j<N_SAMPLE; j++)
        {
            float h = getHeigth(p);
            vec3 sigma_s = RayleighCoefficient(h) + MieCoefficient(h);
            vec3 sigma_a = OzoneAbsorption(h) + MieAbsorption(h);
            vec3 sigma_t = sigma_s + sigma_a;
            opticalDepth += sigma_t * ds;

            vec3 t1 = TransmittanceToAtmosphere(p, lightDir);
            vec3 s  = Scattering(h, lightDir, worldDir);
            vec3 t2 = fastExp(-opticalDepth);
            
            G_2  += t1 * s * t2 * uniform_phase * ds * 1.0;  
            f_ms += t2 * sigma_s * uniform_phase * ds;

            p += worldDir * ds;
        }
    }

    G_2 *= sphereSolidAngle;
    f_ms *= sphereSolidAngle;
    return 50.0 * G_2 * (1.0 / (1.0 - f_ms));
}

vec3 DrawMultiScatteringLut(){
    float mu_s = texcoord.x * 2.0 - 1.0;
    float r = texcoord.y * atmosphere_h + earth_r;

    float cos_theta = mu_s;
    float sin_theta = sqrt(1.0 - cos_theta * cos_theta);
    vec3 lightDir = vec3(sin_theta, cos_theta, 0);
    vec3 p = vec3(0, r, 0);

    return IntegralMultiScattering(p, lightDir);
}

vec3 GetMultiScattering(float h, vec3 p, vec3 lightDir){
    float cosSunZenithAngle = dot(normalize(p), lightDir);
    vec2 uv = vec2(cosSunZenithAngle * 0.5 + 0.5, mix(1.0, h / atmosphere_h, 1.0));
    if(outScreen(uv)) return vec3(0.0);

    uv = vec2(remap(uv.x, 0.0, 1.0, MS_I.x - 1, MS_I.z + 1),
              remap(uv.y, 0.0, 1.0, MS_I.y - 1, MS_I.w + 1));
    uv /= 512.0;

    return texture(colortex7, uv).rgb * 0.02;
}

void getDensity(float h, out float dRayleigh_c, out float dMie_c){
    dRayleigh_c = getRho(h, H_R);
    dMie_c = getRho(h, H_M);
}

vec3 ExtinctionT2(float h, float dRayleigh_c, float dMie_c,
                inout float dRayleigh, inout float dMie, inout float dOzone){
    dRayleigh += dRayleigh_c;
    dMie += dMie_c;
    dOzone += max(0, 1.0 - (abs(h - ozoneCenter) / ozoneWidth));
    return dRayleigh * RayleighSigma
        + dMie * (MieSigma + MieAbsorptionSigma)
        + dOzone * OzoneAbsorptionSigma;
}

mat2x3 AtmosphericScattering(vec3 worldPos, vec3 worldDirO, vec3 lightDir, vec3 I, float mieAmount, const int N_SAMPLES){
    float dis = length(worldPos);
    vec3 worldDir = normalize(worldPos);
    float ds = dis / float(N_SAMPLES);

    float dRayleigh = 0.0;
    float dMie = 0.0;
    float dRayleigh_c = 0.0;
    float dMie_c = 0.0;
    float dOzone = 0.0;

    vec3 p = earthPos;
    float h = getHeigth(p);

    vec3 t1 = TransmittanceToAtmosphere(p, lightDir);
    vec3 G_ALL = GetMultiScattering(h, p, lightDir);

    getDensity(h, dRayleigh_c, dMie_c);
    vec3 f_in_R_prev    = t1 * dRayleigh_c;
    vec3 f_in_M_prev    = t1 * dMie_c;
    vec3 f_multi_R_prev = G_ALL * dRayleigh_c;
    vec3 f_multi_M_prev = G_ALL * dMie_c;

    vec3 inScatteringR = vec3(0.0);
    vec3 inScatteringM = vec3(0.0);
    vec3 multiScatteringR = vec3(0.0);
    vec3 multiScatteringM = vec3(0.0);

    for (int i = 0; i < N_SAMPLES; i++) {
        p += worldDir * ds;
        h = getHeigth(p);

        t1 = TransmittanceToAtmosphere(p, lightDir);
        G_ALL = GetMultiScattering(h, p, lightDir);

        getDensity(h, dRayleigh_c, dMie_c);
        vec3 opticalDepthT2 = ExtinctionT2(h, dRayleigh_c, dMie_c, dRayleigh, dMie, dOzone) * ds;
        vec3 t2 = fastExp(-opticalDepthT2);

        vec3 f_in_R_new    = (t1 * t2) * dRayleigh_c;
        vec3 f_in_M_new    = (t1 * t2) * dMie_c;
        vec3 f_multi_R_new = (G_ALL * t2) * dRayleigh_c;
        vec3 f_multi_M_new = (G_ALL * t2) * dMie_c;

        inScatteringR    += 0.5 * (f_in_R_prev    + f_in_R_new);
        inScatteringM    += 0.5 * (f_in_M_prev    + f_in_M_new);
        multiScatteringR += 0.5 * (f_multi_R_prev + f_multi_R_new);
        multiScatteringM += 0.5 * (f_multi_M_prev + f_multi_M_new);

        f_in_R_prev    = f_in_R_new;
        f_in_M_prev    = f_in_M_new;
        f_multi_R_prev = f_multi_R_new;
        f_multi_M_prev = f_multi_M_new;
    }

    float cos_theta = dot(worldDirO, lightDir);
    vec3 inScattering    = inScatteringR * RayleighSigma * RayleiPhase(cos_theta)
                        + inScatteringM * MieSigma      * MiePhase(cos_theta);
    vec3 multiScattering = multiScatteringR * RayleighSigma
                        + multiScatteringM * MieSigma * mieAmount * MIE_STRENGTHNESS;

    inScattering = mix(inScattering, vec3(getLuminance(inScattering)), rainStrength * 0.33);

    return mat2x3(inScattering * ds * I, multiScattering * ds * I);
}



void Extinction(float h,
                inout float dRayleigh, inout float dMie, inout float dOzone){
    dRayleigh += getRho(h, H_R);
    dMie += getRho(h, H_M);
    dOzone += max(0, 1.0 - abs(h - ozoneCenter) / ozoneWidth);
}

vec3 Transmittance1(vec3 startPos, vec3 endPos, float stepCount){
    vec3 direction = normalize(endPos - startPos);
    float dist = length(endPos - startPos);
    float stepSize = dist / stepCount;
    vec3 stepVec = direction * stepSize;

    vec3 p = startPos + 0.5 * stepVec;

    float dRayleigh = 0.0;
    float dMie = 0.0;
    float dOzone = 0.0;

    for(int i = 0; i < stepCount; i++){
        float h = getHeigth(p);
        Extinction(h, dRayleigh, dMie, dOzone);
        p += stepVec;
    }
    vec3 opticalDepth = dRayleigh * RayleighSigma
                + dMie * (MieSigma + MieAbsorptionSigma)
                + dOzone * OzoneAbsorptionSigma;
    vec3 t = exp(-opticalDepth * stepSize);

    return t;
}

vec3 getSunColor(){
    vec3 sunColor = texelFetch(colortex7, sunColorUV, 0).rgb;
    return sunColor;
}

vec3 getSkyColor(){
    vec3 skyColor = texelFetch(colortex7, skyColorUV, 0).rgb;
    return skyColor;
}

/*
vec3 AtmosphericScattering(vec3 worldDir, vec3 lightDir){
    const int N_SAMPLE = 64;

    float d_p2a = RaySphereIntersection(earthPos, worldDir, vec3(0.0), earth_r + atmosphere_h).y;
    // float d_p2e = RaySphereIntersection(earthPos, worldDir, vec3(0.0), earth_r).x;
    float dis = d_p2a;
    // if(d_p2a < 0.0) return BLACK; 
    // if(d_p2e > 0.0) dis = min(dis, d_p2e);

    float ds = dis / float(N_SAMPLE);
    vec3 p = earthPos + worldDir * ds * 0.0;

    vec3 color = BLACK;
    vec3 opticalDepth = BLACK;

    for(int i = 0; i < N_SAMPLE; i++){
        float h = getHeigth(p);

        vec3 t1 = TransmittanceToAtmosphere(p, lightDir);
        
        vec3 s = Scattering(h, worldDir, lightDir);
        
        vec3 extinction = RayleighCoefficient(h) + MieCoefficient(h) +
                        OzoneAbsorption(h) + MieAbsorption(h);
        opticalDepth += extinction * ds;
        vec3 t2 = fastExp(-opticalDepth);

        vec3 inScattering = t1 * s * t2 * ds;
        color += inScattering;

        vec3 multiScattering = GetMultiScattering(h, p, lightDir) 
                                * (RayleighCoefficient(h) + MieCoefficient(h));
        color += multiScattering * t2 * ds;

        p += worldDir * ds;
    }

	return color * vec3(1.0) * 8.0;
}
*/


const vec3 UP_DIRECTION = vec3(0.0, 1.0, 0.0);
const float PLANET_RADIUS = 6371.0;
const float ATMOSPHERE_HEIGHT = 80.0;
const float SUN_INTENSITY = 22.0;

const vec3 BETA_RAYLEIGH = vec3(5.8e-3, 13.5e-3, 33.1e-3);
const float BETA_MIE = 4.0e-3;
// const float MIE_G = 0.80;

const float SUN_DISK_SIZE = 0.999;
const float SUN_DISK_INTENSITY = 5.0;
const float EXPOSURE1 = 1.5;
const float SUNSET_FACTOR = 2.5;
const float HORIZON_EDGE_FACTOR = 0.1;

float miePhaseFunction(float cosTheta, float g) {
    float g2 = g * g;
    float numerator = (1.0 - g2);
    float denominator = 1.0 + g2 - 2.0 * g * cosTheta;
    denominator = max(denominator, 1e-4);
    denominator = pow(denominator, 1.5);
    return (1.0 / (4.0 * 3.14159265)) * (numerator / denominator);
}

float rayleighPhaseFunction(float cosTheta) {
    float factor = (3.0 / (16.0 * 3.14159265));
    return factor * (1.0 + cosTheta * cosTheta);
}

float getOpticalDepth(vec3 worldPos, vec3 worldDir, float radius, float atmHeight, float scaleHeight) {
    float h = 0.0; // max(0.0, length(worldPos) - radius);
    float cosTheta = dot(worldDir, normalize(worldPos));

    float zenithAngle = acos(max(0.001, dot(worldDir, UP_DIRECTION)));
    return scaleHeight * exp(-(h / scaleHeight)) * (1.0 / max(0.001, cos(zenithAngle)));
    // return scaleHeight / max(0.001, dot(worldDir, UP_DIRECTION));
}
vec3 getASkyColor(in vec3 worldDir, in vec3 sunDir) {
    const float rayleighScatterCoeff = 0.0025;
    const float mieScatterCoeff = 0.0010;
    const float sunIntensity = 20.0;
    const vec3 rayleighColor = vec3(0.18867780436772762, 0.4978442963618773, 1.0);
    const vec3 mieColor = vec3(1.0);
    const float rayleighHeight = 8000.0;
    const float mieHeight = 1200.0;
    const float mieG = 0.8;
    const float earthRadius = 6371000.0;
    const float atmosphereHeight = 100000.0;
    
    float sunCosTheta = dot(worldDir, sunDir);
    
    float zenithCosTheta = worldDir.y;
    float zenithAngle = acos(zenithCosTheta);
    
    float rayleighOpticalDepth = rayleighScatterCoeff * exp(-zenithAngle / rayleighHeight);
    float mieOpticalDepth = mieScatterCoeff * exp(-zenithAngle / mieHeight);
    
    float sunHeight = sunDir.y;
    
    vec3 sunsetColor = vec3(1.0, 0.3, 0.1);
    float sunsetFactor = 1.0 - smoothstep(-0.15, 0.15, sunHeight);
    
    float g2 = mieG * mieG;
    float miePhase = 1.5 * (1.0 - g2) / (2.0 + g2) * 
                    (1.0 + sunCosTheta * sunCosTheta) / 
                    pow(1.0 + g2 - 2.0 * mieG * sunCosTheta, 1.5);
    
    float rayleighPhase = 0.75 * (1.0 + sunCosTheta * sunCosTheta);
    
    vec3 rayleighScatter = rayleighColor * rayleighPhase * rayleighOpticalDepth;
    vec3 mieScatter = mieColor * miePhase * mieOpticalDepth;
    
    float sunSpot = smoothstep(0.9998, 0.9999, sunCosTheta);
    vec3 sunColor = vec3(1.0, 1.0, 0.95) * sunSpot * (1.0 - sunsetFactor * 0.98);
    
    vec3 horizonColor = mix(rayleighColor, sunsetColor, sunsetFactor);
    float horizonInfluence = smoothstep(0.0, 0.4, 1.0 - zenithCosTheta);
    
    vec3 scatterColor = horizonColor * rayleighScatter + mieColor * mieScatter;
    
    float sunHeightFactor = smoothstep(-0.1, 0.3, sunHeight);
    scatterColor = mix(scatterColor * (0.3 + 0.7 * sunHeightFactor), scatterColor, sunHeightFactor);
    
    scatterColor += sunColor * sunIntensity;
    
    scatterColor = mix(scatterColor, horizonColor, horizonInfluence * sunsetFactor);
    
    vec3 sunriseColor = vec3(0.7, 0.3, 0.9);
    float sunriseFactorZenith = smoothstep(-0.2, 0.0, sunHeight);
    float sunriseFactorHorizon = (1.0 - sunriseFactorZenith) * horizonInfluence;
    scatterColor = mix(scatterColor, sunriseColor, sunriseFactorHorizon * 0.3);
    
    scatterColor = 1.0 - exp(-scatterColor * 1.5);
    
    return scatterColor;
}