/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   http://lammps.sandia.gov, Sandia National Laboratories
   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract
   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
   certain rights in this software.  This software is distributed under 
   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.
------------------------------------------------------------------------- */

#include "string.h"
#include "compute_hf_atom.h"
#include "atom.h"
#include "modify.h"
#include "comm.h"
#include "force.h"
#include "memory.h"
#include "error.h"

#include "pair.h"
#include "bond.h"
#include "angle.h"
#include "dihedral.h"
#include "improper.h"

#include "neighbor.h"
#include "neigh_request.h"
#include "neigh_list.h"

//#include "stdio.h"

using namespace LAMMPS_NS;

#define INVOKED_PERATOM 4


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

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

	peratom_flag = 1;
	size_peratom = 0;

	peatomflag = 1;
	timeflag = 1;
	comm_reverse = 1;

	if (narg == 3) {
		pairflag = 1;
		bondflag = angleflag = dihedralflag = improperflag = 1;
	} else {
		pairflag = 0;
		bondflag = angleflag = dihedralflag = improperflag = 0;
		int iarg = 3;
		while (iarg < narg) {
			if (strcmp(arg[iarg],"pair") == 0) pairflag = 1;
			else if (strcmp(arg[iarg],"bond") == 0) bondflag = 1;
			else if (strcmp(arg[iarg],"angle") == 0) angleflag = 1;
			else if (strcmp(arg[iarg],"dihedral") == 0) dihedralflag = 1;
			else if (strcmp(arg[iarg],"improper") == 0) improperflag = 1;
			else error->all("Illegal compute pe/atom command");
			iarg++;
		}
	}

	nmax = 0;
	hf = NULL;
}

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

ComputeHFAtom::~ComputeHFAtom()
{
  memory->sfree(hf);
}

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

void ComputeHFAtom::init()
{
  int count = 0;
  for (int i = 0; i < modify->ncompute; i++)
    if (strcmp(modify->compute[i]->style,"hf/atom") == 0) count++;
  if (count > 1 && comm->me == 0)
    error->warning("More than one compute hf/atom");

  // pairs

  if (force->pair == NULL)
    error->all("Compute group/group requires pair style be defined");
  pair = force->pair;
  cutsq = force->pair->cutsq;

  // need an occasional half neighbor list

  int irequest = neighbor->request((void *) this);
  neighbor->requests[irequest]->pair = 0;
  neighbor->requests[irequest]->compute = 1;
  neighbor->requests[irequest]->occasional = 1;
}

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

void ComputeHFAtom::compute_peratom()
{
  invoked |= INVOKED_PERATOM;

  // grow hf array if necessary

  if ((atom->nlocal > nmax)||(atom->nmax > nmax)) {
    memory->sfree(hf);
    nmax = atom->nmax;
    hf = (double *) memory->smalloc(nmax*sizeof(double),"compute/hf/atom:hf");
    scalar_atom = hf;
  }

  // compute kinetic energy for each atom in group

  double mvv2e = force->mvv2e;
  double **v = atom->v;
  double *mass = atom->mass;
  double *rmass = atom->rmass;
  int *mask = atom->mask;
  int *type = atom->type;
  int nlocal = atom->nlocal;

  // ntotal includes ghosts if either newton flag is set
  int ntotal = nlocal;
  if (force->newton) ntotal += atom->nghost;

  // clear local energy array
  for (int i = 0; i < ntotal; i++) hf[i] = 0.0;

  if (mass)
    for (int i = 0; i < nlocal; i++) {
      if (mask[i] & groupbit) {
			hf[i] = 0.5 * mvv2e * mass[type[i]] * (v[i][0]*v[i][0] + v[i][1]*v[i][1] + v[i][2]*v[i][2]);
      } else hf[i] = 0.0;
    }

  else
    for (int i = 0; i < nlocal; i++) {
      if (mask[i] & groupbit) {
			hf[i] = 0.5 * mvv2e * rmass[i] * (v[i][0]*v[i][0] + v[i][1]*v[i][1] + v[i][2]*v[i][2]);
      } else hf[i] = 0.0;
    }

  // npair includes ghosts if either newton flag is set
  //   b/c some bonds/dihedrals call pair::ev_tally with pairwise info
  // nbond includes ghosts if newton_bond is set

  int npair = nlocal;
  int nbond = nlocal;
  if (force->newton) npair += atom->nghost;
  if (force->newton_bond) nbond += atom->nghost;

  // add in per-atom contributions from each force
  
  if (pairflag && force->pair) {
    double *eatom = force->pair->eatom;
    for (int i = 0; i < npair; i++) {
	 	hf[i] += eatom[i];
	 }
  }

  if (bondflag && force->bond) {
    double *eatom = force->bond->eatom;
    for (int i = 0; i < nbond; i++) hf[i] += eatom[i];
  }

  if (angleflag && force->angle) {
    double *eatom = force->angle->eatom;
    for (int i = 0; i < nbond; i++) hf[i] += eatom[i];
  }

  if (dihedralflag && force->dihedral) {
    double *eatom = force->dihedral->eatom;
    for (int i = 0; i < nbond; i++) hf[i] += eatom[i];
  }

  if (improperflag && force->improper) {
    double *eatom = force->improper->eatom;
    for (int i = 0; i < nbond; i++) hf[i] += eatom[i];
  }

	// communicate ghost energy between neighbor procs

	if (force->newton) comm->reverse_comm_compute(this);

	// zero energy of atoms not in group
	// only do this after comm since ghost contributions must be included

	for (int i = 0; i < nlocal; i++)
		if (!(mask[i] & groupbit)) hf[i] = 0.0;

	int		inum, jnum, i, j, itype, jtype;
	double	xtmp, ytmp, ztmp, delx, dely, delz;
	int		*ilist,*jlist,*numneigh,**firstneigh;
	double	**x = atom->x;
	int		nall = nlocal + atom->nghost;
	double	rsq,eng,fpair,factor_coul,factor_lj;
	double	*special_coul = force->special_coul;
	double	*special_lj = force->special_lj;
	int		newton_pair = force->newton_pair;

	// invoke half neighbor list (will copy or build if necessary)

	neighbor->build_one(list->index);

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

	double	one[4],all[4];
	one[0] = one[1] = one[2] = one[3] = 0.0;

	for (int 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];
		printf("inum=%d ii=%d i=%d r=(% f,% f,% f) itype=%d jnum=%d nall=%d ",
			inum,ii,i,xtmp,ytmp,ztmp,itype,jnum,nall);

		double	esum = 0;
		for (int jj = 0; jj < jnum; jj++) {
			j = jlist[jj];
			if (j < nall) 
				factor_coul = factor_lj = 1.0;
			else {
				factor_coul = special_coul[j/nall];
				factor_lj = special_lj[j/nall];
				//printf("j/nall=%d factor_coul=%f factor_lj=%f jj=%d j=%d\n",j/nall,factor_coul,factor_lj,jj,j);
				j %= nall;
			}
			delx = xtmp - x[j][0];
			dely = ytmp - x[j][1];
			delz = ztmp - x[j][2];
			rsq = delx*delx + dely*dely + delz*delz;
			jtype = type[j];
			//printf("\t jj=%d j=%d del=(% f,% f,% f) jtype=%d\n",jj,j,delx,dely,delz,jtype);
			//printf("jj=%d j=%d jtype=%d rsq=%f cutsq[itype][jtype]=%f\n",jj,j,jtype,rsq,cutsq[itype][jtype]);
			if (rsq < cutsq[itype][jtype]) {
				eng = pair->single(i,j,itype,jtype,rsq,factor_coul,factor_lj,fpair);
				//printf("newton_pair=%d j=%d nlocal=%d\n", newton_pair, j, nlocal);
				if (newton_pair || j < nlocal) {
					one[0] += eng;
					esum += eng;
					one[1] += delx*fpair;
					one[2] += dely*fpair;
					one[3] += delz*fpair;

					// energy computed twice so tally half amount
					// only tally force if I own igroup atom

				} else {
					one[0] += 0.5*eng;
					esum += 0.5*eng;
					one[1] += delx*fpair;
					one[2] += dely*fpair;
					one[3] += delz*fpair;
				}
			}
		}
		//printf("ESUM=%f one[0]=%f\n", esum, one[0]);
		printf("ESUM=%f\n", esum);
	}
	MPI_Allreduce(one,all,4,MPI_DOUBLE,MPI_SUM,world);

}

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

int ComputeHFAtom::pack_reverse_comm(int n, int first, double *buf)
{
  int i,m,last;

  m = 0;
  last = first + n;
  for (i = first; i < last; i++) buf[m++] = hf[i];
  return 1;
}

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

void ComputeHFAtom::unpack_reverse_comm(int n, int *list, double *buf)
{
  int i,j,m;

  m = 0;
  for (i = 0; i < n; i++) {
    j = list[i];
    hf[j] += buf[m++];
  }
}

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

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

/* ----------------------------------------------------------------------
   memory usage of local atom-based array
------------------------------------------------------------------------- */

double ComputeHFAtom::memory_usage()
{
  double bytes = nmax * sizeof(double);
  return bytes;
}