#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MPI_SAFE_CALL(call) do{int ierr=call;if (ierr != MPI_SUCCESS) return 1;}while(0)

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

  // get the number of processes
  int np;
  MPI_SAFE_CALL(MPI_Comm_size(MPI_COMM_WORLD, &np));

  // get my ID
  int my_id;
  MPI_SAFE_CALL(MPI_Comm_rank(MPI_COMM_WORLD, &my_id));

  // default values
  long arg_rep=10;
  long arg_size=4096;
  long arg_verbose=0;

  // parse args
  for (int i=1; i<argc; i=i+2){
    if (i>=(argc-1)) {
        fprintf(stderr, "USE: %s [--repetitions VALUE] [--size VALUE(bytes)] [--verbose 1|0]\n", argv[0]);
	return 1;
    }
    if      (strcmp("--repetitions", argv[i]) == 0) arg_rep     = atol(argv[i+1]);
    else if (strcmp("--size",        argv[i]) == 0) arg_size    = atol(argv[i+1]);
    else if (strcmp("--verbose",     argv[i]) == 0) arg_verbose = atol(argv[i+1]);
    else {
        fprintf(stderr, "USE: %s [--repetitions VALUE] [--size VALUE(bytes)] [--verbose 1|0]\n", argv[0]);
	return 1;
    }
  }

  // validate args
  const long nrep = arg_rep;
  if (nrep < 1) return 1;
  const long msg_size = arg_size;
  if (msg_size < 1) return 1;
  const long verbose=(arg_verbose) ? 1 : 0;

  const int sendcount = msg_size;
  const int recvcount = sendcount;

  // maximum number of inflight messages per sender process
  const int max_req=100;
  size_t msize = ((size_t)sendcount)*((size_t)max_req);

  if (my_id == 0) {
    printf("verbose = %ld\n", verbose);
    printf("nrep = %ld\n", nrep);
    printf("sendcount = %d\n", sendcount);
    printf("recvcount = %d\n", recvcount);
  }


  char* const sendrecvbuf = malloc(msize);
  memset(sendrecvbuf, 0, msize);


  char my_hostname[(MPI_MAX_PROCESSOR_NAME/16+1)*16];
  int len;
  MPI_SAFE_CALL(MPI_Get_processor_name(my_hostname, &len));
  if ((len >= 0) && (len <= MPI_MAX_PROCESSOR_NAME)) {
    my_hostname[len]='\0';
  }
  else {
    my_hostname[0]='\0';
  }

  MPI_Status status;
  long ncalls=0;
  int ireq=0;
  MPI_Request request[max_req];
  const int recv_ratio=4;
  if (verbose) {printf("Hello from host %s, process: %d / %d\n", my_hostname, my_id, np); fflush(stdout);}

  MPI_SAFE_CALL(MPI_Barrier(MPI_COMM_WORLD));
  const double t0 = MPI_Wtime();
  for (long i=0; i<nrep; ++i) {
    if (verbose) {printf("proc %d repetition %ld\n", my_id, i); fflush(stdout);}
    if (my_id%recv_ratio == 0) {
        for (int iproc=0; iproc<np; ++iproc) if (iproc%recv_ratio != 0) {
	    ++ncalls;
            MPI_SAFE_CALL(MPI_Recv(sendrecvbuf, recvcount, MPI_BYTE, MPI_ANY_SOURCE, 0, MPI_COMM_WORLD, &status));
            if (verbose) {printf("proc %d MPI_Recv(%d bytes) from %d\n", my_id, recvcount, status.MPI_SOURCE); fflush(stdout);}
            // MPI_SAFE_CALL(MPI_Recv(sendrecvbuf, recvcount, MPI_BYTE, iproc, 0, MPI_COMM_WORLD, &status));
        }
    }
    else { 
        for (int iproc=0; iproc<np; ++iproc) if  (iproc%recv_ratio == 0) {
            //MPI_SAFE_CALL(MPI_Send(sendrecvbuf, sendcount, MPI_BYTE, iproc, 0, MPI_COMM_WORLD));
            MPI_SAFE_CALL(MPI_Isend(&sendrecvbuf[ireq*sendcount], sendcount, MPI_BYTE, iproc, 0, MPI_COMM_WORLD, &request[ireq]));
            if (verbose) {printf("proc %d MPI_Isend (%d bytes) to %d\n", my_id, sendcount, iproc); fflush(stdout);}
	    ++ireq;
            if (ireq>=max_req) {
                if (verbose) {printf("proc %d MPI_Waitall\n", my_id); fflush(stdout);}
	        MPI_SAFE_CALL(MPI_Waitall(max_req, request, MPI_STATUSES_IGNORE));
		ireq=0;
	    }
        }
    }
  }
  if (ireq>0) {
    if (verbose) {printf("proc %d MPI_Waitall\n", my_id); fflush(stdout);}
    MPI_SAFE_CALL(MPI_Waitall(ireq, request, MPI_STATUSES_IGNORE));
    ireq=0;
  }
  MPI_SAFE_CALL(MPI_Barrier(MPI_COMM_WORLD));
  if ((my_id == 0) && (ncalls>0)){
  const double t1 = MPI_Wtime();
  const double wtime = t1-t0;
  const double latency = 1.0e6*wtime / ncalls;
  const double mbps = 1.0e-6*recvcount*ncalls/wtime;
    printf("wtime = %.3lf\n", wtime);
    printf("Bandwidth per receiver process = %.3lf MB/s\n", mbps);
    printf("MPI_Recv latency = %.3lf us\n", latency);
  }


  MPI_SAFE_CALL(MPI_Finalize());
  return 0;
}