#include <psemek/phys/engine_2d.hpp>

#include <psemek/util/heterogeneous_container.hpp>

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

#include <psemek/util/profiler.hpp>

#include <psemek/log/log.hpp>

#include <optional>

namespace psemek::phys2d
{

	namespace
	{

		float shape_area(ball const & b)
		{
			return geom::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;
		}

		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
		{
			geom::point<float, 2> position;
			geom::vector<float, 2> penetration;
		};

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

		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 = geom::length(d);
			if (l > b1.radius + b2.radius)
				return std::nullopt;

			auto const p = geom::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)
		{
			geom::vector ex{std::cos(s2.rotation), std::sin(s2.rotation)};
			geom::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 = geom::dot(r, ex);
			float y = geom::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, geom::vector{-vx - b1.radius, 0.f}};
					}
					else
					{
						return collision{s1.position, geom::vector{vx + b1.radius, 0.f}};
					}
				}
				else
				{
					if (y > 0.f)
					{
						return collision{s1.position, geom::vector{0.f, -vy - b1.radius}};
					}
					else
					{
						return collision{s1.position, geom::vector{0.f, vy + b1.radius}};
					}
				}
			}
			else
			{
				float cx = geom::clamp(x, {-w, w});
				float cy = geom::clamp(y, {-h, h});

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

				float l = geom::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)
		{
			geom::vector a{ b2.width / 2.f, b2.height / 2.f};
			geom::vector b{-b2.width / 2.f, b2.height / 2.f};

			auto const rot = geom::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;

			geom::vector<float, 2> d;

			float va = geom::dot(a, b1.normal);
			float vb = geom::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)
		{
			geom::vector<float, 2> ex1{ std::cos(s1.rotation), std::sin(s1.rotation)};
			geom::vector<float, 2> ey1{-std::sin(s1.rotation), std::cos(s1.rotation)};

			geom::vector<float, 2> ex2{ std::cos(s2.rotation), std::sin(s2.rotation)};
			geom::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;

			float dxx = geom::dot(ex1, ex2);
			float dxy = geom::dot(ex1, ey2);
			float dyx = geom::dot(ey1, ex2);
			float dyy = geom::dot(ey1, ey2);

			// SAT collision detector, hand-optimized for boxes

			geom::vector<float, 2> normal{0.f, 0.f};
			float penetration = -std::numeric_limits<float>::infinity();
			geom::point<float, 2> point{0.f, 0.f};

			auto side = [&](geom::vector<float, 2> const & n, geom::point<float, 2> const & a, float vx, float vy, geom::point<float, 2> const & c, geom::vector<float, 2> const & ex, geom::vector<float, 2> const & ey, bool invert)
			{
				float vc = geom::dot((c - a), n);

				float v = vc - std::abs(vx) - std::abs(vy);

				if (v > penetration)
				{
					penetration = v;
					point = c - geom::sign(vx) * ex - geom::sign(vy) * ey;
					normal = invert ? -n : n;
				}
			};

			// dot(n, c +/- ex +/- ey) = dot(n, c) +/- dot(n, ex) +/- dot(n, ey)

			// box2 vs box1 +x
			side(ex1, s1.position + ex1 * w1, dxx * w2, dxy * h2, s2.position, ex2 * w2, ey2 * h2, false);
			if (penetration > 0.f) return std::nullopt;
			// box2 vs box1 -x
			side(-ex1, s1.position - ex1 * w1, -dxx * w2, -dxy * h2, s2.position, ex2 * w2, ey2 * h2, false);
			if (penetration > 0.f) return std::nullopt;
			// box2 vs box1 +y
			side(ey1, s1.position + ey1 * h1, dyx * w2, dyy * h2, s2.position, ex2 * w2, ey2 * h2, false);
			if (penetration > 0.f) return std::nullopt;
			// box2 vs box1 -y
			side(-ey1, s1.position - ey1 * h1, -dyx * w2, -dyy * h2, s2.position, ex2 * w2, ey2 * h2, false);
			if (penetration > 0.f) return std::nullopt;

			// box1 vs box2 +x
			side(ex2, s2.position + ex2 * w2, dxx * w1, dyx * h1, s1.position, ex1 * w1, ey1 * h1, true);
			if (penetration > 0.f) return std::nullopt;
			// box1 vs box2 -x
			side(-ex2, s2.position - ex2 * w2, -dxx * w1, -dyx * h1, s1.position, ex1 * w1, ey1 * h1, true);
			if (penetration > 0.f) return std::nullopt;
			// box1 vs box2 +y
			side(ey2, s2.position + ey2 * h2, dxy * w1, dyy * h1, s1.position, ex1 * w1, ey1 * h1, true);
			if (penetration > 0.f) return std::nullopt;
			// box1 vs box2 -y
			side(-ey2, s2.position - ey2 * h2, -dxy * w1, -dyy * h1, s1.position, ex1 * w1, ey1 * h1, true);
			if (penetration > 0.f) return std::nullopt;

			return collision{point, -normal * penetration};
		}

	}

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

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

		struct object_info
		{
			engine::shape_handle shape;
			engine::material_handle material;
			float inv_mass;
			float inv_inertia;
		};

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

			std::vector<geom::vector<float, 2>> position_change;
			std::vector<geom::vector<float, 2>> velocity_change;
			std::vector<float> angular_velocity_change;
		};

		std::vector<group> groups;

		std::optional<geom::vector<float, 2>> gravity;

		void update(float dt);

		void apply_gravity(float dt);
		void integrate(float dt);
		bool resolve_collisions();

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

	void engine::impl::update(float dt)
	{
		apply_gravity(dt);
		integrate(dt);
		resolve_collisions();
//		log_energy();
	}

	void engine::impl::apply_gravity(float dt)
	{
		if (gravity)
		{
			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;
			}
		}
	}

	bool engine::impl::resolve_collisions()
	{
		bool had_collision = false;

		for (std::size_t gi = 0; gi < groups.size(); ++gi)
		{
			for (std::size_t gj = gi; gj < groups.size(); ++gj)
			{
				for (std::size_t i = 0; i < groups[gi].infos.size(); ++i)
				{
					for (std::size_t j = (gi == gj) ? i + 1 : 0; j < groups[gj].infos.size(); ++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) continue;

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

						// TODO: handle inv_mass == 0

						auto const n = geom::normalized(c->penetration);

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

						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 + geom::ort(ri) * dsi.angular_velocity;
						auto uj = dsj.velocity + geom::ort(rj) * dsj.angular_velocity;

						auto const u = uj - ui;

						if (geom::dot(u, n) > 0.f) continue;

						had_collision = true;

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

						geom::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] += ri[1] * ri[1] * infoi.inv_inertia;
						K[1][1] += ri[0] * ri[0] * infoi.inv_inertia;
						K[0][1] -= ri[0] * ri[1] * infoi.inv_inertia;

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

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

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

//						float const e = (mati.elasticity + matj.elasticity) / 2.f;
						float const e = std::sqrt(mati.elasticity * matj.elasticity);

//						float const f = (mati.friction + matj.friction) / 2.f;
						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 = geom::length(J - n * geom::dot(J, n));
						if (q > mu * geom::dot(J, n))
						{
							float k = (mu * (1.f + e) * geom::dot(J1, n)) / (geom::length(J2 - n * geom::dot(J2, n)) - mu * geom::dot(n, J2 - J1));
							J = (1.f + e) * J1 + k * (J2 - J1);
						}

						groups[gi].dynamic_states[i].velocity -= J * infoi.inv_mass;
						groups[gj].dynamic_states[j].velocity += J * infoj.inv_mass;

						groups[gi].dynamic_states[i].angular_velocity -= geom::det(ri, J) * infoi.inv_inertia;
						groups[gj].dynamic_states[j].angular_velocity += geom::det(rj, J) * infoj.inv_inertia;

						groups[gi].static_states[i].position -= c->penetration * infoi.inv_mass / (infoi.inv_mass + infoj.inv_mass);
						groups[gj].static_states[j].position += c->penetration * infoj.inv_mass / (infoi.inv_mass + infoj.inv_mass);
					}
				}
			}
		}

		return !had_collision;
	}

	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;

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

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

		return Eg + Ek + Er;
	}

	void engine::impl::log_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;

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

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

		log::info() << "G " << Eg << "   K " << Ek << "   R " << Er;
	}

	engine::engine()
		: pimpl_{std::make_unique<struct 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));
	}

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

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

	void 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.static_states.push_back(ss);
		g.dynamic_states.push_back(ds);
	}

	void engine::remove_object(group_handle group, std::size_t index)
	{
		auto & g = impl().groups[group];
		g.infos.erase(g.infos.begin() + index);
		g.static_states.erase(g.static_states.begin() + index);
		g.dynamic_states.erase(g.dynamic_states.begin() + index);
	}

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

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

}