/*
  mpi example for matrix vector product  A*x
  
  -- cutting the matrix along rows in swap manner
  -- dim(A) may not be an integral multiple of np 
*/
#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#define N  1024

void printmatrix(double *pa, int m, int n)
{
	int i, j;
	for(i=0; i<m; i++)
	{
		for(j=0; j<n; j++)
			printf("%6.2f", *(pa+i*n+j));
		printf("\n");
	}
}

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

  	MPI_Init(&argc, &argv);
  	MPI_Comm_rank(MPI_COMM_WORLD, &myid);
  	MPI_Comm_size(MPI_COMM_WORLD, &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, 2); 
	}

  	t = MPI_Wtime(); // time

	// set initial value 
  	for(j=0; j<N; j++) 
	  	px[j] = 1; // j+1;
		    
  	for(i=0,istep=myid; i<NL; i++,istep+=np) 
  	{
    	pyl[i] = 0.0;
    	for(j=0; j<N; j++) 
      		pA[i][j] = 1.0/(double)(istep+j+1);
  	}
    printf("myid=%d, A[0][0]=%.6f, A[1][0]=%.6f\n", myid, pA[0][0], pA[1][0]);

	if (myid==0 && N<10)
	{
		printf("myid=%d, local A:\n", myid);
		printmatrix(&pA[0][0], NL, N);
		printf("loca x: \n");
		printmatrix(px,1,N);
	}
//
//  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,istep=0; istep<N; istep+=np,i++)
    	 	py[istep] = pyl[i]; 
    	// printmatrix(py,1,N);
    	
		for (int i_np=1; i_np<np; i_np++)
		{
			MPI_Recv(&NL, 1, MPI_INT, i_np, 111, MPI_COMM_WORLD, &status);
		 	MPI_Recv(pyl, NL, MPI_DOUBLE, i_np, 112, MPI_COMM_WORLD, &status);
		 	for (i=0,istep=i_np; istep<N; istep+=np,i++)
    			py[istep] = pyl[i];
            // printmatrix(py,1,N);
    	}
  	}
  	else
  	{
  		MPI_Send(&NL, 1, MPI_INT, 0, 111, MPI_COMM_WORLD);
  		MPI_Send(pyl, NL, MPI_DOUBLE, 0, 112, MPI_COMM_WORLD);
	}
 
  	t = MPI_Wtime() - t; // time

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

  	MPI_Finalize();  	
  	return 0;
}