#include "math.h"
#include "string.h"
#include "compute_TEST.h"
#include "atom.h"
#include "atom_vec.h"
#include "update.h"
#include "force.h"
#include "pair.h"
#include "modify.h"
#include "group.h"
#include "neighbor.h"
#include "neigh_list.h"
#include "neigh_request.h"
#include "error.h"

using namespace LAMMPS_NS;

enum{DUMMY0,INVOKED_SCALAR,INVOKED_VECTOR,DUMMMY3,INVOKED_PERATOM};

/* ---------------------------------------------------------------------- */

ComputeTEST::ComputeTEST(LAMMPS *lmp, int narg, char **arg) : 
  Compute(lmp, narg, arg)
{
  if (narg != 3) error->all("Illegal compute TEST command");

  vector_flag = 1;
  size_vector = 3;
  extvector = 1;

  int n = strlen(arg[3]) + 1;
  vector = new double[3];
}

/* ---------------------------------------------------------------------- */

ComputeTEST::~ComputeTEST()
{ 
  delete [] vector;
}

/* ---------------------------------------------------------------------- */

void ComputeTEST::init()
{
  // error checks
  if (atom->avec->ghost_velocity == 0)
    error->all("Compute TEST requires ghost atoms store velocity");

  cutsq = force->pair->cutsq;
}

/* ---------------------------------------------------------------------- */

void ComputeTEST::init_list(int id, NeighList *ptr)
{
  list = ptr;
}

/* ---------------------------------------------------------------------- */

void ComputeTEST::compute_vector()
{
  int i,j,ii,jj,kk,inum,jnum,knum,itype,jtype,ktype;
  double xtmp,ytmp,ztmp,delxij,delyij,delzij,delxik,delyik,delzik;
  double vel;
  int *ilist,*jlist,*klist,*numneigh,**firstneigh;

  invoked_vector = update->ntimestep;

  double **x = atom->x; 
  double **v = atom->v; 
  int *type = atom->type;
  int *mask = atom->mask;
  int nlocal = atom->nlocal;
  int nall = nlocal + atom->nghost;
  

  neighbor->build_one(list->index);

  inum = list->inum;
  ilist = list->ilist;
  numneigh = list->numneigh;
  firstneigh = list->firstneigh;

  double A[3] = {0.0,0.0,0.0};

  for (ii = 0; ii < inum; ii++) {
    i = ilist[ii];
    xtmp = x[i][0];
    ytmp = x[i][1];
    ztmp = x[i][2];
    itype = type[i];
    jlist = firstneigh[i];
    jnum = numneigh[i];

    vel=sqrt(v[i][0]*v[i][0] + v[i][1]*v[i][1] + v[i][2]*v[i][2]);

    for (jj = 0; jj < jnum; jj++) {
      j = jlist[jj];

      delxij = xtmp - x[j][0];
      delyij = ytmp - x[j][1];
      delzij = ztmp - x[j][2];
      rsq = delxij*delxij + delyij*delyij + delzij*delzij;
      jtype = type[j];

      if (rsq < cutsq[itype][jtype]) {
	
	for (kk = 0; kk < jnum; kk++) {
   	  k = jlist[kk];
        
          if (!(mask[i] & groupbit) && !(mask[j] & groupbit) && !(mask[k] & groupbit)) continue;
	
          delxik = xtmp - x[k][0];
          delyik = ytmp - x[k][1];
          delzik = ztmp - x[k][2];
	  rsq = delxik*delxik + delyik*delyik + delzik*delzik;
	  ktype = type[k];

          if (rsq < cutsq[itype][ktype]) {

	    A[0] += (delxij+delxik)*vel;
	    A[1] += (delyij+delxik)*vel;
	    A[2] += (delzij+delxik)*vel;
         }
        }
      }                  
    }                  
  }

  MPI_Allreduce(A,vector,3,MPI_DOUBLE,MPI_SUM,world);
}