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

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

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

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

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

@group(0) @binding(5) 
var<uniform> gridMin: vec3<f32>;

@group(0) @binding(6)
var<storage, read> startIndices: array<u32>; // start index of particles in each cell

@group(0) @binding(7)
var<storage, read> endIndices: array<u32>;   // end index (exclusive) of particles in each cell

@group(0) @binding(8)
var<uniform> collisionRadius: f32;

@group(0) @binding(9)
var<uniform> deltaTimeSeconds: f32;

@group(0) @binding(10) var<uniform> gridMax: vec3<f32>;

// particleIndex = hashSortedIndices[ startIndices[ i ] ]

fn getHash(gridCoordinate: vec3<i32>, cellCount: u32) -> u32 {

let maxIndex = i32(cellCount) - 1;

	let x = max(0, min(gridCoordinate.x, maxIndex));
	let y = max(0, min(gridCoordinate.y, maxIndex));
	let z = max(0, min(gridCoordinate.z, maxIndex));

	return u32(x + y * i32(cellCount) + z * i32(cellCount) * i32(cellCount));
}


@compute @workgroup_size(256)
fn computeMain(@builtin(global_invocation_id) globalInvocationId: vec3<u32>) {

	let index 					= globalInvocationId.x;

	let particleIndex 			= hashSortedIndices[ index ];

	if ( particleIndex >= arrayLength(&positions) ) {
		return;
	}

	var currentPosition = positions[ particleIndex ];

	let cellSize = (gridMax - gridMin) / f32(cellCount);  // 2.0 / 16 = 0.125

	let relativePos = currentPosition - gridMin;          // currentPosition + 1

	let gridCoord = vec3<i32>(floor(relativePos / cellSize));  // relativePos divided by cellSize, then floored

	let hash = getHash(gridCoord, cellCount);


	var currentVelocity = velocities[particleIndex];

	var push 			= vec3<f32>(0.0);

	var count 			= 0u;

	let collisionRadiusSquared = collisionRadius * collisionRadius;

	for (var dz = -1; dz <= 1; dz = dz + 1) {
	    for (var dy = -1; dy <= 1; dy = dy + 1) {
	        for (var dx = -1; dx <= 1; dx = dx + 1) {


	            let neighborCell = gridCoord + vec3<i32>(dx, dy, dz);



	            // Compute hash of neighbor cell, with clamping to valid range inside getHash()
	            let neighborHash = getHash(neighborCell, cellCount);

	            let startIndex = startIndices[neighborHash];
	            let endIndex = endIndices[neighborHash];

	            for (var i = startIndex; i < endIndex; i = i + 1u) {
	                let otherIndex = hashSortedIndices[i];
	                if (otherIndex == particleIndex) {
	                    continue;
	                }

	                let otherPosition = positions[otherIndex];
	                let offset = currentPosition - otherPosition;
	                let distSquared = dot(offset, offset);

	                if (distSquared < collisionRadiusSquared && distSquared > 0.00001) {
	                    let distance = sqrt(distSquared);
	                    let direction = offset / distance;
	                    let overlap = collisionRadius - distance;

	                    push += direction * overlap;
	                    count += 1u;
	                }
	            }
	        }
	    }
	}

	if ( count > 0u ) {

		let averagePush = push / f32(count);

		currentPosition += averagePush * .9;

		currentVelocity += averagePush * 3.0;

let pushDir = normalize(averagePush);

// Project current velocity onto push direction
let velAlongPush = dot(currentVelocity, pushDir);

// Damping factor (energy loss on collision)
let dampingFactor = 0.25;

// Reduce velocity along push direction
let velAlongPushDamped = velAlongPush * dampingFactor;

// Velocity perpendicular to push direction remains unchanged
let velPerp = currentVelocity - velAlongPush * pushDir;

// Combine damped velocity components
currentVelocity = velPerp + velAlongPushDamped * pushDir;
	

	}

	let deltaTimeClamped 	= min( deltaTimeSeconds, 0.01 );


	let gridExtent = vec3<f32>(f32(cellCount)) * cellSize;
	//let gridMax = gridMin + gridExtent;

	// Enforce hardcoded bounding box from -1 to +1 on all axes
if (currentPosition.x < gridMin.x) {
	currentPosition.x = gridMin.x;
	currentVelocity.x = abs(currentVelocity.x) * 0.2;
} else if (currentPosition.x > gridMax.x) {
	currentPosition.x = gridMax.x;
	currentVelocity.x = -abs(currentVelocity.x) * 0.2;
}

if (currentPosition.y < gridMin.y) {
	currentPosition.y = gridMin.y;
	currentVelocity.y = abs(currentVelocity.y) * 0.2;
} else if (currentPosition.y > gridMax.y) {
	currentPosition.y = gridMax.y;
	currentVelocity.y = -abs(currentVelocity.y) * 0.2;
}

if (currentPosition.z < gridMin.z) {
	currentPosition.z = gridMin.z;
	currentVelocity.z = abs(currentVelocity.z) * 0.2;
} else if (currentPosition.z > gridMax.z) {
	currentPosition.z = gridMax.z;
	currentVelocity.z = -abs(currentVelocity.z) * 0.2;
}


	if (currentPosition.y < -1.0) {
		currentPosition.y = -1.0;
		currentVelocity.y *= -0.2;
	}


	currentPosition += currentVelocity * deltaTimeClamped;

	positions[ particleIndex ] = currentPosition;

	velocities[ particleIndex ] = currentVelocity;

}