/* "matrix transpose" using MPI_Alltoall
 * 
 * Input:
 *     matrix_A: the original input matrix distributed in block row form
 *
 * Output:
 *     matrix_AT: the original matrix transposed and distributed in block row form
 *
 * See Chap 10, pp. 242 in PPMPI.
 *
 * Written by:
 * -- 
 * John L. Weatherwax                2006-01-29
 * 
 * email: wax@alum.mit.edu
 * 
 * Please send comments and especially bug reports to the
 * above email address.
 * 
 *-----
 */
#include <stdio.h>
#include <math.h>
#include "mpi.h"
#define MAX_DIM 100

typedef float MATRIX_T[MAX_DIM][MAX_DIM];

void Read_matrix(char* prompt, MATRIX_T A_local, int n, int my_rank, int p);
void Print_matrix(char* title, MATRIX_T A_local, int n, 
         int my_rank, int p);
void Transpose_A(int n_bar, int p, int my_rank, MATRIX_T A_local, MATRIX_T AT_local);

int
main(int argc, char* argv[]) {

    int        p,my_rank,n,n_bar,n_rem,n_tol,n_local,new_n_local;
    MATRIX_T   A_local,AT_local;
    float      *keys_local,*new_keys_local,max_key=-1.0;

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

    if (my_rank == 0) {
        printf("Enter n dimension of the global matrix:\n");
        scanf("%d", &n);
	if( n % p != 0 ){ /* Do I NEED to call MPI_Finalize() here? */ 
	  fprintf(stderr,"n = %d; p = %d\n",n,p); 
	  fprintf(stderr,"p must divide n evenly ... exiting\n"); 
	  exit(1); 
	}
    }
    MPI_Bcast(&n, 1, MPI_INT, 0, MPI_COMM_WORLD);
    n_bar = n / p; 
    n_tol = n % p; 
    /* Compute the number of elements in each local array: */ 
    /* n_local = ( my_rank!=p-1 ? n_bar : n_bar+n_tol ); */ 

    Read_matrix("Enter the matrix", A_local, n, my_rank, p);
    Print_matrix("The original matrix:", A_local, n, my_rank, p); 

    Transpose_A(n_bar,p,my_rank,A_local,AT_local); 
    Print_matrix("The transposed matrix:", AT_local, n, my_rank, p); 

    MPI_Finalize();
}  /* main */

/* Takes the transpose of A using calls to MPI_Alltoall() */ 
void Transpose_A(int n_bar, int p, int my_rank, MATRIX_T A_local, MATRIX_T AT_local){

  int ii_row,i1,i2,col_loc;
  float tmp_store[MAX_DIM];

  /* Send each row into the corresponding column of AT_local */
  for( ii_row=0; ii_row < n_bar; ii_row++ ){

    /* Exchange the actual matrix data into tmp_store[] */ 
    MPI_Alltoall(&(A_local[ii_row][0]),n_bar,MPI_FLOAT,
		 tmp_store            ,n_bar,MPI_FLOAT,
		 MPI_COMM_WORLD);
    /* Copy data from tmp_store into AT_local */ 
    for( i1=0; i1 < p; i1++ ){
      col_loc = ii_row + i1*n_bar;
      for( i2=0; i2 < n_bar; i2++ ){ 
	AT_local[i2][col_loc] = tmp_store[i2+i1*n_bar];
      }
    }
  } 
}


/*********************************************************************/
/* Reads a matrix using a block row distribution                     */ 
void Read_matrix(
         char*     prompt   /* in  */,
         MATRIX_T  A_local  /* out */,
         int       n        /* in  */,
         int       my_rank  /* in  */,
         int       p        /* in  */) {

    int       i, j;
    MATRIX_T  temp;
    int       n_bar, n_rem, n_local;
    MPI_Status status;
 
    n_bar = n / p;
    n_rem = n % p;
    /* n_local = ( my_rank != p-1 ? n_bar : n_bar+n_rem ); */
    n_local = n_bar; 

    /* Fill dummy entries in temp with zeroes */
    for (i = 0; i < MAX_DIM; i++)
        for (j = n; j < MAX_DIM; j++)
            temp[i][j] = -1.0;

    if (my_rank == 0) {
        printf("%s\n", prompt);
        for (i = 0; i < n; i++)
            for (j = 0; j < n; j++)
                scanf("%f",&temp[i][j]);
    }
    /* Distribute the elements of A by block rows... */
    MPI_Scatter(temp, n_local*MAX_DIM, MPI_FLOAT, A_local,
        n_local*MAX_DIM, MPI_FLOAT, 0, MPI_COMM_WORLD);

}  /* Read_matrix */

/*********************************************************************/
/* Prints a matrix that is distributed by block rows                 */ 
void Print_matrix(
         char*     title      /* in */,
         MATRIX_T  A_local    /* in */,
         int       n          /* in */,
         int       my_rank    /* in */,
         int       p          /* in */) {

    int       i, j;
    MATRIX_T  temp;
    int       n_bar, rem;
    MPI_Status status; 

    n_bar = n/p;
    rem = n % p; 

    MPI_Gather(A_local, n_bar*MAX_DIM, MPI_FLOAT, 
	       temp,    n_bar*MAX_DIM, MPI_FLOAT, 
	       0, MPI_COMM_WORLD);

    if (my_rank == 0) {
      printf("%s\n", title);
      for (i = 0; i < n; i++) {
	for (j = 0; j < n; j++)
	  printf("%10.6f ", temp[i][j]);
	printf("\n");
      }
    }
} /* Print_matrix */