#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>

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

int main(int argc, char **argv)
{
  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)
	{
	  printf("\n\t Usage: <executable> <number processes along X> \n");
	  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)
	{
	  printf("\n\t ntasks % cartesian_grid_x != 0 ... aborting ...\n");
	  fflush(stdout);
	  MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
	  exit(EXIT_FAILURE);
	}
    }
  
  const int dims[SIZE] = {cartesian_grid_x, cartesian_grid_y};
  const int periods[SIZE] = {0, 0};
  const 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)
	{
	  printf("\n\t Task: %d\n", task);
	  printf("\t\t coords[%d, %d]", coords[X], coords[Y]);
	  printf("\n\t\t nbrright: %d - nbrleft  : %d", nbrright, nbrleft);
	  printf("\n\t\t nbrtop  : %d - nbrbottom: %d \n\n", nbrtop, nbrbottom);
	  fflush(stdout);
	}
    }

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