@group(0) @binding(0)
var<storage, read_write> compare: array<f32>;

@group(0) @binding(1)
var<storage, read_write> indices: array<u32>;

@group(0) @binding(2)
var<uniform> k: u32;  // current stage size (power of two)

@group(0) @binding(3)
var<uniform> j: u32;  // current subsequence size

@group(0) @binding(4)
var<uniform> totalCount: u32;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
	let idx = global_id.x;


	let ixj = idx ^ j;


	if (idx >= totalCount || ixj >= totalCount) {
		return;
	}

	

	if (ixj > idx && ixj < totalCount) {
		let ascending = (idx & k) == 0u;

		let dist_idx = compare[idx];
		let dist_ixj = compare[ixj];

		var swap = false;

		if (ascending) {
		    if (dist_idx < dist_ixj) {
		        swap = true;
		    }
		} else {
		    if (dist_idx > dist_ixj) {
		        swap = true;
		    }
		}

		if (swap) {

			let tempDist = compare[idx];
			let tempDist2 = compare[ixj];

			let tempIndex = indices[idx];
			let tempIndex2 = indices[ixj];

			compare[idx] = tempDist2;
			compare[ixj] = tempDist;

			indices[idx] = tempIndex2;
			indices[ixj] = tempIndex;
			
		} 
	}
}