/* ----------------------------------------------------------------------
   Daniel Schwen
   Helper compute to store potential energies of individual
   pair styles in pair hybrid/weighted 
------------------------------------------------------------------------- */

#include "math.h"
#include "string.h"
#include "stdlib.h"
#include "compute_storepe.h"
#include "atom.h"
#include "update.h"
#include "modify.h"
#include "neighbor.h"
#include "neigh_list.h"
#include "neigh_request.h"
#include "force.h"
#include "pair.h"
#include "comm.h"
#include "memory.h"
#include "error.h"

using namespace LAMMPS_NS;

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

ComputeStorePE::ComputeStorePE(LAMMPS *lmp, int narg, char **arg) :
  Compute(lmp, narg, arg)
{
  if (narg != 4) error->all(FLERR,"Illegal compute store/pe command");

  vector_flag = 1;
  size_vector = atoi(arg[3]);
  vector = new double[size_vector];
  sendvector = new double[size_vector];

  for( int i=0; i<size_vector; ++i ) vector[i] = 0.0;
}

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

ComputeStorePE::~ComputeStorePE()
{
  delete[] vector;
  delete[] sendvector;
}

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

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

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

double ComputeStorePE::memory_usage()
{
  double bytes = size_vector * sizeof(double);
  return bytes;
}

/* ----------------------------------------------------------------------
   gather compute vector data from other nodes
------------------------------------------------------------------------- */

void ComputeStorePE::compute_vector()
{
  invoked_vector = update->ntimestep;
  for( int i=0; i<size_vector; ++i ) sendvector[i] = vector[i];
  MPI_Allreduce(sendvector,vector,size_vector,MPI_DOUBLE,MPI_SUM,world);
}