/*
  mpi example for matrix matrix product
  
  -- row-column partition 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  8 // 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;
  	int i, j, k, ii, ij, ik, src, dest;
  	int istep, jstep;
  	double t, *pa, *pb, *pc;
  	MPI_Status status;

  	MPI_Init(&argc, &argv);
  	MPI_Comm_rank(MPI_COMM_WORLD, &myid);
  	MPI_Comm_size(MPI_COMM_WORLD, &np);
  	
  	int NLmax, NL[np], Nr;
  	
  	Nr = N % np;  	
  	if (Nr==0) 
	{
  		NLmax = N/np;
  		for (i=0; i<np; i++)
  			NL[i] = NLmax;
  	}
	else
	{
		NLmax = N/np + 1;  		
  	    for (i=0; i<Nr; i++)
		  	NL[i] = NLmax;
		for (i=Nr; i<np; i++)
			NL[i] = NLmax-1;
	}

  	pa = malloc(NL[myid]*N*sizeof(double));
  	pb = malloc(N*NLmax*sizeof(double));
  	pc = malloc(NL[myid]*N*sizeof(double));
  	if (!pa | !pb | !pc) 
	{ 
    	printf("Out of memory!\n"); 
    	MPI_Abort(MPI_COMM_WORLD, 2); 
  	}

  	t = MPI_Wtime(); // time

	// set initial value 		
  	for(i=0,istep=myid; i<NL[myid]; i++,istep+=np) 
	{
    	for(j=0; j<N; j++) 
		{
      		*(pa+i*N+j) = 1.0/(double)(istep+j+1);
      		*(pc+i*N+j) = 0.0;
    	}
  	}
  	for(i=0; i<N; i++) 
    	for(j=0,jstep=myid; j<NL[myid]; j++,jstep+=np) // j<NLmax
      		*(pb+i*NL[myid]+j) = (double)(i+jstep+1);

	if (myid==1)
	{
		printf("myid=%d, local A:\n", myid);
		printmatrix(pa, NL[myid], N);
		printf("loca B: \n");
		printmatrix(pb, N, NL[myid]);
	}

// ********************************** //
// computations
  	for (k=0; k<np; k++)
  	{
		// compute C_{ij}
        for(ii=0; ii<NL[myid]; ii++)
    		for(ij=0,jstep=(myid+k)%np; ij<NL[(myid+k)%np]; ij++,jstep+=np)
      			for(ik=0; ik<N; ik++)
          			*(pc+ii*N+jstep) += *(pa+ii*N+ik)* *(pb+ik*NL[(myid+k)%np]+ij);

//  send B to the left processor and 
//  get new B from the right processor
    	src  = ((myid+1+np) % np);
    	dest = ((myid-1+np) % np);
    	if ( k != np-1)
      		MPI_Sendrecv_replace(pb, N*NLmax, MPI_DOUBLE, dest, 111, 
                           src, 111, MPI_COMM_WORLD, &status);
  	}

  	t = MPI_Wtime() - t; // time

  	printf("myid=%d, C[0][0]=%.2f, C[0][1]=%.2f, Time=%.4e\n", myid, *pc, *(pc+1), t);

  	free(pa); free(pb); free(pc);
  	
  	MPI_Finalize();  	  	
	return 0;
}