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

namespace nbody {
	using namespace std;

	MpiSimulation::MpiSimulation(int& argc, char**& argv) {
		Body b[2];
		int bodyBlocklengths[7] = {1, 3, 3, 3, 1, 1, 1};
		MPI_Datatype bodyDatatypes[7] = {MPI_LB, MPI_DOUBLE, MPI_DOUBLE, MPI_DOUBLE, MPI_DOUBLE, MPI_INT, MPI_UB};
		MPI_Aint bodyDisplacements[7];
		Box box[2];
		int boxBlocklengths[4] = {1, 3, 3, 1};
		MPI_Datatype boxDatatypes[4] = {MPI_LB, MPI_DOUBLE, MPI_DOUBLE, MPI_UB};
		MPI_Aint boxDisplacements[4];
		MPI_Aint baseAdress;

		this->bodyType = MPI_DATATYPE_NULL;
		this->boxType = MPI_DATATYPE_NULL;
		MPI_Get_address(&(b[0]), &baseAdress);
		bodyDisplacements[0] = baseAdress - baseAdress;
		MPI_Get_address(&(b[0].position[0]), &bodyDisplacements[1]);
		bodyDisplacements[1] -= baseAdress;
		MPI_Get_address(&(b[0].velocity[0]), &bodyDisplacements[2]);
		bodyDisplacements[2] -= baseAdress;
		MPI_Get_address(&(b[0].acceleration[0]), &bodyDisplacements[3]);
		bodyDisplacements[3] -= baseAdress;
		MPI_Get_address(&(b[0].mass), &bodyDisplacements[4]);
		bodyDisplacements[4] -= baseAdress;
		MPI_Get_address(&(b[0].refinement), &bodyDisplacements[5]);
		bodyDisplacements[5] -= baseAdress;
		MPI_Get_address(&(b[1]), &bodyDisplacements[6]);
		bodyDisplacements[6] -= baseAdress;
		MPI_Type_create_struct(7, bodyBlocklengths, bodyDisplacements, bodyDatatypes, &this->bodyType);
		MPI_Type_commit(&this->bodyType);

		MPI_Get_address(&(box[0]), &baseAdress);
		boxDisplacements[0] = baseAdress - baseAdress;
		MPI_Get_address(&(box[0].min[0]), &boxDisplacements[1]);
		boxDisplacements[1] -= baseAdress;
		MPI_Get_address(&(box[0].max[0]), &boxDisplacements[2]);
		boxDisplacements[2] -= baseAdress;
		MPI_Get_address(&(box[1]), &boxDisplacements[3]);
		boxDisplacements[3] -= baseAdress;
		MPI_Type_create_struct(4, boxBlocklengths, boxDisplacements, boxDatatypes, &this->boxType);
		MPI_Type_commit(&this->boxType);

		MPI_Comm_size(MPI_COMM_WORLD, &this->mpiSize);
		MPI_Comm_rank(MPI_COMM_WORLD, &this->mpiRank);
		this->domains = new Box[this->mpiSize];
		this->correctState = true;

		if (argc == 2) {
			this->correctState = this->readInputData(string(argv[1]));
		} else {
			this->correctState = false;
		}
		MPI_Bcast(&this->correctState, 1, MPI_INT, 0, MPI_COMM_WORLD);
		if (!this->correctState) {
			cerr << "Error occurred: terminating ..." << endl;
			MPI_Type_free(&this->bodyType);
			MPI_Type_free(&this->boxType);
			MPI_Finalize();
			exit(-1);
		}
	}

	bool MpiSimulation::readInputData(string filename) {
		if (this->mpiRank == 0) {
			this->bodies = Tree::dubinskiParse(filename);
			if (this->bodies.empty()) {
				return false;
			}
		}
		return true;
	}

	void MpiSimulation::cleanup() {
		while (!this->comms.empty()) {
			this->comms.back().cleanup();
			this->comms.pop_back();
		}
		MPI_Type_free(&this->bodyType);
		MPI_Type_free(&this->boxType);
	}

	bool MpiSimulation::stateCorrect() {
		return this->correctState != 0;
	}

	MpiSimulation::~MpiSimulation() {
		delete[] this->domains;
	}

	int MpiSimulation::getNumberOfProcesses() {
		return this->mpiSize;
	}

	int MpiSimulation::getProcessId() {
		return this->mpiRank;
	}

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

	void MpiSimulation::distributeBodies() {
		if (this->mpiRank == 0) {
			vector<Node> nodes;
			Box bb;

			nodes.push_back(Node(NULL));
			nodes.front().setBodies(this->bodies);
			bb = nodes.front().getBB();
			bb.extendForBodies(this->bodies);
			nodes.front().setBB(bb);
			while (nodes.size() < (unsigned int) this->mpiSize) {
				int mostBodiesIndex = 0;

				for (unsigned int i = 0; i < nodes.size(); i++) {
					if (nodes[i].getBodies().size() > nodes[mostBodiesIndex].getBodies().size()) {
						mostBodiesIndex = i;
					}
				}
				vector<Box> subdomains = nodes[mostBodiesIndex].getBB().splitLongestSide();
				vector<Body> buf = nodes[mostBodiesIndex].getBodies();
				Node n(NULL);
				Box bb;

				n.setBodies(subdomains[0].extractBodies(buf));
				bb.extendForBodies(n.getBodies());
				n.setBB(bb);
				nodes.insert(nodes.begin() + mostBodiesIndex, n);
				n = Node(NULL);
				n.setBodies(subdomains[1].extractBodies(buf));
				bb.extendForBodies(n.getBodies());
				n.setBB(bb);
				nodes.insert(nodes.begin() + mostBodiesIndex, n);
				nodes.erase(nodes.begin() + mostBodiesIndex + 2);
			}
			this->bodies = nodes[0].getBodies();

			for (unsigned int i = 1; i < nodes.size(); i++) {
				this->comms.push_back(MpiBodyComm(&this->bodyType));
				this->comms.back().sendUnblocking(i, nodes[i].getBodies());
			}
		} else {
			this->comms.push_back(MpiBodyComm(&this->bodyType));
			this->comms[0].recvBlocking(0, this->bodies);
		}
		this->tree.build(this->bodies);
		this->domains[this->mpiRank] = this->tree.getRootBB();
		//this->tree.isCorrect();
	}

	void MpiSimulation::distributeDomains() {
		MPI_Allgather(&this->domains[this->mpiRank], 1, this->boxType, &this->domains[0], 1, this->boxType, MPI_COMM_WORLD);
	}

	void MpiSimulation::distributeLETs() {
		//send out locally essential trees (local bodies needed by remote simulations)
		for (int i = 0; i < this->mpiSize; i++) {
			if (i != this->mpiRank) {
				vector<Body> refinements = this->tree.copyRefinements(this->domains[i]);
				vector<MpiBodyComm>::iterator it = this->comms.begin();

				while (it != this->comms.end() && !it->sendUnblocking(i, refinements)) {
					it++;
				}
				if (it == this->comms.end()) {
					this->comms.push_back(MpiBodyComm(&this->bodyType));
					this->comms.back().sendUnblocking(i, refinements);
				}
			}
		}
		//receive bodies and integrate them into local tree for simulation
		int received = 0;
		while (received < this->mpiSize - 1) {
			vector<Body> refinements;

			//integrate bodies in order of arrival to do communication/computation overlapping
			this->comms[0].recvBlocking(MPI_ANY_SOURCE, refinements);
			this->tree.mergeLET(refinements);
			//this->tree.getRootBB().print();
			received++;
		}
	}

	void MpiSimulation::runStep() {
		Box overallDomain;

		this->distributeDomains();
		//this->domains[this->mpiRank].print();
		for (int i = 0; i < this->mpiSize; i++) {
			overallDomain.extend(this->domains[i]);
		}
		/*
		if (this->mpiRank == 0) {
			overallDomain.print();
		}
		*/
		this->tree.rebuild(overallDomain);
		this->distributeLETs();
		this->tree.computeForces();
		this->tree.advance();
		this->domains[this->mpiRank] = this->tree.getLocalBB();
		if (!this->tree.isCorrect()) {
			cout << "tree not correct" << endl;
		}
	}

}
