#include <mpi.h>
#include <cuda.h>
#include <cuda_runtime.h>
//#include <iostream>
#include <cublas_v2.h>
#include <vector>
#include <numeric>
#include <time.h>
#include <sys/time.h>


#define CUDA(x)           \
{                               \
  cudaError_t err = x;          \
  if (err != cudaSuccess)       \
  {                             \
    printf("!!! CUDA ERROR: \"%s\" at file %s, line %d !!!\n", cudaGetErrorString(err), __FILE__, __LINE__);\
    exit(1);                    \
  }                             \
}



#define CUBLAS(x)                         \
{                                               \
  cublasStatus_t stat = x;                      \
  if (stat != CUBLAS_STATUS_SUCCESS)            \
  {                                             \
     printf ("CUBLAS failed at file %s, line %d, aborting\n", __FILE__, __LINE__); \
     exit(2);                                   \
  }                                             \
} 

#define TIME(x) if (mpi_rank == 0) {CUDA(cudaStreamSynchronize(0));  printf("[%.4f] : %s\n", wallclock(), #x);} 
//#define TIME(x)  if (mpi_rank <= 1) printf("[%d: %.4f] : %s\n", mpi_rank, wallclock(), #x); 

double wallclock()
{
  const double mega = 1e-6;
  struct timeval tp;
  int n = gettimeofday (&tp, NULL);
  if (n != 0)
    printf("**** gettimeofday failed with %d ****\n", n);

  double time = (double) (tp.tv_sec);
  double mic  = (double) (tp.tv_usec);
  time += mic * mega;

  static bool first = true;
  static double t0 = time;
  double result = time - t0;
  t0 = time;
  if (first)
  {
    first = false;
    result = 0;
  }

  return result;
}


double my_random()
{
  return  static_cast<double>(rand()) / RAND_MAX;
}




int main(int argc, char* argv[])
{
  const int N = (argc == 1) ? 600000000 : atoi(argv[1]);

  MPI_Init(&argc, &argv);
 
  int mpi_size, mpi_rank;

  
  MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
  MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);

  TIME("start");

  CUDA( cudaDeviceSynchronize() );    TIME("cudaDeviceSynchronize =   CUDA initialization");
 
  if (mpi_rank == 0)
  {
    int count = -1;
    CUDA( cudaGetDeviceCount(&count) );
    printf("\t\tRunning %d MPI processes; found %d CUDA-capable devices\n", mpi_size, count);
    if (count < mpi_size)
    {
      printf("\t\tError: too many MPI processes for %d devices, aborting\n", count);
//      exit(1);
    }
    printf("\t\tProcessing %d doubles (%.2f GB)\n",  N, N*sizeof(double)/1e9); 
  }
  TIME("Checking the number of devices");  

 
  std::vector<int> buf_size(mpi_size);
  int s = 0;
  for (int i = 0; i < mpi_size; i++)
  {
     buf_size[i] = N/mpi_size;
     s += buf_size[i];
  }
  buf_size[mpi_size - 1] += N - s;

  std::vector<double> v(buf_size[mpi_rank]);  TIME("Array allocation on host");
  srand(time(0)+ 10*mpi_rank);
  for(int i = 0; i < v.size(); ++i)
  {
    v[i] = my_random();
  }
  TIME("Array initialization on host");

  double h_partial_result = std::accumulate(v.begin(), v.end(), 0.0);
  double h_result = -1;
  MPI_Reduce(&h_partial_result, &h_result, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
  TIME("Calculating the sum on host");

  if (mpi_rank == 0)
     printf("\t\tSum on the CPU: %.16g\n", h_result);
  
  
  //  CUDA + CUBLAS
  double * d_v = 0;
  double device_partial_result = 0;
  double d_result = 0;
  cublasHandle_t handle = 0;
  
  CUDA( cudaSetDevice(mpi_rank) );                                                   TIME("cudaSetDevice");
  CUDA( cudaMalloc((void**)&d_v, sizeof(double) * v.size()) );                       TIME("cudaMalloc");
  CUDA( cudaMemcpy(d_v, &v[0], sizeof(double) * v.size(), cudaMemcpyHostToDevice) ); TIME("cudaMemcpy");
  CUBLAS( cublasCreate(&handle) );                                                   TIME("cublasCreate");
  CUBLAS( cublasDasum(handle, v.size(), d_v, 1, &device_partial_result) );           TIME("cublasDasum (calculates the sum on the GPU)");
  CUBLAS( cublasDestroy(handle) );                                                   TIME("cublasDestroy");
  CUDA( cudaFree(d_v) );                                                              TIME("cudaFree"); 
  MPI_Reduce(&device_partial_result, &d_result, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);  
  if (mpi_rank == 0)
     printf("\t\tSum on the GPU: %.16g\n", d_result);
  

  MPI_Finalize();
}