// mpi example for matrix vector product
//
//   compute the matrix-vector production
//   by cutting the matrix along rows
//
#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#define N  1024

int main(int argc, char *argv[])
{
  	int myid, np, err, NL, Nk, Nr;
  	int i, j, lr;
  	double t, (*pA)[N], *px, *pyl, *py;

  	MPI_Init(&argc, &argv);
  	MPI_Comm_rank(MPI_COMM_WORLD, &myid);
  	MPI_Comm_size(MPI_COMM_WORLD, &np);

  	int rc[np], disp[np];

  	Nk = N / np; Nr = N % np;
  	if (myid < Nr)  NL = Nk + 1; else NL = Nk;

  	pA  = malloc(NL*N*sizeof(double));
  	px  = malloc(N*sizeof(double));
  	pyl = malloc(NL*sizeof(double));
  	py  = malloc(N*sizeof(double));
  	if (!pA | !px | !pyl | !py) 
  	{  printf("Out of memory!\n"); MPI_Abort(MPI_COMM_WORLD, 1); }

  	t = MPI_Wtime(); // time

	// set initial value 
  	for(j=0; j<N; j++)  px[j] = 1.0;
  
  	if (myid < Nr ) 
	  	lr = myid*Nk + myid;
  	else 
	  	lr = myid*Nk + Nr;

  	for(i=0; i<NL; i++) 
  	{
    	pyl[i] = 0.0;
    	for(j=0; j<N; j++) 
      		pA[i][j] = 1.0/(double)(i+lr+j+1);
  	}
	//  printf("myid=%d, A[0][0]=%f\n", myid, pA[0][0]);

	// computations 
  	for(i=0; i<NL; i++) 
    	for(j=0; j<N; j++)
	  		pyl[i] = pyl[i] + pA[i][j]*px[j];

  	if (myid==0)
  	{
    	for (i=0; i<Nr; i++) 	rc[i] = Nk + 1; 
    	for (i=Nr; i<np; i++) 	rc[i] = Nk; 
    	for (disp[0]=0, i=0; i<np-1; i++)  disp[i+1] = disp[i] + rc[i];
  	}
  	MPI_Gatherv(pyl, NL, MPI_DOUBLE, py, rc, disp, MPI_DOUBLE, 0, MPI_COMM_WORLD);
 
  	t = MPI_Wtime() - t; // time

  	if (myid == 0)
  	printf("y[0]=%.6f, y[N-1]=%.6f, Time=%.4e\n", py[0], py[N-1], t);

  	free(pA); free(px); free(pyl); free(py);

  	MPI_Finalize();
  	return 0;
}