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

   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 "stdlib.h"

#include "string.h"

#include "fix_store_coord.h"

#include "atom.h"

#include "domain.h"

#include "group.h"

#include "memory.h"

#include "error.h"

using namespace LAMMPS_NS;

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

FixStoreCoord::FixStoreCoord(LAMMPS *lmp, int narg, char **arg) :

  Fix(lmp, narg, arg)

{

  if (narg < 3) error->all("Illegal fix store/coord command");

  restart_peratom = 1;

  peratom_flag = 1;

  size_peratom_cols = 3;

  peratom_freq = 1;

  // optional args

  int comflag = 0;

  int iarg = 3;

  while (iarg < narg) {

    if (strcmp(arg[iarg],"com") == 0) {

      if (iarg+2 > narg) error->all("Illegal fix store/coord command");

      if (strcmp(arg[iarg+1],"no") == 0) comflag = 0;

      else if (strcmp(arg[iarg+1],"yes") == 0) comflag = 1;

      else error->all("Illegal fix store/coord command");

      iarg += 2;

    } else error->all("Illegal fix store/coord command");

  }

  // perform initial allocation of atom-based array

  // register with Atom class

  xoriginal = NULL;

  grow_arrays(atom->nmax);

  atom->add_callback(0);

  atom->add_callback(1);

  // cm = original center of mass

  double cm[3];

  if (comflag) {

    double masstotal = group->mass(igroup);

    group->xcm(igroup,masstotal,cm);

  }

  // xoriginal = initial unwrapped positions of atoms

  // relative to center of mass if comflag is set

  double **x = atom->x;

  int *mask = atom->mask;

  int *image = atom->image;

  int nlocal = atom->nlocal;

  for (int i = 0; i < nlocal; i++) {

    if (mask[i] & groupbit) {

      domain->unmap(x[i],image[i],xoriginal[i]);

      if (comflag) {

	xoriginal[i][0] -= cm[0];

	xoriginal[i][1] -= cm[1];

	xoriginal[i][2] -= cm[2];

      }

    } else xoriginal[i][0] = xoriginal[i][1] = xoriginal[i][2] = 0.0;

  }

}

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

FixStoreCoord::~FixStoreCoord()

{

  // unregister callbacks to this fix from Atom class

  atom->delete_callback(id,0);

  atom->delete_callback(id,1);

  memory->destroy_2d_double_array(xoriginal);

}

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

int FixStoreCoord::setmask()

{

  int mask = 0;

  return mask;

}

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

   memory usage of local atom-based array

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

double FixStoreCoord::memory_usage()

{

  double bytes = atom->nmax*3 * sizeof(double);

  return bytes;

}

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

   allocate atom-based array

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

void FixStoreCoord::grow_arrays(int nmax)

{

  xoriginal =

    memory->grow_2d_double_array(xoriginal,nmax,3,"fix_msd:xoriginal");

  array_atom = xoriginal;

}

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

   copy values within local atom-based array

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

void FixStoreCoord::copy_arrays(int i, int j)

{

  xoriginal[j][0] = xoriginal[i][0];

  xoriginal[j][1] = xoriginal[i][1];

  xoriginal[j][2] = xoriginal[i][2];

}

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

   pack values in local atom-based array for exchange with another proc

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

int FixStoreCoord::pack_exchange(int i, double *buf)

{

  buf[0] = xoriginal[i][0];

  buf[1] = xoriginal[i][1];

  buf[2] = xoriginal[i][2];

  return 3;

}

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

   unpack values in local atom-based array from exchange with another proc

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

int FixStoreCoord::unpack_exchange(int nlocal, double *buf)

{

  xoriginal[nlocal][0] = buf[0];

  xoriginal[nlocal][1] = buf[1];

  xoriginal[nlocal][2] = buf[2];

  return 3;

}

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

   pack values in local atom-based arrays for restart file

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

int FixStoreCoord::pack_restart(int i, double *buf)

{

  buf[0] = 4;

  buf[1] = xoriginal[i][0];

  buf[2] = xoriginal[i][1];

  buf[3] = xoriginal[i][2];

  return 4;

}

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

   unpack values from atom->extra array to restart the fix

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

void FixStoreCoord::unpack_restart(int nlocal, int nth)

{

  double **extra = atom->extra;

  // skip to Nth set of extra values

  int m = 0;

  for (int i = 0; i < nth; i++) m += static_cast<int> (extra[nlocal][m]);

  m++;

  xoriginal[nlocal][0] = extra[nlocal][m++];

  xoriginal[nlocal][1] = extra[nlocal][m++];

  xoriginal[nlocal][2] = extra[nlocal][m++];

}

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

   maxsize of any atom's restart data

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

int FixStoreCoord::maxsize_restart()

{

  return 4;

}

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

   size of atom nlocal's restart data

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

int FixStoreCoord::size_restart(int nlocal)

{

  return 4;

}

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

   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.

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

#ifndef FIX_STORE_COORD_H

#define FIX_STORE_COORD_H

#include "fix.h"

namespace LAMMPS_NS {

class FixStoreCoord : public Fix {

 public:

  FixStoreCoord(class LAMMPS *, int, char **);

  ~FixStoreCoord();

  int setmask();

  double memory_usage();

  void grow_arrays(int);

  void copy_arrays(int, int);

  int pack_exchange(int, double *);

  int unpack_exchange(int, double *);

  int pack_restart(int, double *);

  void unpack_restart(int, int);

  int size_restart(int);

  int maxsize_restart();

 private:

  double **xoriginal;         // original coords of atoms

};

}

#endif

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

   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 "fix_store_force.h"

#include "atom.h"

#include "update.h"

#include "group.h"

#include "respa.h"

#include "memory.h"

#include "error.h"

using namespace LAMMPS_NS;

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

FixStoreForce::FixStoreForce(LAMMPS *lmp, int narg, char **arg) :

  Fix(lmp, narg, arg)

{

  if (narg < 3) error->all("Illegal fix store/coord command");

  peratom_flag = 1;

  size_peratom_cols = 3;

  peratom_freq = 1;

  nmax = atom->nmax;

  foriginal = memory->create_2d_double_array(nmax,3,"store/force:foriginal");

  array_atom = foriginal;

  // zero the array since dump may access it on timestep 0

  // zero the array since a variable may access it before first run

  int nlocal = atom->nlocal;

  for (int i = 0; i < nlocal; i++)

    foriginal[i][0] = foriginal[i][1] = foriginal[i][2] = 0.0;

}

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

FixStoreForce::~FixStoreForce()

{

  memory->destroy_2d_double_array(foriginal);

}

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

int FixStoreForce::setmask()

{

  int mask = 0;

  mask |= POST_FORCE;

  mask |= POST_FORCE_RESPA;

  mask |= MIN_POST_FORCE;

  return mask;

}

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

void FixStoreForce::init()

{

  if (strcmp(update->integrate_style,"respa") == 0)

    nlevels_respa = ((Respa *) update->integrate)->nlevels;

}

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

void FixStoreForce::setup(int vflag)

{

  if (strcmp(update->integrate_style,"verlet") == 0)

    post_force(vflag);

  else {

    ((Respa *) update->integrate)->copy_flevel_f(nlevels_respa-1);

    post_force_respa(vflag,nlevels_respa-1,0);

    ((Respa *) update->integrate)->copy_f_flevel(nlevels_respa-1);

  }

}

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

void FixStoreForce::min_setup(int vflag)

{

  post_force(vflag);

}

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

void FixStoreForce::post_force(int vflag)

{

  if (atom->nlocal > nmax) {

    nmax = atom->nmax;

    memory->destroy_2d_double_array(foriginal);

    foriginal = memory->create_2d_double_array(nmax,3,"store/force:foriginal");

    array_atom = foriginal;

  }

  double **f = atom->f;

  int *mask = atom->mask;

  int nlocal = atom->nlocal;

  for (int i = 0; i < nlocal; i++)

    if (mask[i] & groupbit) {

      foriginal[i][0] = f[i][0];

      foriginal[i][1] = f[i][1];

      foriginal[i][2] = f[i][2];

    } else foriginal[i][0] = foriginal[i][1] = foriginal[i][2] = 0.0;

}

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

void FixStoreForce::post_force_respa(int vflag, int ilevel, int iloop)

{

  if (ilevel == nlevels_respa-1) post_force(vflag);

}

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

void FixStoreForce::min_post_force(int vflag)

{

  post_force(vflag);

}

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

   memory usage of local atom-based array

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

double FixStoreForce::memory_usage()

{

  double bytes = atom->nmax*3 * sizeof(double);

  return bytes;

}

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

   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.

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

#ifndef FIX_STORE_FORCE_H

#define FIX_STORE_FORCE_H

#include "fix.h"

namespace LAMMPS_NS {

class FixStoreForce : public Fix {

 public:

  FixStoreForce(class LAMMPS *, int, char **);

  ~FixStoreForce();

  int setmask();

  void init();

  void setup(int);

  void min_setup(int);

  void post_force(int);

  void post_force_respa(int, int, int);

  void min_post_force(int);

  double memory_usage();

 private:

  int nlevels_respa;

  int nmax;

  double **foriginal;         // stored force on atoms

};

}

#endif