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

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

#include "Configuration.hpp"
#include "MpiComm.hpp"
#include "MpiEnvironment.hpp"
#include "MpiSharedMemWin.hpp"
#include "MpiTypes.hpp"
#include "Particle.hpp"

template <typename T> void printModel(const 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;
	}
}

void printInfoIntranode(const MpiComm& intranodeComm, int worldRank) {
	if (intranodeComm.isMaster()) {
		std::cout << "\nrank[" << worldRank << "]: Size of intranode_comm: " << intranodeComm.size() << "\n";
	}
}

void printInfoInternode(const MpiComm& internodeComm, int worldRank) {
	if (internodeComm.comm() != MPI_COMM_NULL) {
		if (internodeComm.isMaster()) {
			std::cout << "\nrank[" << worldRank << "]: Size of internode_comm: " << internodeComm.size() << "\n";
		}
	}
}

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

	// Create sub-communicators of MPI_COMM_WORLD containing processes sharing one address space (i.e. living on one
	// node)
	if (env.isMaster()) std::cout << "\nCreating MPI communicators...\n";
	const auto& intranodeComm = MpiComm::fromCommSplitTypeShared();
	// const auto intranodeComm = MpiComm::fromCommSplitColor(env.world_rank() < 2 ? 1 : 2); // for debugging on local
	// machine
	printInfoIntranode(intranodeComm, env.worldRank());

	// Create sub-communicator of "node-masters" (communicator containing exactly one process from each node)
	const auto& color = intranodeComm.isMaster() ? 1 : MPI_UNDEFINED;
	const auto& internode_comm = MpiComm::fromCommSplitColor(color);
	printInfoInternode(internode_comm, env.worldRank());
	// Note: internode_comm.comm() is MPI_COMM_NULL for all "non-node-masters", i.e. for all processes who supplied
	// MPI_UNDEFINED as color
	if (env.isMaster()) std::cout << "\nEND Creating MPI communicators...\n";

	if (intranodeComm.isMaster()) {
		int max{0};
		auto intranode_size = intranodeComm.size();
		MPI_Allreduce(&intranode_size, &max, 1, MPI_INT, MPI_MAX, internode_comm.comm());
		if (max != intranode_size) {
			std::cerr << "Error: all communicators must have the same size\n";
			MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
		}
	}

	// 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& nodeSize = chunkSize * intranodeComm.size();
	const auto& startIdx = env.worldRank() * chunkSize;

	// create shared memory window
	std::size_t localWindowCount{0};
	if (intranodeComm.isMaster()) { localWindowCount = config.NoParticles; }
	if (env.isMaster())
		std::cout << "\nCreating shared memory region (" << localWindowCount * sizeof(Particle) << " bytes)...\n";
	MpiSharedMemWin<Particle> memoryWindow{localWindowCount, intranodeComm.comm()};

	// wraps memory window
	using modelType = boost::multi_array_ref<Particle, 1>;
	modelType model{memoryWindow.basePtr(), boost::extents[localWindowCount] };
	// calculate local chunk boundaries
	auto localChunk = model[boost::indices[modelType::index_range(startIdx, startIdx + chunkSize)]];

	/// begin init model

	// generate initial particles for local chunk
	std::random_device rndDev;
	std::mt19937 rnd_eng{rndDev()};
	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
	intranodeComm.barrier();
	if (intranodeComm.isMaster()) {
		MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, memoryWindow.basePtr(), nodeSize, particleType.type(),
		              internode_comm.comm());
	}
	intranodeComm.barrier();

	if (env.worldRank() == 0) {
		std::cout << "Initialization done:\n";
		printModel(memoryWindow);
	}
	/// end init model

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

		// calc forces, update velocity
		for (auto& p1 : localChunk) { // update only local chunk
			const auto& force = std::accumulate(
			    std::begin(memoryWindow), std::end(memoryWindow), 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;
		}

		intranodeComm.barrier();

		// update location
		for (auto& p1 : localChunk) {
			p1.Location += p1.Velocity * delta_t; //
		}

		// broadcast local chunk, receive all other chunks
		intranodeComm.barrier();
		if (intranodeComm.isMaster()) {
			MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, memoryWindow.basePtr(), nodeSize,
			              sparseLocatedParticleType.type(), internode_comm.comm());
		}
		intranodeComm.barrier();
	}

	if (env.worldRank() == 0) {
		std::cout << "End Simulation\n";
		printModel(memoryWindow);

		std::cout << "Execution time:"
		          << std::chrono::duration<double>{std::chrono::system_clock::now() - starttime}.count() << "s\n";
	}

	return EXIT_SUCCESS;
}
