#include <glm/geometric.hpp>
#include "soft_body.hpp"
#include "mesh.hpp"
#include "constraints.hpp"
#include "tetgen.h"
#include "timer.hpp"
#include <omp.h>
#include <iostream>

SoftBody::SoftBody(Mesh* mesh, float compliance) : compliance(compliance) {
	tetgenbehavior behavior;
	behavior.parse_commandline(std::string("pYQfeza0.1").data());

	tetgenio in;
	in.numberofpoints = static_cast<int>(mesh->vertices.size());
	in.pointlist = new REAL[mesh->vertices.size() * 3];
	in.numberoffacets = static_cast<int>(mesh->faces.size());
	in.facetlist = new tetgenio::facet[mesh->faces.size()];
	in.numberofpointattributes = 3;
	in.pointattributelist = new REAL[mesh->vertices.size() * 3];

	for (size_t i = 0; i < mesh->vertices.size(); i++){
		in.pointlist[i * 3 + 0] = mesh->vertices[i].position.x;
		in.pointlist[i * 3 + 1] = mesh->vertices[i].position.y;
		in.pointlist[i * 3 + 2] = mesh->vertices[i].position.z;
		in.pointattributelist[i * 3 + 0] = (mesh->vertices[i].position.x + 1) / 2;
		in.pointattributelist[i * 3 + 1] = (mesh->vertices[i].position.y + 1) / 2;
		in.pointattributelist[i * 3 + 2] = (mesh->vertices[i].position.z + 1) / 2;
	}

	for (size_t i = 0; i < mesh->faces.size(); i++){
		tetgenio::facet &f = in.facetlist[i];
		tetgenio::polygon p;

		p.numberofvertices = 3;
		p.vertexlist = new int[3];
		p.vertexlist[0] = static_cast<int>(mesh->faces[i].a);
		p.vertexlist[1] = static_cast<int>(mesh->faces[i].b);
		p.vertexlist[2] = static_cast<int>(mesh->faces[i].c);

		f.numberofholes = 0;
		f.holelist = nullptr;
		f.numberofpolygons = 1;
		f.polygonlist = new tetgenio::polygon[1];
		f.polygonlist[0] = p;
	}

	tetgenio out;
	tetrahedralize(&behavior, &in, &out);

	vertices.reserve(out.numberofpoints);
	for (size_t i = 0; i < out.numberofpoints; i++){
		float x = static_cast<float>(out.pointlist[i * 3 + 0]);
		float y = static_cast<float>(out.pointlist[i * 3 + 1]);
		float z = static_cast<float>(out.pointlist[i * 3 + 2]);
		float r = static_cast<float>(out.pointattributelist[i * 3 + 0]);
		float g = static_cast<float>(out.pointattributelist[i * 3 + 1]);
		float b = static_cast<float>(out.pointattributelist[i * 3 + 2]);
		vertices.emplace_back(Vertex({x, y, z}, {r, g, b}));
	}

	faces.reserve(out.numberoftrifaces);
	constraintData.triangles.reserve(out.numberoftrifaces);
	for (size_t i = 0; i < out.numberoftrifaces; i++){
		uint32_t a = out.trifacelist[i * 3 + 0];
		uint32_t b = out.trifacelist[i * 3 + 1];
		uint32_t c = out.trifacelist[i * 3 + 2];
		if (out.trifacemarkerlist[i] != 0)
			faces.emplace_back(Face(a, b, c));
		constraintData.triangles.emplace_back(Triangle(Face(a, b, c), 0));
	}
	faces.shrink_to_fit();

	constraintData.edges.reserve(out.numberofedges);
	for (size_t i = 0; i < out.numberofedges; i++) {
		uint32_t a = out.edgelist[i * 2 + 0];
		uint32_t b = out.edgelist[i * 2 + 1];
		float length = glm::length(vertices[a].position - vertices[b].position);
		constraintData.edges.emplace_back(Edge(a, b, length));
	}

	float density = 1;
	vector<float> masses(vertices.size());

	constraintData.tetrahedra.reserve(out.numberoftetrahedra);
	for (size_t i = 0; i < out.numberoftetrahedra; i++){
		uint32_t a = out.tetrahedronlist[i * 4 + 0];
		uint32_t b = out.tetrahedronlist[i * 4 + 1];
		uint32_t c = out.tetrahedronlist[i * 4 + 2];
		uint32_t d = out.tetrahedronlist[i * 4 + 3];

		glm::vec3 p1 = vertices[a].position;
		glm::vec3 p2 = vertices[b].position;
		glm::vec3 p3 = vertices[c].position;
		glm::vec3 p4 = vertices[d].position;

		float volume = glm::dot(p4 - p1, glm::cross(p2 - p1, p3 - p1)) / 6;
		assert(volume > 0);

		float pointMass = volume * density / 4;

		masses[a] += pointMass;
		masses[b] += pointMass;
		masses[c] += pointMass;
		masses[d] += pointMass;

		constraintData.tetrahedra.emplace_back(Tetrahedron(a, b, c, d, volume));
	}

	for (size_t i = 0; i < masses.size(); i++){
		vertices[i].inverseMass = 1.0f / masses[i];
	}

	splitConstraints();
}

void SoftBody::applyVertexOffset(const glm::vec3 &offset) {
	for (Vertex& vertex : vertices){
		vertex.position += offset;
	}
}

vector<unordered_set<const Constraint *>> createPartitions(const Graph &graph){
	unordered_map<const Constraint *, uint32_t> constraintToColor;
	vector<unordered_set<const Constraint *>> colorToConstraintList;

	for (const auto& [constraint, adjacentList] : graph){

		unordered_set<uint32_t> forbiddenColors;
		for (auto adjacent : adjacentList)
			if (constraintToColor.contains(adjacent))
				forbiddenColors.insert(constraintToColor.at(adjacent));

		uint32_t assignColor;
		if (forbiddenColors.size() == colorToConstraintList.size()) {
			assignColor = forbiddenColors.size();
			colorToConstraintList.emplace_back();
		} else {
			uint32_t minColor;
			uint32_t minColorUseCount = UINT32_MAX;
			for (uint32_t color = 0; color < colorToConstraintList.size(); color++){
				if (colorToConstraintList[color].size() < minColorUseCount && !forbiddenColors.contains(color)){
					minColorUseCount = colorToConstraintList[color].size();
					minColor = color;
				}
			}
			assignColor = minColor;
		}

		constraintToColor[constraint] = assignColor;
		colorToConstraintList[assignColor].insert(constraint);
	}

	return colorToConstraintList;
}

void SoftBody::splitConstraints() {
	omp_set_num_threads(omp_get_num_procs() - 2);

	unordered_map<uint32_t, vector<const Constraint *>> pointToConstraints;

	vector<DistanceConstraint> lengthConstraints;
	vector<VolumeConstraint> volumeConstraints;
	lengthConstraints.reserve(constraintData.edges.size());
	volumeConstraints.reserve(constraintData.tetrahedra.size());

	for (const Edge &edge : constraintData.edges)
		lengthConstraints.push_back(DistanceConstraint(edge));
	for (const Tetrahedron &tetrahedron : constraintData.tetrahedra)
		volumeConstraints.push_back(VolumeConstraint(tetrahedron));

	for (const DistanceConstraint &distanceConstraint : lengthConstraints){
		pointToConstraints[distanceConstraint.a].push_back(&distanceConstraint);
		pointToConstraints[distanceConstraint.b].push_back(&distanceConstraint);
	}

	for (const VolumeConstraint &volumeConstraint : volumeConstraints){
		pointToConstraints[volumeConstraint.a].push_back(&volumeConstraint);
		pointToConstraints[volumeConstraint.b].push_back(&volumeConstraint);
		pointToConstraints[volumeConstraint.c].push_back(&volumeConstraint);
		pointToConstraints[volumeConstraint.d].push_back(&volumeConstraint);
	}

	Graph graph;

	auto setNeighbors = [&pointToConstraints, &graph](const Constraint * constraint, const vector<uint32_t> &vIDs){
		unordered_set<const Constraint *> neighbors;
		for (uint32_t vID : vIDs)
			neighbors.insert(pointToConstraints[vID].begin(), pointToConstraints[vID].end());
		neighbors.erase(constraint);
		#pragma omp critical
		graph[constraint].insert(graph[constraint].end(), neighbors.begin(), neighbors.end());
	};

	#pragma omp parallel for default(none) shared(setNeighbors, lengthConstraints)
	for (const DistanceConstraint &distanceConstraint : lengthConstraints){
		setNeighbors(&distanceConstraint, {
			distanceConstraint.a,
			distanceConstraint.b
		});
	}

	#pragma omp parallel for default(none) shared(setNeighbors, volumeConstraints)
	for (const VolumeConstraint &volumeConstraint : volumeConstraints){
		setNeighbors(&volumeConstraint, {
			volumeConstraint.a,
			volumeConstraint.b,
			volumeConstraint.c,
			volumeConstraint.d,
		});
	}

	vector<unordered_set<const Constraint *>> partitions = createPartitions(graph);
	constraintData.partitionCount = partitions.size();

	reorderConstraintIndices(partitions);
}

void SoftBody::reorderConstraintIndices(const vector<unordered_set<const Constraint *>> &partitions) {
	ConstraintData reordered;
	for (uint32_t partition = 0; partition < partitions.size(); partition++){
		reordered.recordNewPartition();
		for (const Constraint * constraint : partitions[partition])
			constraint->writeData(reordered);
		reordered.writePartitionInformation();
	}
	std::swap(reordered.edgePartitions, constraintData.edgePartitions);
	std::swap(reordered.tetrahedronPartitions, constraintData.tetrahedronPartitions);
	std::swap(reordered.edges, constraintData.edges);
	std::swap(reordered.tetrahedra, constraintData.tetrahedra);
	//std::swap(reordered, constraintData);
	//reordered.triangles.swap(constraintData.triangles);
	//reordered.tetrahedra.swap(constraintData.tetrahedra);
	//std::swap(reordered.partitionCount, constraintData.partitionCount);
}