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

#define G 6.673e-11          // GRAVITATIONAL CONSTANT
#define PI 3.1415926535898   // CONSTANT PI

typedef struct ParticleStruct
{
  double px, py;
  double vx, vy;
  double Mass;
  double ForceAccumX, ForceAccumY;
} Particle;

int NumProcs;
int MyID;
MPI_Datatype MPI_PARTICLE;

void ComputeGravForce(Particle A, Particle B, double *ForceX, double *ForceY)
{
  double ABx = B.px - A.px;
  double ABy = B.py - A.py;
  double DistAB2 = ABx*ABx+ABy*ABy;
  double DistAB = sqrt(DistAB2);
  double DistAB3 = DistAB2 * DistAB;
  double Scale = (G * A.Mass * B.Mass) / DistAB3;
  *ForceX = Scale * ABx;
  *ForceY = Scale * ABy;
}

void ComputeGravForces(Particle *Particles, int NumParticles)
{
  int i,j, k;
  double ForceX, ForceY;
  int DestProc, SourceProc;
  int LeftProc, RightProc;
  MPI_Status Status;

 // ALLOCATE SPACE FOR RECEIVED PARTICLES
  Particle *CommParticles = (Particle *)malloc(sizeof(Particle)*NumParticles);
  for (i=0; i<NumParticles; i++)
  {
    Particles[i].ForceAccumX=0;
    Particles[i].ForceAccumY=0;
  }

  // COMPUTE PROC IDS OF LEFT AND RIGHT NEIGHBORING PROCS
  LeftProc = (MyID+NumProcs-1)%NumProcs;
  RightProc = (MyID+1)%NumProcs;  

  // SEND OUR LOCAL PARTICLES TO NEXT PROCESSOR (RIGHT)
  MPI_Send(Particles,NumParticles,MPI_PARTICLE,RightProc,0,MPI_COMM_WORLD);

  for (k=0; k<NumProcs/2; k++)
  {
    // SEND UPDATED "RECEIVED" PARTICLES TO NEXT PROCESSOR
    if (k>0)  
      MPI_Send(CommParticles,NumParticles,MPI_PARTICLE,RightProc,0,MPI_COMM_WORLD);
 
    // RECEIVE PARTICLES FROM PREV PROCESSOR (LEFT)
    MPI_Recv(CommParticles,NumParticles,MPI_PARTICLE,LeftProc,0,MPI_COMM_WORLD,&Status);

    // COMPUTE FORCES BETWEEN RECEIVED AND LOCAL PARTICLES
    // (UPDATE LOCAL AND RECEIVED PARTICLES)
    for (i=0; i<NumParticles; i++)
      for (j=0; j<NumParticles; j++)
      {
        ComputeGravForce(Particles[i],CommParticles[j],&ForceX,&ForceY);
        Particles[i].ForceAccumX+=ForceX;
        Particles[i].ForceAccumY+=ForceY;
        CommParticles[j].ForceAccumX-=ForceX;
        CommParticles[j].ForceAccumY-=ForceY;
      }
  }    

  // SEND FORCES BACK TO SOURCE
  DestProc = (MyID+NumProcs-(NumProcs/2)) % NumProcs;
  MPI_Send(CommParticles,NumParticles,MPI_PARTICLE,DestProc,0,MPI_COMM_WORLD);

  // RECEIVE PARTICLES FROM PREV PROCESSOR
  SourceProc = (MyID+(NumProcs/2))%NumProcs;
  MPI_Recv(CommParticles,NumParticles,MPI_PARTICLE,SourceProc,0,
              MPI_COMM_WORLD,&Status);

  for (i=0; i<NumParticles; i++)
  {
    Particles[i].ForceAccumX+=CommParticles[i].ForceAccumX;
    Particles[i].ForceAccumY+=CommParticles[i].ForceAccumY;
  }

  // COMPUTE FORCES BETWEEN RECEIVED AND LOCAL PARTICLES
  // (UPDATE LOCAL AND RECEIVED PARTICLES)
  for (i=0; i<(NumParticles-1); i++) {
    for (j=i+1; j<NumParticles; j++)
    {
      ComputeGravForce(Particles[i],Particles[j],&ForceX,&ForceY);
      Particles[i].ForceAccumX+=ForceX;
      Particles[i].ForceAccumY+=ForceY;
      Particles[j].ForceAccumX-=ForceX;
      Particles[j].ForceAccumY-=ForceY;
    }
    printf("\t%f\t%f", Particles[i].px, Particles[i].py);
  }

  free(CommParticles);
}

void UpdateParticles(Particle *Particles, int NumParticles, float DeltaT)
{
  int i;
  double Scale, dVx, dVy;

  for (i=0; i<NumParticles; i++)
  {
    Scale = DeltaT / Particles[i].Mass;
    dVx = Particles[i].ForceAccumX * Scale;
    dVy = Particles[i].ForceAccumY * Scale;
    Particles[i].px += (Particles[i].vx + dVx*0.5) * DeltaT;
    Particles[i].py += (Particles[i].vy + dVy*0.5) * DeltaT;
    Particles[i].vx += dVx;
    Particles[i].vy += dVy;
  }
}

void CheckMomentumConservation(Particle *Particles, int NumParticles)
{
  double TotalMomentumX=0, TotalMomentumY=0;
  int i;
  for (i=0; i<NumParticles; i++)
  {
    TotalMomentumX += (Particles[i].Mass * Particles[i].vx);
    TotalMomentumY += (Particles[i].Mass * Particles[i].vy);
  }  
  printf("[%d/%d]: %g %g\n",MyID,NumProcs,TotalMomentumX,TotalMomentumY);
}

void InitConfig(Particle *Particles, int NumParticles)
{
  int TotalNumParticles = NumParticles*NumProcs;
  int i,k;
  double t1, t2;
  for (i=0; i<NumParticles; i++)
  {
    k = MyID*NumParticles+i;
    t1 = (1.0 + ((double)k) / ((double)TotalNumParticles)) * 0.5;
    t2 = (2.0 * PI * k) / ((double)TotalNumParticles);
    Particles[i].Mass = ((double)(TotalNumParticles-k))/((double)TotalNumParticles);
    Particles[i].px = t1*cos(t2);
    Particles[i].py = t1*sin(t2);
    Particles[i].vx = 0;
    Particles[i].vy = 0;
    Particles[i].Mass = 100000000;
  }
 

}

void RunSimulation(int NumParticles, int NumIterations, double TimeStep)
{
  int k;
  double TotalTime, AvgTimePerIter, TotalNumParticles, InteractionRate;

  Particle *Particles = (Particle *)malloc(sizeof(Particle)*NumParticles);
  clock_t Time = clock();
  
  InitConfig(Particles,NumParticles);

  for (k=0; k<NumIterations; k++)
  {
    ComputeGravForces(Particles,NumParticles);
    UpdateParticles(Particles,NumParticles,TimeStep);
  }

  TotalTime = (clock()-Time)/(double)CLOCKS_PER_SEC;
  AvgTimePerIter = TotalTime / (double)NumIterations;
  TotalNumParticles = NumParticles*NumProcs;
  InteractionRate = (TotalNumParticles*(TotalNumParticles-1)) / AvgTimePerIter;
  printf("%g\n",InteractionRate);

//  CheckMomentumConservation(Particles,NumParticles);

  free(Particles);
}

int main(int argc, char **argv)
{
  int NumParticles, NumIterations;
  double TimeStep;

  MPI_Init(&argc,&argv);
  MPI_Comm_size(MPI_COMM_WORLD,&NumProcs);
  MPI_Comm_rank(MPI_COMM_WORLD,&MyID);  
  MPI_Type_contiguous(7,MPI_DOUBLE,&MPI_PARTICLE);

    NumParticles=1000;
    NumIterations=10;
    TimeStep=1;
    RunSimulation(NumParticles,NumIterations,TimeStep);
  
  MPI_Finalize();
  return 0;
}