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

/* Quick and dirty primality test */
bool is_prime(int num) {
    if (num == 2 || num == 3) { return true; }
    if ((num % 2) == 0 || (num % 3) == 0) { return false; }

    for (int div = 5; div * div <= num; div += 6) {
        if ((num % div) == 0 || ((num + 2) % div) == 0) { return false; }
    }

    return true;
}

int main(int argn, char* argv[]) {

    MPI_Init(nullptr, nullptr);

    int mpi_size;
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);

    // Check if a 2D grid is (more or less) possible
    if ((mpi_size > 2) && is_prime(mpi_size)) {
        // Print warning only once (by rank 0)
        int mpi_rank;
        MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
        if (mpi_rank == 0) {
            std::cout << "Warning: 2D grid degenerates to 1D." << std::endl;
        }
    }

    // Group processes into a cartesian communication topology
    int mpi_dims[2] = {0, 0};
    MPI_Dims_create(mpi_size, 2, mpi_dims);

    // Setup a cartesian topology communicator (NON-cyclic)
    const int mpi_periods[2] = {false, false};
    MPI_Comm mpi_comm_grid;
    MPI_Cart_create(
        MPI_COMM_WORLD,     // Old Communicator
        2,                  // number of dimensions
        mpi_dims,           // grid dimensions
        mpi_periods,        // grid periodicity
        true,               // allow process reordering
        &mpi_comm_grid
    );

    // Get rank with respect to the grid communicator
    int mpi_rank;
    MPI_Comm_rank(mpi_comm_grid, &mpi_rank);

    // Get coordinates with respect to the grid communicator
    int mpi_coords[2];
    MPI_Cart_coords(mpi_comm_grid, mpi_rank, 2, mpi_coords);

    // Get direct neightbours in the communication grid
    struct { int north; int east; int south; int west; } mpi_neighbours;
    // Get X-direction (dim 0) neightbours
    MPI_Cart_shift(
        mpi_comm_grid,              // grid communicator
        0,                          // axis index (0 <-> X)
        1,                          // offset
        &(mpi_neighbours.west),     // negated offset neightbour
        &(mpi_neighbours.east)      // offset neightbour
    );
    // Get Y-direction (dim 1) neightbours
    MPI_Cart_shift(
        mpi_comm_grid,
        1,                          // axis index (1 <-> X)
        1,
        &(mpi_neighbours.south),
        &(mpi_neighbours.north)
    );

    // Print grid communicator location and with direct neighbors (offset +-1)
    std::cout << "Rank " << mpi_rank << ": Coords ("
        << mpi_coords[0] << ", "
        << mpi_coords[1] << ") - Neighbours: ";
    if (mpi_neighbours.west  != MPI_PROC_NULL) { std::cout << mpi_neighbours.west  << ' '; }
    if (mpi_neighbours.east  != MPI_PROC_NULL) { std::cout << mpi_neighbours.east  << ' '; }
    if (mpi_neighbours.south != MPI_PROC_NULL) { std::cout << mpi_neighbours.south << ' '; }
    if (mpi_neighbours.north != MPI_PROC_NULL) { std::cout << mpi_neighbours.north << ' '; }
    std::cout << std::endl;

    MPI_Finalize();

    return 0;
}