/*------------------------------------------*/
/* An Example of N-Body Simulation.         */
/* Modified by Saleh Elmohamed  10/20/98    */
/* NPAC, CPS-615 , Fall 98.                 */
/*------------------------------------------*/
/* Function:                                */
/* A simple version of an n-body code.      */
/* It uses the simple but highly            */
/* suboptimal n^2 algorithm, and does not   */
/* take advantage of symmetry.              */
/* The time integrator is a simple leapfrog */
/* scheme.                                  */
/*------------------------------------------*/

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

extern void   srand48();
extern double drand48();

/* Pipeline version of the algorithm... but the MPE is not supported in
   SunMPI ....... */
/* we really need the velocities as well... */

typedef struct {
  double x, y, z;
  double mass;
} Particle;

/* Save the forces and old velocities */
typedef struct {
  double xold, yold, zold;
  double fx, fy, fz;
} ParticleV;

#define MAX_PARTICLES 4000

/* This is 2-D only */
void ComputeForces(Particle *particles, ParticleV *pv, int npart, 
		   Particle *recvbuf, int rlen, double  *max_f )
{
  int i, j;
  double xi, yi, mi, rx, ry, mj, r, fx, fy;
  double xnew, ynew, rmin;

/* Compute forces (2D only) */
  for (i=0; i<npart; i++) {
    rmin = 100.0;
    xi   = particles[i].x;
    yi   = particles[i].y;
    fx   = 0.0;
    fy   = 0.0;
    for (j=0; j<rlen; j++) {
      rx = xi - recvbuf[j].x;
      ry = yi - recvbuf[j].y;
      mj = recvbuf[j].mass;
      r  = rx * rx + ry * ry;
      /* ignore overlap and same particle */
      if (r == 0.0) continue;
      if (r < rmin) rmin = r;
      /* compute forces */
      r  = r * sqrt(r);
      fx += mj * rx / r;
      fy += mj * ry / r;
    }
    pv[i].fx -= fx;
    pv[i].fy -= fy;
    /* Compute a rough estimate of (1/m)|df / dx| */
    fx                      = sqrt(fx*fx + fy*fy)/rmin;
    if (fx > *max_f) *max_f = fx;
  }
}

void PrintParticles(Particle *particles, int npart, double t )
{
  int i;

  for (i=0; i<npart; i++) {
    printf( "%lf %lf\n", particles[i].x, particles[i].y );
  }
}

int main(int argc, char **argv )
{
  Particle  particles[MAX_PARTICLES];   /* Particles on LOCAL node */
  ParticleV pv[MAX_PARTICLES];          /* Particle velocity */
  Particle  *recvbuf;
  int       rank, size, npart, i, j;    /* location of local particles */
  int       step, rlen;
  int       totpart,                    /* total number of particles */
  cnt;                                  /* number of times in loop */
  MPI_Datatype particletype;
  double    time;                       /* Computation time */
  double    dt, dt_old;                 /* Integration time step */
  double    t;                          /* Time actually integrated to */
  double    a0, a1, a2;
  void      *pipe;                      /* Opaque structure for communication 
					   pipe */
  int       debug_flag = 1;

  MPI_Init( &argc, &argv );
  MPI_Comm_rank( MPI_COMM_WORLD, &rank );
  MPI_Comm_size( MPI_COMM_WORLD, &size );

  /* Everyone could have a different size ... */
  npart = 2 * size;
  if (argc > 1) 
    npart = atoi(argv[1]) / size;
  if (npart > MAX_PARTICLES) 
    MPI_Abort( MPI_COMM_WORLD, 1 );

  MPI_Allreduce( &npart, &totpart, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD );

  cnt    = 100;

  MPI_Type_contiguous( 4, MPI_DOUBLE, &particletype );
  MPI_Type_commit( &particletype );

/* Generate the initial values */
  srand48( rank * 117 );
  for (i=0; i<npart; i++) {
    particles[i].x    = drand48();
    particles[i].y    = drand48();
    particles[i].z    = drand48();

    /* Normalize the mass with respect to the gravitational constant 
       (multiply by sqrt(G), G = 6.668E-8 dyne-cm^2/gm^2 or 
       6.668E-16 Joule-m^2/kg^2, sqrt(G) = 2.582e-8 
       */

    particles[i].mass = 1 * 2.582e-8;
    /* This is the zero initial velocity case */
    pv[i].xold              = particles[i].x;
    pv[i].yold              = particles[i].y;
    pv[i].zold              = particles[i].z;
    pv[i].fx        = 0;
    pv[i].fy        = 0;
    pv[i].fz        = 0;
  }

  /* Add a particle in the center with much larger mass */
  if (rank == 0) {
    particles[0].x    = 0.5;
    particles[0].y    = 0.5;
    particles[0].z    = 0.5;
    particles[0].mass = 1000 * 2.582e-8;
    pv[0].xold              = particles[0].x;
    pv[0].yold              = particles[0].y;
    pv[0].zold              = particles[0].z;
    pv[0].fx        = 0;
    pv[0].fy        = 0;
    pv[0].fz        = 0;
  }

  if (debug_flag && rank == 0) 
    PrintParticles( particles, npart, 0.0 );

  dt     = 0.001;
  dt_old = 0.001;
  
  /* Create a communication pipe -- NOT SUPPORTED in SunMPI so forget it*/
#if 0
  MPE_Pipe_create( MPI_COMM_WORLD, particletype, npart, &pipe ); 
#endif

  time = MPI_Wtime();
  t    = 0.0;
  while (cnt--) {
    double max_f, dt_est, new_dt, dt_new;
    
    /* Load the initial sendbuffer -- NOT SUPPORTED in SunMPI so forget it*/
#if 0
    MPE_Pipe_start( pipe, particles, npart, 1 );
#endif

    /* integation is a0 * x^+ + a1 * x + a2 * x^- = f / m */
    a0   = 2.0 / (dt * (dt + dt_old));
    a2   = 2.0 / (dt_old * (dt + dt_old));
    a1   = -(a0 + a2);      /* also -2/(dt*dt_old) */

    /* Compute self forces (on same processor) */
    max_f = 0;
    ComputeForces( particles, pv, npart, particles, npart, &max_f );
    
    /* For other particles, compute interaction forces */
    for (step=1; step<size; step++) {

      /* Push pipe, get new data -- NOT SUPPORTED in SunMPI so forget it*/
#if 0
      MPE_Pipe_push( pipe, &recvbuf, &rlen );
#endif

      /* Compute forces */
      ComputeForces( particles, pv, npart, recvbuf, rlen, &max_f );
    }

    /* Once we have the forces, we compute the changes in position */
    for (i=0; i<npart; i++) {
      double xi, yi;
      /* Very, very simple leapfrog time integration.  We use a variable 
         step version to simplify time-step control. */
      xi             = particles[i].x;
      yi             = particles[i].y;
      particles[i].x = (pv[i].fx - a1 * xi - a2 * pv[i].xold) / a0;
      particles[i].y = (pv[i].fy - a1 * yi - a2 * pv[i].yold) / a0;
      pv[i].xold     = xi;
      pv[i].yold     = yi;
      pv[i].fx       = 0;
      pv[i].fy       = 0;
    }

    t += dt;

    if (debug_flag && rank == 0) 
      PrintParticles( particles, npart, t );

    /* Recompute a time step. Stability criteria is roughly 
       2 = dt * sqrt(1/m |df/dx|), or
       2/sqrt(1/m |df/dx|) >= dt.  We leave a little room */
    
    dt_est = 1.0/sqrt(max_f);
    
    /* Set a minimum: */

    if (dt_est < 1.0e-6) dt_est = 1.0e-6;

    MPI_Allreduce( &dt_est, &dt_new, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD );
    
    /* Modify time step */

    if (dt_new < dt) {
      dt_old = dt;
      dt     = dt_new;
      if (debug_flag && rank == 0) 
	printf( "#New time step is %lf\n", dt );
    }
    else if (dt_new > 4.0 * dt) {
      dt_old = dt;
      dt    *= 2.0;
      if (debug_flag && rank == 0) 
	printf( "#New time step is %lf\n", dt );
    }

    /* We could do graphics here (move particles on the display) */
  }
  time = MPI_Wtime() - time;
  if (rank == 0) {
    printf( "#Computed %d particles in %lf seconds\n", totpart, time );
  }

#if 0
  MPE_Pipe_free( &pipe );
  MPI_Type_free( &particletype );
#endif

  MPI_Finalize();
}