#include <mpi.h>
#include <iostream>

#define SIZE 2
#define X    0
#define Y    1

int main(int argc, char **argv)
{
  using namespace std;
  
  int task, ntasks;
  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &ntasks);
  MPI_Comm_rank(MPI_COMM_WORLD, &task);

  if (argc <= 1)
    {
      if (!task)
	{
	  cout << "\n\t Usage: <executable> <number processes along X> \n" << endl;
	  MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
	  exit(EXIT_FAILURE);
	}
    }
  
  const int cartesian_grid_x = (int)strtol(argv[1], NULL, 10);
  const int cartesian_grid_y = ((ntasks % cartesian_grid_x == 0) ? (ntasks / cartesian_grid_x) : -1);
  if (cartesian_grid_y == -1)
    {
      if (!task)
	{
	  cout << "\n\t ntasks % cartesian_grid_x != 0 ... aborting ...\n" << endl;
	  MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
	  exit(EXIT_FAILURE);
	}
    }
  
  static int dims[SIZE] = {cartesian_grid_x, cartesian_grid_y};
  static int periods[SIZE] = {0, 0};
  static int reorder = 0;
  MPI_Comm comm2d;
  MPI_Cart_create(MPI_COMM_WORLD, SIZE, dims, periods, reorder, &comm2d);

  int coords[SIZE];
  MPI_Cart_coords(comm2d, task, SIZE, coords);

  int nbrright, nbrleft;
  MPI_Cart_shift(comm2d, Y, 1, &nbrleft, &nbrright);

  int nbrtop, nbrbottom;
  MPI_Cart_shift(comm2d, X, 1, &nbrbottom, &nbrtop);

  for (int rank=0 ; rank<ntasks; rank++)
    {
      MPI_Barrier(MPI_COMM_WORLD);
      if (rank == task)
	{
	  cout << "\n\t Task: " << task << endl;
	  cout << "\t\t coords[" << coords[X] << "," << coords[Y] << "]" << endl;
	  cout << "\t\t nbrright: " << nbrright << " - nbrleft  : " << nbrleft << endl;
	  cout << "\t\t nbrtop  : " << nbrtop   << " - nbrbottom: " << nbrbottom << endl;
	  cout << endl;
	}
    }

  MPI_Comm_free(&comm2d);
  MPI_Finalize();
  
  return 0;
}