#import bevy_ui::ui_vertex_output::UiVertexOutput;

@group(1) @binding(0) var<uniform> aspect: f32;
@group(1) @binding(1) var<uniform> dt: f32;
@group(1) @binding(2) var<storage, read_write> nodes: array<Node>;

struct Node {
    pos: vec2<f32>,
    velo: vec2<f32>,
    size: f32,
    rgb: vec3<f32>,
}

@compute @workgroup_size(16)
fn update(@builtin(global_invocation_id) id: vec3<u32>) {
    let i = id.x;
    if (i < arrayLength(&nodes)) {
        let pos = nodes[i].pos + (nodes[i].velo * dt);
        nodes[i].velo.x *= 1 - 2 * (floor(pos.x) % 2);
        nodes[i].velo.y *= 1 - 2 * (floor(pos.y) % 2);
        nodes[i].pos.x = abs((1 - pos.x) % 2 - 1);
        nodes[i].pos.y = abs((1 - pos.y) % 2 - 1);
    }
}

@fragment
fn fragment(in: UiVertexOutput) -> @location(0) vec4<f32> {
    var depth = 0.0;
    var rgb = vec3(0.0);
    for (var i = 0u; i < arrayLength(&nodes); i++) {
        let node = nodes[i];
        let x = (node.pos.x - in.uv.x) * aspect;
        let y = node.pos.y - in.uv.y;
        let d = node.size / ((x * x) + (y * y));
        depth += d;
        rgb += node.rgb / d;
    }
    rgb *= depth;
    return vec4(rgb, depth);
}