#version 300 es
precision highp float;

in vec2 quadVertex;
in vec2 particleID;
in float trailIndex;

uniform sampler2D stateTexture;
uniform float time;
uniform vec2 resolution;

out vec3 vColor;
out float vAlpha;
out vec2 vQuadCoord;

const float PI = 3.14159265359;
const float r_min = 0.002;
const float r_max = 0.008;
const float sigmoid_k = 20.0;
const float sigmoid_s0 = 0.15;

float sigmoid(float x) {
    return 1.0 / (1.0 + exp(-sigmoid_k * (x - sigmoid_s0)));
}

float velocityToRadius(vec2 vel) {
    float speed = length(vel);
    return r_min + (r_max - r_min) * sigmoid(speed);
}

float streamFunction(vec2 q, float t) {
    float x = q.x, y = q.y;
    float phi1 = 0.5*t, phi2 = 0.3*t, phi3 = 0.4*t;
    float r = length(q - vec2(0.5));
    return 0.3*sin(2.0*PI*x + phi1)*sin(2.0*PI*y)
         + 0.25*sin(3.0*PI*x)*sin(3.0*PI*y + phi2)
         + 0.2*sin(4.0*PI*r + phi3)
         + 0.15*(x - 0.5)*(y - 0.5)*sin(0.5*t);
}

float computeVorticity(vec2 q, float t) {
    const float h = 0.001;
    float psi_c = streamFunction(q, t);
    float psi_r = streamFunction(vec2(q.x + h, q.y), t);
    float psi_l = streamFunction(vec2(q.x - h, q.y), t);
    float psi_u = streamFunction(vec2(q.x, q.y + h), t);
    float psi_d = streamFunction(vec2(q.x, q.y - h), t);
    
    float d2psi_dx2 = (psi_r - 2.0*psi_c + psi_l) / (h*h);
    float d2psi_dy2 = (psi_u - 2.0*psi_c + psi_d) / (h*h);
    
    return -(d2psi_dx2 + d2psi_dy2);
}

vec3 hsv2rgb(vec3 c) {
    vec4 K = vec4(1.0, 2.0/3.0, 1.0/3.0, 3.0);
    vec3 p = abs(fract(c.xxx + K.xyz) * 6.0 - K.www);
    return c.z * mix(K.xxx, clamp(p - K.xxx, 0.0, 1.0), c.y);
}

void main() {
    vec2 texCoord = (particleID + 0.5) / 16.0;
    vec4 state = texture(stateTexture, texCoord);
    
    vec2 pos = state.rg;
    vec2 vel = state.ba;
    
    float radius = velocityToRadius(vel);
    
    vec2 aspectRatio = vec2(resolution.y / resolution.x, 1.0);
    vec2 screenPos = (pos * 2.0 - 1.0) * aspectRatio;
    vec2 offset = quadVertex * radius * aspectRatio;
    
    gl_Position = vec4(screenPos + offset, 0.0, 1.0);
    
    float vorticity = computeVorticity(pos, time);
    float hue = mod((vorticity + 3.0) / 6.0, 1.0);
    vColor = hsv2rgb(vec3(hue, 0.8, 0.9));
    
    vAlpha = (1.0 - trailIndex / 20.0) * 0.6;
    vQuadCoord = quadVertex;
}