
#include <iostream.h>
#include <math.h>
#include <stdlib.h>

#include "littleMD3.h"


void init_cell(Cell** cell, int xi, int yi, int zi)
{
  // Allocate the cell
  Cell* this_cell = *cell = new Cell;

  // Set some cell values
  this_cell->my_x = xi;
  this_cell->my_y = yi;
  this_cell->my_z = zi;
    
  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;

  this_cell->min_x = params.min_x + (pdx * xi) / params.cell_dim_x;
  this_cell->max_x = params.min_x + (pdx * (xi+1)) / params.cell_dim_x;
  this_cell->min_y = params.min_y + (pdy * yi) / params.cell_dim_y;
  this_cell->max_y = params.min_y + (pdy * (yi+1)) / params.cell_dim_y;
  this_cell->min_z = params.min_z + (pdz * zi) / params.cell_dim_z;
  this_cell->max_z = params.min_z + (pdz * (zi+1)) / params.cell_dim_z;

  // Build up the list of atoms, using a uniform random distribution
  // Since we are doing each cell separately, we don't know exactly
  // how many atoms to put in it.  Therefore, we will allocate a
  // random number such that the total number of atoms on average
  // would come out right.  We'll do this by pretending to assign
  // all the atoms, but only actually assigning a 1/# of cells fraction
  // of them.
  double* x = new double[params.n_atoms];
  double* y = new double[params.n_atoms];
  double* z = new double[params.n_atoms];
  int atom_count = 0;

  const double prob = 1. / (params.cell_dim_x * params.cell_dim_y 
			    * params.cell_dim_z);
  const double cdx = this_cell->max_x - this_cell->min_x;
  const double cdy = this_cell->max_y - this_cell->min_y;
  const double cdz = this_cell->max_z - this_cell->min_z;

  // Give some seed that is unique to this cell
  unsigned short int seed[3];
  seed[0] = (unsigned short)xi;
  seed[1] = (unsigned short)yi;
  seed[2] = (unsigned short)zi;

  int i;
  for(i=0; i < params.n_atoms; i++) {
    if (erand48(seed) < prob) {
      x[atom_count] = this_cell->min_x + (cdx * erand48(seed));
      y[atom_count] = this_cell->min_y + (cdy * erand48(seed));
      z[atom_count] = this_cell->min_z + (cdz * erand48(seed));
      atom_count++;
    }
  }
//    cout << "Assigning " << atom_count << " atoms to cell " 
//	 << xi << "," << yi << "," << zi << endl;

  // Allocate the atom array for the cell
  this_cell->atoms = new Atom[atom_count];
  this_cell->n_atoms = atom_count;

  // Store the positions into the cells.  Also randomly determine
  // a mass and charge, and zero out the velocity and force
  for(i=0;i<atom_count;i++) {
    Atom* this_atom = &(this_cell->atoms[i]);
    
    this_atom->m = 10. * erand48(seed);
    this_atom->q = 5. * erand48(seed) - 2.5;
    this_atom->x = x[i];
    this_atom->y = y[i];
    this_atom->z = z[i];
    this_atom->vx = this_atom->vy = this_atom->vz = 0;
    this_atom->vhx = this_atom->vhy = this_atom->vhz = 0;
    this_atom->fx = this_atom->fy = this_atom->fz = 0;
  }

  delete [] x;
  delete [] y;
  delete [] z;
}


double  cell_self(Cell* this_cell)
{
  double potentialEnergy = 0.0;
  int i,j;
  for(i=0;i<this_cell->n_atoms;i++)
    for(j=i+1;j<this_cell->n_atoms;j++) {
      potentialEnergy += calc_force(&(this_cell->atoms[i]),
				    &(this_cell->atoms[j]));
    }
  //  CkPrintf("%d,%d): self:%f\n", potentialEnergy);
  return potentialEnergy;
}

double cell_neighbor(Cell* cell1, Cell* cell2)
{
  double potentialEnergy = 0.0;
  int i,j;
  for(i=0;i<cell1->n_atoms;i++)
    for(j=0;j<cell2->n_atoms;j++) {
      potentialEnergy += calc_force(&(cell1->atoms[i]),&(cell2->atoms[j]));
    }
  /*  CkPrintf("(%d,%d): nbr PE:%f (%d,%d)\n", 
	   cell1->my_x, cell1->my_y, potentialEnergy, 
	   cell1->n_atoms, cell2->n_atoms); */
  return potentialEnergy;
}

double calc_force(Atom *atom1 , Atom* atom2)
{
  double potentialEnergy = 0.0;
  // Convert values to SI units
  const double x1 = atom1->x * 1E-10;
  const double y1 = atom1->y * 1E-10;
  const double z1 = atom1->z * 1E-10;
  const double q1 = atom1->q * 1.602E-19;

  const double x2 = atom2->x * 1E-10;
  const double y2 = atom2->y * 1E-10;
  const double z2 = atom2->z * 1E-10;
  const double q2 = atom2->q * 1.602E-19;

  const double dx = x2-x1;
  const double dy = y2-y1;
  const double dz = z2-z1;
  const double r2 = dx*dx + dy*dy + dz * dz;

  const double csi = params.cutoff * 1E-10;
  const double cutoff2 = csi * csi;

  if (r2 > cutoff2)  // Outside cutoff, ignore
    return 0.0;

  const double r = sqrt(r2);
  
  const double kq1q2 = 9e9*q1*q2;
  const double f = kq1q2 / (r2*r);
  const double fx = f*dx;
  const double fy = f*dy;
  const double fz = f*dz;

  atom1->fx += fx;
  atom2->fx -= fx;
  atom1->fy += fy;
  atom2->fy -= fy;
  atom1->fz += fz;
  atom2->fz -= fz;

  //  cout << "Force x1=" << x1 << ", " << y1 << endl;
  //  cout << "Force x2=" << x2 << ", " << y2 << endl;
  //  cout << "Force dx=" << dx << ", " << dy << endl;
  //  cout << "Force f=" << fx << ", " << fy << endl;

  potentialEnergy -= kq1q2 / r;
  return potentialEnergy;
}


#ifdef NEW_UPDATE_CELL
double update_cell(Cell* this_cell, int first_step)
#else
double update_cell(Cell* this_cell)
#endif
{
  double kineticEnergy = 0.0;
  // This is a Verlet integrator, which uses the half-step velocity
  // to calculate new positions.  This is a time-reversible integration
  // scheme that has some nice numerical properties
  // The equations are:
  // v(t) = v(t-1/2) + dt/2 * a(t)
  // v(t+1/2) = v(t) + dt/2 * a(t)
  // x(t+1) = x(t) + dt * v(t+1/2)
  // Notice that at the end of this function, you end up with
  // v(t) and x(t+1), so to get x and v for the same point in time,
  // you must save x before it gets updated.

  const double dt = params.step_sz * 1E-15;
  const double dt2 = dt/2;

  int i;
  for(i=0;i<this_cell->n_atoms; i++) {
    Atom* this_atom = &(this_cell->atoms[i]);
    const double mass = this_atom->m * 1.66e-27;
    const double dt2m = dt2/mass;

    const double tax = dt2m * this_atom->fx;
    const double tay = dt2m * this_atom->fy;
    const double taz = dt2m * this_atom->fz;

    //    convert things to SI
    double vx = this_atom->vx * 1E-10/1E-15;
    double vy = this_atom->vy * 1E-10/1E-15;
    double vz = this_atom->vz * 1E-10/1E-15;
    double vhx = this_atom->vhx * 1E-10/1E-15;
    double vhy = this_atom->vhy * 1E-10/1E-15;
    double vhz = this_atom->vhz * 1E-10/1E-15;

    // For the first step, we are given an initial velocity.
    // After that, we must calculate it.
    if (!first_step) {
      vx = vhx + tax;
      vy = vhy + tay;
      vz = vhz + taz;
    }
    //    cout << "F = " << this_atom->fx << ", " << this_atom->fy << endl;
    //    cout << "A = " << tax << ", " << tay << endl;
    //    cout << "V = " << vx << ", " << vy << endl;
    
    kineticEnergy += 0.5 * mass * (vx*vx+vy*vy+vz*vz);

    vhx = vx + tax;
    vhy = vy + tay;
    vhz = vz + taz;
    
    this_atom->x += dt * vhx / 1E-10; 
    this_atom->y += dt * vhy / 1E-10;
    this_atom->z += dt * vhz / 1E-10;

    // Convert other values back from SI and store
    this_atom->vx = vx * 1E-15 / 1E-10;
    this_atom->vy = vy * 1E-15 / 1E-10;
    this_atom->vz = vz * 1E-15 / 1E-10;
    this_atom->vhx = vhx * 1E-15 / 1E-10;
    this_atom->vhy = vhy * 1E-15 / 1E-10;
    this_atom->vhz = vhz * 1E-15 / 1E-10;
  }

  for(i=0;i<this_cell->n_atoms; i++)
    this_cell->atoms[i].fx = 
      this_cell->atoms[i].fy =
      this_cell->atoms[i].fz = 0;

  return kineticEnergy;
}
