#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sched.h>
#include <mpi.h>
#include <omp.h>

/*
  Compile:
    mpicc -fopenmp -O2 -o mpi_openmp_core_info mpi_openmp_core_info.c

  Run (example):
    mpirun -np 4 ./mpi_openmp_core_info 4
    # -> 4 MPI ranks, each spawns 4 OpenMP threads
*/

int main(int argc, char **argv) {
    int required = MPI_THREAD_FUNNELED;
    int provided;
    MPI_Init_thread(&argc, &argv, required, &provided);

    int world_rank, world_size;
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    int num_threads = 1;
    if (argc >= 2) num_threads = atoi(argv[1]);
    if (num_threads <= 0) num_threads = 1;

    pid_t pid = getpid();

    printf("[Rank %d] PID %d starting, total ranks = %d, OpenMP threads = %d\n",
           world_rank, pid, world_size, num_threads);
    fflush(stdout);

    omp_set_num_threads(num_threads);

#pragma omp parallel
    {
        int tid = omp_get_thread_num();
        int cpu = sched_getcpu();

#pragma omp critical
        {
            printf("[Rank %d] PID %d | Thread %d/%d | Running on core %d\n",
                   world_rank, pid, tid, omp_get_num_threads(), cpu);
            fflush(stdout);
        }
    }

    MPI_Barrier(MPI_COMM_WORLD);
    if (world_rank == 0) {
        printf("All ranks finished.\n");
    }

    MPI_Finalize();
    return 0;
}