/* ------------------------------------------------------------------------ */
/* -- Parallel Matrix Operations                                         -- */
/* -- Solution of AX = B, assume solution exists, no 0 on diagonal of A. -- */
/* -- Warning, this program is not full-proof, it was written for        -- */
/* -- simplicity. In case of failed memory allocation, the program       -- */
/* -- behaviour is undefined. In more robuts codes, this sort of         -- */
/* -- situation should be dealt with carefully.                          -- */
/* --                                                                    -- */
/* -- E. van den Berg                                         05/10/2001 -- */
/* ------------------------------------------------------------------------ */

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


// This program was writting in the Centre of Advanced Calculations in
// Engineering Science (ACES).
#define TAG_MATRIX_PARTITION	44261 // Try to calculate hex value :-) 

// Type definitions
typedef struct
{  unsigned int   m, n;  // Rows, cols
   double        *data;  // Data, ordered by row, then by col
   double       **rows;  // Pointers to rows in data
} TMatrix;

// Global variables
int processor_rank  = 0;
int processor_count = 1;

// Function declaractions
TMatrix createMatrix (const unsigned int rows, const unsigned int cols);
void    freeMatrix   (TMatrix *matrix);
int     validMatrix  (TMatrix matrix);
TMatrix initMatrix   (void);
void    sweepMatrix  (TMatrix A, TMatrix B, unsigned int pivotRow, unsigned int pivotCol,
                      double *rowA, double *rowB);
void    printMatrix  (TMatrix A);
int     readMatrix   (char *filename, TMatrix *A);
int     writeMatrix  (char *filename, TMatrix A);


/* ------------------------------------------------------------------------ */
int main (int argc, char *argv[])
/* ------------------------------------------------------------------------ */
{  MPI_Status   status;
   TMatrix      A,B,C,D;
   unsigned int m, p, i, i0;
   unsigned int offset, size[4] = {0};
   double       time0, time1;
   double       d;

   // Initiate all matrices to 0 x 0 matrices.
   A = initMatrix();
   B = initMatrix();
   C = initMatrix();
   D = initMatrix();

   // Always use MPI_Init first
   if (MPI_Init(&argc, &argv) == MPI_SUCCESS) do
   {
      // Determine which number is assigned to this processor
      MPI_Comm_size (MPI_COMM_WORLD, &processor_count);
      MPI_Comm_rank (MPI_COMM_WORLD, &processor_rank );

      if (processor_rank == 0)
      {  
         // Read data, argument 0 = executable, 1 = matrix A, 2 = matrix B, 3 = output filename
         if (argc == 4) do
         {  if (!readMatrix(argv[1], &A)) break;
            if (!readMatrix(argv[2], &B)) break;
            if (!validMatrix(A) || !validMatrix(B) || (A.m != B.m) || (A.m != A.n)) break;
	    C = createMatrix(1,A.n);
	    D = createMatrix(1,B.n);
	    if (!validMatrix(C) || !validMatrix(D)) break;

            size[0] = A.m;
            size[1] = A.n;
            size[2] = B.m;
            size[3] = B.n;
         } while (0);
         else printf ("Usage : %s <matrix filename> <matrix filename> <output filename>\n", argv[0]);

         // Broadcast (send) size of matrix
         MPI_Bcast((void *)&size, 4, MPI_INT, 0, MPI_COMM_WORLD); if (size[0] == 0) break;
         m = A.m / processor_count;
	 i0= 0;
	 
         // Record starting time
         time0 = MPI_Wtime();

         // Send each process it's own part of A and B
	 offset = A.m - m * (processor_count - 1);
	 for (i = 1; i < processor_count; i++)
	 {  MPI_Send((void *)A.rows[offset], m*A.n, MPI_DOUBLE,
	             i, TAG_MATRIX_PARTITION, MPI_COMM_WORLD);
	    MPI_Send((void *)B.rows[offset], m*B.n, MPI_DOUBLE,
	             i, TAG_MATRIX_PARTITION, MPI_COMM_WORLD);
	    offset += m;
	 }
	 
	 // Set Master's local matrix height of A and B to new value
	 A.m -= m * (processor_count - 1);
	 B.m -= m * (processor_count - 1); 
      }
      else
      {
         // Broadcast (receive) size of matrix
         MPI_Bcast((void *)&size, 4, MPI_INT, 0, MPI_COMM_WORLD); if (size[0] == 0) break;

         // Allocate memory for matrices
         m = size[0] / processor_count;
	 i0= size[0] - (processor_count - processor_rank) * m;
	 A = createMatrix(m, size[1]);
	 B = createMatrix(m, size[3]);
	 C = createMatrix(1, size[1]);
	 D = createMatrix(1, size[3]);
     
         // Broadcast (recieve) B matrix
	 MPI_Recv((void *)A.data, m*A.n, MPI_DOUBLE, 0, TAG_MATRIX_PARTITION, MPI_COMM_WORLD, &status);
	 MPI_Recv((void *)B.data, m*B.n, MPI_DOUBLE, 0, TAG_MATRIX_PARTITION, MPI_COMM_WORLD, &status);	 
      }

      
      
      // Sweep matrix
      offset = size[0] - m * (processor_count - 1);
      for (i = 0; i < size[0]; i++)
      {  
         if (i < offset)
	      p = 0;
	 else p = (i - offset) / m + 1;

         if (processor_rank == p)
         {  // Send row
	    MPI_Bcast((void *)A.rows[i - i0], A.n, MPI_DOUBLE, p, MPI_COMM_WORLD);
	    MPI_Bcast((void *)B.rows[i - i0], B.n, MPI_DOUBLE, p, MPI_COMM_WORLD);
	    sweepMatrix(A, B, i - i0, i, A.rows[i - i0], B.rows[i - i0]);
	 }
	 else
	 {  // Receive
	    MPI_Bcast((void *)C.rows[0], C.n, MPI_DOUBLE, p, MPI_COMM_WORLD);
	    MPI_Bcast((void *)D.rows[0], D.n, MPI_DOUBLE, p, MPI_COMM_WORLD);
	    sweepMatrix(A, B, -1, i, C.rows[0], D.rows[0]);
	 }
      }


      // Send back all results to processor 0
      if (processor_rank == 0)
      {
         // Restore original matrix sizes of A and B
	 A.m = size[0];
	 B.m = size[2];
      
         // Receive part of B matrix f (result) rom each process
         offset = A.m - m * (processor_count - 1);
         for (i = 1; i < processor_count; i++)
         {  MPI_Recv((void *)B.rows[offset], m*B.n, MPI_DOUBLE,
                     i, TAG_MATRIX_PARTITION, MPI_COMM_WORLD, &status);
            offset += m;
         }

         // Record finish time
         time1 = MPI_Wtime();
	 
	 // Write result into file
	 writeMatrix(argv[3],B);

         // Print time statistics
         printf ("Total time using [%2d] processors : [%f] seconds\n", processor_count, time1 - time0);
      }
      else
      {	 // Send back result
	 MPI_Send((void *)B.data, B.m*B.n, MPI_DOUBLE, 0, TAG_MATRIX_PARTITION, MPI_COMM_WORLD);
      }
  
   } while (0); // Enable break statement to be used above

   // Free matrix data
   freeMatrix(&A); freeMatrix(&B);

   // Wait for everyone to stop   
   MPI_Barrier(MPI_COMM_WORLD);

   // Always use MPI_Finalize as the last instruction of the program
   MPI_Finalize();

   return 0;
}


/* ------------------------------------------------------------------------ */
TMatrix createMatrix(const unsigned int rows, const unsigned int cols)
/* ------------------------------------------------------------------------ */
{  TMatrix           matrix;
   unsigned long int m, n;
   unsigned int      i;

   m = rows; n = cols;
   matrix.m    = rows;
   matrix.n    = cols;
   matrix.data = (double *) malloc(sizeof(double) * m * n);
   matrix.rows = (double **) malloc(sizeof(double *) * m);
   if (validMatrix(matrix))
   {  matrix.m = rows; 
      matrix.n = cols;
      for (i = 0; i < rows; i++)
      {  matrix.rows[i] = matrix.data + (i * cols);
      }
   }
   else
   {  freeMatrix(&matrix);
   }

   return matrix;
}


/* ------------------------------------------------------------------------ */
void freeMatrix (TMatrix *matrix)
/* ------------------------------------------------------------------------ */
{  if (matrix == NULL) return;

   if (matrix -> data) { free(matrix -> data); matrix -> data = NULL; }
   if (matrix -> rows) { free(matrix -> rows); matrix -> rows = NULL; }
   matrix -> m = 0;
   matrix -> n = 0;
}


/* ------------------------------------------------------------------------ */
int validMatrix (TMatrix matrix)
/* ------------------------------------------------------------------------ */
{  if ((matrix.data == NULL) || (matrix.rows == NULL) ||
       (matrix.m == 0) || (matrix.n == 0))
        return 0;
   else return 1;
}


/* ------------------------------------------------------------------------ */
TMatrix initMatrix()
/* ------------------------------------------------------------------------ */
{  TMatrix matrix;
 
   matrix.m = 0;
   matrix.n = 0;
   matrix.data = NULL;
   matrix.rows = NULL;

   return matrix;
}


/* ------------------------------------------------------------------------ */
void sweepMatrix (TMatrix A, TMatrix B, unsigned int pivotRow, unsigned int pivotCol,
                  double *rowA, double *rowB)
/* ------------------------------------------------------------------------ */
{  unsigned int i, j;
   double d;
   
   // Normalize row
   d = rowA[pivotCol]; if (d == 0) return; // Please replace by proper error handling
   for (j = 0; j < A.n; j++) rowA[j] /= d;
   for (j = 0; j < B.n; j++) rowB[j] /= d;

   // Error checking goes here!
   
   for (i = 0; i < A.m; i++)
   {  if (i == pivotRow) continue;
      d = A.rows[i][pivotCol]; if (d == 0) continue; // Please replace by proper error handling
      for (j = 0; j < A.n; j++) A.rows[i][j] -= d * rowA[j];
      for (j = 0; j < B.n; j++) B.rows[i][j] -= d * rowB[j];
   }
}


/* ------------------------------------------------------------------------ */
void printMatrix(TMatrix A)
/* ------------------------------------------------------------------------ */
{  unsigned int i, j;

   if (validMatrix(A))
   {  for (i = 0; i < A.m; i++)
      {  for (j = 0; j < A.n; j++) printf ("%7.3f ", A.rows[i][j]);
	 printf ("\n");
      }
   }
}


/* ------------------------------------------------------------------------ */
int readMatrix(char *filename, TMatrix *A)
/* ------------------------------------------------------------------------ */
{  FILE *fp;
   unsigned int m, n, i, j;
   float d;
   int result = 0;

   if ((fp = fopen (filename, "r")) == NULL) return 0;

   do
   {
      if (fscanf (fp, "%d%d",  &m, &n) != 2) break;
      if ((m == 0) || (n == 0)) break;
      *A = createMatrix(m,n);
      if (!validMatrix(*A)) break;

      for (i = 0; i < m; i ++)
      {  for (j = 0; j < n; j ++)
         {  if (fscanf (fp, "%f", &d) != 1) break;
            A -> rows[i][j] = d;
         }
         if (j != n) break;
      }
      if (i != m) break;

      result = 1;
   } while (0);

   fclose (fp);

   return result;
}


/* ------------------------------------------------------------------------ */
int writeMatrix(char *filename, TMatrix A)
/* ------------------------------------------------------------------------ */
{  FILE *fp;
   unsigned int i, j;

   if ((fp = fopen (filename, "w")) == NULL) return 0;

   fprintf (fp, "%d %d\n", A.m, A.n);

   for (i = 0; i < A.m; i++)
   {  for (j = 0; j < A.n; j++)
      {  fprintf (fp, "%20.10f ", A.rows[i][j]);
      }
      fprintf (fp, "\n");
   }

   fclose(fp);

   return 1;
}