#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;
};

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


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;
	// vertices[vID].velocity += dt * gravity;
	vertices[vID].position += dt * vertices[vID].velocity;
}

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 = 0.3 / dt / dt;

	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 = 0 / dt / dt;
    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){
	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;
	}
}