#include <psemek/phys/engine_2d.hpp>

#include <psemek/util/heterogeneous_container.hpp>
#include <psemek/prof/profiler.hpp>

#include <psemek/math/constants.hpp>
#include <psemek/math/matrix.hpp>
#include <psemek/math/gauss.hpp>
#include <psemek/math/rotation.hpp>
#include <psemek/math/math.hpp>
#include <psemek/math/interval.hpp>

#include <psemek/log/log.hpp>

#include <optional>
#include <set>

namespace psemek::phys2d
{

	namespace
	{

		float shape_area(ball const & b)
		{
			return math::pi * b.radius * b.radius;
		}

		float shape_inertia(ball const & b)
		{
			return shape_area(b) * b.radius * b.radius / 2.f;
		}

		float shape_area(half_space const &)
		{
			return std::numeric_limits<float>::infinity();
		}

		float shape_inertia(half_space const &)
		{
			return std::numeric_limits<float>::infinity();
		}

		float shape_area(box const & b)
		{
			return b.width * b.height;
		}

		float shape_inertia(box const & b)
		{
			return shape_area(b) * (b.width * b.width + b.height * b.height) / 12.f;
		}

		math::point<float, 2> shape_closest(ball const & b, static_state const & s, math::point<float, 2> const & p)
		{
			auto d = p - s.position;
			auto l = math::length(d);
			if (l <= b.radius)
				return p;
			return s.position + (d / l) * b.radius;
		}

		math::point<float, 2> shape_closest(half_space const & b, static_state const &, math::point<float, 2> const & p)
		{
			auto v = math::dot(b.normal, p - math::point<float, 2>::zero()) - b.value;
			if (v <= 0.f)
				return p;
			return p - b.normal * v;
		}

		math::point<float, 2> shape_closest(box const & b, static_state const & s, math::point<float, 2> const & p)
		{
			math::vector ex{std::cos(s.rotation), std::sin(s.rotation)};
			math::vector ey{-std::sin(s.rotation), std::cos(s.rotation)};

			float w = b.width / 2.f;
			float h = b.height / 2.f;

			auto d = p - s.position;

			float x = math::dot(d, ex);
			float y = math::dot(d, ey);

			x = math::clamp(x, {-w, w});
			y = math::clamp(y, {-h, h});

			return s.position + ex * x + ey * y;
		}

		math::box<float, 2> shape_bbox(ball const & b, static_state const & s)
		{
			return {{{s.position[0] - b.radius, s.position[0] + b.radius}, {s.position[1] - b.radius, s.position[1] + b.radius}}};
		}

		math::box<float, 2> shape_bbox(box const & b, static_state const & s)
		{
			float cc = std::abs(std::cos(s.rotation));
			float ss = std::abs(std::sin(s.rotation));

			float w = b.width / 2.f;
			float h = b.height / 2.f;

			float tx = w * cc + h * ss;
			float ty = w * ss + h * cc;

			return {{{s.position[0] - tx, s.position[0] + tx}, {s.position[1] - ty, s.position[1] + ty}}};
		}

		math::box<float, 2> shape_bbox(half_space const &, static_state const &)
		{
			return math::box<float, 2>::full();
		}

		template <typename Shape>
		struct wrapped_shape
		{
			Shape shape;
			float area;
			float inertia;

			wrapped_shape(Shape const & shape)
				: shape{shape}
				, area(shape_area(shape))
				, inertia(shape_inertia(shape))
			{}
		};

		template <typename Shape>
		wrapped_shape<Shape> wrap(Shape const & shape)
		{
			return {shape};
		}

		// penetration points from object 1 to object 2
		struct collision
		{
			math::point<float, 2> position;
			math::vector<float, 2> penetration;
		};

		std::optional<collision> invert(std::optional<collision> c)
		{
			if (c)
				c->penetration *= -1.f;
			return c;
		}

		template <typename Points1, typename Normals1, typename Points2, typename Normals2>
		std::optional<collision> convex_collision(Points1 const & points1, Normals1 const & normals1, Points2 const & points2, Normals2 const & normals2)
		{
			static constexpr float inf = std::numeric_limits<float>::infinity();

			bool invert = false;
			math::point<float, 2> point{0.f, 0.f};
			float penetration = -inf;
			math::vector<float, 2> normal{0.f, 0.f};

			// Normals of 1st body against points of 2nd body
			{
				auto p = std::begin(points1);
				auto n = std::begin(normals1);

				for (; n != std::end(normals1); ++p, ++n)
				{
					float min = inf;
					math::point<float, 2> minp{0.f, 0.f};
					for (auto const & p2 : points2)
					{
						float v = math::dot(p2 - *p, *n);
						if (v < min)
						{
							min = v;
							minp = p2;
						}
					}

					if (min > 0.f)
						return std::nullopt;

					if (min > penetration)
					{
						point = minp;
						penetration = min;
						normal = *n;
					}
				}
			}

			// Normals of 2nd body against points of 1st body
			{
				auto p = std::begin(points2);
				auto n = std::begin(normals2);

				for (; n != std::end(normals2); ++p, ++n)
				{
					float min = inf;
					math::point<float, 2> minp{0.f, 0.f};
					for (auto const & p1 : points1)
					{
						float v = math::dot(p1 - *p, *n);
						if (v < min)
						{
							min = v;
							minp = p1;
						}
					}

					if (min > 0.f)
						return std::nullopt;

					if (min > penetration)
					{
						invert = true;
						point = minp;
						penetration = min;
						normal = *n;
					}
				}
			}

			return collision{point - normal * penetration / 2.f, (invert ? 1.f : -1.f) * normal * penetration};
		}

		std::optional<collision> shape_collision(ball const & b1, static_state const & s1, ball const & b2, static_state const & s2)
		{
			auto const d = s2.position - s1.position;
			float const l = math::length(d);
			if (l > b1.radius + b2.radius)
				return std::nullopt;

			auto const p = math::lerp(s1.position, s2.position, (b1.radius + l - b2.radius) / 2.f / l);
			return collision{p, d * (b1.radius + b2.radius - l) / l};
		}

		std::optional<collision> shape_collision(ball const & b1, static_state const & s1, half_space const & b2, static_state const &)
		{
			float v = b2.normal[0] * s1.position[0] + b2.normal[1] * s1.position[1] - b2.value - b1.radius;
			if (v > 0.f)
				return std::nullopt;

			return collision{s1.position - b2.normal * b1.radius, v * b2.normal};
		}

		std::optional<collision> shape_collision(ball const & b1, static_state const & s1, box const & b2, static_state const & s2)
		{
			math::vector ex{std::cos(s2.rotation), std::sin(s2.rotation)};
			math::vector ey{-std::sin(s2.rotation), std::cos(s2.rotation)};

			float w = b2.width / 2.f;
			float h = b2.height / 2.f;

			auto r = s1.position - s2.position;

			float x = math::dot(r, ex);
			float y = math::dot(r, ey);

			if (std::abs(x) < w && std::abs(y) < h)
			{
				// ball center is inside the box

				float vx = w - std::abs(x);
				float vy = h - std::abs(y);

				if (vx < vy)
				{
					if (x > 0.f)
					{
						return collision{s1.position, math::vector{-vx - b1.radius, 0.f}};
					}
					else
					{
						return collision{s1.position, math::vector{vx + b1.radius, 0.f}};
					}
				}
				else
				{
					if (y > 0.f)
					{
						return collision{s1.position, math::vector{0.f, -vy - b1.radius}};
					}
					else
					{
						return collision{s1.position, math::vector{0.f, vy + b1.radius}};
					}
				}
			}
			else
			{
				float cx = math::clamp(x, {-w, w});
				float cy = math::clamp(y, {-h, h});

				math::vector n = ex * (cx - x) + ey * (cy - y);

				float l = math::length(n);

				if (l < b1.radius)
				{
					float v = b1.radius - l;
					n = n / l;
					return collision{s1.position + n * (b1.radius - v / 2.f), n * v};
				}
			}


			return std::nullopt;
		}

		std::optional<collision> shape_collision(half_space const & b1, static_state const & s1, ball const & b2, static_state const & s2)
		{
			return invert(shape_collision(b2, s2, b1, s1));
		}

		std::optional<collision> shape_collision(half_space const &, static_state const &, half_space const &, static_state const &)
		{
			return std::nullopt;
		}

		std::optional<collision> shape_collision(half_space const & b1, static_state const &, box const & b2, static_state const & s2)
		{
			math::vector a{ b2.width / 2.f, b2.height / 2.f};
			math::vector b{-b2.width / 2.f, b2.height / 2.f};

			auto const rot = math::plane_rotation<float, 2>(0, 1, s2.rotation);
			a = rot(a);
			b = rot(b);

			float v = b1.normal[0] * s2.position[0] + b1.normal[1] * s2.position[1] - b1.value;

			math::vector<float, 2> d;

			float va = math::dot(a, b1.normal);
			float vb = math::dot(b, b1.normal);
			v -= std::max(std::abs(va), std::abs(vb));
			if (std::abs(va) > std::abs(vb))
			{
				if (va > 0.f)
					d = -a;
				else
					d = a;
			}
			else if (std::abs(vb) > std::abs(va))
			{
				if (vb > 0.f)
					d = -b;
				else
					d = b;
			}
			else
			{
				// multiple points of contact
				auto da = (va > 0.f) ? -a : a;
				auto db = (vb > 0.f) ? -b : b;
				d = da * 0.5f + db * 0.5f;
			}

			if (v > 0.f)
				return std::nullopt;
			return collision{s2.position + d, - v * b1.normal};
		}

		std::optional<collision> shape_collision(box const & b1, static_state const & s1, ball const & b2, static_state const & s2)
		{
			return invert(shape_collision(b2, s2, b1, s1));
		}

		std::optional<collision> shape_collision(box const & b1, static_state const & s1, half_space const & b2, static_state const & s2)
		{
			return invert(shape_collision(b2, s2, b1, s1));
		}

		std::optional<collision> shape_collision(box const & b1, static_state const & s1, box const & b2, static_state const & s2)
		{
			math::vector<float, 2> ex1{ std::cos(s1.rotation), std::sin(s1.rotation)};
			math::vector<float, 2> ey1{-std::sin(s1.rotation), std::cos(s1.rotation)};

			math::vector<float, 2> ex2{ std::cos(s2.rotation), std::sin(s2.rotation)};
			math::vector<float, 2> ey2{-std::sin(s2.rotation), std::cos(s2.rotation)};

			float w1 = b1.width  / 2.f;
			float h1 = b1.height / 2.f;
			float w2 = b2.width  / 2.f;
			float h2 = b2.height / 2.f;

			math::point<float, 2> points1[4] =
			{
				s1.position - ex1 * w1 - ey1 * h1,
				s1.position + ex1 * w1 - ey1 * h1,
				s1.position + ex1 * w1 + ey1 * h1,
				s1.position - ex1 * w1 + ey1 * h1,
			};

			math::vector<float, 2> normals1[4] =
			{
				-ey1,
				ex1,
				ey1,
				-ex1
			};

			math::point<float, 2> points2[4] =
			{
				s2.position - ex2 * w2 - ey2 * h2,
				s2.position + ex2 * w2 - ey2 * h2,
				s2.position + ex2 * w2 + ey2 * h2,
				s2.position - ex2 * w2 + ey2 * h2,
			};

			math::vector<float, 2> normals2[4] =
			{
				-ey2,
				ex2,
				ey2,
				-ex2
			};

			return convex_collision(points1, normals1, points2, normals2);
		}

		struct contact
		{
			std::size_t gi, i, gj, j;
			math::vector<float, 2> ri, rj;
			math::vector<float, 2> penetration;
			math::vector<float, 2> normal;
			float penetration_depth;
			math::vector<float, 2> velocity;
			float velocity_projection;
		};

		struct velocity_comparator
		{
			bool operator()(contact const & c1, contact const & c2) const
			{
				if (c1.velocity_projection == c2.velocity_projection)
					return std::tie(c1.gi, c1.gj, c1.i, c1.j) < std::tie(c2.gi, c2.gj, c2.i, c2.j);
				return c1.velocity_projection < c2.velocity_projection;
			}
		};

		struct penetration_comparator
		{
			bool operator()(contact const & c1, contact const & c2) const
			{
				if (c1.penetration_depth == c2.penetration_depth)
					return std::tie(c1.gi, c1.gj, c1.i, c1.j) < std::tie(c2.gi, c2.gj, c2.i, c2.j);
				return c1.penetration_depth > c2.penetration_depth;
			}
		};

	}

	struct engine::impl
	{
		util::heterogeneous_container<engine::shape_handle
			, wrapped_shape<ball>
			, wrapped_shape<half_space>
			, wrapped_shape<box>
//			, convex_polygon
			> shapes;

		util::flat_list<material, engine::material_handle> materials;

		struct constraint
		{
			std::function<void()> resolve;
			std::vector<std::pair<group_handle, std::size_t>> objects;
		};

		util::flat_list<constraint> constraints;
		std::vector<constraint_handle> constraint_handles;

		struct group
		{
			std::vector<object_info> infos;
			std::vector<std::vector<constraint_handle>> constraints;
			std::vector<static_state> static_states;
			std::vector<static_state> static_states_temp;
			std::vector<dynamic_state> dynamic_states;

			std::vector<std::vector<std::set<contact, velocity_comparator>::iterator>> contacts_by_velocity;
			std::vector<std::vector<std::set<contact, penetration_comparator>::iterator>> contacts_by_penetration;
		};

		std::vector<group> groups;

		math::vector<float, 2> gravity{0.f, 0.f};
		float air_friction = 0.f;

		std::set<contact, velocity_comparator> contacts_by_velocity;
		std::set<contact, penetration_comparator> contacts_by_penetration;

		void explode(math::point<float, 2> const & center, float strength, float attenuation);

		void update(float dt);

		void apply_gravity(float dt);
		void integrate(float dt);
		void resolve_constraints(float dt);
		void gather_collisions();
		void resolve_impulses();
		void resolve_penetrations();
		void apply_air_friction(float dt);

		void log_energy();
		float energy();
	};

	void engine::impl::explode(math::point<float, 2> const & center, float strength, float attenuation)
	{
		for (auto & g : groups)
		{
			for (std::size_t i = 0; i < g.infos.size(); ++i)
			{
				if (g.infos[i].inv_mass == 0.f) continue;

				auto const closest = shapes.visit([&](auto const & s) -> math::point<float, 2> { return shape_closest(s.shape, g.static_states[i], center); }, g.infos[i].shape);

				auto r = closest - center;
				float l = math::length(r);

				auto J = (r / l) * strength / (1.f + l * l * attenuation);

				g.dynamic_states[i].velocity += J * g.infos[i].inv_mass;
				g.dynamic_states[i].angular_velocity += math::det(closest - g.static_states[i].position, J) * g.infos[i].inv_inertia;
			}
		}
	}

	void engine::impl::update(float dt)
	{
		profile_function;

		apply_gravity(dt);
		integrate(dt);
		resolve_constraints(dt);
		gather_collisions();
		resolve_impulses();
		resolve_penetrations();
		apply_air_friction(dt);
	}

	void engine::impl::apply_gravity(float dt)
	{
		if (gravity != math::vector{0.f, 0.f})
		{
			for (auto & g : groups)
			{
				for (std::size_t i = 0; i < g.infos.size(); ++i)
				{
					if (g.infos[i].inv_mass == 0.f) continue;

					g.dynamic_states[i].velocity += gravity * dt;
				}
			}
		}
	}

	void engine::impl::integrate(float dt)
	{
		for (auto & g : groups)
		{
			for (std::size_t i = 0; i < g.infos.size(); ++i)
			{
				g.static_states[i].position += dt * g.dynamic_states[i].velocity;
				g.static_states[i].rotation += dt * g.dynamic_states[i].angular_velocity;
			}
		}
	}

	void engine::impl::resolve_constraints(float dt)
	{
		for (std::size_t gi = 0; gi < groups.size(); ++gi)
			for (std::size_t i = 0; i < groups[gi].infos.size(); ++i)
				groups[gi].static_states_temp[i] = groups[gi].static_states[i];
		(void)dt;

		for (int i = 0; i < 4; ++i)
		{
			for (auto h : constraint_handles)
			{
				constraints[h].resolve();
			}
		}

		for (std::size_t gi = 0; gi < groups.size(); ++gi)
		{
			for (std::size_t i = 0; i < groups[gi].infos.size(); ++i)
			{
				groups[gi].dynamic_states[i].velocity += (groups[gi].static_states[i].position - groups[gi].static_states_temp[i].position) / dt;
				groups[gi].dynamic_states[i].angular_velocity += (groups[gi].static_states[i].rotation - groups[gi].static_states_temp[i].rotation) / dt;
			}
		}
	}

	void engine::impl::gather_collisions()
	{
		profile_function;

		for (std::size_t gi = 0; gi < groups.size(); ++gi)
		{
			groups[gi].contacts_by_velocity.resize(groups[gi].infos.size());
			groups[gi].contacts_by_penetration.resize(groups[gi].infos.size());
			for (std::size_t i = 0; i < groups[gi].infos.size(); ++i)
			{
				groups[gi].contacts_by_velocity[i].clear();
				groups[gi].contacts_by_penetration[i].clear();
			}
		}

		contacts_by_velocity.clear();
		contacts_by_penetration.clear();

		auto test_objects = [&](std::size_t gi, std::size_t i, std::size_t gj, std::size_t j)
		{
			auto shi = groups[gi].infos[i].shape;
			auto shj = groups[gj].infos[j].shape;

			auto & sti = groups[gi].static_states[i];
			auto & stj = groups[gj].static_states[j];

			auto c = shapes.visit([&](auto const & si, auto const & sj){
				return shape_collision(si.shape, sti, sj.shape, stj);
			}, shi, shj);

			if (!c) return;

			if (c->penetration == math::vector<float, 2>::zero()) return;

			auto ri = c->position - sti.position;
			auto rj = c->position - stj.position;

			auto & dsi = groups[gi].dynamic_states[i];
			auto & dsj = groups[gj].dynamic_states[j];

			auto ui = dsi.velocity + math::ort(ri) * dsi.angular_velocity;
			auto uj = dsj.velocity + math::ort(rj) * dsj.angular_velocity;

			auto const u = uj - ui;

			contact ct;
			ct.gi = gi;
			ct.gj = gj;
			ct.i = i;
			ct.j = j;
			ct.ri = ri;
			ct.rj = rj;
			ct.penetration = c->penetration;
			ct.normal = math::normalized(ct.penetration);
			ct.penetration_depth = math::length(ct.penetration);
			ct.velocity = u;
			ct.velocity_projection = math::dot(u, ct.normal);

			auto res = contacts_by_velocity.insert(ct);
			assert(res.second);
			groups[gi].contacts_by_velocity[i].push_back(res.first);
			groups[gj].contacts_by_velocity[j].push_back(res.first);
		};

		struct bbox_data
		{
			math::box<float, 2> box;
			std::size_t g, i;
		};

		std::vector<bbox_data> bboxes;
		for (std::size_t gi = 0; gi < groups.size(); ++gi)
		{
			for (std::size_t i = 0; i < groups[gi].infos.size(); ++i)
			{
				bbox_data data;
				data.box = shapes.visit([&](auto const & s){ return shape_bbox(s.shape, groups[gi].static_states[i]); }, groups[gi].infos[i].shape);
				data.g = gi;
				data.i = i;
				bboxes.push_back(data);
			}
		}

		std::vector<std::size_t> min, max;
		min.reserve(bboxes.size());
		max.reserve(bboxes.size());

		for (std::size_t i = 0; i < bboxes.size(); ++i)
		{
			min.push_back(i);
			max.push_back(i);
		}

		std::sort(min.begin(), min.end(), [&](auto i, auto j){ return bboxes[i].box[0].min < bboxes[j].box[0].min; });
		std::sort(max.begin(), max.end(), [&](auto i, auto j){ return bboxes[i].box[0].max < bboxes[j].box[0].max; });

		std::set<std::size_t> current;

		for (std::size_t mini = 0, maxi = 0; mini < bboxes.size();)
		{
			if (bboxes[min[mini]].box[0].min < bboxes[max[maxi]].box[0].max)
			{
				std::size_t i = min[mini];
				for (auto j : current)
				{
					if (!(bboxes[i].box[1] & bboxes[j].box[1]).empty())
						test_objects(bboxes[i].g, bboxes[i].i, bboxes[j].g, bboxes[j].i);
				}
				current.insert(i);

				++mini;
			}
			else
			{
				current.erase(max[maxi]);
				++maxi;
			}
		}
	}

	void engine::impl::resolve_impulses()
	{
		profile_function;

		std::size_t const iterations = contacts_by_velocity.size() * 2;

		for (std::size_t iteration = 0; iteration < iterations; ++iteration)
		{
			contact const & c = *contacts_by_velocity.begin();

			if (c.velocity_projection > 0.f)
			{
				break;
			}

			auto const & infoi = groups[c.gi].infos[c.i];
			auto const & infoj = groups[c.gj].infos[c.j];

			auto const & mati = materials[infoi.material];
			auto const & matj = materials[infoj.material];

			math::matrix<float, 2, 2> K;
			K[0][0] = infoi.inv_mass + infoj.inv_mass;
			K[1][1] = K[0][0];
			K[0][1] = 0.f;
			K[1][0] = 0.f;

			K[0][0] += c.ri[1] * c.ri[1] * infoi.inv_inertia;
			K[1][1] += c.ri[0] * c.ri[0] * infoi.inv_inertia;
			K[0][1] -= c.ri[0] * c.ri[1] * infoi.inv_inertia;

			K[0][0] += c.rj[1] * c.rj[1] * infoj.inv_inertia;
			K[1][1] += c.rj[0] * c.rj[0] * infoj.inv_inertia;
			K[0][1] -= c.rj[0] * c.rj[1] * infoj.inv_inertia;

			K[1][0] = K[0][1];

			auto const & n = c.normal;

			// Plastic sliding impulse
			// Normal relative velocity -> 0
			// Tangential relative velocity -> unchanged
			auto const J1 = - n * c.velocity_projection / math::dot(n, K * n);
			// Plastic sticking impulse
			// Normal relative velocity -> 0
			// Tangential relative velocity -> 0
			auto const J2 = - *math::solve(K, c.velocity);

			float const e = std::sqrt(mati.elasticity * matj.elasticity);
			float const f = std::sqrt(mati.friction * matj.friction);
			float const mu = f;

			auto J = (1.f + e) * J1 + f * (J2 - J1);

			// Test for friction cone
			float q = math::length(J - n * math::dot(J, n));
			if (q > mu * math::dot(J, n))
			{
				float k = (mu * (1.f + e) * math::dot(J1, n)) / (math::length(J2 - n * math::dot(J2, n)) - mu * math::dot(n, J2 - J1));
				J = (1.f + e) * J1 + k * (J2 - J1);
			}

			auto dvi = -J * infoi.inv_mass;
			auto dvj =  J * infoj.inv_mass;

			auto dai = -math::det(c.ri, J) * infoi.inv_inertia;
			auto daj =  math::det(c.rj, J) * infoj.inv_inertia;

			groups[c.gi].dynamic_states[c.i].velocity += dvi;
			groups[c.gj].dynamic_states[c.j].velocity += dvj;

			groups[c.gi].dynamic_states[c.i].angular_velocity += dai;
			groups[c.gj].dynamic_states[c.j].angular_velocity += daj;

			auto ct = contacts_by_velocity.begin();

			std::size_t cti = -1, ctj = -1;

			auto update_contact = [&](auto old_it, std::size_t G, std::size_t I, std::size_t i, math::vector<float, 2> const & dv, float da)
			{
				contact const & c = *old_it;

				std::size_t OG, OI;
				bool flip;

				if (c.gi == G && c.i == I)
				{
					OG = c.gj;
					OI = c.j;
					flip = false;
				}
				else
				{
					OG = c.gi;
					OI = c.i;
					flip = true;
				}

				std::size_t oi;
				for (oi = 0; oi < groups[OG].contacts_by_velocity[OI].size(); ++oi)
				{
					if (groups[OG].contacts_by_velocity[OI][oi] == old_it)
						break;
				}

				auto node = contacts_by_velocity.extract(old_it);
				contact & cc = node.value();
				cc.velocity += (flip ? 1.f : -1.f) * (dv + math::ort(flip ? cc.rj : cc.ri) * da);
				cc.velocity_projection = math::dot(cc.velocity, cc.normal);

				auto res = contacts_by_velocity.insert(std::move(node));
				assert(res.inserted);
				auto new_it = res.position;

				groups[G].contacts_by_velocity[I][i] = new_it;
				groups[OG].contacts_by_velocity[OI][oi] = new_it;
			};

			for (std::size_t i = 0; i < groups[c.gi].contacts_by_velocity[c.i].size(); ++i)
			{
				auto old_it = groups[c.gi].contacts_by_velocity[c.i][i];
				if (old_it == ct)
				{
					cti = i;
					continue;
				}

				update_contact(old_it, c.gi, c.i, i, dvi, dai);
			}

			for (std::size_t i = 0; i < groups[c.gj].contacts_by_velocity[c.j].size(); ++i)
			{
				auto old_it = groups[c.gj].contacts_by_velocity[c.j][i];
				if (old_it == ct)
				{
					ctj = i;
					continue;
				}

				update_contact(old_it, c.gj, c.j, i, dvj, daj);
			}

			{
				auto node = contacts_by_velocity.extract(ct);
				contact & cc = node.value();
				cc.velocity -= dvi + math::ort(cc.ri) * dai;
				cc.velocity += dvj + math::ort(cc.rj) * daj;
				cc.velocity_projection = math::dot(cc.velocity, cc.normal);

				auto res = contacts_by_velocity.insert(std::move(node));
				assert(res.inserted);
				auto new_it = res.position;

				auto & ccc = *new_it;

				groups[ccc.gi].contacts_by_velocity[ccc.i][cti] = new_it;
				groups[ccc.gj].contacts_by_velocity[ccc.j][ctj] = new_it;
			}
		}
	}

	void engine::impl::resolve_penetrations()
	{
		profile_function;

		for (auto it = contacts_by_velocity.begin(); it != contacts_by_velocity.end();)
		{
			auto jt = it++;
			auto node = contacts_by_velocity.extract(jt);
			auto res = contacts_by_penetration.insert(std::move(node));
			assert(res.inserted);
			auto new_it = res.position;
			contact const & c = *new_it;
			groups[c.gi].contacts_by_penetration[c.i].push_back(new_it);
			groups[c.gj].contacts_by_penetration[c.j].push_back(new_it);
		}

		std::size_t const iterations = contacts_by_penetration.size() * 2;

		for (std::size_t iteration = 0; iteration < iterations; ++iteration)
		{
			contact const & c = *contacts_by_penetration.begin();

			if (c.penetration_depth < 0.f)
				break;

			auto const & infoi = groups[c.gi].infos[c.i];
			auto const & infoj = groups[c.gj].infos[c.j];

			auto di = - c.penetration * infoi.inv_mass / (infoi.inv_mass + infoj.inv_mass);
			auto dj =   c.penetration * infoj.inv_mass / (infoi.inv_mass + infoj.inv_mass);

			groups[c.gi].static_states[c.i].position += di;
			groups[c.gj].static_states[c.j].position += dj;

			auto ct = contacts_by_penetration.begin();

			std::size_t cti = -1, ctj = -1;

			auto update_contact = [&](auto old_it, std::size_t G, std::size_t I, std::size_t i, math::vector<float, 2> const & d)
			{
				contact const & c = *old_it;

				std::size_t OG, OI;
				bool flip;

				if (c.gi == G && c.i == I)
				{
					OG = c.gj;
					OI = c.j;
					flip = false;
				}
				else
				{
					OG = c.gi;
					OI = c.i;
					flip = true;
				}

				std::size_t oi;
				for (oi = 0; oi < groups[OG].contacts_by_penetration[OI].size(); ++oi)
				{
					if (groups[OG].contacts_by_penetration[OI][oi] == old_it)
						break;
				}

				auto node = contacts_by_penetration.extract(old_it);
				contact & cc = node.value();
				cc.penetration += (flip ? -1.f : 1.f) * d;
				cc.penetration_depth = math::dot(cc.penetration, cc.normal);

				auto res = contacts_by_penetration.insert(std::move(node));
				assert(res.inserted);
				auto new_it = res.position;

				groups[G].contacts_by_penetration[I][i] = new_it;
				groups[OG].contacts_by_penetration[OI][oi] = new_it;
			};

			for (std::size_t i = 0; i < groups[c.gi].contacts_by_penetration[c.i].size(); ++i)
			{
				auto old_it = groups[c.gi].contacts_by_penetration[c.i][i];
				if (old_it == ct)
				{
					cti = i;
					continue;
				}

				update_contact(old_it, c.gi, c.i, i, di);
			}

			for (std::size_t i = 0; i < groups[c.gj].contacts_by_penetration[c.j].size(); ++i)
			{
				auto old_it = groups[c.gj].contacts_by_penetration[c.j][i];
				if (old_it == ct)
				{
					ctj = i;
					continue;
				}

				update_contact(old_it, c.gj, c.j, i, dj);
			}

			{
				auto node = contacts_by_penetration.extract(ct);
				contact & cc = node.value();
				cc.penetration += di;
				cc.penetration -= dj;
				cc.penetration_depth = math::dot(cc.penetration, cc.normal);

				auto res = contacts_by_penetration.insert(std::move(node));
				assert(res.inserted);
				auto new_it = res.position;

				auto & ccc = *new_it;

				groups[ccc.gi].contacts_by_penetration[ccc.i][cti] = new_it;
				groups[ccc.gj].contacts_by_penetration[ccc.j][ctj] = new_it;
			}
		}
	}

	void engine::impl::apply_air_friction(float dt)
	{
		if (air_friction != 0.f)
		{
			float const factor = 1.f - std::min(1.f, air_friction * dt);

			for (auto & g : groups)
			{
				for (std::size_t i = 0; i < g.infos.size(); ++i)
				{
					if (g.infos[i].inv_mass == 0.f) continue;

					g.dynamic_states[i].velocity *= factor;
					g.dynamic_states[i].angular_velocity *= factor;
				}
			}
		}
	}

	float engine::impl::energy()
	{
		float Eg = 0.f;
		float Er = 0.f;
		float Ek = 0.f;

		for (auto & g : groups)
		{
			for (std::size_t i = 0; i < g.infos.size(); ++i)
			{
				if (g.infos[i].inv_mass == 0.f) continue;

				Eg -= math::dot(gravity, g.static_states[i].position - math::point<float, 2>::zero()) / g.infos[i].inv_mass;

				Ek += math::length_sqr(g.dynamic_states[i].velocity) / g.infos[i].inv_mass / 2.f;
				Er += math::sqr(g.dynamic_states[i].angular_velocity) / g.infos[i].inv_inertia / 2.f;
			}
		}

		return Eg + Ek + Er;
	}

	engine::engine()
		: pimpl_{make_impl()}
	{}

	engine::~engine() = default;

	engine::shape_handle engine::add_shape(ball const & b)
	{
		return impl().shapes.insert(wrap(b));
	}

	engine::shape_handle engine::add_shape(half_space const & s)
	{
		return impl().shapes.insert(wrap(s));
	}

	engine::shape_handle engine::add_shape(box const & b)
	{
		return impl().shapes.insert(wrap(b));
	}

	std::variant<ball const *, half_space const *, box const *> engine::get_shape(shape_handle handle) const
	{
		using result_type = std::variant<ball const *, half_space const *, box const *>;
		return impl().shapes.visit([](auto const & w) -> result_type {
			return &(w.shape);
		}, handle);
	}

	void engine::remove_shape(shape_handle handle)
	{
		impl().shapes.erase(handle);
	}

	engine::material_handle engine::add_material(material const & m)
	{
		return impl().materials.insert(m);
	}

	material const & engine::get_material(material_handle handle) const
	{
		return impl().materials[handle];
	}

	material & engine::get_material(material_handle handle)
	{
		return impl().materials[handle];
	}

	void engine::remove_material(material_handle handle)
	{
		impl().materials.erase(handle);
	}

	engine::group_handle engine::create_group()
	{
		auto h = static_cast<group_handle>(impl().groups.size());
		impl().groups.emplace_back();
		return h;
	}

	void engine::remove_group(group_handle handle)
	{
		auto & g = impl().groups[handle];
		g.infos.clear();
		g.static_states.clear();
		g.dynamic_states.clear();
	}

	std::size_t engine::group_size(group_handle handle) const
	{
		return impl().groups[handle].infos.size();
	}

	static_state const * engine::group_static_state(group_handle handle) const
	{
		return impl().groups[handle].static_states.data();
	}

	dynamic_state const * engine::group_dynamic_state(group_handle handle) const
	{
		return impl().groups[handle].dynamic_states.data();
	}

	static_state * engine::group_static_state(group_handle handle)
	{
		return impl().groups[handle].static_states.data();
	}

	dynamic_state * engine::group_dynamic_state(group_handle handle)
	{
		return impl().groups[handle].dynamic_states.data();
	}

	std::size_t engine::add_object(group_handle group, shape_handle shape, material_handle material, static_state const & ss, dynamic_state const & ds)
	{
		float const density = impl().materials[material].density;
		float const area = impl().shapes.visit([](auto const & s){ return s.area; }, shape);
		float const inertia = impl().shapes.visit([](auto const & s){ return s.inertia; }, shape);

		auto & g = impl().groups[group];
		g.infos.push_back({shape, material, 1.f / (area * density), 1.f / (inertia * density)});
		g.constraints.emplace_back();
		g.static_states.push_back(ss);
		g.static_states_temp.emplace_back();
		g.dynamic_states.push_back(ds);

		return g.infos.size() - 1;
	}

	void engine::remove_object(group_handle group, std::size_t index)
	{
		auto & g = impl().groups[group];

		for (auto h : g.constraints[index])
			remove_constraint(h);

		g.infos.erase(g.infos.begin() + index);
		g.constraints.erase(g.constraints.begin() + index);
		g.static_states.erase(g.static_states.begin() + index);
		g.static_states_temp.erase(g.static_states_temp.begin() + index);
		g.dynamic_states.erase(g.dynamic_states.begin() + index);
	}

	engine::object_info const & engine::get_object(group_handle group, std::size_t index)
	{
		return impl().groups[group].infos[index];
	}

	engine::constraint_handle engine::add_constraint(group_handle g, std::size_t index, fixed_constraint const & c)
	{
		impl::constraint constraint;
		constraint.objects.push_back({g, index});
		constraint.resolve = [=, this]{
			auto & ss = impl().groups[g].static_states[index];

			math::vector<float, 2> t;
			t[0] = c.target[0] - ss.position[0] - c.point[0] * std::cos(ss.rotation) + c.point[1] * std::sin(ss.rotation);
			t[1] = c.target[1] - ss.position[1] - c.point[0] * std::sin(ss.rotation) - c.point[1] * std::cos(ss.rotation);

			math::matrix<float, 2, 3> J;
			J[0][0] = -1.f;
			J[0][1] = 0.f;
			J[0][2] = c.point[0] * std::sin(ss.rotation) + c.point[1] * std::cos(ss.rotation);
			J[1][0] = 0.f;
			J[1][1] = -1.f;
			J[1][2] = - c.point[0] * std::cos(ss.rotation) + c.point[1] * std::sin(ss.rotation);

			auto delta = math::least_squares(J, -t);

			if (delta)
			{
				ss.position[0] += (*delta)[0];
				ss.position[1] += (*delta)[1];
				ss.rotation += (*delta)[2];
			}
		};

		auto handle = impl().constraints.insert(std::move(constraint));
		impl().constraint_handles.push_back(handle);
		impl().groups[g].constraints[index].push_back(handle);
		return handle;
	}

	engine::constraint_handle engine::add_constraint(group_handle g1, std::size_t index1, group_handle g2, std::size_t index2, joint_constraint const & c)
	{
		impl::constraint constraint;
		constraint.objects.push_back({g1, index1});
		constraint.objects.push_back({g2, index2});
		constraint.resolve = [=, this]{
			auto & ss1 = impl().groups[g1].static_states[index1];
			auto & ss2 = impl().groups[g2].static_states[index2];

			math::vector<float, 2> t;
			t[0] = ss2.position[0] + c.point2[0] * std::cos(ss2.rotation) - c.point2[1] * std::sin(ss2.rotation) - ss1.position[0] - c.point1[0] * std::cos(ss1.rotation) + c.point1[1] * std::sin(ss1.rotation);
			t[1] = ss2.position[1] + c.point2[0] * std::sin(ss2.rotation) + c.point2[1] * std::cos(ss2.rotation) - ss1.position[1] - c.point1[0] * std::sin(ss1.rotation) - c.point1[1] * std::cos(ss1.rotation);

			math::matrix<float, 2, 6> J;
			J[0][0] = -1.f;
			J[0][1] = 0.f;
			J[0][2] = c.point1[0] * std::sin(ss1.rotation) + c.point1[1] * std::cos(ss1.rotation);
			J[0][3] = 1.f;
			J[0][4] = 0.f;
			J[0][5] = - c.point2[0] * std::sin(ss2.rotation) - c.point2[1] * std::cos(ss2.rotation);
			J[1][0] = 0.f;
			J[1][1] = -1.f;
			J[1][2] = - c.point1[0] * std::cos(ss1.rotation) + c.point1[1] * std::sin(ss1.rotation);
			J[1][3] = 0.f;
			J[1][4] = 1.f;
			J[1][5] = c.point2[0] * std::cos(ss2.rotation) - c.point2[1] * std::sin(ss2.rotation);

			auto delta = math::least_squares(J, -t);

			if (delta)
			{
				ss1.position[0] += (*delta)[0];
				ss1.position[1] += (*delta)[1];
				ss1.rotation += (*delta)[2];
				ss2.position[0] += (*delta)[3];
				ss2.position[1] += (*delta)[4];
				ss2.rotation += (*delta)[5];
			}
		};

		auto handle = impl().constraints.insert(std::move(constraint));
		impl().constraint_handles.push_back(handle);
		impl().groups[g1].constraints[index1].push_back(handle);
		impl().groups[g2].constraints[index2].push_back(handle);
		return handle;
	}

	void engine::remove_constraint(constraint_handle handle)
	{
		auto const & constraint = impl().constraints[handle];
		for (auto [g, i] : constraint.objects)
		{
			auto & cs = impl().groups[g].constraints[i];
			if (auto it = std::find(cs.begin(), cs.end(), handle); it != cs.end())
				cs.erase(it);
		}
		impl().constraints.erase(handle);
		impl().constraint_handles.erase(std::find(impl().constraint_handles.begin(), impl().constraint_handles.end(), handle));
	}

	void engine::set_gravity(math::vector<float, 2> const & g)
	{
		impl().gravity = g;
	}

	void engine::set_air_friction(float friction)
	{
		impl().air_friction = friction;
	}

	void engine::explode(math::point<float, 2> const & center, float strength, float attenuation)
	{
		impl().explode(center, strength, attenuation);
	}

	void engine::update(float dt)
	{
		impl().update(dt);
	}

}
