#version 300 es
precision highp float;

in vec2 vTexCoord;
out vec4 fragColor;

uniform sampler2D stateTexture;
uniform float time;
uniform float dt;

const float PI = 3.14159265359;
const float alpha = 0.1;
const float beta = 0.05;
const float k = 2.0;

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 dPsi_dx(vec2 q, float t) {
    const float h = 0.001;
    return (streamFunction(vec2(q.x + h, q.y), t) - 
            streamFunction(vec2(q.x - h, q.y), t)) / (2.0 * h);
}

float dPsi_dy(vec2 q, float t) {
    const float h = 0.001;
    return (streamFunction(vec2(q.x, q.y + h), t) - 
            streamFunction(vec2(q.x, q.y - h), t)) / (2.0 * h);
}

vec2 derivative(vec2 q, float t) {
    vec2 vel = vec2(dPsi_dy(q, t), -dPsi_dx(q, t));
    vec2 center_offset = q - vec2(0.5);
    vec2 F_boundary = -4.0 * k * vec2(
        pow(center_offset.x, 3.0),
        pow(center_offset.y, 3.0)
    );
    return (1.0 - alpha) * vel + beta * F_boundary;
}

void main() {
    vec4 state = texture(stateTexture, vTexCoord);
    vec2 pos = state.rg;
    
    // RK4 integration
    vec2 k1 = derivative(pos, time);
    vec2 k2 = derivative(pos + 0.5*dt*k1, time + 0.5*dt);
    vec2 k3 = derivative(pos + 0.5*dt*k2, time + 0.5*dt);
    vec2 k4 = derivative(pos + dt*k3, time + dt);
    
    vec2 newPos = pos + (dt / 6.0) * (k1 + 2.0*k2 + 2.0*k3 + k4);
    vec2 vel = (k1 + 2.0*k2 + 2.0*k3 + k4) / 6.0;
    
    fragColor = vec4(newPos, vel);
}