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

/* 

Exercise 6.10.4, see Chap 6, p. 110 in PPMPI.
by John Weathewax

Here we implement a sparse matrix transpose from processor 0 to processor 1 
and print.

*/ 

#define MAX_ORDER 100

/* Create a structure for sparse matrix rows */ 
typedef struct {
  int               nnz; 
  int              *col;
  float            *val; 
} sparse_row_t; 

/* Create a data type for sparse matrices */  
typedef sparse_row_t sparse_matrix_t[MAX_ORDER];

/* declarations: */ 
void Find_Row_Elements(int col,sparse_matrix_t local_A,int n, 
		       int *n_rows,int rows[],float vals[]);
void Pack_A_Sparse_Column(int n_rows,int rows[],float vals[],
			  char buffer[],int *position);
void UnPack_A_Sparse_Column(int *n_rows,int rows[],float vals[],
			    char buffer[]);

main( int argc, char* argv[] ){
    int                  my_rank;
    int                  p;
    sparse_matrix_t      local_A; 
    int                  n;
    int                  i,j,l,nnz,found;

    /* declarations: */ 
    void Create_Sparse_Matrix(sparse_matrix_t local_A,int *n,int my_rank); 
    void Send_Transpose(sparse_matrix_t local_A, int n,int my_rank); 

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

    if( my_rank==0 ){ /* error checking... */ 
      if( p != 2 ){
	printf("Number of processors must be 2.\n" ); 
	MPI_Finalize(); 
	exit(1); 
      }
    }
    Create_Sparse_Matrix(local_A,&n,my_rank); 
    MPI_Bcast(&n,1,MPI_INT,0,MPI_COMM_WORLD); 

    Send_Transpose(local_A,n,my_rank);

    /* 
       -- print matrix A^T. --
       Note: It is possibly not valid to print from process 1.
       This could be a problem but it works on my system so I won't 
       go to the trouble of passing a message to process 0 and printing 
       from there.
    */ 
    if( my_rank==1 ){ 
      for( i=0; i < n; i++ ){
	nnz = local_A[i].nnz; 
	printf("nnz = %d: ",nnz);
	for( j=0; j < n; j++ ){
	  /* Find if column j has an element: -1 => no column exists */ 
	  found=-1;
	  for( l=0; l < nnz; l++ ) if( local_A[i].col[l] == j ) found = l; 
	  if( found != -1 ){
	    printf("%f ",local_A[i].val[found]);
	  }else{
	    printf("%f ",0.0); 
	  }
	}
	printf("\n"); 
      }
    }

    MPI_Finalize(); 
}



/**********************************************************************/
/* 
   Packs all nonzero elements into a data structure for sending
   to process==1.  Note: this will ONLY be called by process 0.
*/ 
void Pack_A_Sparse_Column(int n_rows,int rows[],float vals[],
			  char buffer[],int *position){

  (*position)=0; 
  MPI_Pack(&n_rows,1,MPI_INT,buffer,MAX_ORDER,position,MPI_COMM_WORLD); 
  MPI_Pack(rows,n_rows,MPI_INT,buffer,MAX_ORDER,position,MPI_COMM_WORLD); 
  MPI_Pack(vals,n_rows,MPI_FLOAT,buffer,MAX_ORDER,position,MPI_COMM_WORLD); 
}



/**********************************************************************/
/* 
   Unpacks all nonzero elements into a data structure after recieved
   with process==1.  Note: this will ONLY be called by process 0.
*/ 
void UnPack_A_Sparse_Column(int *n_rows,int rows[],float vals[],
			    char buffer[]){

  int position=0; 
  MPI_Unpack(buffer,MAX_ORDER,&position,n_rows,1,MPI_INT,MPI_COMM_WORLD); 
  MPI_Unpack(buffer,MAX_ORDER,&position,rows,(*n_rows),MPI_INT,MPI_COMM_WORLD); 
  MPI_Unpack(buffer,MAX_ORDER,&position,vals,(*n_rows),MPI_FLOAT,MPI_COMM_WORLD); 

}



/**********************************************************************/
void Send_Transpose(
         sparse_matrix_t     local_A  /* in/out */, 
         int                       n  /* in  */, 
         int                 my_rank  /* in  */) {

    int                  i,j,l,nnz;
    /* n_rows, rows & vals record the nonzero rows given an column */ 
    int                  rows[MAX_ORDER],n_rows; 
    float                vals[MAX_ORDER];
    /* packed data buffer */ 
    char                 buffer[MAX_ORDER];
    int                  position; 
    MPI_Status           status; 

    if( my_rank==0 ){ /* I am sending...pack data to send */

      for( j=0; j < n; j++ ){ /* loop over each column of local_A */ 
	Find_Row_Elements(j,local_A,n,&n_rows,rows,vals);
	Pack_A_Sparse_Column(n_rows,rows,vals,buffer,&position); 
	MPI_Send(buffer,position,MPI_PACKED,1,0,MPI_COMM_WORLD); 
      }

    }else{ /* I am recieving...unpack and store */ 

      for( i=0; i < n; i++ ){ /* loop over each row of local_A */ 
	MPI_Recv(buffer,MAX_ORDER,MPI_PACKED,0,0,MPI_COMM_WORLD,&status);
	UnPack_A_Sparse_Column(&n_rows,rows,vals,buffer);
	nnz = n_rows; 
	local_A[i].nnz = nnz; 
	local_A[i].col = (int*) malloc(nnz*sizeof(int));
	local_A[i].val = (float*) malloc(nnz*sizeof(float)); 
	/* Fill in arrays col & val: */ 
	for( l=0; l < nnz; l++ ){
	  local_A[i].col[l] = rows[l]; 
	  local_A[i].val[l] = vals[l]; 
	}
      }
    }
}



/**********************************************************************/
/* 
   Finds the nonzero row elements given a column "col". 
*/ 
void Find_Row_Elements(int col,sparse_matrix_t local_A,int n, 
		       int *n_rows,int rows[],float vals[]){

  int i,j,nnz,indx; 

  /* zeros the "rows" array: */
  (*n_rows)=0; 

  /* 
     loop over each row and see if we have an element of that 
     row in column "col" 
  */ 
  for( i=0; i < n; i++ ){
    nnz = local_A[i].nnz;
    if( nnz == 0 ) continue; /* this row has NO elements... */ 
    for( indx=0; indx < nnz; indx++ ){
      if( local_A[i].col[indx] == col ){
	(*n_rows)++;
	rows[(*n_rows)-1] = i;                    /* "-1" is to perform C's zero indexing */
	vals[(*n_rows)-1] = local_A[i].val[indx]; 
      }
    }
  }
}



/**********************************************************************/
void Create_Sparse_Matrix(sparse_matrix_t local_A,int *n,int my_rank){

  /* Note: error checking is ignored in this simple code... 
     matrix is (in matlab format) = [ 1 2 3; 0 4 0; 5 0 0 ]
     with transpose               = [ 1 0 5; 2 4 0; 3 0 0 ]
     Note: Any rows of all zeros will have local_A[?].nnz=0;
  */ 
  int nnz; 
  int i,j,l,found;

  /* Create a sparse matrix in process==0 only: */ 
  if( my_rank==0 ){
    (*n) = 3; 
    /* Fill row 0... */ 
    nnz=3; 
    local_A[0].nnz = nnz;
    local_A[0].col = (int*) malloc(nnz*sizeof(int));
    local_A[0].col[0] = 0; local_A[0].col[1] = 1;  local_A[0].col[2] = 2; 
    local_A[0].val = (float*) malloc(nnz*sizeof(float)); 
    local_A[0].val[0] = 1.0; local_A[0].val[1] = 2.0;  local_A[0].val[2] = 3.0;
    /* Fill row 1... */ 
    nnz=1; 
    local_A[1].nnz = nnz; 
    local_A[1].col = (int*) malloc(nnz*sizeof(int));
    local_A[1].col[0] = 1; 
    local_A[1].val = (float*) malloc(nnz*sizeof(float)); 
    local_A[1].val[0] = 4.0;
    /* Fill row 2... */ 
    nnz=1; 
    local_A[2].nnz = nnz;
    local_A[2].col = (int*) malloc(nnz*sizeof(int));
    local_A[2].col[0] = 0;
    local_A[2].val = (float*) malloc(nnz*sizeof(float)); 
    local_A[2].val[0] = 5.0;

    /* For debugging: print A out... */ 
    for( i=0; i < (*n); i++ ){
      nnz = local_A[i].nnz;
      printf("nnz = %d: ",nnz);
      for( j=0; j < (*n); j++ ){
	/* Find if column j has an element: -1 => no column exists */ 
	found=-1;
	for( l=0; l < nnz; l++ ) if( local_A[i].col[l] == j ) found = l; 
	if( found != -1 ){
	  printf("%f ",local_A[i].val[found]);
	}else{
	  printf("%f ",0.0); 
	}
      }
      printf("\n"); 
    }

    printf("\n");

  }

}