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

#include <boost/multi_array.hpp>
#include <mpi.h>
#include <omp.h>

#include "Configuration.hpp"
#include "MpiEnvironment.hpp"
#include "MpiTypes.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; }
	}
}

// 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
	MpiEnvironment env{argc, argv, MPI_THREAD_FUNNELED};
	// MPI_THREAD_FUNNELED is needed for OpenMP to work.
	// MPI calls must be made only from master thread

	// define mpi datatypes
	const auto& configType = GenConfigurationType();
	const auto& particleType = GenParticleType();
	const auto& sparseLocatedParticleType = GenSparseLocatedParticleType();

	// create configuration
	Configuration config;
	if (env.isMaster()) {
		config = parseArgs(argc, argv);
		if (config.NoParticles % env.worldSize() != 0) {
			std::cerr << "Error: number of particles can not divided by number of ranks\n";
			MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
		}
	}
	MPI_Bcast(&config, 1, configType.type(), 0, MPI_COMM_WORLD);

	const auto& chunkSize = config.NoParticles / env.worldSize();
	const auto& startIdx = env.worldRank() * 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 rndDev;
	std::mt19937 rndEng{rndDev()};
	std::uniform_real_distribution<double> dist{-1.0, 1.0};

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

	// broadcast local chunk, receive all other chunks
	MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, model.data(), chunkSize, particleType.type(), MPI_COMM_WORLD);

	if (env.worldRank() == 0) {
#pragma omp parallel
		{
			const auto& nthreads = omp_get_num_threads();
			const auto& tid = omp_get_thread_num();
			if (tid == 0) { std::cout << "omp_num_threads: " << nthreads << '\n'; }
		}
		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
		MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, model.data(), chunkSize, sparseLocatedParticleType.type(),
		              MPI_COMM_WORLD);
	}

	if (env.worldRank() == 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";
	}

	return EXIT_SUCCESS;
}