#include <iostream>
#include <mpi.h>
#include <string.h>

#include "DKSBase.h"

using namespace std;

void printData(int *data, int N, int nprocs, const char *message = "") {
	if (strcmp(message, "") != 0)
		cout << message;
		
	for (int i = 0; i < nprocs; i++) {
		for (int j = 0; j < N; j++)
			cout << data[i*N + j] << "\t";
		cout << endl;
	}
}

void initData(int *data, int N, int rank) {
	for (int i = 0; i < N; i++)
		data[i] = (rank+1);
}

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

	int ierr;
	int rank, nprocs;
		
	MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
    
    cout << "Rank " << (rank+1) << " from " << nprocs << endl;

	int n = 8;
	int sizen = sizeof(int)*n;
	int sizeall = sizeof(int)*n*nprocs;
	
	int *hdata_in = new int[n];
	int *hdata_out = new int[n];
	initData(hdata_in, n, rank);
	cout << "In data for process " << rank+1 << ":\t";
	printData(hdata_in, n, 1);
	
	
	DKSBase base = DKSBase();
	base.setAPI("Cuda", 4);
	base.setDevice("-gpu", 4);
	base.initDevice();
	
	
	if (rank == 0) {
	
		int *hdata_out_all = new int[nprocs*n];
		void* mem_ptr;
		mem_ptr = base.allocateMemory<int>(nprocs*n, ierr);
		
		MPI_Gather(hdata_in, n, MPI_INT, mem_ptr, n, MPI_INT, 0, MPI_COMM_WORLD);
		
		base.readData<int>(mem_ptr, hdata_out_all, n*nprocs);
		
		MPI_Scatter(mem_ptr, n, MPI_INT, hdata_out, n, MPI_INT, 0, MPI_COMM_WORLD);
		
		base.freeMemory<int>(mem_ptr, n*nprocs);
		
		printData(hdata_out_all, n, nprocs, "Out data 1:\n");
		cout << "Scatter data for proces: " << rank + 1 << ": \t";
		printData(hdata_in, n, 1);
	} else {
		
		MPI_Gather(hdata_in, n, MPI_INT, NULL, NULL, NULL, 0, MPI_COMM_WORLD);
		
		MPI_Scatter(NULL, NULL, NULL, hdata_out, n, MPI_INT, 0, MPI_COMM_WORLD);
		
		cout << "Scatter data for proces: " << rank + 1 << ": \t";
		printData(hdata_in, n, 1);
		
	}


	MPI_Finalize();
	return 0;
}