#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include "bcast-linear.c"
#include "bcast-flattree.c"
#include "bcast-binomial-tree.c"
#include "bcast-scatter-rdb-allgather.c"
#include "bcast-scatter-LR-allgather.c"
#include "bcast-knary.c"

extern int bcast_linear_segment_size;
extern int bcast_knary_segment_size;
    
#define BUFSIZE 10000000
int MAX_ITER=100;
char buffer[BUFSIZE];
int debug_level;

int main (int argc, char *argv[]) {
    int size;
    int count;
    int my_rank, num_procs;
    int tag = 0;
    MPI_Status status;              
    double start, end, elapsed;
    double seconds, bps, avg_latency;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);

    /*
    if (2 != num_procs) {
        fprintf (stderr, "mptest: this program requires exactly two processes.\n");
        exit(1);
    }
    */       
    if (1 < argc) {
        size = atoi(argv[1]);
    } else {
        size = 100;
    }

    if (2 < argc) MAX_ITER = atoi(argv[2]);

    if ((0 > size) || (BUFSIZE < size)) {
        fprintf (stderr, "mpitest: buffer size must be between 0 and %d.\n", BUFSIZE);
        exit (1);
    }

    MPI_Barrier(MPI_COMM_WORLD);
    if (0 != my_rank) {
        for (count = 0; count < MAX_ITER; count++) {
#ifdef LINEAR
          bcast_linear_segment_size = LINEAR;
	  bcast_linear(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(BINARY)
          bcast_knary_segment_size = BINARY;
	  bcast_knary(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(FLATTREE)
	  bcast_flattree(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(BINOMIALTREE)
	  bcast_binomial_tree(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(SCATTERALLGATHER_RDB)
	  bcast_scatter_rdb_allgather(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(SCATTERALLGATHER_LR)
	  bcast_scatter_LR_allgather(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#else
            MPI_Bcast(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#endif
	    MPI_Barrier(MPI_COMM_WORLD);
        }
    } 
    else if (my_rank == 0) 
    {
      start = MPI_Wtime();
      for (count = 0; count < MAX_ITER; count++) {
#ifdef LINEAR
          bcast_linear_segment_size = LINEAR;
	  bcast_linear(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(BINARY)
          bcast_knary_segment_size = BINARY;
	  bcast_knary(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(FLATTREE)
	  bcast_flattree(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(BINOMIALTREE)
	  bcast_binomial_tree(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(SCATTERALLGATHER_RDB)
	  bcast_scatter_rdb_allgather(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#elif defined(SCATTERALLGATHER_LR)
	  bcast_scatter_LR_allgather(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#else
            MPI_Bcast(buffer, size, MPI_CHAR, 0, MPI_COMM_WORLD);
#endif
	    MPI_Barrier(MPI_COMM_WORLD);
      }
      end = MPI_Wtime();
      elapsed = end - start;
      seconds = 1.0 * elapsed;
      bps = (2.0 * 8.0 * size * MAX_ITER) / seconds;
      avg_latency = seconds / MAX_ITER;

#ifdef LINEAR
      printf("Bcast with linear tree pipeline with segment size %d\n", LINEAR);
#elif defined(BINARY)
      printf("Bcast with binary tree pipeline with segment size %d\n", BINARY);
#elif defined(FLATTREE)
      printf("Bcast with flat tree.\n");
#elif defined(BINOMIALTREE)
      printf("Bcast with binomial tree.\n");
#elif defined(SCATTERALLGATHER_RDB)
      printf("Bcast with scatter-allgather(RDB).\n");
#elif defined(SCATTERALLGATHER_LR)
      printf("Bcast with scatter-allgather(LR).\n");
#else
      printf("Bcast with the default MPI_Bcast.\n");

#endif
      printf ("  Message size = %d bytes.\n", size);
      printf ("  Iterations = %d.\n", MAX_ITER);
      printf ("  Elapsed time = %g sec.\n", seconds);
      printf ("  Average per invocation time = %g sec.\n", avg_latency);
      printf ("  Average bandwidth = %g bps.\n", bps);        
    }
    MPI_Finalize();
    return 0;
}