#include "application.hpp"
#include "vulkan/swapchain.hpp"
#include "vulkan/pipeline.hpp"
#include "vulkan/instance.hpp"
#include "vulkan/buffer.hpp"
#include "vulkan/command_pool.hpp"
#include "vulkan/image.hpp"
#include "vulkan/synchronization.hpp"
#include "vulkan/descriptor_pool.hpp"
#include "camera.hpp"
#include "input.hpp"
#include "soft_body.hpp"
#include "mesh.hpp"
#include "constraints.hpp"
#include "timer.hpp"
#include "grabber.hpp"

Application::Application() {
	createSyncObjects();
	swapchain = make_unique<Swapchain>();
	descriptorPool = make_unique<DescriptorPool>();
	graphicsPipeline = unique_ptr<GraphicsPipeline>(new GraphicsPipeline("shaders/vert.spv", "shaders/frag.spv",
											swapchain->renderPass,
											{descriptorPool->layouts[DescriptorSet::WORLD]}));

	VkDeviceSize bufferSize = sizeof(CameraUniformData);

	cameraUniformBuffer = make_unique<Buffer>(bufferSize,
											  VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
											  VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE,
							   VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT | VMA_ALLOCATION_CREATE_MAPPED_BIT);
	cameraUniformBuffer->setName("Camera");
	camera = make_unique<Camera>(swapchain->extent);

	struct GrabInformation {
		float originalInverseMass;
		uint32_t vID;
		float distanceToFace;
		bool grabbing;
	} initialGrabInformation {};
	initialGrabInformation.distanceToFace = 1e20;
	grabBuffer = make_unique<Buffer>(sizeof(GrabInformation), &initialGrabInformation, sizeof(GrabInformation),
									 VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE, 0);
	grabBuffer->setName("Grab");
	grabber = make_unique<Grabber>();

	createMeshBuffers();

	SizeInformation sizeInformation {};
	sizeInformation.vertexCount = vertexBuffers[0]->size / sizeof(Vertex);
	sizeInformation.faceCount = faceBuffer->size / sizeof(Face);

	sizeInformationBuffer = make_unique<Buffer>(
			sizeof(SizeInformation), &sizeInformation, sizeof(sizeInformation),
			VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
			VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE, 0);
	sizeInformationBuffer->setName("Sizes");

	properties.gravity = {0, -9.81, 0};
	properties.k = 10;
	properties.dt = 1.f / 60.f;

	propertiesBuffer = make_unique<Buffer>(
			sizeof(Properties), &properties, sizeof(properties),
			VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
			VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE, 0);
	propertiesBuffer->setName("Simulation properties");


	descriptorPool->bindBuffer(*cameraUniformBuffer, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, DescriptorSet::WORLD, 0);

	descriptorPool->bindBuffer(*vertexBuffers[1 - currentDrawVertexBuffer], VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, DescriptorSet::MESH, 0);
	descriptorPool->bindBuffer(*faceBuffer, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, DescriptorSet::MESH, 1);
	descriptorPool->bindBuffer(*edgeBuffer, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, DescriptorSet::MESH, 2);
	descriptorPool->bindBuffer(*triangleBuffer, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, DescriptorSet::MESH, 3);
	descriptorPool->bindBuffer(*tetrahedronBuffer, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, DescriptorSet::MESH, 4);
	descriptorPool->bindBuffer(*sizeInformationBuffer, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, DescriptorSet::MESH, 5);

	descriptorPool->bindBuffer(*propertiesBuffer, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, DescriptorSet::SIMULATION, 0);
	descriptorPool->bindBuffer(*grabBuffer, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, DescriptorSet::SIMULATION, 1);

	createComputePipelines();

	printVmaStats();
}

#include <future>
#include <chrono>

using namespace std::chrono;

void Application::mainLoop() {
	std::future compute = std::async(std::launch::async, [this](){
		while (!glfwWindowShouldClose(Instance::instance->window)){
			auto t1 = system_clock::now();
			update();
			auto t2 = system_clock::now();

			microseconds updateDuration = duration_cast<microseconds>(t2 - t1);

			microseconds sleepDuration(static_cast<int64_t>(properties.dt * 1000 * 1000));
			std::this_thread::sleep_for(sleepDuration - updateDuration);
		}
	});

	while (!glfwWindowShouldClose(Instance::instance->window)){
		glfwPollEvents();
		drawFrame();
	}
	compute.wait();
	vkDeviceWaitIdle(Instance::GetDevice());
}

Application::~Application() {

}

void Application::createSyncObjects() {
	imageAvailable = make_unique<Semaphore>();
	renderFinished = make_unique<Semaphore>();
	computeSemaphore = make_unique<Semaphore>();
	transferFinished = make_unique<Semaphore>();
	renderFence = make_unique<Fence>(true);
	computeFence = make_unique<Fence>(true);
	transferFence = make_unique<Fence>(true);
}

void Application::createMeshBuffers() {
	Mesh sphere("models/icosphere.ply");
	Mesh bunny("models/bunny_medium.ply");

	auto body = std::make_unique<SoftBody>(&sphere, 1.f / 50);

	for (size_t i = 0; i < 10; i++){
		auto copy = std::make_unique<SoftBody>(*body.get());
		copy->applyVertexOffset({i * 2, 0, 0});
		softBodies.push_back(std::move(copy));
	}

	body = std::make_unique<SoftBody>(&bunny, 1.f / 3);
	for (size_t i = 0; i < 10; i++){
		auto copy = std::make_unique<SoftBody>(*body.get());
		copy->applyVertexOffset({i * 2, 0, 2});
		softBodies.push_back(std::move(copy));
	}

	vector<Vertex> vertices;
	vector<Face> faces;

	for (std::unique_ptr<SoftBody> &currentSoftBody : softBodies){
		currentSoftBody->firstIndex = faces.size() * 3;
		currentSoftBody->vertexOffset = static_cast<int32_t>(vertices.size());

		int32_t vertexOffset = currentSoftBody->vertexOffset;

		for (auto &face : currentSoftBody->faces){
			face.a += vertexOffset;
			face.b += vertexOffset;
			face.c += vertexOffset;
		}

		for (auto &edge : currentSoftBody->constraintData.edges){
			edge.a += vertexOffset;
			edge.b += vertexOffset;
		}

		for (auto &triangle : currentSoftBody->constraintData.triangles){
			triangle.a += vertexOffset;
			triangle.b += vertexOffset;
			triangle.c += vertexOffset;
		}

		for (auto &tetrahedron : currentSoftBody->constraintData.tetrahedra){
			tetrahedron.a += vertexOffset;
			tetrahedron.b += vertexOffset;
			tetrahedron.c += vertexOffset;
			tetrahedron.d += vertexOffset;
		}

		vertices.insert(vertices.end(), currentSoftBody->vertices.begin(), currentSoftBody->vertices.end());
		faces.insert(faces.end(), currentSoftBody->faces.begin(), currentSoftBody->faces.end());

		constraintData.partitionCount = std::max(constraintData.partitionCount, currentSoftBody->constraintData.partitionCount);

		auto combine = [&currentSoftBody, this] (
				auto &globalIndices, auto &bodyIndices,
				vector<ConstraintData::Partition> &globalPartitions, vector<ConstraintData::Partition> &bodyPartitions){
			if (globalPartitions.size() < currentSoftBody->constraintData.partitionCount)
				globalPartitions.resize(currentSoftBody->constraintData.partitionCount);

			uint32_t offsetAdded = 0;

			for (uint32_t partition = 0; partition < constraintData.partitionCount; partition++){

				ConstraintData::Partition &globalPartition = globalPartitions[partition];
				globalPartition.offset += offsetAdded;

				if (partition < bodyPartitions.size()){
					const ConstraintData::Partition &bodyPartition = bodyPartitions[partition];

					auto dst = globalIndices.begin() + globalPartition.offset;
					auto srcStart = bodyIndices.begin() + bodyPartition.offset;
					uint32_t count = bodyPartition.size;
					globalIndices.insert(dst, srcStart, srcStart + count);

					globalPartition.size += count;

					offsetAdded += count;
				}
			}
		};

		combine(constraintData.edges, currentSoftBody->constraintData.edges,
				constraintData.edgePartitions, currentSoftBody->constraintData.edgePartitions);

		combine(constraintData.tetrahedra, currentSoftBody->constraintData.tetrahedra,
				constraintData.tetrahedronPartitions, currentSoftBody->constraintData.tetrahedronPartitions);

		constraintData.triangles.insert(constraintData.triangles.end(), currentSoftBody->constraintData.triangles.begin(), currentSoftBody->constraintData.triangles.end());
		constraintData.tetrahedra.insert(constraintData.tetrahedra.end(), currentSoftBody->constraintData.tetrahedra.begin(), currentSoftBody->constraintData.tetrahedra.end());
	}

	printf("Vertices: %zu\nFaces: %zu\nEdges: %zu\nTriangles: %zu\nTetrahedra: %zu\nTotal Constraints: %zu\n",
		   vertices.size(), faces.size(), constraintData.edges.size(), constraintData.triangles.size(), constraintData.tetrahedra.size(),
		   constraintData.edges.size() + constraintData.tetrahedra.size());

	printf("Partitions: %u\n", constraintData.partitionCount);

	class SimulationBuffer : public Buffer {
	public:
		SimulationBuffer(void* data, VkDeviceSize size, VkBufferUsageFlags additionalUsageFlags=0)
				: Buffer(size, data, size, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | additionalUsageFlags, VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE, 0) {}
	};

	vertexBuffers[0] = make_unique<SimulationBuffer>(vertices.data(), vertices.size() * sizeof(Vertex),
												  VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
	vertexBuffers[1] = make_unique<SimulationBuffer>(vertices.data(), vertices.size() * sizeof(Vertex),
												  VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
	faceBuffer = make_unique<SimulationBuffer>(faces.data(), faces.size() * sizeof(Face), VK_BUFFER_USAGE_INDEX_BUFFER_BIT);
	edgeBuffer = make_unique<SimulationBuffer>(constraintData.edges.data(), constraintData.edges.size() * sizeof(Edge));
	triangleBuffer = make_unique<SimulationBuffer>(constraintData.triangles.data(), constraintData.triangles.size() * sizeof(Triangle));
	tetrahedronBuffer = make_unique<SimulationBuffer>(constraintData.tetrahedra.data(), constraintData.tetrahedra.size() * sizeof(Tetrahedron));

	vertexBuffers[0]->setName("Vertices 0");
	vertexBuffers[1]->setName("Vertices 1");
	faceBuffer->setName("Faces");
	edgeBuffer->setName("Edges");
	triangleBuffer->setName("Triangles");
	tetrahedronBuffer->setName("Tetrahedra");
}

void Application::createComputePipelines() {
	vector<VkDescriptorSetLayout> layouts;
	vector<VkPushConstantRange> pushRanges;

	{
		layouts.push_back(descriptorPool->layouts[DescriptorSet::MESH]);
		layouts.push_back(descriptorPool->layouts[DescriptorSet::WORLD]);
		layouts.push_back(descriptorPool->layouts[DescriptorSet::SIMULATION]);
		pushRanges.push_back({
									 .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
									 .offset = 0,
									 .size = sizeof(GrabPushConstants)
							 });
		grabPipeline = unique_ptr<ComputePipeline>(new ComputePipeline("shaders/grab.spv", layouts, pushRanges));
	}

	layouts.clear();
	pushRanges.clear();

	{
		layouts.push_back(descriptorPool->layouts[DescriptorSet::MESH]);
		layouts.push_back(descriptorPool->layouts[DescriptorSet::SIMULATION]);
		pushRanges.push_back({
									 .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
									 .offset = 0,
									 .size = sizeof(PBDPushData)
							 });
		pbdPipeline = unique_ptr<ComputePipeline>(new ComputePipeline("shaders/pbd.spv", layouts, pushRanges));
	}

	layouts.clear();
	pushRanges.clear();

	{
		layouts.push_back(descriptorPool->layouts[DescriptorSet::MESH]);
		pushRanges.push_back({
									 .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
									 .offset = 0,
									 .size = sizeof(uint32_t)
							 });
		normalPipeline = unique_ptr<ComputePipeline>(new ComputePipeline("shaders/normal.spv", layouts, pushRanges));
	}

}


void Application::updateUniformBuffer() {
	static float elapsed = 0;

	elapsed += 0.007;

	CameraUniformData ubo {};
	ubo.view = camera->view();
	ubo.projection = camera->projection();
	ubo.projection[1][1] *= -1;
	ubo.viewport = camera->viewport();

	memcpy(cameraUniformBuffer->allocationInfo.pMappedData, &ubo, sizeof(CameraUniformData));

	vmaFlushAllocation(Instance::GetAllocator(),
					   cameraUniformBuffer->allocation,
					   cameraUniformBuffer->allocationInfo.offset,
					   cameraUniformBuffer->allocationInfo.size);
}

void Application::drawFrame() {
	vkWaitForFences(Instance::GetDevice(), 1, &renderFence->handle, VK_TRUE, UINT64_MAX);

	uint32_t imageIndex;
	VkResult result = vkAcquireNextImageKHR(Instance::GetDevice(), swapchain->handle, UINT64_MAX, imageAvailable->handle, VK_NULL_HANDLE, &imageIndex);
	if (result == VK_ERROR_OUT_OF_DATE_KHR){
		swapchain->recreateSwapchain();
		return;
	}

	vkResetFences(Instance::GetDevice(), 1, &renderFence->handle);

	camera->update(0.017);
	updateUniformBuffer();

	VkCommandBuffer cmdBuffer = Instance::instance->renderingCommandPool->buffers[0];
	{
		vkResetCommandBuffer(cmdBuffer, 0);

		VkCommandBufferBeginInfo beginInfo{};
		beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;

		vkBeginCommandBuffer(cmdBuffer, &beginInfo);

		VkBufferMemoryBarrier vertexBufferBarrier{};
		vertexBufferBarrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
		vertexBufferBarrier.size = vertexBuffers[currentDrawVertexBuffer]->size;
		vertexBufferBarrier.offset = 0;
		vertexBufferBarrier.buffer = vertexBuffers[currentDrawVertexBuffer]->handle;
		vertexBufferBarrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_SHADER_READ_BIT;
		vertexBufferBarrier.dstAccessMask = VK_ACCESS_VERTEX_ATTRIBUTE_READ_BIT;

		vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_VERTEX_INPUT_BIT, 0, 0,
							 nullptr, 1, &vertexBufferBarrier, 0, nullptr);

		VkBufferMemoryBarrier uniformBufferBarrier {};
		uniformBufferBarrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
		uniformBufferBarrier.size = cameraUniformBuffer->size;
		uniformBufferBarrier.offset = 0;
		uniformBufferBarrier.buffer = cameraUniformBuffer->handle;
		uniformBufferBarrier.srcAccessMask = VK_ACCESS_HOST_WRITE_BIT;
		uniformBufferBarrier.dstAccessMask = VK_ACCESS_UNIFORM_READ_BIT;

		vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_HOST_BIT, VK_PIPELINE_STAGE_VERTEX_SHADER_BIT | VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0,
							 0, nullptr, 1, &uniformBufferBarrier, 0, nullptr);

		VkRenderPassBeginInfo renderPassInfo{};
		renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
		renderPassInfo.renderPass = swapchain->renderPass;
		renderPassInfo.framebuffer = swapchain->frameBuffers[imageIndex];
		renderPassInfo.renderArea.offset = {0, 0};
		renderPassInfo.renderArea.extent = swapchain->extent;

		VkClearValue clearValues[2]{};
		clearValues[0].color = {{0, 0, 0, 1}};
		clearValues[1].depthStencil = {1.0f, 0};

		renderPassInfo.clearValueCount = 2;
		renderPassInfo.pClearValues = clearValues;

		vkCmdBeginRenderPass(cmdBuffer, &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);

		VkViewport viewport{};
		viewport.x = 0;
		viewport.y = 0;
		viewport.width = static_cast<float>(swapchain->extent.width);
		viewport.height = static_cast<float>(swapchain->extent.height);
		viewport.minDepth = 0;
		viewport.maxDepth = 1;
		vkCmdSetViewport(cmdBuffer, 0, 1, &viewport);

		VkRect2D scissor{};
		scissor.offset = {0, 0};
		scissor.extent = swapchain->extent;
		vkCmdSetScissor(cmdBuffer, 0, 1, &scissor);

		recordDrawCommands(cmdBuffer);

		vkCmdEndRenderPass(cmdBuffer);
		vkEndCommandBuffer(cmdBuffer);
	}

	VkSubmitInfo submitInfo {};
	submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

	VkSemaphore waitSemaphores[] = {imageAvailable->handle};
	VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
	submitInfo.waitSemaphoreCount = 1;
	submitInfo.pWaitSemaphores = waitSemaphores;
	submitInfo.pWaitDstStageMask = waitStages;
	submitInfo.commandBufferCount = 1;
	submitInfo.pCommandBuffers = &cmdBuffer;

	VkSemaphore signalSemaphores[] = {renderFinished->handle};
	submitInfo.signalSemaphoreCount = 1;
	submitInfo.pSignalSemaphores = signalSemaphores;

	submitMutex.lock();
	vkQueueSubmit(Instance::instance->graphicsAndPresentQueue, 1, &submitInfo, renderFence->handle);
	submitMutex.unlock();

	VkPresentInfoKHR presentInfo {};
	presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
	presentInfo.waitSemaphoreCount = 1;
	presentInfo.pWaitSemaphores = signalSemaphores;

	VkSwapchainKHR swapchains[] = {swapchain->handle};
	presentInfo.swapchainCount = 1;
	presentInfo.pSwapchains = swapchains;
	presentInfo.pImageIndices = &imageIndex;

	submitMutex.lock();
	result = vkQueuePresentKHR(Instance::instance->graphicsAndPresentQueue, &presentInfo);
	submitMutex.unlock();

	if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || Instance::instance->windowResized){
		Instance::instance->windowResized = false;
		swapchain->recreateSwapchain();
	}
}

void Application::recordDrawCommands(VkCommandBuffer cmdBuffer) {
	vkCmdBindPipeline(cmdBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, graphicsPipeline->handle);

	VkBuffer buffers[] = {vertexBuffers[currentDrawVertexBuffer]->handle};
	VkDeviceSize offsets[] = {0};
	vkCmdBindVertexBuffers(cmdBuffer, 0, 1, buffers, offsets);
	vkCmdBindIndexBuffer(cmdBuffer, faceBuffer->handle, 0, VK_INDEX_TYPE_UINT32);
	vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, graphicsPipeline->layout, 0, 1, &descriptorPool->sets[DescriptorSet::WORLD], 0, nullptr);

	for (const auto &softBody : softBodies){
		vkCmdDrawIndexed(cmdBuffer, softBody->faces.size() * 3, 1, softBody->firstIndex, 0, 0);
	}
}

void Application::update() {
	vkWaitForFences(Instance::GetDevice(), 1, &transferFence->handle, VK_TRUE, UINT64_MAX);
	vkResetFences(Instance::GetDevice(), 1, &transferFence->handle);

	currentDrawVertexBuffer = 1 - currentDrawVertexBuffer;
	descriptorPool->bindBuffer(*vertexBuffers[1 - currentDrawVertexBuffer], VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, DescriptorSet::MESH, 0);

	VkCommandBufferBeginInfo beginInfo {};
	beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
	beginInfo.flags = 0;

	VkCommandBuffer cmdBuffer = Instance::instance->computeCommandPool->buffers[0];
	{
		vkResetCommandBuffer(cmdBuffer, 0);
		vkBeginCommandBuffer(cmdBuffer, &beginInfo);

		recordGrabCommands(cmdBuffer);
		recordPBDCommands(cmdBuffer);
		recordNormalCommands(cmdBuffer);

		vkEndCommandBuffer(cmdBuffer);

		VkSubmitInfo submit {};
		submit.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
		submit.commandBufferCount = 1;
		submit.pCommandBuffers = &cmdBuffer;
		submit.signalSemaphoreCount = 1;
		submit.pSignalSemaphores = &computeSemaphore->handle;

		submitMutex.lock();
		vkQueueSubmit(Instance::instance->computeAndTransferQueue, 1, &submit, nullptr);
		submitMutex.unlock();
	}

	cmdBuffer = Instance::instance->computeCommandPool->buffers[1];
	vkResetCommandBuffer(cmdBuffer, 0);
	{
		vkBeginCommandBuffer(cmdBuffer, &beginInfo);

		VkBufferCopy copyRegion {};
		copyRegion.size = vertexBuffers[0]->size;
		copyRegion.srcOffset = 0;
		copyRegion.dstOffset = 0;

		vkCmdCopyBuffer(cmdBuffer, vertexBuffers[1 - currentDrawVertexBuffer]->handle, vertexBuffers[currentDrawVertexBuffer]->handle, 1, &copyRegion);

		vkEndCommandBuffer(cmdBuffer);

		VkPipelineStageFlags waitStage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;

		VkSubmitInfo submit {};
		submit.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
		submit.commandBufferCount = 1;
		submit.pCommandBuffers = &cmdBuffer;
		submit.waitSemaphoreCount = 1;
		submit.pWaitSemaphores = &computeSemaphore->handle;
		submit.pWaitDstStageMask = &waitStage;

		submitMutex.lock();
		vkQueueSubmit(Instance::instance->computeAndTransferQueue, 1, &submit, transferFence->handle);
		submitMutex.unlock();
	}
}

uint32_t Application::GetGroupCount(uint32_t threads, uint32_t blockSize) {
	return (threads - 1) / blockSize + 1;
}

void Application::recordGrabCommands(VkCommandBuffer cmdBuffer) {
	VkMemoryBarrier barrier {};
	barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
	barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
	barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;

	// TODO maybe add buffermemorybarrier for camera uniform, because it can be changed from main drawing thread

	vkCmdBindPipeline(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, grabPipeline->handle);
	vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, grabPipeline->layout, 0, 1, &descriptorPool->sets[DescriptorSet::MESH], 0, nullptr);
	vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, grabPipeline->layout, 1, 1, &descriptorPool->sets[DescriptorSet::WORLD], 0, nullptr);
	vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, grabPipeline->layout, 2, 1, &descriptorPool->sets[DescriptorSet::SIMULATION], 0, nullptr);

	GrabPushConstants pushConstants {};

	if (grabber->started()){
		pushConstants.state = 0;
		pushConstants.screenPosition = grabber->previousCursorPosition;
		vkCmdPushConstants(cmdBuffer, grabPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(GrabPushConstants), &pushConstants);

		uint32_t faceInvocations = GetGroupCount(faceBuffer->size / sizeof(Face), BLOCK_SIZE_GRAB);

		vkCmdDispatch(cmdBuffer, faceInvocations, 1, 1);

		vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier,
							 0, nullptr, 0, nullptr);

		pushConstants.state = 1;
		vkCmdPushConstants(cmdBuffer, grabPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(GrabPushConstants), &pushConstants);
		vkCmdDispatch(cmdBuffer, 1, 1, 1);
	}

	vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier,
						 0, nullptr, 0, nullptr);

	glm::vec2 screenDelta;
	if (grabber->moved(screenDelta)){
		pushConstants.state = 2;
		pushConstants.screenDelta = screenDelta;
		vkCmdPushConstants(cmdBuffer, grabPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(GrabPushConstants), &pushConstants);
		vkCmdDispatch(cmdBuffer, 1, 1, 1);
	}

	vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier,
						 0, nullptr, 0, nullptr);

	if (grabber->stopped()){
		pushConstants.state = 3;
		vkCmdPushConstants(cmdBuffer, grabPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(GrabPushConstants), &pushConstants);
		vkCmdDispatch(cmdBuffer, 1, 1, 1);
	}

	vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier,
						 0, nullptr, 0, nullptr);
}

void Application::recordPBDCommands(VkCommandBuffer cmdBuffer) {
	VkMemoryBarrier barrier {};
	barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
	barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
	barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;

	uint32_t vertexGroupCount = GetGroupCount(vertexBuffers[1 - currentDrawVertexBuffer]->size / sizeof(Vertex), BLOCK_SIZE_PBD);

	vkCmdBindPipeline(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pbdPipeline->handle);
	vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pbdPipeline->layout, 0, 1, &descriptorPool->sets[DescriptorSet::MESH], 0, nullptr);
	vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pbdPipeline->layout, 1, 1, &descriptorPool->sets[DescriptorSet::SIMULATION], 0, nullptr);

	uint32_t state;

	for (size_t i = 0; i < properties.k; i++){
		state = 0;
		vkCmdPushConstants(cmdBuffer, pbdPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &state);
		vkCmdDispatch(cmdBuffer, vertexGroupCount, 1, 1);
		vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);

		state = 1;
		vkCmdPushConstants(cmdBuffer, pbdPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &state);

		for (uint32_t partition = 0; partition < constraintData.partitionCount; partition++){
			auto edgePartition = constraintData.edgePartitions[partition];
			auto tetrahedronPartition = constraintData.tetrahedronPartitions[partition];

			if (edgePartition.size + tetrahedronPartition.size == 0)
				continue;

			ConstraintData::Partition partitions[2] = {edgePartition, tetrahedronPartition};

			vkCmdPushConstants(cmdBuffer, pbdPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT,
							   offsetof(PBDPushData, edgePartition),
							   sizeof(partitions), partitions);

			uint32_t invocations = GetGroupCount(edgePartition.size + tetrahedronPartition.size, BLOCK_SIZE_PBD);
			vkCmdDispatch(cmdBuffer, invocations, 1, 1);

			vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);
		}

		state = 2;
		vkCmdPushConstants(cmdBuffer, pbdPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &state);
		vkCmdDispatch(cmdBuffer, vertexGroupCount, 1, 1);
		vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);
	}
}

void Application::recordNormalCommands(VkCommandBuffer cmdBuffer) {
	VkMemoryBarrier barrier {};
	barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
	barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
	barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;

	uint32_t vertexGroupCount = GetGroupCount(vertexBuffers[1 - currentDrawVertexBuffer]->size / sizeof(Vertex), BLOCK_SIZE_NORMAL);
	uint32_t faceGroupCount = GetGroupCount(faceBuffer->size / sizeof(Face), BLOCK_SIZE_NORMAL);

	vkCmdBindPipeline(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, normalPipeline->handle);
	vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, normalPipeline->layout, 0, 1, &descriptorPool->sets[DescriptorSet::MESH], 0, nullptr);

	uint32_t state = 0;
	vkCmdPushConstants(cmdBuffer, normalPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &state);
	vkCmdDispatch(cmdBuffer, vertexGroupCount, 1, 1);
	vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);

	state = 1;
	vkCmdPushConstants(cmdBuffer, normalPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &state);
	vkCmdDispatch(cmdBuffer, faceGroupCount, 1, 1);
	vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);

	state = 2;
	vkCmdPushConstants(cmdBuffer, normalPipeline->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(uint32_t), &state);
	vkCmdDispatch(cmdBuffer, vertexGroupCount, 1, 1);
	vkCmdPipelineBarrier(cmdBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);
}