/* simulation of N-body problem (gravitational forces)  - Chuck Liang

compile with (on deepdish cluster):

mpecc grav3d.c -o grav3d

Please run only on 1 processor until I figure out how to fix the
X11 forwarding (os/network problem).

*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>
#include "mpe_graphics.h"
#include "mpe.h"

#include "threed.c"

/* Physics stuff:

  Given two objects with masses M1 and M2, the gravitationl force between
  these objects is F = (G*M1*M2)/(d*d), where d is the distance between the
  center of gravity of the two objects.  Newton also says that F=MA,
  so A1 = G*M2/(d*d) is the accelartion caused by the pull of M2 on M1, and
  A2 = G*M1/(d*d) is that caused by M1 on M2.
 
  The gravitational constant G is 6.673e-11 meters^3/(kg*sec^2) in our
  universe (yours may be different :)

  Acceleration is given in meters/sec^2

  Distances are in meters and mass in kg.
*/

/* universal constants */

// screen coordinates will be in units of 10e6 meters 
//#define UDIST 1
#define UDIST 1.0e8

// time unit between frames - in seconds
//#define UTIME 1
//#define UTIME 60
#define UTIME 360
//#define UTIME 3600
//#define UTIME 86400
//#define UTIME 2629800
//#define UTIME 31557600
//#define UTIME 31557600.0e3

// wait for mouse click between frames?
#define SYNCH 0
#define TRACE 0

#define ITERS 10000000

// window width, height:
#define SW 1024
#define SH 768
#define SD 512


// wrapround?
#define WRAPAROUND 1

// number of particles
#define PSIZE 200

// gravitational constant 
#define G 6.673e-11

// mass of earth in kg, radius in meters, mass of sun
#define MEARTH 5.9742e24
#define REARTH 6.4e6
#define MSUN 1.99e30

// body will be a vector of 6 doubles: xcord, ycord, mass, xvel, yvel, radius.
// radius is used only for drawing (unit in pixels)

#define XCORD(AV) AV[0]
#define YCORD(AV) AV[1]
#define ZCORD(AV) AV[2]
#define MASS(AV) AV[3]
#define XVEL(AV) AV[4]
#define YVEL(AV) AV[5]
#define ZVEL(AV) AV[6]
#define RADIUS(AV) AV[7]

// not used for now: if radius is <0, particle has been destroyed/merged
#define DESTROYED(AV) (AV[7]<0)

double BODY[PSIZE][8];
MPE_XGraph graph;

int minb, maxb;  // partition process is responsible for
MPE_Color mycolor;

/**** grid  ****/
#define DROW 24
#define DCOL 32
#define DSPN 12
int grid[DCOL][DROW][DSPN];  // contains indices into BODY array

// calculate effect of force on A and B bodies.
// returns if bodies have collided (not used now)
//(i is local, j maybe remote)
int calcforce(int i, int j)
{
  double *A, *B;
  double d2, dx, dy, dist, sx, sy, sm, sz, dz;
  double accAX, accAY, accAZ, deeper; // accelerations
  A = BODY[i]; B = BODY[j];
  dx = XCORD(B)-XCORD(A);  dy = YCORD(B)-YCORD(A);
  dz = ZCORD(B)-ZCORD(A);
  d2 = (dx*dx) + (dy*dy) + (dz*dz); 
  if (d2==0) d2=1e-100;
  dist = sqrt(d2);

  // now determine if bodies have collide:
  if (ZCORD(A)>ZCORD(B)) deeper = ZCORD(A);  else deeper=ZCORD(B);
  if (dist < UDIST*(dscale(RADIUS(A) + RADIUS(B), deeper/UDIST)))
    { // destroy smaller i 

      //      printf("collusion!\n");

      /*
      if (i<j)
	{
   	  RADIUS(A) = -1;  // signals A is destroyed
          XVEL(A) = YVEL(A) = MASS(A) = 0;
	}
      else 
	{ // merge j into i
	  MASS(A) += MASS(B);
	  RADIUS(A) += RADIUS(B);
	}
      */
      sm = MASS(A)+MASS(B);
     // adjust velocity vector: smaller mass into larger one:
      XVEL(A) = (XVEL(B)*MASS(B)/sm)+(XVEL(A)*MASS(A)/sm);
      YVEL(A) = (YVEL(B)*MASS(B)/sm)+(YVEL(A)*MASS(A)/sm);
      ZVEL(A) = (ZVEL(B)*MASS(B)/sm)+(ZVEL(A)*MASS(A)/sm);
      return 1;
    } // collusion

  dx = dx/dist;  dy = dy/dist; // scale into unit vector
  dz = dz/dist;
  // calculate acceleration vectors
  accAX = dx * (G*MASS(B))/d2; 
  accAY = dy * (G*MASS(B))/d2; 
  accAZ = dz * (G*MASS(B))/d2;
  XVEL(A) += accAX * UTIME;
  YVEL(A) += accAY * UTIME;
  ZVEL(A) += accAZ * UTIME;
  return 0;
}


// update positions, draw
void update()
{
  int i, j, x, y, k, z, oldx, oldy, oldz;
  for(i=minb;i<=maxb;i++)
    { if (!DESTROYED(BODY[i])) {
      x = ((int) ((XCORD(BODY[i]) / UDIST)+0.5)) % XDIM;
      y = ((int) ((YCORD(BODY[i]) / UDIST)+0.5)) % YDIM;
      z = ((int) ((ZCORD(BODY[i]) / UDIST)+0.5)) % ZDIM;
      if (x<0) x = WRAPAROUND * (XDIM + x);
      if (y<0) y = WRAPAROUND * (YDIM + y);
      if (z<0) z = WRAPAROUND * (ZDIM + z);
      // draw to screen (but don't update yet)
      oldx = x; oldy = y; oldz = z; // record previous coordinates
      // erase old circle:
      fillcircle(graph,x,y,z,RADIUS(BODY[i]),MPE_WHITE);

      // calculate new coordinates for object (after calcforce above)
      XCORD(BODY[i]) += XVEL(BODY[i]) * UTIME;
      YCORD(BODY[i]) += YVEL(BODY[i]) * UTIME;
      ZCORD(BODY[i]) += ZVEL(BODY[i]) * UTIME;
      // scale to screen coordinates
      x = ((int) ((XCORD(BODY[i]) / UDIST)+0.5)) % XDIM;
      y = ((int) ((YCORD(BODY[i]) / UDIST)+0.5)) % YDIM;
      z = ((int) ((ZCORD(BODY[i]) / UDIST)+0.5)) % ZDIM;
      if (x<0) x = WRAPAROUND * (XDIM + x); // wraparound?
      if (y<0) y = WRAPAROUND * (YDIM + y);
      if (z<0) z = WRAPAROUND * (ZDIM + z);
      // draw to screen (but don't update yet)
      fillcircle(graph,x,y,z,RADIUS(BODY[i]),mycolor);
      // draw line between points
      if (TRACE) drawline(graph,oldx,oldy,oldz,x,y,z,MPE_BLUE); }
    }
  //  MPE_Update(graph);  - only rank 0 should draw
} // update


void initialize(int rank)
{ int i, r, c, d, dx, dy, dz;
  switch (rank % 8)
    {
      case 0 : mycolor = MPE_RED; break;
      case 1 : mycolor = MPE_BLUE; break;
      case 2 : mycolor = MPE_GREEN; break;
      case 3 : mycolor = MPE_BROWN; break;
      case 4 : mycolor = MPE_YELLOW; break;
      case 5 : mycolor = MPE_ORANGE; break;
      case 6 : mycolor = MPE_GRAY; break;
      default : mycolor = MPE_BLACK; break;
    }

  // initial mesh data
  if (rank == 0)
    {
      /*
      for (i=0;i<PSIZE;i++)
	{ r = i / (DCOL*DSPN);  c = (i-(r*DCOL*DSPN)) / DSPN;  
  	  d = i % DSPN; // depth
  	  XCORD(BODY[i]) = c*UDIST*DCOL;
          YCORD(BODY[i]) = r*UDIST*DROW;
          ZCORD(BODY[i]) = d*UDIST*DSPN;
          RADIUS(BODY[i]) = 1; 
          MASS(BODY[i]) = MEARTH * RADIUS(BODY[i]); // yikes!
	  grid[c][r] = i;
	  dx = c - (DCOL/2);  dy = c - (DROW/2);
          XVEL(BODY[i]) = 2000000  / (1+dx*dx);
          YVEL(BODY[i]) = 50000 / (1+dy*dy);
	}
    } // randomize
      */

      for(i=0;i<PSIZE;i++)
	{
	  XCORD(BODY[i]) = (random() % XDIM) * UDIST;
  	  YCORD(BODY[i]) = (random() % YDIM) * UDIST;
	  ZCORD(BODY[i]) = (random() % ZDIM) * UDIST;
	  RADIUS(BODY[i]) = 1 + (random() % 10);
          MASS(BODY[i]) = MEARTH * RADIUS(BODY[i]); // yikes!
	  XVEL(BODY[i]) = YVEL(BODY[i]) = ZVEL(BODY[i]) = 0;
	}
  // sun
      //    i = grid[DCOL/2][DROW/2];
    i = 0;
    MASS(BODY[i]) = MSUN;
    XCORD(BODY[i])=UDIST*XDIM/2;  
    YCORD(BODY[i])=UDIST*YDIM/2;  ZCORD(BODY[i])=UDIST*(ZDIM/2);
    RADIUS(BODY[i]) = 30;
    XVEL(BODY[i]) = 0;  YVEL(BODY[i])=0; ZVEL(BODY[i])=0;
    /*
    //earth
    i = 1;
    MASS(BODY[i]) = MEARTH;
    XCORD(BODY[i])=UDIST*XDIM/4;  
    YCORD(BODY[i])=UDIST*YDIM/4;  ZCORD(BODY[i])=UDIST*1;
    RADIUS(BODY[i]) = 20;
    XVEL(BODY[i]) = 50;  YVEL(BODY[i])=0; ZVEL(BODY[i])=0;
    */
    } // initial data 

  MPI_Bcast(BODY,8*PSIZE,MPI_DOUBLE,0,MPI_COMM_WORLD);
} // initialize

void drawgrid(MPE_Color);
void simulate(int rank)
{ 
 int i, j, k, x, y, c, x2, y2;
 // draw bodies in initial positions
 update();
 MPI_Barrier(MPI_COMM_WORLD); // initial synch
 for (k=0;k<ITERS;k++)
  {
    //    if (SYNCH && (rank==0)) drawgrid(MPE_WHITE); 
    for(i=minb;i<=maxb;i++)
      for(j=0;j<PSIZE;j++)
	{
	  if (i!=j && !DESTROYED(BODY[i]) && !DESTROYED(BODY[j]))
	    { 
	        c = calcforce(i,j);
	    } // if
	} // for j,i -- done calculating vel vectors
    update();

    // exchange info
    MPI_Allgather(BODY+minb,8*(maxb-minb+1),MPI_DOUBLE,
    		  BODY,8*(maxb-minb+1),MPI_DOUBLE,MPI_COMM_WORLD);
    // MPI book is wrong! the send and recv counts must be the same!

    //    if (SYNCH && rank==0) drawgrid(MPE_GREEN); 

    if (SYNCH) // wait for mouse press on root
      { if (rank==0) MPE_Get_mouse_press( graph, &x, &y, &j );
        MPI_Barrier(MPI_COMM_WORLD);
      }
    if (rank==0) MPE_Update(graph);  // draw next frame
  } // for k - outermost iterations
} // simulate

int main(int argc, char **argv)
{
  int rank, size, i, j;
  char display[128];  
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  if (argc>1) srandom(rank+atoi(argv[1])); // seed random number generator

  if (argc>1) tiltR = toRadians(atoi(argv[1]));

//  MPE_Open_graphics(&graph,MPI_COMM_WORLD,(char*)0,0,0,1024,768,
//  		    MPE_GRAPH_INDEPENDENT);

  minb = rank * (PSIZE/size); // assuming no remainder
  maxb = minb + (PSIZE/size) - 1;
  tiltR=toRadians(220);  ZR = 0.5;  //SR = 0.8;
  XDIM = 800;  YDIM = 640;  ZDIM = 400;
  xshift = (int)(ZDIM * cos(tiltR));
  yshift = (int)(ZDIM * sin(tiltR));
  
  sprintf(display,"%s","10.1.0.2:10.0");
    

  MPE_Open_graphics(&graph,MPI_COMM_WORLD,(char*)0,1,1,
		    XDIM+(int)abs(xshift)+10,YDIM+(int)abs(yshift)+10, 
		    MPE_GRAPH_INDEPENDENT);



  drawbox(graph,1,1,1,XDIM,YDIM,ZDIM,MPE_GREEN);

  initialize(rank);
  simulate(rank);

  MPI_Finalize();
  exit(0);
}