#version 430 core

struct particle_t {
    float pos_x;
    float pos_y;
    float dir_x;
    float dir_y;
};

struct cell_t {
    float pheromone;
    float result;
    int population;
};

uniform uint map_width;
uniform uint map_height;

uniform uint population;

uniform float speed;
uniform float steering;
uniform float sensing_distance;
uniform float sensing_angle;

uniform float dt;

layout(std430, binding = 1) coherent buffer particles_buffer
{
    particle_t particles[];
};

layout(std430, binding = 2) coherent buffer map_buffer
{
    cell_t map[];
};

vec2 rotate(vec2 v, float angle){
    float c = cos(angle);
    float s = sin(angle);
    return vec2(v.x * c - v.y * s, v.x * s + v.y * c);
}

vec2 ahead(vec2 at, vec2 dir, float d){
    return at + d * dir;
}

vec2 wrap_coord(vec2 v){
    return mod(v, vec2(map_width, map_height));
}

uint cell_id(vec2 v){
    ivec2 map_pos = ivec2(floor(v));
    return uint(map_pos.y * map_width + map_pos.x);
}

layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
void main(){
    uint particle_id = uint(gl_GlobalInvocationID.x);
    if(particle_id >= population){
        return;
    }

    // Update particle position
    particle_t p = particles[particle_id];
    vec2 pos = vec2(p.pos_x, p.pos_y);
    vec2 dir = vec2(p.dir_x, p.dir_y);
    pos = wrap_coord(pos + dt * speed * dir);
    particles[particle_id].pos_x = pos.x;
    particles[particle_id].pos_y = pos.y;
    
    // Count particle
    uint p_cell_id = cell_id(pos);
    atomicAdd(map[p_cell_id].population, 1);

    // Sense and turn
    vec2 look_ahead = wrap_coord(ahead(pos, dir, sensing_distance));
    vec2 look_left  = wrap_coord(ahead(pos, rotate(dir,  sensing_angle), sensing_distance));
    vec2 look_right = wrap_coord(ahead(pos, rotate(dir, -sensing_angle), sensing_distance));

    float pheromone_ahead = map[cell_id(look_ahead)].pheromone;
    float pheromone_left  = map[cell_id(look_left )].pheromone;
    float pheromone_right = map[cell_id(look_right)].pheromone;

    if(pheromone_left > pheromone_ahead && pheromone_left > pheromone_right){
        dir = rotate(dir,  steering * dt);
    }else if(pheromone_right > pheromone_ahead && pheromone_right > pheromone_left){
        dir = rotate(dir, -steering * dt);
    }
    particles[particle_id].dir_x = dir.x;
    particles[particle_id].dir_y = dir.y;
}