#include <numeric>
#include <iostream>
#include <vector>
#include <cstdlib>
#include <cmath>
#include <time.h>
#include <sys/time.h>
#include <sstream>
#include <string>
#include "mpi.h"

const double ARmass = 39.94800000; //A.U.s
const double ARsigma = 3.40500000; // Angstroms
const double AReps   = 119.800000; // Kelvins
const double CellDim = 12.0000000; // Angstroms
int NPartPerCell = 10;
void swap (int *,int *);
void sort(int[],int[],int,int);

using namespace std;
// Class for keeping track of the properties for a particle
class Particle{
public:
  double x;		// position in x axis
  double y;		// position in y axis
  double fx;		// total forces on x axis
  double fy;		// total forces on y axis
  double ax;		// acceleration on x axis
  double ay;		// acceleration on y axis
  
  double vx;		// velocity on x axis
  double vy;		// velocity on y axis
  // Default constructor
  Particle() 
    : x(0.0),y(0.0),fx(0.0),fy(0.0),ax(0.0),
      ay(0.0),vx(0.0),vy(0.0){
  }
  
  double update(){
    // We are using a 1.0 fs timestep, this is converted
    double DT = 0.000911633;
    ax = fx/ARmass;
    ay = fy/ARmass;
    vx += ax*0.5*DT;
    vy += ay*0.5*DT;
    x += DT*vx;
    y += DT*vy;
    fx = 0.0;
    fy = 0.0;
    return 0.5*ARmass*(vx*vx+vy*vy);
  }
};
 

class Cell {
public:
  vector<Particle> particles;
  // Buffer to hold coordinates recieved
  double **remote_particles;
  int absx,absy;
  MPI_Request reqr[8],reqs[8];
  int nreqr, nreqs;
  // Buffer to hold coordinates to send
  double *PartCoordsBuff;
  Cell(int x, int y, int nParticles) : particles(nParticles){
    absx = x;
    absy = y;
    //  We will be filling the cells, making sure than
    //  No atom is less than 2 Angstroms from another
    double rcelldim = double(CellDim);
    double centerx = rcelldim*double(absx) + 0.5*rcelldim;
    double centery = rcelldim*double(absy) + 0.5*rcelldim;
   
    // Randomly initialize particles to be some distance 
    // from the center of the square
    // place first atom in cell
    particles[0].x = centerx + ((drand48()-0.5)*(CellDim-2));
    particles[0].y = centery + ((drand48()-0.5)*(CellDim-2));
    
    double r,rx,ry;
    for (int i = 1; i < particles.size(); i++) {
      r = 0;
      while(r < 2.7) {   // limit to how close atoms can get together
	r = 2.70001;
	rx = centerx + ((drand48()-0.5)*(CellDim-2)); 
	ry = centery + ((drand48()-0.5)*(CellDim-2));
      	//loop over all other current atoms
	for(int j = 0; j<i; j++){
	  double rxj = rx - particles[j].x;
	  double ryj = ry - particles[j].y;
	  double rt = rxj*rxj+ryj*ryj;
	  if(rt < r) r=rt;
	}
      }
      particles[i].x = rx;
      particles[i].y = ry;
    }
    
    // Now that we have our particles, we need to allocate the proper spaces
    // For this example, we will actually allocate space for all of the particles,
    // needed, which is a waste of space, but simplifies.  The first element of each 
    // array will be the number of particles sent.
    remote_particles = (double**)malloc(sizeof(double*)*8);
    for(int i=0;i<8;i++)
      remote_particles[i] = (double *)calloc(sizeof(double),NPartPerCell*2+1);
    PartCoordsBuff = (double*)malloc(sizeof(double)*((particles.size()*2)+1));
  }
  
  void Communicate(int*** CellMap, int Dimension){
    // Pack the message
    for(int i=0;i<particles.size();i++){
      PartCoordsBuff[1+i*2] = particles[i].x;
      PartCoordsBuff[1+i*2+1] = particles[i].y;
    }
    
    // Assuming a good MPI layer, communication between two of the same 
    // processor should be pretty efficient.  Not screening greatly simplifies
    // the code
    
    // i+1,j
    int ircount =0;
    if(absx!=Dimension-1){
      int tag = (absx+1)*Dimension+absy;
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
      		MPI_DOUBLE,CellMap[absx+1][absy][0],
      		tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    //i-1,j
    if(absx!= 0){
      int tag = (absx-1)*Dimension+absy;
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
		MPI_DOUBLE,CellMap[absx-1][absy][0],
		tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    // i,j+1
    if(absy!=Dimension-1){
      int tag = (absx)*Dimension+(absy+1);
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
		MPI_DOUBLE,CellMap[absx][absy+1][0],
		tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    // i,j-1
    if(absy !=0){
      int tag = absx*Dimension+(absy-1);
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
		MPI_DOUBLE,CellMap[absx][absy-1][0],
		tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    // i+1,j+1
    if(absx != Dimension-1 && absy != Dimension-1){
      int tag = (absx+1)*Dimension + (absy+1);
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
		MPI_DOUBLE,CellMap[absx+1][absy+1][0],
		tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    //i+1,j-1
    if(absx != Dimension-1 && absy != 0){ 
      int tag = (absx+1)*Dimension + (absy-1);
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
		MPI_DOUBLE,CellMap[absx+1][absy-1][0],
		tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    //i-1,j+1
    if(absx != 0 && absy != Dimension-1){
      int tag = (absx-1)*Dimension + (absy+1);
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
		MPI_DOUBLE,CellMap[absx-1][absy+1][0],
		tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    //i-1,j-1
    if(absx !=0 && absy !=0){
      int tag = (absx-1)*Dimension + (absy-1);
      MPI_Irecv(remote_particles[ircount],NPartPerCell*2+1,
		MPI_DOUBLE,CellMap[absx-1][absy-1][0],
	       	tag,MPI_COMM_WORLD,&reqr[ircount]);
      ircount++;
    }
    
    // Now the sending
    int iscount = 0;

    //i+1,j   
    if(absx != Dimension-1){
      int tag = absx*Dimension + absy;
      MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx+1][absy][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    // i-1,j
    if(absx != 0){
      int tag = absx*Dimension + absy;
      MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx-1][absy][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    // i,j+1
    if(absy != Dimension-1){
      int tag = (absx)*Dimension + (absy);
      MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx][absy+1][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    // i,j-1
    if(absy != 0){
      int tag = (absx)*Dimension + (absy);
      MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx][absy-1][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    // i+1,j+1
    if(absx != Dimension-1 && absy != Dimension-1){
      int tag = (absx)*Dimension + (absy);
       MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx+1][absy+1][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    //i+1,j-1
    if(absx!=Dimension-1 && absy!=0){
      int tag = (absx)*Dimension + (absy);
      MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx+1][absy-1][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    //i-1,j+1
    if(absx!=0 && absy!= Dimension-1){
      int tag = (absx)*Dimension + (absy);
      MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx-1][absy+1][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    //i-1,j-1
    if(absx!=0 && absy!=0){
      int tag = (absx)*Dimension + (absy);
      MPI_Isend(PartCoordsBuff,NPartPerCell*2+1,MPI_DOUBLE,CellMap[absx-1][absy-1][0],
		tag,MPI_COMM_WORLD,&reqs[iscount]);
      iscount++;
    }
    nreqr = ircount;
    nreqs = iscount;
  }
  void PostWaits(void){
    MPI_Status statr[8];
    MPI_Status stats[8];
    MPI_Waitall(nreqs,reqs,stats);
    MPI_Waitall(nreqr,reqr,statr);
  }
};
    
//Force calculation
//------------------------------------
// Overloaded to handle new case
double interact(Particle &atom1, double atom2x, double atom2y){
  double rx,ry,rz,r,fx,fy,fz,f;
  double sigma6,sigma12;
  
  // computing base values
  rx = atom1.x - atom2x;
  ry = atom1.y - atom2y;
  r = rx*rx + ry*ry;

  if(r < 0.000001)
    return 0.0;
  
  r=sqrt(r);
  double sigma2 = (ARsigma/r)*(ARsigma/r);
  sigma6 = sigma2*sigma2*sigma2;
  sigma12 = sigma6*sigma6;
  f = ((sigma12-0.5*sigma6)*48.0*AReps)/r;
  fx = f * rx;
  fy = f * ry;
  // updating particle properties
  atom1.fx += fx;
  atom1.fy += fy;
  return 4.0*AReps*(sigma12-sigma6);
}

double interact(Particle &atom1, Particle &atom2){
  return interact(atom1,atom2.x,atom2.y);
}

int ComputeAtomsPerCell(int ***CellMap, 
			int NRows,int NCols, 
			int NParts){
//  Uncomment Lines for a parallel version
  int MY_PE;
  int max = 0;
  MPI_Comm_rank(MPI_COMM_WORLD,&MY_PE);
  if(MY_PE == 0){
    double *rCellSizes = (double*)malloc(sizeof(double)*NRows*NCols);
    double ransum = 0.0;
    
    for(int i=0;i<NRows;i++){
      for(int j=0;j<NCols;j++){
	double ran = drand48();
	rCellSizes[i*NRows + j] = ran;
	ransum += ran;
      }
    }
    int molsum = 0;
    for(int i=0;i<NRows;i++){
      for(int j=0;j<NCols;j++){
	double ran = rCellSizes[i*NRows+j]/ransum;
	CellMap[i][j][3] = int(NParts*ran)+1;
	molsum += CellMap[i][j][3];
	if(CellMap[i][j][3] > max) max = CellMap[i][j][3];
      }
    }
    free(rCellSizes);
    while(molsum < NParts){
      for(int i=0;i<NRows;i++){
	if(molsum >=NParts){
	  break;
	}
	for(int j=0;j<NCols;j++){
	  CellMap[i][j][3]++;
	  molsum++;
	  if(CellMap[i][j][3] > max) max = CellMap[i][j][3];
	  if(molsum >= NParts) {
	    break;
	  }
	}
      }
    }
    while(molsum > NParts){
      for(int i=0;i<NRows;i++){
	if(molsum <= NParts){
	  break;
	}
	for(int j=0;j<NCols;j++){
	  if(CellMap[i][j][3] > 1){
	    CellMap[i][j][3]--;
	    molsum--;
	  }
	  if(CellMap[i][j][3] > max) max = CellMap[i][j][3];
	  if(molsum <=NParts) {
	    break;
	  }
	}
      }
    }
  }
  MPI_Bcast(&max,1,MPI_INT,0,MPI_COMM_WORLD);
  for(int i=0;i<NRows;i++){
    for(int j=0;j<NCols;j++){
      MPI_Bcast(&CellMap[i][j][3],1,MPI_INT,0,MPI_COMM_WORLD);
    }
  }
  return max;
}

int main(int argc,char* argv[]){
  
  MPI_Init(&argc,&argv);
  int MY_PE,NCPUS;
  MPI_Comm_rank(MPI_COMM_WORLD,&MY_PE);
  MPI_Comm_size(MPI_COMM_WORLD,&NCPUS);

  // Get Command Line Arguments
  if(argc!=3){
    if(MY_PE == 0)
      cerr << "\nIncorrect syntax: should be two arguments\n";
    MPI_Finalize();
    exit(2);
  }

  // The Number of CPUs must be a Perfect Square
  int SQncpus = int(sqrt(double(NCPUS)));
  if(NCPUS % SQncpus != 0) {
    if(MY_PE == 0)
      cerr << "\nThe Number of cores must be a perfect square (e.g. 4,16,64, etc...)\n";
    MPI_Finalize();
    exit(3);
  }
  
  
  int NumParticles = atoi(argv[1]);
  int NumIterations= atoi(argv[2]);
  
  int Dimension    = int(sqrt(NumParticles/double(NPartPerCell)));
  int TotalCells   = Dimension*Dimension;
  int NCellsPerCPU = 0;
  int MaxCellSize;
      
  if(TotalCells < NCPUS){
    if(MY_PE == 0)
      cerr << "\nThe Number of Cells is less than the number of CPUS: \n"
	   << "NCells = " << TotalCells << " NCPUS = "<<NCPUS;
    MPI_Finalize();
    exit(5);
  }
  
  // Create a map for the which processor each cell is assigned
  // This map is replicated, but small

  // First Assign the processors
  // Need to rename Cell Map
  int ***CellCpus = (int***)malloc(sizeof(int**)*Dimension);
  for(int i=0;i<Dimension;i++){
    CellCpus[i] = (int**)malloc(sizeof(int*)*Dimension);
    for(int j=0;j<Dimension;j++)
      CellCpus[i][j] = (int*)calloc(sizeof(int),4); // 0 MPI_Rank the Cell Belongs
                                                    // 1 2 are the local x,y (Not really needed)
                                                    // 3 is the number of particles in the cell
  }	
  
  // OK now that we know which cell is on which cpu, we need to allocate the main space
  //This is replicated, so that it gives exactly the same number of particles as the
  //serial version
  MaxCellSize = ComputeAtomsPerCell(CellCpus,Dimension,Dimension,NumParticles);
  
  if(MY_PE == 0){
    cout << "\nThe Total Number of Cells is " << TotalCells
	 << " With a maximum of " << MaxCellSize << " particles per cell,"
	 << "\nand "  << NumParticles << " particles total in system\n";
    fflush(stdout);
  }

  // Now we need to reload the balance
  // We want to distribute them so that the number of atoms per processor is
  // as even as we can get it.

  // To map the work on a processor, we need a new array the tells us
  // the Cells absolute coordinates
  int icpu = -1;
  int isum = 0;
    
  int *cc = (int*)calloc(sizeof(int),NCPUS); // holds how many atoms
  int *ci = (int*)calloc(sizeof(int),NCPUS); // holds the corresponding Cpus
  
  // Since I made sure that the number of Cells is at least the number of CPUS
  // this should be valid
  // This assigns at least one cell to each processor, so none are empty
  icpu =-1;
  int iatoms=0;
  for(int i=0;i<NCPUS;i++) ci[i] = i;

  int starti,startj;
  int ii,jj;
  for(ii=0;ii<Dimension;ii++){
    for(jj=0;jj<Dimension;jj++){  
      icpu++;
      if(icpu > NCPUS-1){
	break;
      }
      CellCpus[ii][jj][0] = icpu;
      cc[icpu] += CellCpus[ii][jj][3];
      isum++;
    }
    if(icpu > NCPUS-1) {
      break;
    }
  }
  starti=ii;
  startj=jj;
  // Now quick sort the cc and ci arrays so that I can put cells in the lowest
  
  while(isum < TotalCells){
    sort(cc,ci,0,NCPUS);
    for(int i=0;i<NCPUS;i++)
    icpu = -1;
    for(ii=starti; ii<Dimension;ii++){
      for(jj=startj;jj<Dimension;jj++){
	icpu++;
	if(icpu > NCPUS-1 || isum >= TotalCells){
	  break;
	}
	CellCpus[ii][jj][0] = ci[icpu];
	cc[icpu] += CellCpus[ii][jj][3];
	isum ++;
      }
      if(icpu > NCPUS-1 || isum >= TotalCells){
	break;
      }
      startj =0;
    }
    starti=ii;
    startj=jj;
  }
  free(cc);free(ci);

  int LocalCPUnCells = 0;
  vector <int> LocalCPUCells;
  
  for(int i=0;i<Dimension;i++){
    for(int j=0;j<Dimension;j++){
      if(MY_PE == CellCpus[i][j][0]){
	LocalCPUnCells++;
	LocalCPUCells.push_back(i);
	LocalCPUCells.push_back(j);
      }
    }
  }
  
  // Allocate Cell, space for Particles couched in the Cell creator
  // Local space, only enough for what is on the CPU
  
  vector< Cell > cells;
  for(int i=0;i<LocalCPUnCells;i++){
    int absx = LocalCPUCells[i*2];
    int absy = LocalCPUCells[i*2+1];
    cells.push_back(Cell(absx,absy,CellCpus[absx][absy][3]));
  }

  MPI_Barrier(MPI_COMM_WORLD);
  double TimeStart = MPI_Wtime();
  for (int t = 0; t < NumIterations; t++) { // For each timestep
    double TotPot1=0.0;
    double TotPot2=0.0;
    double TotKin=0.0;
      
    // Start the communications going
    // We can communicate whilst all of the cells are computing the 
    // interactions with themselves.

    for(int cx1 = 0; cx1<LocalCPUnCells; cx1++){
      int absx = cells[cx1].absx;
      int absy = cells[cx1].absy;
      cells[cx1].Communicate(CellCpus,Dimension);
    }
    
    // Computing the interactions of the cells with the particles inside.
    for (int cx1 = 0; cx1 < LocalCPUnCells; cx1++) {
      // Consider interactions between particles within the cell
      for (int i = 0; i < cells[cx1].particles.size(); i++) {
	for (int j = 0; j < cells[cx1].particles.size(); j++) {
	  if(i!=j)
	    TotPot1 += interact(cells[cx1].particles[i],
				cells[cx1].particles[j]);
	}
      }
    }
     // We have to wait here for the comm to finish 
    for(int cx1=0;cx1 < LocalCPUnCells;cx1++)
      cells[cx1].PostWaits();
    
    // Consider each other cell...
    for(int cx1=0;cx1 < LocalCPUnCells;cx1++){
      for (int cx2 = 0; cx2 < cells[cx1].nreqr; cx2++) {
	// Consider all interactions between particles
	int nparts=cells[cx1].remote_particles[cx2][0];
	for (int i = 0; i < cells[cx1].particles.size(); i++) {
	  for (int j = 0; j < nparts; j++) {
	    TotPot2 += interact(cells[cx1].particles[i],
				cells[cx1].remote_particles[cx2][j*2+1],
				cells[cx1].remote_particles[cx2][j*2+2]);
	  }
	}
      }
    }
    
    // End iteration over cells    
    // Apply the accumulated forces; update accelerations, velocities, positions
    for (int cx1 = 0; cx1 < LocalCPUnCells; cx1++) {
      for(int i=0; i < cells[cx1].particles.size(); i++) {
	TotKin += cells[cx1].particles[i].update();
      }
    }
     double TotPotSum = 0.0;
     double TotKinSum = 0.0;
     double LSum = TotPot1+TotPot2;
    
     MPI_Reduce(&LSum,&TotPotSum,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
     MPI_Reduce(&TotKin,&TotKinSum,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
    if(MY_PE == 0)
      printf("\nIteration#%d with Total Energy %e per Atom",
	     t,(TotKinSum+TotPotSum)/NumParticles);
  } // iterate time steps
  double TimeEnd = MPI_Wtime();
  if(MY_PE == 0){
    cout << "\nTime for " << NumIterations << " is "<< TimeEnd-TimeStart;
    fflush(stdout);
  }
    
  for(int i=0;i<Dimension;i++){
    for(int j=0;j<Dimension;j++){
      free(CellCpus[i][j]);
    }
    free(CellCpus[i]);
  }
  free(CellCpus);

  MPI_Finalize();
  return 0;
}

void swap(int *a, int *b){
  int t=*a; 
  *a=*b; 
  *b=t;
}

void sort(int arr[], int arr2[], int beg, int end){
  if (end > beg + 1){
    int piv = arr[beg], l = beg + 1, r = end;
    while (l < r){
      if (arr[l] <= piv)
        l++;
      else{
	r--;
        swap(&arr[l], &arr[r]);
	swap(&arr2[l], &arr2[r]);
      }
    }
    l--;
    swap(&arr[l], &arr[beg]);
    swap(&arr2[l],&arr2[beg]);
    sort(arr, arr2, beg, l);
    sort(arr, arr2, r, end);
  }
}