#include <iostream>
#include <sys/time.h>
#include <unistd.h>     /* gethostname() */
#include <stdlib.h>     /* rand() */
#include <mpi.h>
using namespace std;

#define TRIALS 100

double sendAndTime(int length);
void receive(int length);

int main(int argc, char * argv[])
{
    int my_rank;
    double max_bytes_per_sec = 0.0;
    int max_bps_pkt_size = 0;
    char hostname[256];

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

    if (my_rank == 0)
    {
        MPI_Recv(&hostname[0], 256, MPI_CHAR, MPI_ANY_SOURCE,
            MPI_ANY_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        cout << "Remote host name: " << hostname << endl;
    }
    else
    {
        gethostname(&hostname[0], 255);
        MPI_Send(&hostname[0], 256, MPI_CHAR, 0, 42, MPI_COMM_WORLD);
    }

    for (int length = 100; length < 100000; length += 100)
    {
        if (my_rank == 0)
        {
            cout << "Testing packet length " << length << endl;
            double rtt = sendAndTime(length);
            double bytes_per_sec = (length * TRIALS) / rtt;
            if (bytes_per_sec > max_bytes_per_sec)
            {
                max_bps_pkt_size = length;
                max_bytes_per_sec = bytes_per_sec;
            }
        }
        else
        {
            receive(length);
        }
    }

    if (my_rank == 0)
    {
        cout << "Maximum bytes per second (round trip) was " << max_bytes_per_sec << endl;
        cout << "This occurred with a packet size of " << max_bps_pkt_size << endl;
    }
    else
    {
        cout << "Slave exiting." << endl;
    }

    MPI_Finalize();

    return 0;
}

double sendAndTime(int length)
{
    char buff[length];
    for (int i = 0; i < length; i++)
    {
        buff[i] = rand() & 0xFF;
    }

    struct timeval before, after;
    gettimeofday(&before, NULL);        /* Start timing */

    for (int i = 0; i < TRIALS; i++)
    {
        int ret = MPI_Send(&buff, length, MPI_CHAR, 0, 42, MPI_COMM_WORLD);
//        cout << "MPI_Send() returned " << ret << endl;

        MPI_Recv(&buff, length, MPI_CHAR,
            MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }

    gettimeofday(&after, NULL);         /* Stop timing */

    double time_before = before.tv_sec + before.tv_usec / 1000000.0;
    double time_after = after.tv_sec + after.tv_usec / 1000000.0;
    double diff = time_after - time_before;
    return diff;
}

void receive(int length)
{
    char buff[length];

    for (int i = 0; i < TRIALS; i++)
    {
        MPI_Recv(&buff, length, MPI_CHAR,
            MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

        MPI_Send(&buff, length, MPI_CHAR, 0, 42, MPI_COMM_WORLD);
    }
}