#include <stdio.h>
#include <iostream.h>

#include "littleMD3.h"

SimParams params;
Cell**** cells;  // Its a 3-D array of pointers to Cell
double pot_energy;
double kin_energy;
bool first_step;

int main(int, char**)
{
  init();

  int step;
  for (step=0; step < params.steps; step++) {
    calc_forces();
    update_coordinates();
    cout << "Step " << step 
	 << " PE = " << pot_energy 
	 << " KE = " << kin_energy 
	 << " Energy = " << kin_energy + pot_energy << endl;

    pot_energy = 0;
    kin_energy = 0;
    first_step = false;
  }
  done();
}

void setParamsDebug()
{
  params.steps = 20;
  params.n_atoms = 5000;
  params.min_x = 0;
  params.min_y = 0;
  params.min_z = 0;
  params.max_x = 40;
  params.max_y = 40;
  params.max_z = 20;
  params.cutoff = 15;
  params.margin = 2;
  params.step_sz = 1.;
}

void setParamsReal()
{
  params.steps = 10;
  params.n_atoms = 5000;
  params.min_x = 0;
  params.min_y = 0;
  params.min_z = 0;
  params.max_x = 300;
  params.max_y = 300;
  params.max_z = 300;
  params.cutoff = 15;
  params.margin = 2;
  params.step_sz = 1.;
}

void init()
{
  // Set some startup parameters
  // These values can be customized

  setParamsDebug();
  //setParamsReal();

  first_step = true;
  pot_energy = 0;
  kin_energy = 0;
  
  const double pdx = params.max_x - params.min_x;
  const double pdy = params.max_y - params.min_y;
  const double pdz = params.max_z - params.min_z;

  params.cell_dim_x = 1+ (int) (pdx / (params.cutoff + params.margin));
  params.cell_dim_y = 1+ (int) (pdy / (params.cutoff + params.margin));
  params.cell_dim_z = 1+ (int) (pdz / (params.cutoff + params.margin));

  // Allocate a 3-D array of cell pointers using dynamic memory, by
  // allocating two auxiliary arrays of pointers and then arrays of
  // pointers to the cells themselves

  cells = new Cell***[params.cell_dim_x];
  int i,j;
  for(i=0; i < params.cell_dim_x;i++) {
    cells[i] = new Cell**[params.cell_dim_y];
    for(j=0; j < params.cell_dim_y; j++) {
      cells[i][j] = new Cell*[params.cell_dim_z];
    }
  }

  // Call init_cell() for each cell*
  int xi, yi, zi;
  for(xi=0; xi < params.cell_dim_x; xi++) 
    for(yi=0; yi < params.cell_dim_y; yi++)
      for(zi=0; zi < params.cell_dim_z; zi++)
	init_cell(&(cells[xi][yi][zi]),xi,yi,zi);
}

void done()
{
  int i,j;
  for(i=0;i<params.cell_dim_x; i++) {
    for(j=0;j<params.cell_dim_y; j++)
      delete [] cells[i][j];
    delete [] cells[i];
  }

  delete [] cells;
  printf("All done\n");
}

void calc_forces()
{
  // Loop through the cells and calculate the forces
  // These two loops (outermost) are the (only) ones you should parallelize

  int xi, yi, zi;
  for(xi=0; xi<params.cell_dim_x; xi++)
    for(yi=0; yi < params.cell_dim_y; yi++)
      for(zi=0; zi < params.cell_dim_z; zi++) {
	// First, do interactions with my own atoms
	pot_energy += cell_self(cells[xi][yi][zi]);

	//printf("(%d,%d,%d) self pe: %e\n", xi,yi, zi, pot_energy);      

	// Now do my interaction with neighbors with cell coordinates
	// of neighbors. The other neighbors need not be done,
	// because they were caught earlier by Newton's law.
	// We also need to check the bounds of our cell array

	const int neighbors[13][3] = { {-1,-1,-1}, {-1,-1, 0}, {-1,-1, 1},
				       {-1, 0,-1}, {-1, 0, 0}, {-1, 0, 1},
				       {-1, 1,-1}, {-1, 1, 0}, {-1, 1, 1},
				       { 0,-1,-1}, { 0,-1, 0}, { 0,-1, 1},
				       { 0, 0,-1} };
	int nbor;
	for(nbor=0; nbor < 13; nbor++) {
	  const int xn = xi + neighbors[nbor][0];
	  const int yn = yi + neighbors[nbor][1];
	  const int zn = zi + neighbors[nbor][2];

	  if (xn < 0 || xn >= params.cell_dim_x)
	    continue;

	  if (yn < 0 || yn >= params.cell_dim_y)
	    continue;

	  if (zn < 0 || zn >= params.cell_dim_z)
	    continue;

	  pot_energy += cell_neighbor(cells[xi][yi][zi],cells[xn][yn][zn]);
	  //	  printf("(%d,%d,%d): nbor %d %e in calcf\n",
	  //		 xi, yi, zi, nbor, pot_energy);
	}
      }
}

void update_coordinates()
{
  int xi, yi, zi;
  for(xi=0; xi<params.cell_dim_x; xi++)
    for(yi=0; yi < params.cell_dim_y; yi++)
      for(zi=0; zi < params.cell_dim_z; zi++) {
#ifdef NEW_UPDATE_CELL
	kin_energy += update_cell(cells[xi][yi][zi], first_step);
#else
	kin_energy += update_cell(cells[xi][yi][zi]);
#endif
	//	printf("(%d,%d, %d): finished update_cell: %e\n",
	//	       xi, yi, zi, kin_energy);
      }
}


