#include <psemek/phys/engine_3d.hpp>

#include <psemek/geom/gauss.hpp>
#include <psemek/geom/rotation.hpp>
#include <psemek/geom/box.hpp>

#include <psemek/util/flat_list.hpp>
#include <psemek/util/function.hpp>
#include <psemek/util/not_implemented.hpp>
#include <psemek/util/unused.hpp>

namespace psemek::phys3d
{

	namespace
	{

		constexpr float inf = std::numeric_limits<float>::infinity();

		struct shape
		{
			engine::any_shape shape;
			float volume;
			geom::vector<float, 3> inertia;
		};

		struct object
		{
			engine::shape_handle shape;
			engine::material_handle material;
			object_state state;

			float inv_mass;
			geom::vector<float, 3> inertia;

			geom::matrix<float, 3, 3> global_inv_inertia = {};
		};

		float volume(ball const & s)
		{
			return 4.f / 3.f * geom::pi * s.radius * s.radius * s.radius;
		}

		float volume(box const & s)
		{
			return s.dimensions[0] * s.dimensions[1] * s.dimensions[2];
		}

		float volume(half_space const &)
		{
			return inf;
		}

		geom::vector<float, 3> inertia(ball const & s)
		{
			float const I = 0.4f * s.radius * s.radius;
			return {I, I, I};
		}

		geom::vector<float, 3> inertia(box const & s)
		{
			float const x = s.dimensions[0] * s.dimensions[0] / 12.f;
			float const y = s.dimensions[1] * s.dimensions[1] / 12.f;
			float const z = s.dimensions[2] * s.dimensions[2] / 12.f;

			return {y + z, x + z, x + y};
		}

		geom::vector<float, 3> inertia(half_space const &)
		{
			return {inf, inf, inf};
		}

		struct collision_data
		{
			geom::vector<float, 3> point1, point2;
			geom::vector<float, 3> normal;
			// depth should be equal to dot(position 2 + point 2 - position1 - point1, normal)
			float depth;
		};

		collision_data reversed(collision_data d)
		{
			std::swap(d.point1, d.point2);
			d.normal = -d.normal;
			return d;
		}

		std::optional<collision_data> reversed(std::optional<collision_data> const & d)
		{
			if (d)
				return reversed(*d);
			return std::nullopt;
		}

		std::optional<collision_data> collision(object_state const & state1, ball const & shape1, object_state const & state2, ball const & shape2)
		{
			auto r = state2.position - state1.position;
			auto l = geom::length(r);
			if (l >= shape1.radius + shape2.radius)
				return std::nullopt;

			auto n = r / l;
			auto p1 =  n * shape1.radius;
			auto p2 = -n * shape2.radius;

			return collision_data{p1, p2, n, l - shape1.radius - shape2.radius};
		}

		std::optional<collision_data> collision(object_state const & state1, ball const & shape1, object_state const & state2, box const & shape2)
		{
			unused(state1);
			unused(shape1);
			unused(state2);
			unused(shape2);
			util::not_implemented();
		}

		std::optional<collision_data> collision(object_state const & state1, ball const & shape1, object_state const & /* state2 */, half_space const & shape2)
		{
			auto v = geom::dot(shape2.normal, state1.position - state1.position.zero()) - shape2.value;

			if (v >= shape1.radius)
				return std::nullopt;

			return collision_data{- shape2.normal * shape1.radius, {0.f, 0.f, 0.f}, -shape2.normal, v - shape1.radius};
		}

		std::optional<collision_data> collision(object_state const & state1, box const & shape1, object_state const & state2, ball const & shape2)
		{
			return reversed(collision(state2, shape2, state1, shape1));
		}

		std::optional<collision_data> collision(object_state const & state1, box const & shape1, object_state const & state2, box const & shape2)
		{
			float const C = 1000.f;
			float const slop = -0.001f;
			(void)C;

			geom::vector<float, 3> normals[15];

			normals[0] = geom::rotate(state1.rotation, {1.f, 0.f, 0.f});
			normals[1] = geom::rotate(state1.rotation, {0.f, 1.f, 0.f});
			normals[2] = geom::rotate(state1.rotation, {0.f, 0.f, 1.f});

			normals[3] = geom::rotate(state2.rotation, {1.f, 0.f, 0.f});
			normals[4] = geom::rotate(state2.rotation, {0.f, 1.f, 0.f});
			normals[5] = geom::rotate(state2.rotation, {0.f, 0.f, 1.f});

			normals[ 6] = geom::cross(normals[0], normals[3]);
			normals[ 7] = geom::cross(normals[0], normals[4]);
			normals[ 8] = geom::cross(normals[0], normals[5]);
			normals[ 9] = geom::cross(normals[1], normals[3]);
			normals[10] = geom::cross(normals[1], normals[4]);
			normals[11] = geom::cross(normals[1], normals[5]);
			normals[12] = geom::cross(normals[2], normals[3]);
			normals[13] = geom::cross(normals[2], normals[4]);
			normals[14] = geom::cross(normals[2], normals[5]);

			auto tx1 = normals[0] * shape1.dimensions[0] / 2.f;
			auto ty1 = normals[1] * shape1.dimensions[1] / 2.f;
			auto tz1 = normals[2] * shape1.dimensions[2] / 2.f;

			auto tx2 = normals[3] * shape2.dimensions[0] / 2.f;
			auto ty2 = normals[4] * shape2.dimensions[1] / 2.f;
			auto tz2 = normals[5] * shape2.dimensions[2] / 2.f;

			geom::vector<float, 3> max_n{0.f, 0.f, 0.f};
			float max_v = -inf;
			geom::vector<float, 3> max_p1{0.f, 0.f, 0.f};
			geom::vector<float, 3> max_p2{0.f, 0.f, 0.f};

			static const auto zero = geom::point<float, 3>::zero();

			for (auto & n : normals)
			{
				auto nl = geom::length(n);
				if (nl == 0.f) continue;
				n /= nl;

				geom::interval<float> i1, i2;

				geom::vector<float, 3> min_pp1, max_pp1;
				geom::vector<float, 3> min_pp2, max_pp2;

				for (float x : {-1.f, 1.f})
				{
					for (float y : {-1.f, 1.f})
					{
						for (float z : {-1.f, 1.f})
						{
							auto p = tx1 * x + ty1 * y + tz1 * z;
							float v = geom::dot(n, state1.position + p - zero);
							if (v < i1.min)
							{
								i1.min = v;
								min_pp1 = p;
							}
							if (v > i1.max)
							{
								i1.max = v;
								max_pp1 = p;
							}
						}
					}
				}

				for (float x : {-1.f, 1.f})
				{
					for (float y : {-1.f, 1.f})
					{
						for (float z : {-1.f, 1.f})
						{
							auto p = tx2 * x + ty2 * y + tz2 * z;
							float v = geom::dot(n, state2.position + p - zero);
							if (v < i2.min)
							{
								i2.min = v;
								min_pp2 = p;
							}
							if (v > i2.max)
							{
								i2.max = v;
								max_pp2 = p;
							}
						}
					}
				}

				auto i = (i1 & i2);

				float v = -i.length();

				if (v > max_v)
				{
					max_v = v;

					if (i1.center() < i2.center())
					{
						max_n = n;
						max_p1 = max_pp1;
						max_p2 = min_pp2;
					}
					else
					{
						max_n = -n;
						max_p1 = min_pp1;
						max_p2 = max_pp2;
					}
				}

				if (max_v >= slop)
					break;
			}

			if (max_v >= slop)
				return std::nullopt;

			return collision_data{max_p1, max_p2, max_n, max_v - slop};
		}

		std::optional<collision_data> collision(object_state const & state1, box const & shape1, object_state const & /* state2 */, half_space const & shape2)
		{
			float min_v = inf;

			for (float tx : {-1.f, 1.f})
			{
				for (float ty : {-1.f, 1.f})
				{
					for (float tz : {-1.f, 1.f})
					{
						geom::vector<float, 3> t{tx, ty, tz};
						t = geom::rotate(state1.rotation, geom::pointwise_mult(t, shape1.dimensions / 2.f));

						auto p = state1.position + t;

						auto v = geom::dot(shape2.normal, p - p.zero()) - shape2.value;

						if (v < min_v)
							min_v = v;
					}
				}
			}

			float const C = 1000.f;
			float const slop = -0.001f;

			geom::vector<float, 3> sum_r{0.f, 0.f, 0.f};
			float sum_w = 0.f;

			for (float tx : {-1.f, 1.f})
			{
				for (float ty : {-1.f, 1.f})
				{
					for (float tz : {-1.f, 1.f})
					{
						geom::vector<float, 3> t{tx, ty, tz};
						t = geom::rotate(state1.rotation, geom::pointwise_mult(t, shape1.dimensions / 2.f));

						auto p = state1.position + t;

						auto v = geom::dot(shape2.normal, p - p.zero()) - shape2.value;

						float w = std::exp(- C * (v - min_v));

						sum_w += w;
						sum_r += w * t;
					}
				}
			}

			if (min_v >= slop)
				return std::nullopt;

			return collision_data{sum_r / sum_w, {0.f, 0.f, 0.f}, -shape2.normal, min_v - slop};
		}

		std::optional<collision_data> collision(object_state const & state1, half_space const & shape1, object_state const & state2, ball const & shape2)
		{
			return reversed(collision(state2, shape2, state1, shape1));
		}

		std::optional<collision_data> collision(object_state const & state1, half_space const & shape1, object_state const & state2, box const & shape2)
		{
			return reversed(collision(state2, shape2, state1, shape1));
		}

		std::optional<collision_data> collision(object_state const & /* state1 */, half_space const & /* shape1 */, object_state const & /* state2 */, half_space const & /* shape2 */)
		{
			return std::nullopt;
		}

		std::optional<collision_data> collision(object_state const & state1, engine::any_shape const & shape1, object_state const & state2, engine::any_shape const & shape2)
		{
			return std::visit([&](auto const & s1, auto const & s2){ return collision(state1, s1, state2, s2); }, shape1, shape2);
		}

	}

	struct engine::impl
	{
		util::flat_list<shape, shape_handle> shapes;
		util::flat_list<material, material_handle> materials;
		util::flat_list<object, object_handle> objects;
		std::vector<bool> object_mask;

		geom::vector<float, 3> gravity{0.f, 0.f, 0.f};
		float collision_bias = 0.25f;

		template <typename Shape>
		shape_handle add_shape(Shape const & s);

		bool is_object(object_handle h) const;

		using solver = util::function<void()>;

		template <typename Value, typename Jacobian, typename Limits, std::size_t N>
		solver make_solver(Value && value_fn, Jacobian && jacobian_fn, Limits && limits_fn, std::array<object_handle, N> const & objects);

		solver make_collision_solver(object_handle h1, object_handle h2, collision_data const & c, float dt);

		void update(float dt);
	};

	template <typename Shape>
	engine::shape_handle engine::impl::add_shape(Shape const & s)
	{
		return shapes.insert({s, volume(s), inertia(s)});
	}

	bool engine::impl::is_object(object_handle h) const
	{
		return h < objects.capacity() && object_mask[h];
	}

	template <typename Value, typename Jacobian, typename Limits, std::size_t N>
	engine::impl::solver engine::impl::make_solver(Value && value_fn, Jacobian && jacobian_fn, Limits && limits_fn, std::array<object_handle, N> const & handles)
	{
		return [this, value_fn = std::move(value_fn), jacobian_fn = std::move(jacobian_fn), limits_fn = std::move(limits_fn), handles]{

			auto f = value_fn();
			auto J = jacobian_fn();
			auto limits = limits_fn();

			auto M_inv_J_T = geom::transpose(J);

			geom::vector<float, 6 * N> v;
			for (std::size_t i = 0; i < N; ++i)
			{
				object const & obj = objects[handles[i]];
				v[6 * i + 0] = obj.state.velocity[0];
				v[6 * i + 1] = obj.state.velocity[1];
				v[6 * i + 2] = obj.state.velocity[2];
				v[6 * i + 3] = obj.state.angular_velocity[0];
				v[6 * i + 4] = obj.state.angular_velocity[1];
				v[6 * i + 5] = obj.state.angular_velocity[2];

				for (std::size_t j = 0; j < M_inv_J_T.columns(); ++j)
				{
					M_inv_J_T[6 * i + 0][j] *= obj.inv_mass;
					M_inv_J_T[6 * i + 1][j] *= obj.inv_mass;
					M_inv_J_T[6 * i + 2][j] *= obj.inv_mass;

					geom::vector<float, 3> r;
					r[0] = M_inv_J_T[6 * i + 3][j];
					r[1] = M_inv_J_T[6 * i + 4][j];
					r[2] = M_inv_J_T[6 * i + 5][j];

					r = obj.global_inv_inertia * r;

					M_inv_J_T[6 * i + 3][j] = r[0];
					M_inv_J_T[6 * i + 4][j] = r[1];
					M_inv_J_T[6 * i + 5][j] = r[2];
				}
			}

			auto lambda = geom::solve(J * M_inv_J_T, - J * v - f);

			if (!lambda)
				return;

			for (std::size_t i = 0; i < lambda->dimension(); ++i)
				(*lambda)[i] = geom::clamp((*lambda)[i], limits[i]);

			auto dv = M_inv_J_T * (*lambda);

			for (std::size_t i = 0; i < N; ++i)
			{
				object & obj = objects[handles[i]];
				obj.state.velocity[0] += dv[6 * i + 0];
				obj.state.velocity[1] += dv[6 * i + 1];
				obj.state.velocity[2] += dv[6 * i + 2];
				obj.state.angular_velocity[0] += dv[6 * i + 3];
				obj.state.angular_velocity[1] += dv[6 * i + 4];
				obj.state.angular_velocity[2] += dv[6 * i + 5];
			}
		};
	}

	engine::impl::solver engine::impl::make_collision_solver(object_handle h1, object_handle h2, collision_data const & c, float dt)
	{
		auto tx = geom::ort(c.normal);
		auto ty = geom::cross(c.normal, tx);

		auto value_fn = [this, c, dt, h1, h2]{

			auto const & s1 = objects[h1].state;
			auto const & s2 = objects[h2].state;

			auto v1 = s1.velocity + geom::cross(s1.angular_velocity, c.point1);
			auto v2 = s2.velocity + geom::cross(s2.angular_velocity, c.point2);

			auto dv = v2 - v1;

			auto const & m1 = materials[objects[h1].material];
			auto const & m2 = materials[objects[h2].material];

			float b = (m1.elasticity + m2.elasticity) / 2.f;

//			float f = std::sqrt(m1.friction * m2.friction);
//			float fc = std::exp(- f * dt);

			return geom::vector{
				std::min(c.depth * collision_bias / dt, geom::dot(c.normal, dv) * b),
				0.f,
				0.f,
//				- fc * geom::dot(tx, dv),
//				- fc * geom::dot(ty, dv),
			};
		};

		auto jacobian_fn = [c, tx, ty]{
			// f = dot(n, p2 - p1)
			// df = dot(n, v2 + cross(omega2, r2) - v1 - cross(omega1, r1))

			auto t1 = geom::cross(c.point1, c.normal);
			auto t2 = geom::cross(c.point2, c.normal);

			auto tx1 = geom::cross(c.point1, tx);
			auto tx2 = geom::cross(c.point2, tx);

			auto ty1 = geom::cross(c.point1, ty);
			auto ty2 = geom::cross(c.point2, ty);

			geom::matrix<float, 3, 12> J;

			J[0][0] = -c.normal[0];
			J[0][1] = -c.normal[1];
			J[0][2] = -c.normal[2];

			J[0][3] = -t1[0];
			J[0][4] = -t1[1];
			J[0][5] = -t1[2];

			J[0][6] = c.normal[0];
			J[0][7] = c.normal[1];
			J[0][8] = c.normal[2];

			J[0][9]  = t2[0];
			J[0][10] = t2[1];
			J[0][11] = t2[2];

			J[1][0] = -tx[0];
			J[1][1] = -tx[1];
			J[1][2] = -tx[2];

			J[1][3] = -tx1[0];
			J[1][4] = -tx1[1];
			J[1][5] = -tx1[2];

			J[1][6] = tx[0];
			J[1][7] = tx[1];
			J[1][8] = tx[2];

			J[1][9]  = tx2[0];
			J[1][10] = tx2[1];
			J[1][11] = tx2[2];

			J[2][0] = -ty[0];
			J[2][1] = -ty[1];
			J[2][2] = -ty[2];

			J[2][3] = -ty1[0];
			J[2][4] = -ty1[1];
			J[2][5] = -ty1[2];

			J[2][6] = ty[0];
			J[2][7] = ty[1];
			J[2][8] = ty[2];

			J[2][9]  = ty2[0];
			J[2][10] = ty2[1];
			J[2][11] = ty2[2];

			return J;
		};

		auto limits_fn = [this, h1, h2, dt]{

			auto const & m1 = materials[objects[h1].material];
			auto const & m2 = materials[objects[h2].material];

			float f = std::sqrt(m1.friction * m2.friction) * dt * 0;
//			float fc = std::exp(- f * dt);

			return geom::box<float, 3>{{
				{0.f, inf},
//				{-inf, inf},
//				{-inf, inf},
				{-f, f},
				{-f, f},
//				{0.f, 0.f},
//				{0.f, 0.f},
			}};
		};

		return make_solver(value_fn, jacobian_fn, limits_fn, std::array{h1, h2});
	}

	void engine::impl::update(float dt)
	{
		// Update inertia
		for (object_handle h = 0; h < objects.capacity(); ++h)
		{
			if (!is_object(h)) continue;

			auto & o = objects[h];
			auto R = geom::quaternion_rotation<float>(o.state.rotation).linear_matrix();
			if (auto i = geom::inverse(R * geom::matrix<float, 3, 3>::diagonal(o.inertia) * geom::transpose(R)); i && o.inv_mass != 0.f)
				o.global_inv_inertia = *i;
			else
				o.global_inv_inertia = o.global_inv_inertia.zero();
		}

		// Apply forces
		for (auto & o : objects)
		{
			if (o.inv_mass != 0.f)
			{
				o.state.velocity += gravity * dt;
			}
		}

		// Gather constraints
		std::vector<util::function<void()>> solvers;
		for (object_handle h1 = 0; h1 < objects.capacity(); ++h1)
		{
			if (!is_object(h1)) continue;
			for (object_handle h2 = h1 + 1; h2 < objects.capacity(); ++h2)
			{
				if (!is_object(h2)) continue;

				if (auto c = collision(objects[h1].state, shapes[objects[h1].shape].shape, objects[h2].state, shapes[objects[h2].shape].shape))
					solvers.push_back(make_collision_solver(h1, h2, *c, dt));
			}
		}

		// Apply constraints
		for (auto & s : solvers)
			s();

		// Integrate velocity
		for (auto & o : objects)
		{
			o.state.position += o.state.velocity * dt;
			o.state.rotation = geom::normalized(geom::exp(geom::quaternion<float>::vector(o.state.angular_velocity * dt)) * o.state.rotation);
		}
	}

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

	engine::~engine() = default;

	engine::shape_handle engine::add_shape(ball const & s)
	{
		return impl().add_shape(s);
	}

	engine::shape_handle engine::add_shape(box const & s)
	{
		return impl().add_shape(s);
	}

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

	engine::any_shape_ptr engine::get_shape(shape_handle h) const
	{
		return std::visit([](auto const & x) -> any_shape_ptr { return &x; }, impl().shapes[h].shape);
	}

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

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

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

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

	engine::object_handle engine::add_object(shape_handle shape, material_handle material, object_state const & s)
	{
		float mass = impl().shapes[shape].volume * impl().materials[material].density;
		auto inertia = impl().shapes[shape].inertia * mass;

		auto h = impl().objects.insert(object{shape, material, s, 1.f / mass, inertia});
		if (impl().object_mask.size() <= h)
			impl().object_mask.resize(impl().objects.capacity(), false);
		impl().object_mask[h] = true;
		return h;
	}

	engine::shape_handle engine::get_object_shape(object_handle h) const
	{
		return impl().objects[h].shape;
	}

	engine::material_handle engine::get_object_material(object_handle h) const
	{
		return impl().objects[h].material;
	}

	object_state const & engine::get_object_state(object_handle h) const
	{
		return impl().objects[h].state;
	}

	object_state & engine::get_object_state(object_handle h)
	{
		return impl().objects[h].state;
	}

	float engine::get_object_mass(object_handle h) const
	{
		return 1.f / impl().objects[h].inv_mass;
	}

	geom::matrix<float, 3, 3> engine::get_object_inertia(object_handle h) const
	{
		return *geom::inverse(impl().objects[h].global_inv_inertia);
	}

	void engine::remove_object(object_handle h)
	{
		impl().object_mask[h] = false;
		impl().objects.erase(h);
	}

	engine::object_handle engine::object_count() const
	{
		return impl().objects.capacity();
	}

	bool engine::is_object(object_handle h) const
	{
		return impl().is_object(h);
	}

	void engine::add_impulse(object_handle h, geom::vector<float, 3> const & impulse)
	{
		impl().objects[h].state.velocity += impulse * impl().objects[h].inv_mass;
	}

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

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

}