#include <iostream>
#include <chrono>
#include <vector>
#include <numeric>
#include <cstring>

#define BOOST_MPI_HOMOGENEOUS

#include <boost/mpi.hpp>

typedef double scalar_type;
typedef std::chrono::duration<double, std::ratio<1>> seconds;

BOOST_CLASS_TRACKING(std::vector<scalar_type>, boost::serialization::track_never)
BOOST_CLASS_IMPLEMENTATION(std::vector<scalar_type>, boost::serialization::object_serializable)

#define N_REP 100
#define MAX_SIZE (1<<20)

int main(int argc, char** argv) {
    boost::mpi::environment env(argc, argv);
    boost::mpi::communicator comm;
    
    if (comm.rank() == 0) {
        std::cout << "------------" << std::endl;
        std::cout << "Plain array:" << std::endl;
        std::cout << "------------" << std::endl;
        std::cout << "N\tTime [sec]\t# bytes per second" << std::endl;
    }
    for (unsigned long n = 2; n <= MAX_SIZE; n <<= 1) {
        double elapsed = 0;
        for (int i = 0; i < N_REP; i++) {
            if (comm.rank() == 1) {
                scalar_type* data = new scalar_type[n];
                
                // Touch the memory before benchmarking the send
                std::memset(data, 0, n*sizeof(scalar_type));
                
                comm.barrier();
                comm.send(0,0, data, n);
            }
            else if (comm.rank() == 0) {
                scalar_type* data = new scalar_type[n];
                
                // Touch the memory before benchmarking the send
                std::memset(data, 0, n*sizeof(scalar_type));
                
                comm.barrier();
                auto start = std::chrono::high_resolution_clock::now();
                comm.recv(1,0, data, n);
                auto stop = std::chrono::high_resolution_clock::now();
                elapsed += std::chrono::duration_cast<seconds>(stop - start).count();
            }
        }
        if (comm.rank() == 0) {
            std::cout << n << "\t" << elapsed/N_REP << "\t" << N_REP*n*sizeof(scalar_type)/elapsed << std::endl;
        }
    }
    
    if (comm.rank() == 0) {
        std::cout << std::endl;
        std::cout << "------------" << std::endl;
        std::cout << "std::vector:" << std::endl;
        std::cout << "------------" << std::endl;
        std::cout << "N\tTime [sec]\t# bytes per second" << std::endl;
    }
    for (unsigned long n = 2; n <= MAX_SIZE; n <<= 1) {
        double elapsed = 0;
        for (int i = 0; i < N_REP; i++) {
            if (comm.rank() == 1) {
                std::vector<scalar_type> data(n);
                
                comm.barrier();
                comm.send(0,0,data);
            }
            else if (comm.rank() == 0) {
                std::vector<scalar_type> data;
                
                comm.barrier();
                auto start = std::chrono::high_resolution_clock::now();
                comm.recv(1,0, data);
                auto stop = std::chrono::high_resolution_clock::now();
                elapsed += std::chrono::duration_cast<seconds>(stop - start).count();
            }
        }
        if (comm.rank() == 0) {
            std::cout << n << "\t" << elapsed/N_REP << "\t" << N_REP*n*sizeof(scalar_type)/elapsed << std::endl;
        }
    }
    
    if (comm.rank() == 0) {
        std::cout << std::endl;
        std::cout << "----------------------------------------------------" << std::endl;
        std::cout << "std::vector with split sending of skeleton and data:" << std::endl;
        std::cout << "----------------------------------------------------" << std::endl;
        std::cout << "N\tTime [sec]\t# bytes per second" << std::endl;
    }
    for (unsigned long n = 2; n <= MAX_SIZE; n <<= 1) {
        double elapsed = 0;
        for (int i = 0; i < N_REP; i++) {
            if (comm.rank() == 1) {
                std::vector<scalar_type> data(n);
                
                comm.barrier();
                comm.send(0,0, boost::mpi::skeleton(data));
                comm.send(0,0, boost::mpi::get_content(data));
            }
            else if (comm.rank() == 0) {
                std::vector<scalar_type> data;
                
                comm.barrier();
                auto start = std::chrono::high_resolution_clock::now();
                comm.recv(1,0, boost::mpi::skeleton(data));
                comm.recv(1,0, boost::mpi::get_content(data));
                auto stop = std::chrono::high_resolution_clock::now();
                elapsed += std::chrono::duration_cast<seconds>(stop - start).count();
            }
        }
        if (comm.rank() == 0) {
            std::cout << n << "\t" << elapsed/N_REP << "\t" << N_REP*n*sizeof(scalar_type)/elapsed << std::endl;
        }
    }
    
    if (comm.rank() == 0) {
        std::cout << std::endl;
        std::cout << "--------------------------------" << std::endl;
        std::cout << "std::vector sent as plain array:" << std::endl;
        std::cout << "--------------------------------" << std::endl;
        std::cout << "N\tTime [sec]\t# bytes per second" << std::endl;
    }
    for (unsigned long n = 2; n <= MAX_SIZE; n <<= 1) {
        double elapsed = 0;
        for (int i = 0; i < N_REP; i++) {
            if (comm.rank() == 1) {
                std::vector<scalar_type> data(n);
                
                comm.barrier();
                std::size_t s = data.size();
                comm.send(0,0, s);
                comm.send(0,0, data.data(), s);
            }
            else if (comm.rank() == 0) {
                std::vector<scalar_type> data;
                
                comm.barrier();
                auto start = std::chrono::high_resolution_clock::now();
                std::size_t s;
                comm.recv(1,0,s);
                data.resize(s);
                comm.recv(1,0, data.data(), data.size());
                auto stop = std::chrono::high_resolution_clock::now();
                elapsed += std::chrono::duration_cast<seconds>(stop - start).count();
            }
        }
        if (comm.rank() == 0) {
            std::cout << n << "\t" << elapsed/N_REP << "\t" << N_REP*n*sizeof(scalar_type)/elapsed << std::endl;
        }
    }
    
    return 0;
}