struct Settings {
  agentCount: u32,
  padding0: u32,
  padding1: u32,
  padding2: u32,
};

struct Counters {
  aliveAgentCount: atomic<u32>,
};

@group(1) @binding(0) var<uniform> settings: Settings;
@group(1) @binding(2) var<storage, read_write> counters: Counters;
@group(1) @binding(3) var<storage, read_write> compactedAgents: array<Agent>;

var<workgroup> workgroupAliveCount: atomic<u32>;
var<workgroup> workgroupCompactedOffset: u32;
var<workgroup> workgroupCopyCount: u32;

@compute @workgroup_size(64)
fn main(
  @builtin(global_invocation_id) global_id: vec3<u32>,
  @builtin(local_invocation_id) local_id: vec3<u32>
) {
  let id = get_id(global_id);

  if local_id.x == 0u {
    atomicStore(&workgroupAliveCount, 0u);
  }

  workgroupBarrier();

  var localCompactedIndex = 0u;
  if id < settings.agentCount {
    let agent = agents[id];
    if agent.colorIndex >= 0.0 {
      localCompactedIndex = atomicAdd(&workgroupAliveCount, 1u);
    }
  }

  workgroupBarrier();

  if local_id.x == 0u {
    let groupAliveCount = atomicLoad(&workgroupAliveCount);
    if groupAliveCount > 0u {
      workgroupCompactedOffset = atomicAdd(&counters.aliveAgentCount, groupAliveCount);
    } else {
      workgroupCompactedOffset = 0u;
    }
  }

  workgroupBarrier();

  if id < settings.agentCount {
    let agent = agents[id];
    if agent.colorIndex >= 0.0 {
      compactedAgents[workgroupCompactedOffset + localCompactedIndex] = agent;
    }
  }
}

@compute @workgroup_size(64)
fn copyCompactedAgents(
  @builtin(global_invocation_id) global_id: vec3<u32>,
  @builtin(local_invocation_id) local_id: vec3<u32>
) {
  let id = get_id(global_id);

  if local_id.x == 0u {
    workgroupCopyCount = atomicLoad(&counters.aliveAgentCount);
  }

  workgroupBarrier();

  if id >= workgroupCopyCount {
    return;
  }

  agents[id] = compactedAgents[id];
}