#include <iostream>
#include <sstream>
#include <mpi.h>
#include <sys/times.h>

using namespace std;

char* cmdname;
void usage() {
   cerr << "Usage: " << cmdname << " [bufsize [count]]" << endl;
   MPI_Abort(MPI_COMM_WORLD, 1);
}

int setup_chain(int nofprocesses) {
   int* indices = new int[nofprocesses];
   for (int i = 0; i < nofprocesses; ++i) {
      indices[i] = i;
   }
   // shuffle indices but keep 0 at pos 0
   for (int i = 1; i < nofprocesses-1; ++i) {
      int j = i + lrand48() % (nofprocesses - i);
      if (i != j) {
	 int tmp = indices[i]; indices[i] = indices[j]; indices[j] = tmp;
      }
   }
   for (int i = 1; i < nofprocesses; ++i) {
      int next = i + 1; if (next == nofprocesses) next = 0;
      MPI_Send(&indices[next], 1, MPI_INT, indices[i], 0, MPI_COMM_WORLD);
   }
   int next = indices[1];
   delete[] indices;
   return next;
}

void chase_packets(int bufsize, int count, int next, int rank) {
   unsigned char* buf = new unsigned char[bufsize];
   MPI_Status status;
   if (rank == 0) {
      while (count-- > 0) {
	 MPI_Send(buf, bufsize, MPI_UNSIGNED_CHAR, next, 0, MPI_COMM_WORLD);
	 MPI_Recv(buf, bufsize, MPI_UNSIGNED_CHAR,
	    MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
      }
   } else {
      while (count-- > 0) {
	 MPI_Recv(buf, bufsize, MPI_UNSIGNED_CHAR,
	    MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
	 MPI_Send(buf, bufsize, MPI_UNSIGNED_CHAR, next, 0, MPI_COMM_WORLD);
      }
   }
}

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

   int nofprocesses; MPI_Comm_size(MPI_COMM_WORLD, &nofprocesses);
   int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);

   int bufsize = 32;
   int count = 32;
   if (rank == 0) {
      cmdname = *argv++; --argc;
      if (argc > 0) {
	 istringstream arg(*argv++); --argc;
	 if (!(arg >> bufsize) || bufsize <= 0) usage();
      }
      if (argc > 0) {
	 istringstream arg(*argv++); --argc;
	 if (!(arg >> count) || count <= 0) usage();
      }
      if (argc > 0) usage();
   }
   MPI_Bcast(&bufsize, 1, MPI_INT, 0, MPI_COMM_WORLD);
   MPI_Bcast(&count, 1, MPI_INT, 0, MPI_COMM_WORLD);

   int next;
   if (rank == 0) {
      next = setup_chain(nofprocesses);
   } else {
      MPI_Status status;
      MPI_Recv(&next, 1, MPI_INT, MPI_ANY_SOURCE, MPI_ANY_TAG,
	 MPI_COMM_WORLD, &status);
   }

   struct timeval tbuf1; gettimeofday(&tbuf1, 0);
   chase_packets(bufsize, count, next, rank);
   struct timeval tbuf2; gettimeofday(&tbuf2, 0);
   long int usec = (tbuf2.tv_sec - tbuf1.tv_sec) * 1000000 +
                   (tbuf2.tv_usec - tbuf1.tv_usec);

   MPI_Finalize();

   if (rank == 0) {
      double avgTimeInMicrosecs =
	 (double) usec / (double) count / nofprocesses;
      cout << "avg message time in us: " << avgTimeInMicrosecs << endl;
   }
}