#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"

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(UniformBufferObject);

	uniformBuffer = 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);

	descriptorPool->bindBuffer(*uniformBuffer, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, DescriptorSet::WORLD, 0);
	camera = make_unique<Camera>(swapchain->extent);

	createMeshBuffers();

	SizeInformation sizeInformation {};
	sizeInformation.vertexCount = vertexBuffer->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);

	properties.gravity = {0, -9.81, 0};
	properties.k = 10;
	properties.dt = 1.f / 60.f / static_cast<float>(properties.k);

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

	descriptorPool->bindBuffer(*vertexBuffer, 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);

	createComputePipelines();

	char* stats;
	vmaBuildStatsString(Instance::GetAllocator(), &stats, VK_TRUE);
	// printf("%s", stats);
	vmaFreeStatsString(Instance::GetAllocator(), stats);
}

void Application::mainLoop() {
	while (!glfwWindowShouldClose(Instance::instance->window)){
		glfwPollEvents();
		update();
		drawFrame();
	}
	vkDeviceWaitIdle(Instance::GetDevice());
}

Application::~Application() {

}

void Application::createSyncObjects() {
	imageAvailable = make_unique<Semaphore>();
	renderFinished = make_unique<Semaphore>();
	computeFinished = make_unique<Semaphore>();
	renderInFlight = make_unique<Fence>(true);
	computeInFlight = 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 < 5; 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) {}
	};

	vertexBuffer = make_unique<SimulationBuffer>(vertices.data(), vertices.size() * sizeof(Vertex), VK_BUFFER_USAGE_VERTEX_BUFFER_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));
}

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

	{
		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;

	UniformBufferObject ubo {};
	ubo.model = glm::mat4(1);
	ubo.view = camera->view();
	ubo.projection = camera->projection();
	ubo.projection[1][1] *= -1;

	memcpy(uniformBuffer->allocationInfo.pMappedData, &ubo, sizeof(UniformBufferObject));
}

void Application::recordGraphicsCommandBuffer(uint32_t imageIndex) {
	VkCommandBuffer cmdBuffer = Instance::instance->commandPool->graphicsBuffer;

	VkCommandBufferBeginInfo beginInfo {};
	beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;

	vkBeginCommandBuffer(cmdBuffer, &beginInfo);

	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();

	vkCmdEndRenderPass(cmdBuffer);
	vkEndCommandBuffer(cmdBuffer);
}

void Application::drawFrame() {
	vkWaitForFences(Instance::GetDevice(), 1, &renderInFlight->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, &renderInFlight->handle);

	vkResetCommandBuffer(Instance::instance->commandPool->graphicsBuffer, 0);
	recordGraphicsCommandBuffer(imageIndex);

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

	VkSubmitInfo submitInfo {};
	submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

	VkSemaphore waitSemaphores[] = {imageAvailable->handle, computeFinished->handle};
	VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT};
	submitInfo.waitSemaphoreCount = 2;
	submitInfo.pWaitSemaphores = waitSemaphores;
	submitInfo.pWaitDstStageMask = waitStages;
	submitInfo.commandBufferCount = 1;
	submitInfo.pCommandBuffers = &Instance::instance->commandPool->graphicsBuffer;

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

	vkQueueSubmit(Instance::instance->graphicsQueue, 1, &submitInfo, renderInFlight->handle);

	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;

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

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


	VkCommandBuffer cmdBuffer = Instance::instance->commandPool->computeBuffer;
	vkResetCommandBuffer(cmdBuffer, 0);

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

	vkBeginCommandBuffer(cmdBuffer, &beginInfo);
	recordComputeCommands(cmdBuffer);
	vkEndCommandBuffer(cmdBuffer);

	VkSubmitInfo submit {};
	submit.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
	submit.commandBufferCount = 1;
	submit.pCommandBuffers = &Instance::instance->commandPool->computeBuffer;
	submit.signalSemaphoreCount = 1;
	submit.pSignalSemaphores = &computeFinished->handle;

	vkQueueSubmit(Instance::instance->computeQueue, 1, &submit, computeInFlight->handle);
}

void Application::recordDrawCommands() {
	VkCommandBuffer cmdBuffer = Instance::instance->commandPool->graphicsBuffer;

	vkCmdBindPipeline(cmdBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, graphicsPipeline->handle);

	VkBuffer buffers[] = {vertexBuffer->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::recordComputeCommands(VkCommandBuffer cmdBuffer) {
#define BlOCK_SIZE 32

	auto getGroupCount = [](uint32_t threads, uint32_t blockSize){
		return (threads - 1) / blockSize + 1;
	};

	uint32_t vertexGroupCount = getGroupCount(vertexBuffer->size / sizeof(Vertex), BlOCK_SIZE);
	uint32_t faceGroupCount = getGroupCount(faceBuffer->size / sizeof(Face), BlOCK_SIZE);

	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;

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

	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);


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