#include <stdio.h>
#include "mpi.h"
#include <stdlib.h>





/*       *****************************************************
         * Block Matrix Multiply on a 2*2 grid of processors *
         * PE 0: 00,  PE1: 01, PE2: 10, PE3: 11              *
         *****************************************************
*/


/* block size */
#define b 8

/* arr is an int* */
void printblock(int arr[b][b]) {
 int i,j;
 for(i=0;i<b;i++){
  for(j=0;j<b;j++){
   printf(" %4d", arr[i][j] );
  }
 printf("\n      ");
 }
}

/* A, B, C are int* */
void multblock(int C[b][b], int A[b][b],  int B[b][b])  {
 int i,j,k;
 for(i=0;i<b;i++){
  for(j=0;j<b;j++){
   for(k=0;k<b;k++) C[i][j] += A[i][k]*B[k][j];
  }
 }
}

main(int argc, char **argv)    {
/* no args
   the data is created in the 4 PEs
   and printed by PE0 
*/
    int my_id;		/* process id */
    int p;		/* number of processes */
    char* message;	/* storage for the message */
    int i,j,k, ioff, joff;

    int A00[b][b],  A01[b][b],  A10[b][b],  A11[b][b];
    int B00[b][b],  B01[b][b],  B10[b][b],  B11[b][b];
    int C00[b][b],  C01[b][b],  C10[b][b],  C11[b][b];

    MPI_Status status;	/* return status for receive */

    MPI_Init( &argc, &argv );
    MPI_Comm_rank( MPI_COMM_WORLD, &my_id );
    MPI_Comm_size( MPI_COMM_WORLD, &p );

    MPI_Barrier(MPI_COMM_WORLD); 
    switch(my_id)  {

    case 0:
      printf("PE0: Init\n");
      /* Initialize A00, BOO  and C00 */
       for(i=0;i<b;i++){
        for(j=0,joff=b;j<b;j++,joff++){
         A00[i][j] = i+j;  B00[i][j] = i-j;   C00[i][j] =0;
       }}

      /* Row Exchange 0 <--> 1 */
      printf("PE0: <--> PE1: Row Exchange\n");
       MPI_Recv( (int *)A01, b*b, MPI_INT, 1, 1, MPI_COMM_WORLD, &status); 
       MPI_Send( (int *)A00, b*b, MPI_INT, 1, 2, MPI_COMM_WORLD);
      /* Col Exchange 0 <--> 2 */
      printf("PE0: <--> PE2: Col Exchange\n");
       MPI_Recv( (int *)B10, b*b, MPI_INT, 2, 3, MPI_COMM_WORLD, &status); 
       MPI_Send( (int *)B00, b*b, MPI_INT, 2, 4, MPI_COMM_WORLD);

      /* Block Multiply C00 = A00*B00 + A01*B10 */
       multblock(C00,A00,B00); multblock(C00,A01,B10);
      /* Gather */
      printf("PE0: Gather C01 <-- PE1\n");
       MPI_Recv( (int *)C01, b*b, MPI_INT, 1, 5, MPI_COMM_WORLD, &status); 
      printf("PE0: Gather C10 <-- PE2\n");
       MPI_Recv( (int *)C10, b*b, MPI_INT, 2, 6, MPI_COMM_WORLD, &status); 
      printf("PE0: Gather C11 <-- PE3\n");
       MPI_Recv( (int *)C11, b*b, MPI_INT, 3, 7, MPI_COMM_WORLD, &status); 
      /* Print */
      printf("PE0: Print blocks\n");
       printf("\n C00: ");  printblock(C00);
       printf("\n C01: ");  printblock(C01);
       printf("\n C10: ");  printblock(C10);
       printf("\n C11: ");  printblock(C11);
       printf("\n");
       break;

    case 1:
      printf("PE1: Init\n");
      /* Initialize A01, B01 and C01 */
       for(i=0,ioff=b;i<b;i++,ioff++){
        for(j=0,joff=b;j<b;j++,joff++){
         A01[i][j] = i+joff; B01[i][j] = i-joff; C01[i][j] = 0;
	}}
      /* Row Exchange 0 <--> 1 */
      printf("PE1: <--> PE0: Row Exchange\n");
       MPI_Send( (int *)A01, b*b, MPI_INT, 0, 1, MPI_COMM_WORLD);
       MPI_Recv( (int *)A00, b*b, MPI_INT, 0, 2, MPI_COMM_WORLD, &status); 
      /* Col Exchange 1 <--> 3 */
      printf("PE1: <--> PE3: Col Exchange\n");
       MPI_Recv( (int *)B11, b*b, MPI_INT, 3, 3, MPI_COMM_WORLD, &status); 
       MPI_Send( (int *)B01, b*b, MPI_INT, 3, 4, MPI_COMM_WORLD);
      /* Block Multiply C01 = A00*B01 + A01*B11 */
       multblock(C01,A00,B01); multblock(C01,A01,B11);
      /* Gather */
      printf("PE1: Gather C01 --> PE0\n");
       MPI_Send( (int *)C01, b*b, MPI_INT, 0, 5, MPI_COMM_WORLD);
       break;

    case 2:
      printf("PE2: Init\n");
      /* Initialize A10, B10 and C10 */
       for(i=0,ioff=b;i<b;i++,ioff++){
        for(j=0,joff=b;j<b;j++,joff++){
         A10[i][j] = ioff+j; B10[i][j] = ioff-j; C10[i][j] = 0;
	}}
      /* Row Exchange 2 <--> 3 */
      printf("PE2: <--> PE3: Row Exchange\n");
       MPI_Recv( (int *)A11, b*b, MPI_INT, 3, 1, MPI_COMM_WORLD, &status); 
       MPI_Send( (int *)A10, b*b, MPI_INT, 3, 2, MPI_COMM_WORLD);
      /* Col Exchange 0 <--> 2 */
      printf("PE2: <--> PE0: Col Exchange\n");
       MPI_Send( (int *)B10, b*b, MPI_INT, 0, 3, MPI_COMM_WORLD);
       MPI_Recv( (int *)B00, b*b, MPI_INT, 0, 4, MPI_COMM_WORLD, &status); 
      /* Block Multiply C10 = A10*B00 + A11*B10 */
       multblock(C10,A10,B00); multblock(C10,A11,B10);
      /* Gather */
      printf("PE2: Gather C10 --> PE0\n");
       MPI_Send( (int *)C10, b*b, MPI_INT, 0, 6, MPI_COMM_WORLD);
       break;

    case 3:
      printf("PE3: Init\n");
      /* Initialize A11, B11 and C11 */
       for(i=0,ioff=b;i<b;i++,ioff++){
        for(j=0,joff=b;j<b;j++,joff++){
         A11[i][j] = ioff + joff; B11[i][j] = ioff - joff; C11[i][j] = 0;
	}}
      /* Row Exchange 2 <--> 3 */
      printf("PE3: <--> PE2: Row Exchange\n");
       MPI_Send( (int *)A11, b*b, MPI_INT, 2, 1, MPI_COMM_WORLD);
       MPI_Recv( (int *)A10, b*b, MPI_INT, 2, 2, MPI_COMM_WORLD, &status); 
      /* Col Exchange 1 <--> 3 */
      printf("PE3: <--> PE1: Col Exchange\n");
       MPI_Send( (int *)B11, b*b, MPI_INT, 1, 3, MPI_COMM_WORLD);
       MPI_Recv( (int *)B01, b*b, MPI_INT, 1, 4, MPI_COMM_WORLD, &status); 
      /* Block Multiply C11 = A10*B01 + A11*B11 */
       multblock(C11,A10,B01); multblock(C11,A11,B11);
      /* Gather */
      printf("PE3: Gather C11 --> PE0\n");
       MPI_Send( (int *)C11, b*b, MPI_INT, 0, 7, MPI_COMM_WORLD);
       break;
    }	    
EXIT:
    MPI_Finalize();
}