#version 450

layout (local_size_x = 32) in;

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

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

struct Tetrahedron {
	uint a;
	uint b;
	uint c;
	uint d;
	float restVolume;
	float compliance;
};


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

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

layout (std430, set = 0, binding = 4) buffer TetrahedronBuffer {
	Tetrahedron tetrahedra[];
};

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;
	if (vertices[vID].w == 0){
		return;
	}
	vertices[vID].velocity += dt / k * gravity;
	vertices[vID].position += dt / k * vertices[vID].velocity;

	float dist = vertices[vID].position.y + 5;
	if (dist < 0){
		vec3 normal = vec3(0, 1, 0);
		vertices[vID].position -= dist * normal;

		// contact friction
		{
			float mStatic = 1;
			float mKinetic = 1;

			vec3 a = vertices[vID].position - vertices[vID].prevPosition;
			vec3 D = a - dot(a, normal) / dot(normal, normal) * normal;

			vec3 add = -D;
			if (length(D) >= mStatic * -dist){
				add *= min(mKinetic * -dist / length(D), 1);
			}
			vertices[vID].position += add;
		}
	}
}

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

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

	float alpha = edge.compliance / (dt / k) / (dt / k);

	float s = -(currentLength - edge.restLength) / (v1.w + v2.w + alpha);

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

	vec3 delta1 = s * v1.w * g1;
	vec3 delta2 = s * v2.w * g2;

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

void solveTetrahedron(uint tetID){
	Tetrahedron tetrahedron = tetrahedra[tetID];

	Vertex va = vertices[tetrahedron.a];
	Vertex vb = vertices[tetrahedron.b];
	Vertex vc = vertices[tetrahedron.c];
	Vertex vd = vertices[tetrahedron.d];

	vec3 a = va.position;
	vec3 b = vb.position;
	vec3 c = vc.position;
	vec3 d = vd.position;

	float volumeError = dot(d - a, cross(b - a, c - a)) / 6 - tetrahedron.restVolume;

	vec3 ga = cross(d - b, c - b) / 6;
	vec3 gb = cross(c - a, d - a) / 6;
	vec3 gc = cross(d - a, b - a) / 6;
	vec3 gd = cross(b - a, c - a) / 6;

	float w =
	va.w * dot(ga, ga) +
	vb.w * dot(gb, gb) +
	vc.w * dot(gc, gc) +
	vd.w * dot(gd, gd);

	if (w == 0) return;

	float alpha = tetrahedron.compliance / (dt / k) / (dt / k);

	float s = -volumeError / (w + alpha);

	vec3 addA = s * ga * va.w;
	vec3 addB = s * gb * vb.w;
	vec3 addC = s * gc * vc.w;
	vec3 addD = s * gd * vd.w;

	vertices[tetrahedron.a].position += addA;
	vertices[tetrahedron.b].position += addB;
	vertices[tetrahedron.c].position += addC;
	vertices[tetrahedron.d].position += addD;
}

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

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);
			} else if (id - edgePartition.size < tetrahedronPartition.size){
				solveTetrahedron(id - edgePartition.size + tetrahedronPartition.offset);
			}
			break;
		case 2:
			if (id < vertexCount){
				postSolve(id);
			}
			break;
	}
}