/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2023, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include <mrpt/core/round.h>
#include <mrpt/gui/CDisplayWindow3D.h>
#include <mrpt/opengl/CGridPlaneXY.h>
#include <mrpt/opengl/CSphere.h>
#include <mrpt/random.h>
#include <mrpt/system/CTicTac.h>
#include <mrpt/system/os.h>

#include <chrono>
#include <iostream>
#include <thread>

using namespace mrpt;
using namespace std;
using namespace mrpt::gui;
using namespace mrpt::random;
using namespace mrpt::opengl;

const size_t N_MASSES = 1500;

const double BOX = 1000;
const double V0 = 300;
const double MASS_MIN = log(300.0), MASS_MAX = log(5000.0);
const double M2R = 0.5;
const double LARGEST_STEP = 0.0005;
const double G = 300;
const double COLLIS_LOSS = 0.99;

struct TMass
{
	TMass() : obj3d() {}
	double x{0}, y{0}, z{0};
	double vx{0}, vy{0}, vz{0};
	double mass{1};
	double radius;
	opengl::CSphere::Ptr obj3d;
};

void simulateGravity(vector<TMass>& objs, double At);

// ------------------------------------------------------
//				GravityDemo
// ------------------------------------------------------
void GravityDemo()
{
	CDisplayWindow3D win(
		"MRPT example: 3D gravity simulator- JLBC 2008", 1000, 700);

	getRandomGenerator().randomize();

	win.setCameraElevationDeg(50.0f);
	win.setCameraZoom(1000);
	win.getDefaultViewport()->setViewportClipDistances(1, 10000);

	Scene::Ptr& theScene = win.get3DSceneAndLock();

	// Modify the scene:
	// ------------------------------------------------------
	{
		opengl::CGridPlaneXY::Ptr obj =
			opengl::CGridPlaneXY::Create(-2000, 2000, -2000, 2000, 0, 100);
		obj->setColor(0.3f, 0.3f, 0.3f);
		theScene->insert(obj);
	}

	// Create the masses:
	// ----------------------------------------------------

	vector<TMass> masses(N_MASSES);

	// Init at random poses & create opengl objects:
	for (size_t i = 0; i < N_MASSES; i++)
	{
		masses[i].x = getRandomGenerator().drawUniform(-BOX, BOX);
		masses[i].y = getRandomGenerator().drawUniform(-BOX, BOX);
		masses[i].z = getRandomGenerator().drawUniform(-BOX, BOX) / 10;

		double a = atan2(masses[i].y, masses[i].x);

		masses[i].vx = -V0 * sin(a) +
			getRandomGenerator().drawUniform(-V0 * 0.05, V0 * 0.05);
		masses[i].vy = V0 * cos(a) +
			getRandomGenerator().drawUniform(-V0 * 0.05, V0 * 0.05);
		masses[i].vz = 0;  // getRandomGenerator().drawUniform(-V0,V0);

		masses[i].mass =
			exp(getRandomGenerator().drawUniform(MASS_MIN, MASS_MAX));
		opengl::CSphere::Ptr& obj = masses[i].obj3d = opengl::CSphere::Create();

		obj->setColor(
			getRandomGenerator().drawUniform<float>(0.1, 0.9),
			getRandomGenerator().drawUniform<float>(0.1, 0.9),
			getRandomGenerator().drawUniform<float>(0.1, 0.9));

		masses[i].radius = M2R * pow(masses[i].mass, 1.0 / 3.0);
		obj->setRadius(mrpt::d2f(masses[i].radius));  // Guess why ^(1/3) ;-)

		obj->setLocation(masses[i].x, masses[i].y, masses[i].z);
		theScene->insert(obj);
	}

	// IMPORTANT!!! IF NOT UNLOCKED, THE WINDOW WILL NOT BE UPDATED!
	win.unlockAccess3DScene();

	mrpt::system::CTicTac tictac;
	tictac.Tic();

	double t0 = tictac.Tac();

	while (!mrpt::system::os::kbhit() && win.isOpen())
	{
		// Move the scene:
		double t1 = tictac.Tac();
		double At = t1 - t0;
		t0 = t1;

		// Simulate a At, possibly in several small steps:
		// Update the 3D scene:
		win.get3DSceneAndLock();

		size_t n_steps = mrpt::round(ceil(At / LARGEST_STEP) + 1);
		double At_steps = At / n_steps;
		n_steps = min(n_steps, size_t(3));
		for (size_t j = 0; j < n_steps; j++)
			simulateGravity(masses, At_steps);

		for (size_t i = 0; i < masses.size(); i++)
		{
			opengl::CSphere::Ptr& obj = masses[i].obj3d;
			obj->setLocation(masses[i].x, masses[i].y, masses[i].z);
		}
		// IMPORTANT!!! IF NOT UNLOCKED, THE WINDOW WILL NOT BE UPDATED!
		win.unlockAccess3DScene();

		// Update window:
		win.forceRepaint();
		std::this_thread::sleep_for(1ms);
	};
}

struct TForce
{
	TForce() { f[0] = f[1] = f[2] = 0; }
	double f[3];
};

void simulateGravity(vector<TMass>& objs, double At)
{
	const size_t N = objs.size();

	vector<TForce> forces(N);

	// Index in the array must be larger than its content!!
	vector<pair<size_t, double>> lstMass_i_joins_j(
		N, pair<size_t, double>(string::npos, 10000.0));

	for (size_t i = 0; i < (N - 1); i++)
	{
		const double Ri = objs[i].radius;

		// Compute overall gravity force:
		for (size_t j = i + 1; j < N; j++)
		{
			double Ax = objs[j].x - objs[i].x;
			double Ay = objs[j].y - objs[i].y;
			double Az = objs[j].z - objs[i].z;
			double D2 = square(Ax) + square(Ay) + square(Az);

			double D = sqrt(D2);
			if (D == 0) continue;

			const double Rj = objs[j].radius;

			if (D < (Ri + Rj))	// Collission!!
			{
				// Index in the array must be larger than its content!!
				if (D < lstMass_i_joins_j[j].second)
				{
					lstMass_i_joins_j[j].first = i;
					lstMass_i_joins_j[j].second = D;
				}
			}
			else
			{
				double K =
					G * objs[i].mass * objs[j].mass / square(max(D, 1.0));
				double D_1 = 1.0 / D;
				Ax *= D_1;
				Ay *= D_1;
				Az *= D_1;

				const double fx = Ax * K;
				const double fy = Ay * K;
				const double fz = Az * K;

				forces[i].f[0] += fx;
				forces[i].f[1] += fy;
				forces[i].f[2] += fz;

				forces[j].f[0] -= fx;
				forces[j].f[1] -= fy;
				forces[j].f[2] -= fz;
			}
		}
	}

	// F = m a
	for (size_t i = 0; i < N; i++)
	{
		const double M_1 = 1.0 / objs[i].mass;

		forces[i].f[0] *= M_1;
		forces[i].f[1] *= M_1;
		forces[i].f[2] *= M_1;

		objs[i].vx += forces[i].f[0] * At;
		objs[i].vy += forces[i].f[1] * At;
		objs[i].vz += forces[i].f[2] * At;

		objs[i].x += objs[i].vx * At;
		objs[i].y += objs[i].vy * At;
		objs[i].z += objs[i].vz * At;
	}

	//	return;
	// Join masses that collided:
	for (int i = N - 1; i >= 0; i--)
	{
		const size_t newObj = lstMass_i_joins_j[i].first;
		if (newObj == string::npos) continue;

		const double Mi = objs[i].mass;
		const double Mj = objs[newObj].mass;
		const double newMass = Mi + Mj;
		const double newMass_1 = 1.0 / newMass;

		objs[newObj].vx =
			COLLIS_LOSS * newMass_1 * (Mj * objs[newObj].vx + Mi * objs[i].vx);
		objs[newObj].vy =
			COLLIS_LOSS * newMass_1 * (Mj * objs[newObj].vy + Mi * objs[i].vy);
		objs[newObj].vz =
			COLLIS_LOSS * newMass_1 * (Mj * objs[newObj].vz + Mi * objs[i].vz);

		objs[newObj].x = newMass_1 * (Mj * objs[newObj].x + Mi * objs[i].x);
		objs[newObj].y = newMass_1 * (Mj * objs[newObj].y + Mi * objs[i].y);
		objs[newObj].z = newMass_1 * (Mj * objs[newObj].z + Mi * objs[i].z);

		objs[newObj].mass = newMass;
		objs[newObj].radius = M2R * pow(newMass, 1.0 / 3.0);
		objs[newObj].obj3d->setRadius(mrpt::d2f(objs[newObj].radius));

		objs[i].obj3d->setVisibility(false);  // Hide sphere
		objs.erase(objs.begin() + i);
	}
}

// ------------------------------------------------------
//						MAIN
// ------------------------------------------------------
int main()
{
	try
	{
		GravityDemo();
		return 0;
	}
	catch (const std::exception& e)
	{
		std::cerr << "MRPT error: " << mrpt::exception_to_str(e) << std::endl;
		return -1;
	}
}
