/*   Metropolis_CS.c                                                      */
/*   Demonstration of integration by random numbers                       */
/*        Client-Server                                                   */
/*                                                                        */
/*   Metropolis Algorithm                                                 */
/*   Gaussian Random Numbers distribution                                 */ 
/*                                                                        */

/*   Syntax:                                                              */
/*   mpiexec -np 5 ./Metropolis_CS -p  >  data_file                       */

                                         /* Michel Vallieres */

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

#define NMAX 1000  /* trial size of randomarrays x & y */


/* function to integrate */
double func( double x, double y )
{
  return M_PI *  ( x*x + y*y );
}



/* probability profile of the random numbers */
double weight( double x, double y )
{
  return exp( - ( x*x + y*y ) );
}



/* Metropolis algorithm to generate Random       */
/* Points in the plane with probability profile  */
/* weight(x,y)                                   */
void Metropolis ( double x[], double y[], double (*weight)(double,double),  
                  double delta, int N, int *accepted_count )
{
  double xt, yt, xx, yy, eta, ratio;
  int local_count, count;
 
  count = 0;                                           /* accepted count */

  for( local_count=0 ; local_count < N-1 ; local_count++ )
    { 
       xx = (double)rand()/(double)RAND_MAX;           /* random # in 0,1 */
       yy = (double)rand()/(double)RAND_MAX;
       xt = x[local_count] + delta * ( 2*xx - 1.0 );   /* in range -1 to 1 */
       yt = y[local_count] + delta * ( 2*yy - 1.0 );   /* trial step */
                                                       /* probability ratio */
       x[local_count+1] = x[local_count];              /* new step = old step */
       y[local_count+1] = y[local_count];
       ratio = weight( xt, yt ) / weight( x[local_count] , y[local_count] );   
       eta = (double)rand()/(double)RAND_MAX;          /* random # in 0-1  */
       if ( ratio > eta )                              /* Metropolis selection */
         {
            count++;
            x[local_count+1]  = xt;                    /* selected step */
            y[local_count+1]  = yt;
         }

    }
  *accepted_count = *accepted_count + count;
}




/* code executed by node 0 */
/* control of calculation & final Monte Carlo sums */
void node_zero_code( int numprocs )
{
  double f, sumf, sumf2, integral, sigma, store[2];
  int    N_average, N_input, leftover, N_node, N_total, count, first;
  int    node, kill;
  double x0, y0;
  MPI_Status recv_status;

  x0 = 0.5;                                     /* initial position in plane*/
  y0 = 0.6; 
                                                /* transmit to server node */
  MPI_Ssend( &x0, 1, MPI_DOUBLE, numprocs-1, 2123, MPI_COMM_WORLD );
  MPI_Ssend( &y0, 1, MPI_DOUBLE, numprocs-1, 3123, MPI_COMM_WORLD );

  f = func( x0, y0 );                           /* Monte-Carlo sums */
  sumf = f;
  sumf2 = f*f;

  N_total = 1;                                  /* Monte-Carlo counters */
  first = 1;

                                                /* accuracy improvement */
  fprintf( stderr, "\n Enter # of random numbers to use ( 0 to quit ): " );
  while ( scanf( "%d", &N_input ) != EOF && N_input > 0 )
     {
                                                /* adjust N_input to # nodes */
       N_average =  N_input/(numprocs-2);
       leftover = N_input % (numprocs-2);
                                                /* add to total */
       N_total = N_total + N_input - first;
                                                /* send to slave nodes */
       for ( node=1 ; node<numprocs-1 ; node++ )
	 {
	   N_node = N_average;
	   if ( node <= leftover )
	     N_node++;

          if ( first == 0 ) N_node++;          /* adjust count on MC steps */
          first = 0;                           /* on subsequent calls      */

	                                        /* send # to slave nodes */
	   MPI_Ssend( &N_node, 1, MPI_INT, node, 51, MPI_COMM_WORLD );
	 }

                                                /* receive partial sums */
       for ( node=1 ; node<numprocs-1 ; node++ )
	 {
	   MPI_Recv( store, 2,  MPI_DOUBLE, MPI_ANY_SOURCE, 423, 
		     MPI_COMM_WORLD, &recv_status );
	   sumf = sumf + store[0];
	   sumf2 = sumf2 + store[1];
	 }

                                                /* final Monte Carlo integral */
                                                /* x0 & y0 */ 
       integral = sumf/(double)N_total;
       sigma = sqrt( ( sumf2/(double)N_total - integral*integral ) / 
                                   (double)N_total );
       fprintf( stderr, " Integral = %f  +/- sigma = %f \n", integral, sigma );

                                                /* more trials? */
       fprintf( stderr, "\n Enter # of random number to use ( 0 to quit ): " );

     } 

                                               /* kill slave nodes & server */
  kill = -1;
  for ( node=1 ; node<numprocs-1 ; node++ )
    MPI_Ssend( &kill, 1, MPI_INT, node, 51, MPI_COMM_WORLD );

  MPI_Ssend( &kill, 1, MPI_INT, numprocs-1, 123, MPI_COMM_WORLD );

                                                /* receive count (sent before suicide)*/
  MPI_Recv( &count, 1, MPI_INT, numprocs-1, 301, MPI_COMM_WORLD, &recv_status );

                                                /* Acceptance ratio - should be > ~0.5 */
  fprintf( stderr, "\n Acceptance ratio: %d %d  %f\n\n",
	   N_total, count, (double)count/(double)N_total );

                                                /* end all */
  MPI_Finalize();
  exit(0);
                   
}




/* worker code: requests & receives random numbers and  */
/* calculates MC partial sums                           */
void worker_nodes_code( int numprocs )
{
  static double *x, *y;                         /* static: permanent storage */
  static int current_x_y_size;
  double f, sumf, sumf2, store[2];
  int    N, ir;
  MPI_Status recv_status;
                                                /* current size of x & y */
                                                /* arrays - reserve memory */
  current_x_y_size = NMAX;
  x = (double *)malloc( current_x_y_size * sizeof(double) );
  y = (double *)malloc( current_x_y_size * sizeof(double) );

                                                /* wait for # of MC steps */
  N = 1;
  while ( N > 0 )
    { 
                                                /* receive N from node 0 */
     MPI_Recv( &N, 1, MPI_INT, 0, 51, MPI_COMM_WORLD, &recv_status );

                                                /* suicide point */
     if ( N < 0 )
       {
         MPI_Finalize();
         exit(0);
        }
                                                /* request MC steps */
     MPI_Ssend( &N, 1, MPI_INT, numprocs-1, 123, MPI_COMM_WORLD );

                                                /* space for x & y arrays */
     if ( N > current_x_y_size ||
            current_x_y_size == 0 )             /* enough ?               */
       {                                        /* reserve more           */
         free ( x );
         free ( y );
         current_x_y_size = N + 100;
         x = (double *)malloc( current_x_y_size * sizeof(double) );
         y = (double *)malloc( current_x_y_size * sizeof(double) );
       }

                                                /* receive random #s */
      MPI_Recv( x, N, MPI_DOUBLE, numprocs-1, 201, 
		MPI_COMM_WORLD, &recv_status );
      MPI_Recv( y, N, MPI_DOUBLE, numprocs-1, 202, 
		MPI_COMM_WORLD, &recv_status );

                                                /* Monte-Carlo integral */
      sumf = 0.0;
      sumf2 = 0.0;
      for ( ir=1 ; ir<N ; ir++ )
	{
	  f =  func( x[ir], y[ir] );
	  sumf = sumf + f;
	  sumf2 = sumf2 + f*f;
	}
      store[0] = sumf;
      store[1] = sumf2;
      MPI_Ssend( store, 2, MPI_DOUBLE, 0, 423, MPI_COMM_WORLD );
    
    }

}




/* server code: answers to requests for Random Numbers */
void server_code( int plot_control )
{
  double delta;
  int    N, ir;
  static int accepted_count;
  static int current_x_y_size;
  static double *x, *y;
  int    to;
  MPI_Status recv_status;
  double x00, y00;                              /* carry over & initial */
                                                /* x, y position        */
 
                                                /* accepted trial steps count */
  accepted_count = 1;
                                                /* initial position in plane*/
  MPI_Recv( &x00, 1, MPI_DOUBLE, 0, 2123, MPI_COMM_WORLD, &recv_status );
  MPI_Recv( &y00, 1, MPI_DOUBLE, 0, 3123, MPI_COMM_WORLD, &recv_status );
 
                                                /* current size of x & y */
                                                /* arrays - reserve memory */
  current_x_y_size = NMAX;
  x = (double *)malloc( current_x_y_size * sizeof(double) );
  y = (double *)malloc( current_x_y_size * sizeof(double) );

  delta = 0.2;                                  /* Metropolis delta - */
                                                /* size of jump       */
  N = 1;
  while ( N > 0 )                               /* loop over requests */
                                                /* from nodes */
    {
                                                /* receive request */
      MPI_Recv( &N, 1, MPI_INT, MPI_ANY_SOURCE, 123, 
		MPI_COMM_WORLD, &recv_status );

                                                /* in case of early suicide */
      if ( N <= 0 ) break;

                                                /* space for x & y arrays */
      if ( N > current_x_y_size ||
               current_x_y_size == 0 )          /* enough ?               */
        {                                       /* reserve more           */
          free ( x );
          free ( y );
          current_x_y_size = N + 100;
          x = (double *)malloc( current_x_y_size * sizeof(double) );
          y = (double *)malloc( current_x_y_size * sizeof(double) );
         }

                                                /* Metropolis Algorithm */
      x[0] = x00;
      y[0] = y00;
      Metropolis ( x, y, weight, delta, N , &accepted_count );

                                                /* plot info */
      if ( plot_control == 1 )
        for ( ir=1 ; ir<N ; ir++ )
	  printf( " %f %f %f\n", x[ir], y[ir], func( x[ir], y[ir] ) );

      x00 = x[N-1];
      y00 = y[N-1];
                                                /* origin of request */
      to = recv_status.MPI_SOURCE;
                                                /* send random numbers */
      MPI_Ssend( x, N, MPI_DOUBLE, to, 201, MPI_COMM_WORLD ); 
      MPI_Ssend( y, N, MPI_DOUBLE, to, 202, MPI_COMM_WORLD );
    }
                                               /* suicide signal ( N<0 ) */
                                               /* has been received      */
    
                                               /* send count */
  MPI_Ssend( &accepted_count, 1, MPI_INT, 0, 301, MPI_COMM_WORLD );

                                               /* suicide */
  MPI_Finalize();
  exit(0);
}



/* demo program to illustrate Monte-Carlo integration using */
/* Gaussian random numbers in the plane */
int main( int argc, char *argv[] )
{
  int    ir, N, N_total, N_input, N_node, N_average, leftover;
  int    plot_control, count;
  int    myid, numprocs;

                                                /* join the MPI virtual machine */
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &myid);
  MPI_Comm_size(MPI_COMM_WORLD, &numprocs);
                                                /**** MASTER NODE ****/
                                                /* node zero code */
                                                /* overall organization of code */
  if ( myid == 0 )
       node_zero_code( numprocs );

                                                /**** SLAVE NODES ****/
  else if ( myid > 0 && myid < numprocs-1 )
       worker_nodes_code( numprocs );
                                                /*****  SERVER ****/
  else
    {
                                                /* plot control         */
       plot_control = 0;                        /* only server needs it */
       if ( argc == 2 )                        
         if ( strcmp( argv[1],  "-p") == 0  )
           plot_control = 1;
                                                /* call the server code */
       server_code( plot_control );
    }

}