#version 450

layout (local_size_x = 32) in;

#include "structs.comp"

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

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


layout (push_constant) uniform PushConstants {
	uint state;
};

void atomicAddVec3(uint vID, vec3 add){
	for (int i = 0; i < 3; i++) {
		uint expected_memory = vertices[vID].normal[i];
		float floatInput = uintBitsToFloat(vertices[vID].normal[i]) + add[i];
		uint actual_content = atomicCompSwap(vertices[vID].normal[i], expected_memory, floatBitsToUint(floatInput));
		while (actual_content != expected_memory) {
			expected_memory = actual_content;
			floatInput = uintBitsToFloat(expected_memory) + add[i];
			actual_content = atomicCompSwap(vertices[vID].normal[i], expected_memory, floatBitsToUint(floatInput));
		}
	}
}

void reset(uint vID){
	vertices[vID].normal = floatBitsToUint(vec3(0, 0, 0));
}

void accumulate(uint fID){
	Face f = faces[fID];

	Vertex v1 = vertices[f.a];
	Vertex v2 = vertices[f.b];
	Vertex v3 = vertices[f.c];

	vec3 weightedNormal = cross(v3.position - v1.position, v2.position - v1.position);
	atomicAddVec3(f.a, weightedNormal);
	atomicAddVec3(f.b, weightedNormal);
	atomicAddVec3(f.c, weightedNormal);
}

void norm(uint vID){
	vertices[vID].normal = floatBitsToUint(normalize(uintBitsToFloat(vertices[vID].normal)));
}

void main() {
	uint id = gl_GlobalInvocationID.x;
	switch (state){
		case 0:
			if (id < vertexCount){
				reset(id);
			}
			break;
		case 1:
			if (id < faceCount){
				accumulate(id);
			}
			break;
		case 2:
			if (id < vertexCount){
				norm(id);
			}
			break;
	}
}