//
// program to add numbers (read in from file, or generated
//   randomly)
// illustrates use of collective-communication routines
//
// command-line argument:  
//   name of file containing numbers (doubles)
//   or "-n N" to randomly generate N doubles
//
// this is version 2, using MPI_Scatter to scatter input data
//   among processes
// it's faster than version 1, but still not very fast!
//
#include <iostream>
#include <cstdlib>		// has exit(), etc.
#include <cmath>		// has ceil()
#include <fstream>
#include <vector>
#include <algorithm>		// has copy()
#include <mpi++.h>		// MPI header file

// function declarations -- see definitions, following main,
//   for comments
void errorExit(const char msg[], const bool print);
double * getFileInput(const char filename[], const bool root,
		      const int nprocs, int * count);
double * getRandomInput(const char inputCount[], const bool root,
			const int nprocs, int * count);
int roundToMultiple(const int n, const int m);

// main program

int main(int argc, char *argv[]) {

  double *nums = NULL;		// will contain data to sort
  int count = 0;		// will contain count of data to sort
  const int rootID = 0;		// "root" process for collective comm.

  // initialize 
  if (MPI_Init(&argc, &argv) != MPI_SUCCESS) 
    errorExit("MPI initialization error\n", true);
  int nprocs;
  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  int myID;
  MPI_Comm_rank(MPI_COMM_WORLD, &myID);

  // process command-line arguments and get input
  if (argc < 2)
      errorExit("usage:  add-numbers inputfile | -n N\n", 
		myID == rootID);
  if (argc == 2)
    nums = getFileInput(argv[1], myID == rootID, nprocs, &count);
				// calls errorExit if problems
  else
    nums = getRandomInput(argv[2], myID == rootID, nprocs, &count);
				// calls errorExit if problems

  // process rootID now has input data in "nums"
  // count is a multiple of nprocs 

  // start timing 
  double startTime = MPI_Wtime();

  // broadcast size
  MPI_Bcast(&count, 1, MPI_INT, rootID, MPI_COMM_WORLD);

  // in each process, allocate space for "its" count/nprocs numbers
  int numsPerProcess = count/nprocs;
  double * partOfNums = new double[numsPerProcess];
    
  // now scatter data
  MPI_Scatter(nums, numsPerProcess, MPI_DOUBLE,
	      partOfNums, numsPerProcess, MPI_DOUBLE, 
	      rootID, MPI_COMM_WORLD);

  // each process adds up 1/nprocs of the numbers
  double partialSum = 0;
  for (int i = 0; i < numsPerProcess; ++i)
    partialSum += partOfNums[i];
  cerr << "partial sum in process " << myID << " = " << partialSum
       << endl;

  // now "reduce" results to get whole sum
  double sum;
  MPI_Reduce(&partialSum, &sum, 1, MPI_DOUBLE, MPI_SUM, rootID,
             MPI_COMM_WORLD);

  // end timing
  double endTime = MPI_Wtime();

  // print results
  if (myID == rootID) {
    cout << "sum = " << sum << endl;
    cout << "nprocs = " << nprocs << endl;
    cout << "problem size = " << count << endl;
    cout << "time = " << endTime - startTime << " seconds\n";
  }
  
  // clean up and exit
  delete [] nums;
  MPI_Finalize();
  return EXIT_SUCCESS;
}

// function definitions 

// function to bail out with error message
// post:  "msg" has been printed to standard error,
//          if "print" is true
//        MPI_Finalize and exit have been called, to
//          clean up and bail out
void errorExit(const char msg[], const bool print) {
  if (print)
    cerr << msg;
  MPI_Finalize();
  exit(EXIT_FAILURE);
}

// function to get input from file
// pre:  filename is the first command-line argument (after pgm name)
// post:  if file "filename" does not exist, errorExit has been
//          called
//        if file "filename" exists:  if "root", *count contains 
//          number of items read and return value points to
//          dynamically-allocated array containing input data
//        *count is a multiple of nprocs (actual data is padded with
//          zeroes if necessary to make this true)
double * getFileInput(const char filename[], const bool root, 
		      const int nprocs, int * count) {
  ifstream infile;
  infile.open(filename);
  if (!infile)
    errorExit("input file not found\n", root);
  if (root) {
      double temp;
      vector<double> vTemp;
      while (infile >> temp)
	vTemp.push_back(temp);
      infile.close();
      *count = roundToMultiple(vTemp.size(), nprocs);
      double * numsToReturn = new double[*count];
      copy(vTemp.begin(), vTemp.end(), numsToReturn);
      for (int i = vTemp.size(); i < *count; ++i)
	numsToReturn[i] = 0;
      return numsToReturn;
    }
    else {
      infile.close();
      return NULL;
    }
}

// function to generate input randomly
// pre:  inputCount is the second command-line argument (after pgm
//         name)
// post:  if "inputCount" is not a positive integer, errorExit has been
//          called
//        if "inputCount" is a positive integer:  if "root", *count 
//          contains number of input items and return value points to 
//          dynamically-allocated array containing input data
//          (generated using rand())
//        *count is a multiple of nprocs (actual data is padded with
//          zeroes if necessary to make this true)
double * getRandomInput(const char inputCount[], const bool root,
			const int nprocs, int * count) {
  *count = 0;
  *count = atoi(inputCount);
  if (*count <= 0)
    errorExit("usage:  add-numbers inputfile | -n N\n", root);
  *count = roundToMultiple(*count, nprocs);
  if (root) {
      double * numsToReturn = new double[*count];
      for (int i = 0; i < *count; ++i)
	numsToReturn[i] = (double) rand() / (double) (RAND_MAX);
      return numsToReturn;
    }
  else
    return NULL;
}

// function to round up to nearest multiple
// pre:  n, m are positive integers
// post:  return value is smallest x >= n such that x is
//          a multiple of m
int roundToMultiple(const int n, const int m) {
  return m * (int) ceil((double) n / (double) m);
}