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

#define MAX(a,b) ((a)>(b)?(a):(b))

void get_matrix( int n, int is, int ie, double a[ie-is+1][n]) {
  int i, j;
  for ( i=is; i<=ie; i++ ) {
    for ( j=0; j<n; j++ ) {
      a[i-is][j] = 2*(i+1)-3*(j+1);
    }
  }
}

void print_block( int n, int is, int ie, double a[ie-is+1][n]) {
  int i, j;
  for ( i=is; i<=ie; i++ ) {
    printf("row %2.2d: ",i);
    for ( j=0; j<n; j++ ) {
      printf("%6.2lf ",a[i-is][j]);
    }
    printf("\n");
  }
  fflush(stdout);
}


#define N 10
#define SCALE 1       // ex. 3
//#define SCALE 2     // ex. 4

int main( int argc, char *argv[] ) {
  MPI_Status stat;
  int np, me, tag=313, tag_print=314, tmp, i, j, is, ie, nloc;
  
  MPI_Init(&argc,&argv);
  MPI_Comm_size(MPI_COMM_WORLD,&np);
  MPI_Comm_rank(MPI_COMM_WORLD,&me);
  //
  int id[np], nd[np];   // id[i] start index of block owned by i:th task 
                        // nd[i] number of entries in i:th block
  nloc = N/np;
  for ( i=0; i < np; i++ ) {
    if ( i < N%np )
      nd[i] = nloc+1;
    else
      nd[i] = nloc;
  }
  id[0] = 0;
  for ( i=0; i<np-1; i++ )
    id[i+1] = id[i] + nd[i]; 
  is = id[me];
  ie = is + nd[me] - 1;
  //  if ( me==np-1) 
  //  ie=N-1;
  // else
  //  ie = id[me+1]-1;
  printf("P%d takes care of rows %d to %d\n",me,is,ie);
  //
  double a[ie-is+1][N], d[N];
  //
  get_matrix(N,is,ie,a);

  //
  //
  if ( SCALE==1 ) {
    //
    //  d(i) = a(i,i)
    //
    double buf[ie-is+1];
    for ( i=is; i<=ie; i++ ) 
      buf[i-is] = a[i-is][i];   // buf contains diagonal entries of my block
    //
    // Gather to d[:] diagonal elements of a[:][:] owned by different tasks
    // stored in buf.
    // (As they are not necessary of same size we use allgatherv instead
    // of allgather):
    //
    MPI_Allgatherv( buf, nd[me], MPI_DOUBLE, d, nd, id, MPI_DOUBLE, MPI_COMM_WORLD);
  } else {
    //
    // d(j) = max_i |a(i,j)|
    //
    //
    // First maximum over my block:
    //
    for ( j=0; j<N; j++ ) {
      d[j] = 0.0;
      for ( i=is; i<=ie; i++ ) 
	d[j] = MAX( d[j], fabs(a[i-is][j]) );
    }
    //
    // Then maximum over max of each block:
    //
    MPI_Allreduce( MPI_IN_PLACE, d, N, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD );
  }
  //
  // "Balance" block of matrix owned by me:
  //
  for ( i=is; i<=ie; i++ ) {
    for ( j=0; j<N; j++ )
      a[i-is][j] = a[i-is][j]/d[j];
  }
  //
  // Print balanced matrix in proper order:
  // (We try to do it blockwise. Needs synchronization of different tasks..
  //  this is NOT a bullet-proof solution)
  if ( me==0 ) {
    print_block( N, is, ie, a );
    for ( i=1; i<np; i++ ) {
      MPI_Send( &tmp, 1, MPI_INT, i, tag_print, MPI_COMM_WORLD );
      MPI_Recv( &tmp, 1, MPI_INT, i, tag_print, MPI_COMM_WORLD, &stat );
    }
  } else {
    MPI_Recv( &tmp,1, MPI_INT, 0, tag_print, MPI_COMM_WORLD, &stat );
    print_block( N, is, ie, a );
    MPI_Send( &tmp, 1, MPI_INT, 0, tag_print, MPI_COMM_WORLD );
  }
  //
  MPI_Finalize();
}