#include <iostream>
#include <cmath>
#include <cstring>
#include <Body.hpp>
#include <Box.hpp>
#include <Tree.hpp>
#include <Node.hpp>
#include <algorithm>
#include <iostream>
#include <BarnesHutTree.hpp>
#include "MpiSimulation.hpp"

namespace nbody {
	MpiSimulation::MpiSimulation(const std::string& inputFile) {
		//create MPI datatypes for bodies and domain boxes
		int vec3Blocklengths[3] = { 1,1,1 };
		MPI_Datatype vec3Datatypes[3] = { MPI_DOUBLE, MPI_DOUBLE , MPI_DOUBLE };
		MPI_Aint vec3Offsets[3]{
			offsetof(Vec3, X),
			offsetof(Vec3, Y),
			offsetof(Vec3, Z)
		};
		MPI_Type_create_struct(3, vec3Blocklengths, vec3Offsets, vec3Datatypes, &vec3Type);
		MPI_Type_commit(&vec3Type);

		int bodyBlocklengths[6] = { 1, 1, 1, 1, 1, 1 };
		MPI_Datatype bodyDatatypes[6] = { MPI_UINT64_T, vec3Type, vec3Type, vec3Type, MPI_DOUBLE, MPI_CXX_BOOL };
		MPI_Aint bodyOffsets[6];
		bodyOffsets[0] = offsetof(Body, id);
		bodyOffsets[1] = offsetof(Body, position);
		bodyOffsets[2] = offsetof(Body, velocity);
		bodyOffsets[3] = offsetof(Body, acceleration);
		bodyOffsets[4] = offsetof(Body, mass);
		bodyOffsets[5] = offsetof(Body, refinement);
		MPI_Type_create_struct(6, bodyBlocklengths, bodyOffsets, bodyDatatypes, &bodyType);
		MPI_Type_commit(&bodyType);

		int boxBlocklengths[2] = { 1, 1 };
		MPI_Datatype boxDatatypes[2] = { vec3Type, vec3Type };
		MPI_Aint boxOffsets[2];
		boxOffsets[0] = offsetof(Box, min);
		boxOffsets[1] = offsetof(Box, max);
		MPI_Type_create_struct(2, boxBlocklengths, boxOffsets, boxDatatypes, &boxType);
		MPI_Type_commit(&boxType);

		//get number of processes and own process id
		MPI_Comm_size(MPI_COMM_WORLD, &parallelSize);
		MPI_Comm_rank(MPI_COMM_WORLD, &parallelRank);
		domains.reserve(parallelSize);
		correctState = true;

		//parse input data
		if (!inputFile.empty()) {
			correctState = readInputData(inputFile);
		}
		else {
			correctState = false;
		}
		//broadcast current state and terminate if input file cannot be read
		int mpiCorrectState = static_cast<int>(correctState);
		MPI_Bcast(&mpiCorrectState, 1, MPI_INT, 0, MPI_COMM_WORLD);
		if (!correctState) {
			std::cerr << "Error occurred: terminating ...\n";
			MPI_Type_free(&bodyType);
			MPI_Type_free(&boxType);
			MPI_Abort(MPI_COMM_WORLD, -1);
			abort();
		}
	}

	bool MpiSimulation::stateCorrect() {
		return correctState;
	}

	MpiSimulation::~MpiSimulation() {
		flushSendStore();
		//cleanup MPI types
		MPI_Type_free(&bodyType);
		MPI_Type_free(&boxType);
		MPI_Type_free(&vec3Type);
		delete this->tree;
		tree = nullptr;
	}

	std::size_t MpiSimulation::getNumberOfProcesses() const {
		return parallelSize;
	}

	std::size_t MpiSimulation::getProcessId() const {
		return parallelRank;
	}

	MPI_Datatype* MpiSimulation::getDatatype() {
		return &bodyType;
	}

	//mpi send wrapper
	void MpiSimulation::send(std::vector<Body>&& bodies, int target) {
		const auto bodySize = bodies.size();
		SendStore store{ std::move(bodies) };
		MPI_Isend(store.bodies.data(), bodySize, bodyType, target, 0, MPI_COMM_WORLD, &store.request);
		sendStores.push_back(std::move(store));
	}

	//mpi recv wrapper
	int MpiSimulation::recv(std::vector<Body>& bodies, int source) {
		MPI_Status status;
		int count;

		//do blocking recv; any source receive can be done with source == MPI_ANY_SOURCE
		MPI_Probe(source, 0, MPI_COMM_WORLD, &status);
		MPI_Get_count(&status, bodyType, &count);
		bodies.resize(count);
		MPI_Recv(bodies.data(), count, bodyType, status.MPI_SOURCE, 0, MPI_COMM_WORLD, &status);
		//return source to determine message source for any source receives
		return status.MPI_SOURCE;
	}

	//initial body distribution
	void MpiSimulation::distributeBodies() {
		//process 0 distributes bodies, others receive
		if (parallelRank == 0) {
			std::vector<Node> nodes;
			Box bb;
			nodes.push_back(Node(nullptr));
			nodes.front().setBodies(bodies);
			bb.extendForBodies(bodies);
			nodes.front().setBB(bb);
			//determine how to distribute bodies to processes
			//split box with most particles by halfing its longest side
			//until number of boxes equals number of processes
			while (nodes.size() < static_cast<std::size_t>(parallelSize)) {
				std::size_t mostBodiesIndex = 0;

				for (std::size_t i = 1; i < nodes.size(); i++) {
					if (nodes[i].getBodies().size() > nodes[mostBodiesIndex].getBodies().size()) {
						mostBodiesIndex = i;
					}
				}
				std::vector<Box> subdomains = nodes[mostBodiesIndex].getBB().splitLongestSide();
				std::vector<Body> buf = nodes[mostBodiesIndex].getBodies();
				Node n(nullptr);

				n.setBodies(subdomains[0].extractBodies(buf));
				n.setBB(subdomains[0]);
				nodes.insert(std::begin(nodes) + mostBodiesIndex, n);
				n = Node(nullptr);
				n.setBodies(subdomains[1].extractBodies(buf));
				n.setBB(subdomains[1]);
				nodes.insert(std::begin(nodes) + mostBodiesIndex, n);
				nodes.erase(std::begin(nodes) + mostBodiesIndex + 2);
			}
			bodies = nodes[0].getBodies();
			for (std::size_t i = 1; i < nodes.size(); i++) {
				send(nodes[i].getBodies(), i);
			}
		} else {
			recv(bodies, 0);
		}
		flushSendStore();
	}

	void MpiSimulation::distributeDomains(const std::vector<Body>& localBodies) {
		Box localDomain;

		//determine local domain size
		localDomain.extendForBodies(localBodies);

		distributeDomains(localDomain);
	}

	void MpiSimulation::distributeDomains() {
		distributeDomains(bodies);
	}

	//domain distribution, all processes need to know the spatial domains of the others
	void MpiSimulation::distributeDomains(const Box& localDomain) {
		//distribute local domain sizes to all processes through collective MPI operation
		domains[parallelRank] = localDomain;
		MPI_Allgather(&domains[parallelRank], 1, boxType, &domains[0], 1, boxType, MPI_COMM_WORLD);
		overallDomain = localDomain;
		//determine overall domain size
		for (std::size_t i = 0; i < static_cast<std::size_t>(parallelSize); i++) {
			overallDomain.extend(domains[i]);
		}
	}

	void MpiSimulation::flushSendStore() {
		std::vector<MPI_Request> requests;
		std::transform(std::begin(sendStores), std::end(sendStores), std::back_inserter(requests), [](const SendStore& ss) {return ss.request; });
		MPI_Waitall(requests.size(), requests.data(), MPI_STATUSES_IGNORE);
		sendStores.clear();
	}

	//distribute bodies needed by other processes for their local simlation
	void MpiSimulation::distributeLETs() {
		//send out locally essential trees (local bodies needed by remote simulations, determined by remote domain size)
		for (std::size_t i = 0; i < static_cast<std::size_t>(parallelSize); i++) {
			if (i != static_cast<std::size_t>(parallelRank)) {
				auto refinements = tree->copyRefinements(domains[i]);

				send(std::move(refinements), i);
			}
		}

		//receive bodies and integrate them into local tree for simulation
		for (std::size_t i = 0; i < static_cast<std::size_t>(parallelSize - 1); i++) {
			std::vector<Body> refinements;

			//any source receive can be blocking, because we need to wait for data anyhow
			//order is not important, and receiving and merging arriving particles can be overlapped
			recv(refinements, MPI_ANY_SOURCE);
			tree->mergeLET(refinements);
		}
		if (!tree->isCorrect()) {
			std::cerr << "wrong tree\n";
		}
		flushSendStore();
	}

	void MpiSimulation::buildTree() {
		tree = new BarnesHutTree(parallelRank);
		tree->build(bodies, overallDomain);
		if (!tree->isCorrect()) {
			std::cerr << "wrong tree\n";
		}
	}

	void MpiSimulation::rebuildTree() {
		//rebuild tree with moved local particles
		tree->rebuild(overallDomain);
	}

	//run a simulation step
	void MpiSimulation::runStep() {
		//tree is already built here

		//distribute local bodies needed by remote processes
		std::cout << "  rank " << parallelRank << ": distribute local particles required for remote simulation ...\n";
		distributeLETs();
		//force computation
		std::cout << "  rank " << parallelRank << ": compute forces ...\n";
		tree->computeForces();
		//advance/move particles and distribute updated domain to other processes
		std::cout << "  rank " << parallelRank << ": move particles and redistribute simulation domains ...\n";
		distributeDomains(tree->advance());
		//rebuild tree with new particle positions
		std::cout << "  rank " << parallelRank << ": rebuild tree with new particle positions ...\n";
		rebuildTree();
		if (!tree->isCorrect()) {
			std::cerr << "wrong tree\n";
		}
	}
} // namespace nbody
