#include <algorithm>
#include <cassert>
#include <chrono>
#include <iomanip>
#include <iostream>
#include <numeric>
#include <random>
#include <vector>

#include <boost/mpi/collectives.hpp>
#include <boost/mpi/communicator.hpp>
#include <boost/mpi/environment.hpp>
#include <boost/multi_array.hpp>

namespace mpi = boost::mpi;

#include "Configuration.hpp"
#include "Particle.hpp"

template <typename T> void printModel(T& model) {
	std::size_t i{0};
	for (const auto& p : model) {
		std::cout << std::setw(7) << i++ << "X:" << std::setprecision(7) << std::setw(12) << p.Location.X << '\n';
		if (i >= 10) break;
	}
}

// Boost MPI does not support allGather with the MPI_IN_PLACE option.
// This is a replacement using Boost MPI Broadcast.
// n Broadcast calls are needed.
//
// This has the same semantics as the call
// MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, buf, count, type, comm);
template <typename T> void allGatherInplace(mpi::communicator& world, T* buf, std::size_t count) {
	for (std::size_t i = 0; i < world.size(); ++i) { mpi::broadcast(world, buf + i * count, count, i); }
}

// use std::for_each(std::execution::par, ... ) in C++17 instead
template <typename ForwardIt, typename UnaryFunction>
void omp_parallel_for_each(ForwardIt first, ForwardIt last, UnaryFunction f) {
	static_assert(std::is_base_of<typename std::iterator_traits<ForwardIt>::iterator_category,
	                              std::random_access_iterator_tag>::value,
	              "The omp_parallel_for_each() function only accepts forward iterators.\n");
	const auto& diff = std::distance(first, last);
	assert(diff < std::numeric_limits<int>::max());
	const auto& size = static_cast<int>(diff);

#pragma omp parallel for // OpenMP 2.0 compability: no iterators, signed loop variable
	for (int i = 0; i < size; ++i) {
		const auto& it = std::next(first, i);
		f(*it);
	}
}

int main(int argc, char* argv[]) {
	const auto& starttime = std::chrono::system_clock::now();
	constexpr double delta_t = 1e-5;

	// init MPI
	mpi::environment env{mpi::threading::level::funneled};
	mpi::communicator world;
	const bool isMaster{world.rank() == 0};

	// no need to define types here

	// create configuration
	Configuration config;
	if (isMaster) {
		config = parseArgs(argc, argv);
		if (config.NoParticles % world.size() != 0) {
			std::cerr << "Error: number of particles can not divided by number of ranks\n";
			world.abort(EXIT_FAILURE);
		}
	}
	mpi::broadcast(world, config, 0);

	const auto& chunkSize = config.NoParticles / world.size();
	const auto& startIdx = world.rank() * chunkSize;

	/// begin init model

	using modelType = boost::multi_array<Particle, 1>;
	modelType model{ boost::extents[config.NoParticles] };
	// calculate local chunk boundaries
	auto localChunk = model[boost::indices[modelType::index_range(startIdx, startIdx + chunkSize)]];

	// generate initial particles for local chunk
	std::random_device rnd_dev;
	std::mt19937 rnd_eng{0};
	std::uniform_real_distribution<double> dist{-1.0, 1.0};

	std::generate(std::begin(localChunk), std::end(localChunk), [&] {
		return Particle{1.0, {dist(rnd_eng), dist(rnd_eng), dist(rnd_eng)}};
	});

	// broadcast local chunk, receive all other chunks
	allGatherInplace(world, model.data(), chunkSize);

	if (world.rank() == 0) {
		std::cout << "Initialization done:" << '\n';
		printModel(model);
	}
	/// end init model

	// main simulation loop
	for (std::size_t step = 0; step < config.NoIterations; ++step) {

		// calc forces, update velocity
		omp_parallel_for_each(std::begin(localChunk), std::end(localChunk), [&](Particle& p1) {
			const auto& force =
			    std::accumulate(std::begin(model), std::end(model), Vec3{}, [&](Vec3 fAcc, const Particle& p2) {
				    const auto& diff = p1.Location - p2.Location;
				    const auto& dist = Norm(diff);
				    if (dist > 1e-8) {
					    const auto& f = (p1.Mass * p2.Mass) / std::pow(dist, 2.0); // f = G * ((m1 * m2) / r^2)
					    const auto& f_direction = diff / dist;
					    return fAcc + (f_direction * f);
				    } else {
					    return fAcc;
				    }
				});
			const auto& acceleration = force / p1.Mass;
			p1.Velocity += acceleration * delta_t;
		});

		// update location
		omp_parallel_for_each(std::begin(localChunk), std::end(localChunk), [&](Particle& p1) {
			p1.Location += p1.Velocity * delta_t; //
		});
		// broadcast local chunk, receive all other chunks
		// boost mpi does not support the definition of "sparse" data types
		allGatherInplace(world, model.data(), chunkSize);
	}

	if (world.rank() == 0) {
		std::cout << "End Simulation\n";
		printModel(model);
		std::cout << "Execution time:"
		          << std::chrono::duration<double>{std::chrono::system_clock::now() - starttime}.count() << "s\n";
	}
	// resources are freed automatically
	return EXIT_SUCCESS;
}