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

int main (int argc, char *argv[])
{
  int num_tasks;            /* Total number of tasks in the program */
  int this_task;            /* The current process's task number */
  int xdim = 1;             /* Mesh x dimension */
  int ydim = 1;             /* Mesh y dimension */
  int reps = 10000;         /* Number of repetitions of the inner loop */
  int msgsize = 0;          /* Message size in bytes */
  void *msgbuf;             /* Buffer from which to send/receive messages */
  double mean_secs = 0.0;   /* Mean per-hop time in seconds */
  double stdev_secs = 0.0;  /* Standard deviation of per-hop time in seconds */
  int numhops;              /* Length of communication critical path */
  struct poptOption cmdlineopts[] = {
    {"width",   'w', POPT_ARG_INT, &xdim,    0, "Mesh width (tasks)"},
    {"height",  'h', POPT_ARG_INT, &ydim,    0, "Mesh height (tasks)"},
    {"reps",    'r', POPT_ARG_INT, &reps,    0, "Number of wavefronts to time"},
    {"msgsize", 'm', POPT_ARG_INT, &msgsize, 0, "Message size (bytes)"},
    POPT_AUTOHELP
    {NULL, 0, 0, NULL, 0}
  };      
  poptContext poptcon;      /* popt option context */
  char lastarg;             /* Character associated with final popt option */
  MPI_Status dummy_status;  /* MPI_STATUS_NULL doesn't exist in MPI 1. */
  int i;

  /* Initialize MPI. */
  MPI_Init (&argc, &argv);
  MPI_Comm_size (MPI_COMM_WORLD, &num_tasks);
  MPI_Comm_rank (MPI_COMM_WORLD, &this_task);

  /* Use popt to parse the command line. */
  poptcon = poptGetContext (NULL, argc, (const char **)argv, cmdlineopts, 0);
  if ((lastarg = poptGetNextOpt(poptcon)) < -1) {
    fprintf (stderr, "%s: %s\n",
	     poptBadOption (poptcon, POPT_BADOPTION_NOALIAS),
	     poptStrerror (lastarg));
    return 1;
  }

  /* Allocate memory for the message buffer. */
  if (!(msgbuf = (void *) malloc (msgsize)))
    abort();

  /* Perform the communication pattern. */
  numhops = xdim + ydim - 1;
  for (i=0; i<reps; i++) {
    double elapsed_secs = MPI_Wtime();

    if (this_task%xdim != 0)
      MPI_Recv (msgbuf, msgsize, MPI_BYTE, this_task-1, 1, MPI_COMM_WORLD, &dummy_status);
    if (this_task >= xdim)
      MPI_Recv (msgbuf, msgsize, MPI_BYTE, this_task-xdim, 1, MPI_COMM_WORLD, &dummy_status);
    if ((this_task+1)%xdim != 0)
      MPI_Send (msgbuf, msgsize, MPI_BYTE, this_task+1, 1, MPI_COMM_WORLD);
    if (this_task < num_tasks-xdim)
      MPI_Send (msgbuf, msgsize, MPI_BYTE, this_task+xdim, 1, MPI_COMM_WORLD);
    if (this_task == 0)
      MPI_Recv (msgbuf, msgsize, MPI_BYTE, num_tasks-1, 1, MPI_COMM_WORLD, &dummy_status);
    else
      if (this_task == num_tasks-1)
	MPI_Send (msgbuf, msgsize, MPI_BYTE, 0, 1, MPI_COMM_WORLD);
    elapsed_secs = (MPI_Wtime() - elapsed_secs) / numhops;
    mean_secs += elapsed_secs;
    stdev_secs += elapsed_secs * elapsed_secs;
  }
  stdev_secs = sqrt ((reps*stdev_secs - mean_secs*mean_secs) / (reps*(reps-1)));
  mean_secs /= reps;

  /* Output the mean and standard deviation of the per-hop time. */
  if (this_task == 0) {
    printf ("Per-hop latency (usecs, mean)\tPer-hop latency (usecs, std. dev.)\n");
    printf ("%.10g\t%.10g\n", mean_secs*1e6, stdev_secs*1e6);
  }

  /* Finish up. */
  free (msgbuf);
  MPI_Finalize();
  return 0;
}