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

#define REPETITIONS 1
#define SAMPLES     100000
#define SKIP        100
#define BUFLEN      (1<<16)

void pingpong(int is_root, MPI_Comm comm, void *buf, int count, int peer)
{
  if (is_root) {
    MPI_Send(buf, count, MPI_BYTE, peer, 0, comm);
    MPI_Recv(buf, count, MPI_BYTE, peer, 0, comm, MPI_STATUS_IGNORE);
  } else {
    MPI_Recv(buf, count, MPI_BYTE, peer, 0, comm, MPI_STATUS_IGNORE);
    MPI_Send(buf, count, MPI_BYTE, peer, 0, comm);
  }
}

static char buf[BUFLEN];

int main(int argc, char *argv[])
{
  MPI_Init(&argc, &argv);

  int size, rank;
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  assert(size <= 2);

  MPI_Comm comm = MPI_COMM_WORLD;
  int is_root = (rank == 0);
  int use_dpm = (size == 1);

  if (use_dpm) {
    MPI_Comm_get_parent(&comm);
    is_root = (comm == MPI_COMM_NULL);
    if (is_root) {
      MPI_Comm_spawn(argv[0], MPI_ARGV_NULL, 1, MPI_INFO_NULL,
                     0, MPI_COMM_SELF, &comm, MPI_ERRCODES_IGNORE);
    }
  }


  for (int r = 0; r < REPETITIONS; r++)
    {
      int peer = use_dpm ? 0 : (is_root ? 1 : 0);

      double tic = MPI_Wtime();
      for (int i = 0; i < SAMPLES+SKIP; i++) {
        if (i == SKIP) tic = MPI_Wtime();
        pingpong(is_root, comm, buf, BUFLEN, peer);
      }
      double toc = MPI_Wtime();

      if (is_root) {
        printf("samples: %d  buflen: %d  latency: %.2f us\n",
               SAMPLES, BUFLEN, (toc-tic)/SAMPLES * 1e6);
      }

    }

  if (comm != MPI_COMM_WORLD)
    MPI_Comm_disconnect(&comm);
  MPI_Finalize();
  return 0;
}