#version 430 core

#define PI 3.14159265359

layout(local_size_x = 8, local_size_y = 8) in;

layout(rgba32f) uniform image2D img_output;

uniform int width;
uniform int height;

uniform int quantum_n;
uniform int quantum_l;
uniform int quantum_m;

uniform float u_phi = 0;
//vec4 base_color = vec4(1);
vec4 base_color = vec4(209, 227, 255, 255) / 255.0f;

const double a0 = 5.29177210903e-11;
const double h = 6.582119569e-16;
const double mq = 9.10938356e-31;
const double e = 8.8541878128e-12;
const double q = 1.60217662e-19;

int factorial(int n){
    int res = 1;
    for(int i = 2; i <= n; i++){
        res *= i;
    }
    return res;
}
double laguerre(double x, int n, double a){
    if(n == 1)
        return 1 + a - x;
    else 
        return 1;
} 
double harmonics(double theta, double phi){
    double res = 1;
    
    if(quantum_m == 0 && quantum_l == 2){
        res *= 0.25;
        res *= sqrt(5.0/PI);
        res *= (3 * pow(cos(float(theta)), 2) - 1);
    }else if(quantum_m == 1 && quantum_l == 2){
        res *= 0.5;
        res *= sqrt(15.0/PI);
        res *= sin(float(theta)) * cos(float(phi)) * cos(float(theta));
    }
    
    return res;
}
double wave_function(double r, double theta, double phi){
    double p1 = pow(float(2/(quantum_n*a0)),3);
    double p21 = factorial(quantum_n-quantum_l-1);
    double p22 = (2 * quantum_n * factorial(quantum_n+quantum_l));
    double root = sqrt(p1 * (p21/p22));

    double exponent = exp(float(-r/(quantum_n*a0)));

    double p3 = pow(float((2 * r)/(quantum_n * a0)), quantum_l);

    double lag = laguerre((2 * r)/(quantum_n * a0), quantum_n - quantum_l - 1, 2 * quantum_l + 1);

    double Y = harmonics(theta, phi);

    double wave = exponent * root * p3 * lag * Y;
    //return root / 2e13 + (quantum_n+quantum_l+quantum_m+ r) * 0.000001f;
    //return wave / 1e14;
    //return (wave * wave) / 2.28e27;
    return (wave * wave) / 1e27;
    //return (sin(float(r/sqrt(2)) * 10) + 1.0) * 0.5 + (quantum_n+quantum_l+quantum_m) * 0.000001f;
}

vec4 heatmap(float x){
    vec4 colors[6];

    colors[0] = vec4(0);
    colors[1] = vec4(58, 20, 78, 255) / 255.0f;
    colors[2] = vec4(148, 50, 99, 255) / 255.0f;
    colors[3] = vec4(249, 114, 13, 255) / 255.0f;
    colors[4] = vec4(255, 198, 76, 255) / 255.0f;
    colors[5] = vec4(1);

    colors[0].w = 0.0f;
    colors[1].w = 0.1f;
    colors[2].w = 0.25f;
    colors[3].w = 0.5f;
    colors[4].w = 0.75f;
    colors[5].w = 1.0f;

    int index = -1;
    for(int i = 1; i < 6; i++){
        if(x <= colors[i].w){
            index = i;
            break;
        }
    }

    vec4 res;
    if(index != -1){
        float range = colors[index].w - colors[index - 1].w;
        res = mix(colors[index - 1], colors[index], (x - colors[index - 1].w)/range);
    }else
        res = colors[5]; 
    res.w = 1;
    return res;
}

void main(){
    ivec2 id = ivec2(gl_GlobalInvocationID.xy);

    vec2 offset = vec2(0.0f, 0.0f);

    float x;
    float y; 

    // map x,y to [-1,1]
    x = ((((float(id.x) - 0.5f + offset.x)/float(width))) - 0.5f) * 2.0f;
    y = ((((float(id.y) - 0.5f + offset.y)/float(height))) - 0.5f) * 2.0f;
        
    double raw_r = sqrt(pow(x, 2) + pow(y, 2));
    double r = raw_r * 35 * a0;
    double theta = acos(float(y/raw_r));
    float phi = u_phi;

    vec4 base_color = vec4(0);

    vec4 max_color = vec4(1);

    // double n_ang ;
    
    // if(x > 0){
    //     n_ang = mod(theta + phi, PI);
    // }else
    //     n_ang = mod(theta - phi, PI);

    float val = 4.0f * float(wave_function(r, theta, phi));
    //val = 1 - exp(-1 * val * 5);

    //vec4 pixel = mix(base_color, max_color, val);
    vec4 pixel = heatmap(val);
    
    imageStore(img_output, id, pixel);
    //imageStore(img_output, id, heatmap(float(raw_r)));
}