#include <iostream>
#include <sstream>
#include <fstream>
#include <chrono>
#include <vector>
#include <algorithm>
#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 1000
#define MAX_SIZE (1<<20)

template<class F> auto time_fn(F&& f) {
  auto start = std::chrono::high_resolution_clock::now();
  f();
  auto stop = std::chrono::high_resolution_clock::now();
  return std::chrono::duration_cast<seconds>(stop - start).count();
}

template<class SENDR, class RECVR>
void test(std::string test_name, SENDR&& sendr, RECVR&& recvr) {
  boost::mpi::communicator comm;
  std::ofstream f;
  if (comm.rank() == 0) {
    f.open (test_name + ".dat");
    std::cout << "------------" << "\n";
    std::cout << test_name << ":"  << "\n";
    std::cout << "------------" << "\n";
    auto h = "#N\tTime [sec]\t# bytes per second";
    std::cout << h << "\n";  f << h << "\n";
  }
  for (std::size_t n = 2; n <= MAX_SIZE; n <<= 1) {
    double elapsed = 0;
    for (int i = 0; i < N_REP; i++) {
      if (comm.rank() == 1) {
        sendr(comm, n);
      } else if (comm.rank() == 0) {
        elapsed += recvr(comm);
      }
    }
    if (comm.rank() == 0) {
      std::stringstream ss;
      ss << n << "\t" << elapsed/N_REP << "\t"
         << N_REP*n*sizeof(scalar_type)/elapsed;
      std::cout << ss.str() << "\n";   f << ss.str() << "\n";
    }
  }
}

int main(int argc, char** argv) {
  boost::mpi::environment env(argc, argv);

  test("plain_array_with_alloc", [&](auto& comm, std::size_t n){
      auto* data = new scalar_type[n];
      comm.send(0, 0, n);
      comm.send(0, 0, data, n);
      delete [] data;
    }, [&](auto& comm) {
      return time_fn([&](){
          std::size_t n;
          comm.recv(1, 0, n);
          auto* data = new scalar_type[n];
          comm.recv(1,0, data, n);
          delete[] data;
        });
    });

   test("plain_array_no_alloc", [&](auto& comm, std::size_t n){
      auto* data = new scalar_type[n];
      comm.send(0, 0, n);
      comm.send(0, 0, data, n);
      delete[] data;
    }, [&](auto& comm) {
       std::size_t n;
       comm.recv(1, 0, n);
       auto* data = new scalar_type[n];
       auto time = time_fn([&](){
           comm.recv(1,0, data, n);
         });
       delete [] data;
       return time;
    });

  test("vector_as_array_with_alloc", [&](auto& comm, std::size_t n){
      std::vector<scalar_type> data(n);
      comm.send(0, 0, data.size());
      comm.send(0, 0, data.data(), data.size());
    }, [&](auto& comm) {
      std::vector<scalar_type> data;
      return time_fn([&](){
          std::size_t n;
          comm.recv(1, 0, n);
          data.resize(n); //< this default constructs elements!
          comm.recv(1,0, data.data(), data.size());
        });
    });

  test("vector_as_array_no_alloc", [&](auto& comm, std::size_t n){
      std::vector<scalar_type> data(n);
      comm.send(0, 0, data.size());
      comm.send(0, 0, data.data(), data.size());
    }, [&](auto& comm) {
      std::vector<scalar_type> data;
      std::size_t n;
      comm.recv(1, 0, n);
      data.reserve(n); //< this default constructs elements!
      return time_fn([&](){
          comm.recv(1,0, data.data(), data.size());
        });
    });

  test("plain_vector", [&](auto& comm, std::size_t n){
      std::vector<scalar_type> data(n);
      comm.send(0, 0, data.size());
      comm.send(0, 0, data);
    }, [&](auto& comm) {
      std::vector<scalar_type> data;
      std::size_t n;
      comm.recv(1, 0, n);
      data.resize(n); //< this default constructs elements!
      return time_fn([&](){
          comm.recv(1,0, data);
        });
    });

  test("vector_inefficient", [&](auto& comm, std::size_t n){
      std::vector<scalar_type> data(n);
      comm.send(0, 0, data.size());
      comm.send(0, 0, data.data(), data.size());
    }, [&](auto& comm) {
      std::vector<scalar_type> data;
      return time_fn([&](){
          std::size_t n;
          comm.recv(1, 0, n);

          auto* input = new scalar_type[n];
          comm.recv(1, 0, input, n);
          std::copy(input, input + n, std::back_inserter(data));
          delete [] input;
        });
    });

  return 0;
}