#version 450

layout (local_size_x = 32) in;

struct Vertex {
	vec3 position;
	vec3 color;
	vec3 normal;
	vec3 velocity;
	vec3 prevPosition;
	float w;
};

struct Edge {
	uint a;
	uint b;
	float restLength;
};




layout (std430, set = 0, binding = 0) buffer VertexBuffer {
	Vertex vertices[];
};

layout (std430, set = 0, binding = 2) buffer EdgeBuffer {
	Edge edges[];
};

layout (std140, set = 0, binding = 5) uniform Sizes {
	uint vertexCount;
	uint faceCount;
};

layout (std140, set = 1, binding = 0) uniform Properties {
	vec3 gravity;
	float dt;
	uint k;
};

struct Partition {
	uint offset;
	uint size;
};

layout (push_constant, std430) uniform PushConstants {
	uint state;
	Partition edgePartition;
	Partition tetrahedronPartition;
};

void preSolve(uint vID){
	vertices[vID].prevPosition = vertices[vID].position;
	// vertices[vID].velocity += dt * gravity;
	vertices[vID].position += dt * vertices[vID].velocity;
}

void solveEdge(uint eID){
	Edge e = edges[eID];
	
	Vertex v1 = vertices[e.a];
	Vertex v2 = vertices[e.b];

	vec3 diff = v1.position - v2.position;
	float currentLength = length(diff);

	float alpha = 0.0001 / (dt * dt);

	float s = (currentLength - e.restLength) / (1 + 1 + alpha);

	vec3 g1 = normalize(diff);
	vec3 g2 = -g1;

	vec3 delta1 = -s * 1 * g1;
	vec3 delta2 = -s * 1 * g2;

	vertices[e.a].position += delta1;
	vertices[e.b].position += delta2;
}

void solveTetrahedron(uint tetID){

}

void postSolve(uint vID){
	vertices[vID].velocity = (vertices[vID].position - vertices[vID].prevPosition) / dt;
}

void main() {
	uint id = gl_GlobalInvocationID.x;
	switch (state){
		case 0:
			if (id < vertexCount){
				preSolve(id);
			}
			break;
		case 1:
			if (id < edgePartition.size){
				solveEdge(id + edgePartition.offset);
			}
			if (id < tetrahedronPartition.size){
				solveTetrahedron(id + tetrahedronPartition.offset);
			}
			break;
		case 2:
			if (id < vertexCount){
				postSolve(id);
			}
			break;
	}
}