#version 430

// Adaptive step size control with PI controller
// Implements optimal step size prediction based on error estimates

layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

// Error estimates from previous steps
layout(std430, binding = 0) buffer ErrorBuffer {
    float error[];
};

// Previous step sizes
layout(std430, binding = 1) buffer StepSizeBuffer {
    float step_size[];
};

// Output: new step sizes
layout(std430, binding = 2) buffer NewStepSizeBuffer {
    float new_step_size[];
};

// Step acceptance flags
layout(std430, binding = 3) buffer AcceptanceBuffer {
    int accepted[];
};

uniform float safety_factor;      // Usually 0.9
uniform float min_factor;         // Minimum step size reduction factor
uniform float max_factor;         // Maximum step size increase factor
uniform float error_tolerance;    // Target error tolerance
uniform float pi_alpha;           // PI controller proportional gain
uniform float pi_beta;            // PI controller integral gain
uniform int num_systems;          // Number of ODE systems to process

void main() {
    uint index = gl_GlobalInvocationID.x;
    
    if (index >= num_systems) return;
    
    float current_error = error[index];
    float current_step = step_size[index];
    
    // Check if step should be accepted
    bool accept_step = current_error <= error_tolerance;
    accepted[index] = accept_step ? 1 : 0;
    
    // Compute new step size using PI controller
    float error_ratio = current_error / error_tolerance;
    
    // Basic step size factor
    float factor = safety_factor * pow(1.0 / error_ratio, pi_alpha);
    
    // Add integral term if we have previous error
    if (index > 0 && error[index - 1] > 0.0) {
        float prev_error_ratio = error[index - 1] / error_tolerance;
        float integral_term = pow(error_ratio / prev_error_ratio, pi_beta);
        factor *= integral_term;
    }
    
    // Clamp the factor to reasonable bounds
    factor = clamp(factor, min_factor, max_factor);
    
    // Special handling for rejected steps
    if (!accept_step) {
        // Be more conservative when rejecting
        factor = min(factor, 0.5);
    }
    
    // Compute new step size
    float new_step = current_step * factor;
    
    // Additional safety checks
    // Prevent step size from becoming too small or too large
    float min_step = current_step * 0.01;  // Don't reduce by more than 100x
    float max_step = current_step * 10.0;  // Don't increase by more than 10x
    
    new_step = clamp(new_step, min_step, max_step);
    
    new_step_size[index] = new_step;
}