#include "mpimanager.h"
//#include "io.h"

#include <iostream>
#include <string>
#include <vector>
#include <map>




MpiManager::MpiManager(int argc, char *argv[])
{
	MPI_Init(&argc, &argv);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank_world);
	MPI_Comm_size(MPI_COMM_WORLD, &nproc_world);
	MPI_Get_processor_name(node_name, &namelen);
}

MpiManager::~MpiManager()
{
	MPI_Barrier(MPI_COMM_WORLD);
	MPI_Finalize();
}

void MpiManager::getNode()
{
	MPI_Status status;
	
	int  temp_rank;
	char temp_name[100];

	if (rank_world==0)
	{
		std::map<std::string, std::vector<int> > nodes;

		nodes[node_name].push_back(rank_world);

		// receive info from procs
		for (int p=1; p<nproc_world; ++p)
		{
			MPI_Recv(&temp_rank, 1, MPI_INT, p, 0, MPI_COMM_WORLD, &status);
			MPI_Recv(temp_name, 100, MPI_CHAR, p, 0, MPI_COMM_WORLD, &status);
			nodes[temp_name].push_back(temp_rank);
		}

		nodes_world = (int)nodes.size();

		// send info to procs
		int mynode = 0;
		for (std::map<std::string, std::vector<int> >::iterator n=nodes.begin(); n!=nodes.end(); ++n)
		{
			for (std::vector<int>::iterator r=n->second.begin(); r!=n->second.end(); ++r)
			{
				if (*r==0)
				{
					nproc_group = (int)n->second.size();
					node = mynode;
					ranks = new int[nproc_group];
					for (int p=0; p< nproc_group; ++p)
					{
						ranks[p] = n->second[p];
					}
				}
				else
				{
					nproc_group = (int)n->second.size();
					MPI_Send(&nodes_world, 1, MPI_INT, *r, 0, MPI_COMM_WORLD);
					MPI_Send(&mynode, 1, MPI_INT, *r, 0, MPI_COMM_WORLD);
					MPI_Send(&nproc_group, 1, MPI_INT, *r, 0, MPI_COMM_WORLD);
					MPI_Send(&n->second[0], nproc_group, MPI_INT, *r, 0, MPI_COMM_WORLD);
				}
			}
			mynode++;
		}	
	}
	else
	{
		// send info to main
		MPI_Send(&rank_world, 1, MPI_INT, 0, 0, MPI_COMM_WORLD);
		MPI_Send(node_name, 100, MPI_CHAR, 0, 0, MPI_COMM_WORLD);

		// receive infofrom main
		MPI_Recv(&nodes_world, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);
		MPI_Recv(&node, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);
		MPI_Recv(&nproc_group, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);
		ranks = new int[nproc_group];
		MPI_Recv(ranks, nproc_group, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);
	}
}

void MpiManager::createGroup()
{
	getNode();

	MPI_Group	group_original;
	MPI_Group	group;

	MPI_Comm_group(MPI_COMM_WORLD, &group_original);

	MPI_Group_incl(group_original, nproc_group, ranks, &group);

	MPI_Comm_create(MPI_COMM_WORLD, group, &comm); 

	MPI_Group_rank (group, &rank_group);
}

void MpiManager::createGroup(int nodes, int nproc_group)
{
	if (nodes*nproc_group != nproc_world)
	{
		std::cout << "ERROR: nodes*ppn invalid!" << std::endl;
		// do something here?
	}

	this->nodes_world = nodes;
	this->nproc_group = nproc_group;

	MPI_Group	group_original;
	MPI_Group	group;

	MPI_Comm_group(MPI_COMM_WORLD, &group_original);

	ranks = new int[nproc_group];

	node = rank_world/nproc_group;
	rank_group = rank_world - node*nproc_group;

	for (int p=0; p<nproc_group; ++p)
	{
		ranks[p] = p + node*nproc_group;
	}

	MPI_Group_incl(group_original, nproc_group, ranks, &group);

	MPI_Comm_create(MPI_COMM_WORLD, group, &comm); 

	MPI_Group_rank (group, &rank_group);
}

//void MpiManager::info(std::ofstream &ofs, std::string type)
//{
//	saveVariable(ofs, "mpi.nproc_world",	nproc_world, type);
//	saveVariable(ofs, "mpi.nodes_world",	nodes_world, type);
//	saveVariable(ofs, "mpi.ensembles",		nproc_group, type);
//	saveVariable(ofs, "mpi.node_name",		node_name, type);
//	saveVariable(ofs, "mpi.rank_group",		rank_group, type);
//	saveVariable(ofs, "mpi.rank_world",		rank_world, type);
//
//	for (int p=0; p< nproc_group; ++p)
//	{
//		char c[100];
//		sprintf(c, "mpi.groupmembers.proc_%02d", p);
//		saveVariable(ofs, c, ranks[p], type);
//	}
//}

void MpiManager::print()
{
	std::cout << "mpi.nproc_world: " << nproc_world << std::endl;
	std::cout << "mpi.nodes_world: " << nodes_world << std::endl;
	std::cout << "mpi.ensembles: " << nproc_group << std::endl;
	std::cout <<  "mpi.node_name: " << node_name << std::endl;
	std::cout << "mpi.rank_group: " << rank_group << std::endl;
	std::cout << "mpi.rank_world: " << rank_world << std::endl;

	for (int p = 0; p < nproc_group; ++p)
	{
		char c[100];
		sprintf(c, "mpi.groupmembers.proc_%02d: ", p);
		std::cout << c << ranks[p] << std::endl;
	}
}

