/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   https://www.lammps.org/, 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 dir1ectory.
------------------------------------------------------------------------- */

#include "compute_pressure_cartesian.h"

#include "atom.h"
#include "citeme.h"
#include "domain.h"
#include "error.h"
#include "force.h"
#include "math_const.h"
#include "memory.h"
#include "modify.h"
#include "neigh_list.h"
#include "neigh_request.h"
#include "neighbor.h"
#include "pair.h"

#include <cmath>
#include <cstring>
#include <mpi.h>

using namespace LAMMPS_NS;
using namespace MathConst;

/*-----------------------------------------------------------------------------------
  Contributing author: Olav Galteland (Norwegian University of Science and Technology)
                        olav.galteland@ntnu.no
------------------------------------------------------------------------------------*/

static const char cite_compute_pressure_cartesian[] =
    "compute pressure/cartesian:\n\n"
    "@article{ikeshoji2003molecular,\n"
    "title={Molecular-level calculation scheme for pressure in inhomogeneous systems of flat and "
    "spherical layers},\n"
    "author={Ikeshoji, Tamio and Hafskjold, Bj{\\o}rn and Furuholt, Hilde},\n"
    "journal={Molecular Simulation},\n"
    "volume={29},\n"
    "number={2},\n"
    "pages={101--109},\n"
    "year={2003},\n"
    "publisher={Taylor & Francis}\n"
    "}\n\n";

/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */

ComputePressureCartesian::ComputePressureCartesian(LAMMPS *lmp, int narg, char **arg) :
    Compute(lmp, narg, arg), dens(NULL), pkxx(NULL), pkyy(NULL), pkzz(NULL), pcxx(NULL), pcyy(NULL),
    pczz(NULL), tdens(NULL), tpkxx(NULL), tpkyy(NULL), tpkzz(NULL), tpcxx(NULL), tpcyy(NULL),
    tpczz(NULL)
{

  if (lmp->citeme) lmp->citeme->add(cite_compute_pressure_cartesian);

  // narg == 5 for one-dimensional and narg == 7 for two-dimensional
  if (narg == 5)
    dims = 1;
  else if (narg == 7)
    dims = 2;
  else
    error->all(FLERR, "Illegal compute pressure/cartesian command. Illegal number of arguments.");

  if (strcmp(arg[3], "x") == 0)
    dir1 = 0;
  else if (strcmp(arg[3], "y") == 0)
    dir1 = 1;
  else if (strcmp(arg[3], "z") == 0)
    dir1 = 2;
  else
    error->all(FLERR, "Illegal compute pressure/cartesian command.");

  dir2 = 0;
  bin_width1 = utils::numeric(FLERR, arg[4], false, lmp);
  bin_width2 = 0.0;
  nbins1 = (int) ((domain->boxhi[dir1] - domain->boxlo[dir1]) / bin_width1);
  nbins2 = 1;
  invV = bin_width1;

  if (bin_width1 <= 0.0)
    error->all(FLERR, "Illegal compute pressure/cartesian command. Bin width must be > 0");
  else if (bin_width1 > domain->boxhi[dir1] - domain->boxlo[dir1])
    error->all(
        FLERR,
        "Illegal compute pressure/cartesian command. Bin width must be smaller than box size.");

  if (dims == 2) {
    if (strcmp(arg[5], "x") == 0)
      dir2 = 0;
    else if (strcmp(arg[5], "y") == 0)
      dir2 = 1;
    else if (strcmp(arg[5], "z") == 0)
      dir2 = 2;
    else
      error->all(FLERR, "Illegal compute pressure/cartesian command.");

    bin_width2 = utils::numeric(FLERR, arg[6], false, lmp);
    nbins2 = (int) ((domain->boxhi[dir2] - domain->boxlo[dir2]) / bin_width2);
    invV *= bin_width2;

    if (bin_width2 <= 0.0)
      error->all(FLERR, "Illegal compute pressure/cartesian command. Bin width must be > 0");
    else if (bin_width2 > domain->boxhi[dir2] - domain->boxlo[dir2])
      error->all(
          FLERR,
          "Illegal compute pressure/cartesian command. Bin width must be smaller than box size.");
  }

  for (int i = 0; i < 3; i++)
    if ((dims == 1 && i != dir1) || (dims == 2 && (i != dir1 && i != dir2)))
      invV *= domain->boxhi[i] - domain->boxlo[i];
  invV = 1.0 / invV;

  array_flag = 1;
  vector_flag = 0;
  extarray = 0;
  size_array_cols = 7 + dims;    // dir1, dir2, number density, pkxx, pkyy, pkzz, pcxx, pcyy, pczz
  size_array_rows = nbins1 * nbins2;

  memory->create(dens, nbins1 * nbins2, "dens");
  memory->create(pkxx, nbins1 * nbins2, "pkxx");
  memory->create(pkyy, nbins1 * nbins2, "pkyy");
  memory->create(pkzz, nbins1 * nbins2, "pkzz");
  memory->create(pcxx, nbins1 * nbins2, "pcxx");
  memory->create(pcyy, nbins1 * nbins2, "pcyy");
  memory->create(pczz, nbins1 * nbins2, "pczz");
  memory->create(tdens, nbins1 * nbins2, "tdens");
  memory->create(tpkxx, nbins1 * nbins2, "tpkxx");
  memory->create(tpkyy, nbins1 * nbins2, "tpkyy");
  memory->create(tpkzz, nbins1 * nbins2, "tpkzz");
  memory->create(tpcxx, nbins1 * nbins2, "tpcxx");
  memory->create(tpcyy, nbins1 * nbins2, "tpcyy");
  memory->create(tpczz, nbins1 * nbins2, "tpczz");
  memory->create(array, size_array_rows, size_array_cols, "pressure:cartesian:output");
}

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

ComputePressureCartesian::~ComputePressureCartesian()
{
  memory->destroy(dens);
  memory->destroy(pkxx);
  memory->destroy(pkyy);
  memory->destroy(pkzz);
  memory->destroy(pcxx);
  memory->destroy(pcyy);
  memory->destroy(pczz);
  memory->destroy(tdens);
  memory->destroy(tpkxx);
  memory->destroy(tpkyy);
  memory->destroy(tpkzz);
  memory->destroy(tpcxx);
  memory->destroy(tpcyy);
  memory->destroy(tpczz);
  memory->destroy(array);
}

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

void ComputePressureCartesian::init()
{
  if (force->pair == NULL)
    error->all(FLERR, "No pair style is defined for compute pressure/cartesian");
  if (force->pair->single_enable == 0)
    error->all(FLERR, "Pair style does not support compute pressure/cartesian");

  // need an occasional half neighbor list.
  int irequest = neighbor->request(this, instance_me);
  neighbor->requests[irequest]->pair = 0;
  neighbor->requests[irequest]->compute = 1;
  neighbor->requests[irequest]->occasional = 1;
}

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

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

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

/* ----------------------------------------------------------------------
   count pairs and compute pair info on this proc
   only count pair once if newton_pair is off
   both atom I,J must be in group
   if flag is set, compute requested info about pair
------------------------------------------------------------------------- */

void ComputePressureCartesian::compute_array()
{
  //invoked_array = update->ntimestep;
  int me;
  MPI_Comm_rank(world, &me);

  int i, j, ii, jj, inum, jnum, itype, jtype;
  int bin, bin1, bin2, bin3;
  tagint itag, jtag;
  double xtmp, ytmp, ztmp, delx, dely, delz;
  double rsq, fpair, factor_coul, factor_lj;
  int *ilist, *jlist, *numneigh, **firstneigh;

  double **x = atom->x;
  double **v = atom->v;
  double *mass = atom->mass;
  tagint *tag = atom->tag;
  int *type = atom->type;
  int *mask = atom->mask;
  int nlocal = atom->nlocal;
  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);

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

  // Zero arrays
  for (bin = 0; bin < nbins1 * nbins2; bin++) {
    tdens[bin] = 0;
    tpkxx[bin] = 0;
    tpkyy[bin] = 0;
    tpkzz[bin] = 0;
    tpcxx[bin] = 0;
    tpcyy[bin] = 0;
    tpczz[bin] = 0;
  }

  // calculate number density and kinetic contribution to pressure
  for (i = 0; i < nlocal; i++) {
    bin1 = (int) (x[i][dir1] / bin_width1) % nbins1;
    bin2 = 0;
    if (dims == 2) bin2 = (int) (x[i][dir2] / bin_width2) % nbins2;

    j = bin1 + bin2 * nbins1;
    tdens[j] += 1;
    tpkxx[j] += mass[type[i]] * v[i][0] * v[i][0];
    tpkyy[j] += mass[type[i]] * v[i][1] * v[i][1];
    tpkzz[j] += mass[type[i]] * v[i][2] * v[i][2];
  }

  // loop over neighbors of my atoms
  // skip if I or J are not in group
  // for newton = 0 and J = ghost atom,
  //   need to insure I,J pair is only output by one proc
  //   use same itag,jtag logic as in Neighbor::neigh_half_nsq()
  // for flag = 0, just count pair interactions within force cutoff
  // for flag = 1, calculate requested output fields

  Pair *pair = force->pair;
  double **cutsq = force->pair->cutsq;

  double risq, rjsq;
  double xi1, xi2, xj1, xj2;

  for (ii = 0; ii < inum; ii++) {
    i = ilist[ii];
    if (!(mask[i] & groupbit)) continue;

    xtmp = x[i][0];
    ytmp = x[i][1];
    ztmp = x[i][2];
    xi1 = x[i][dir1];
    xi2 = x[i][dir2];
    itag = tag[i];
    itype = type[i];
    jlist = firstneigh[i];
    jnum = numneigh[i];

    for (jj = 0; jj < jnum; jj++) {
      j = jlist[jj];
      factor_lj = special_lj[sbmask(j)];
      factor_coul = special_coul[sbmask(j)];
      j &= NEIGHMASK;
      if (!(mask[j] & groupbit)) continue;

      // itag = jtag is possible for long cutoffs that include images of self
      // do calculation only on appropriate processor
      if (newton_pair == 0 && j >= nlocal) {
        jtag = tag[j];
        if (itag > jtag) {
          if ((itag + jtag) % 2 == 0) continue;
        } else if (itag < jtag) {
          if ((itag + jtag) % 2 == 1) continue;
        } else {
          if (x[j][2] < ztmp) continue;
          if (x[j][2] == ztmp) {
            if (x[j][1] < ytmp) continue;
            if (x[j][1] == ytmp && x[j][0] < xtmp) continue;
          }
        }
      }
      xj1 = x[j][dir1];
      xj2 = x[j][dir2];

      delx = x[j][0] - xtmp;
      dely = x[j][1] - ytmp;
      delz = x[j][2] - ztmp;

      rsq = delx * delx + dely * dely + delz * delz;
      jtype = type[j];

      // Check if inside cut-off
      if (rsq >= cutsq[itype][jtype]) continue;
      pair->single(i, j, itype, jtype, rsq, factor_coul, factor_lj, fpair);
      if (dims == 1) compute_pressure_1d(fpair, xi1, xj1, delx, dely, delz);
      if (dims == 2) compute_pressure_2d(fpair, xi1, xi2, xj1, xj2, delx, dely, delz);
    }
  }

  // normalize pressure
  for (bin = 0; bin < nbins1 * nbins2; bin++) {
    tdens[bin] *= invV;
    tpkxx[bin] *= invV;
    tpkyy[bin] *= invV;
    tpkzz[bin] *= invV;
    tpcxx[bin] *= invV;
    tpcyy[bin] *= invV;
    tpczz[bin] *= invV;
  }

  // communicate across processors
  MPI_Allreduce(tdens, dens, nbins1 * nbins2, MPI_DOUBLE, MPI_SUM, world);
  MPI_Allreduce(tpkxx, pkxx, nbins1 * nbins2, MPI_DOUBLE, MPI_SUM, world);
  MPI_Allreduce(tpkyy, pkyy, nbins1 * nbins2, MPI_DOUBLE, MPI_SUM, world);
  MPI_Allreduce(tpkzz, pkzz, nbins1 * nbins2, MPI_DOUBLE, MPI_SUM, world);
  MPI_Allreduce(tpcxx, pcxx, nbins1 * nbins2, MPI_DOUBLE, MPI_SUM, world);
  MPI_Allreduce(tpcyy, pcyy, nbins1 * nbins2, MPI_DOUBLE, MPI_SUM, world);
  MPI_Allreduce(tpczz, pczz, nbins1 * nbins2, MPI_DOUBLE, MPI_SUM, world);

  // populate array to output.
  for (bin = 0; bin < nbins1 * nbins2; bin++) {
    array[bin][0] = (bin % nbins1 + 0.5) * bin_width1;
    if (dims == 2) array[bin][1] = ((int) (bin / nbins1) + 0.5) * bin_width2;
    array[bin][0 + dims] = dens[bin];
    array[bin][1 + dims] = pkxx[bin];
    array[bin][2 + dims] = pkyy[bin];
    array[bin][3 + dims] = pkzz[bin];
    array[bin][4 + dims] = pcxx[bin];
    array[bin][5 + dims] = pcyy[bin];
    array[bin][6 + dims] = pczz[bin];
  }
}

void ComputePressureCartesian::compute_pressure_1d(double fpair, double xi, double xj, double delx,
                                                   double dely, double delz)
{
  int bin_s, bin_e, bin_step, bin, bin_limit;
  double xa, xb, l_sum;

  if (xi < domain->boxlo[dir1])
    xi += (domain->boxhi[dir1] - domain->boxlo[dir1]);
  else if (xi > domain->boxhi[dir1])
    xi -= (domain->boxhi[dir1] - domain->boxlo[dir1]);
  if (xj < domain->boxlo[dir1])
    xj += (domain->boxhi[dir1] - domain->boxlo[dir1]);
  else if (xj > domain->boxhi[dir1])
    xj -= (domain->boxhi[dir1] - domain->boxlo[dir1]);

  // Integrating contour from bin_s to bin_e
  bin_s = (int) lround(((xi - domain->boxlo[dir1]) / bin_width1)) % nbins1;
  bin_e = (int) lround(((xj - domain->boxlo[dir1]) / bin_width1)) % nbins1;

  // If not periodic in dir1
  if (domain->periodicity[dir1] == 0) {
    bin_s = (int) lround(((xi - domain->boxlo[dir1]) / bin_width1));
    bin_e = (int) lround(((xj - domain->boxlo[dir1]) / bin_width1));

    if (bin_e == nbins1) bin_e--;
    if (bin_s == nbins1) bin_s--;
  }

  bin_step = 1;
  if (domain->periodicity[dir1] == 1) {
    if (bin_e - bin_s > 0.5 * nbins1)
      bin_step = -1;
    else if (bin_s - bin_e > 0.5 * nbins1)
      bin_step = 1;
    else if (bin_s > bin_e)
      bin_step = -1;
  } else {
    if (bin_s > bin_e) bin_step = -1;
  }
  if (domain->periodicity[dir1] == 1)
    bin_limit = (bin_e + bin_step) % nbins1 < 0 ? (bin_e + bin_step) % nbins1 + nbins1
                                                : (bin_e + bin_step) % nbins1;
  else
    bin_limit = bin_e + bin_step;

  bin = bin_s;
  // Integrate from bin_s to bin_e with step bin_step.
  while (bin != bin_limit) {

    // Calculating exit and entry point (xa, xb). Checking if inside current bin.
    if (bin == bin_s) {
      if (domain->periodicity[dir1] == 1)
        xa = fmod(xi, domain->boxhi[dir1]) + domain->boxlo[dir1];
      else
        xa = xi;
    } else
      xa = (bin_step == 1) ? bin * bin_width1 : (bin + 1) * bin_width1;
    if (bin == bin_e) {
      if (domain->periodicity[dir1] == 1)
        xb = fmod(xj, domain->boxhi[dir1]) + domain->boxlo[dir1];
      else
        xb = xj;
    } else
      xb = (bin_step == 1) ? (bin + 1) * bin_width1 : bin * bin_width1;

    if (bin < 0 || bin >= nbins1) error->all(FLERR, "ERROR: Bin outside simulation.");

    if (bin_s != bin_e) {
      if (dir1 == 0) {
        tpcxx[bin] += (fpair * delx * delx) * (xb - xa) / delx;
        tpcyy[bin] += (fpair * dely * dely) * (xb - xa) / delx;
        tpczz[bin] += (fpair * delz * delz) * (xb - xa) / delx;
      } else if (dir1 == 1) {
        tpcxx[bin] += (fpair * delx * delx) * (xb - xa) / dely;
        tpcyy[bin] += (fpair * dely * dely) * (xb - xa) / dely;
        tpczz[bin] += (fpair * delz * delz) * (xb - xa) / dely;
      } else if (dir1 == 2) {
        tpcxx[bin] += (fpair * delx * delx) * (xb - xa) / delz;
        tpcyy[bin] += (fpair * dely * dely) * (xb - xa) / delz;
        tpczz[bin] += (fpair * delz * delz) * (xb - xa) / delz;
      }
    }
    // Particle i and j in same bin. Avoiding zero divided by zero.
    else {
      tpcxx[bin] += fpair * delx * delx;
      tpcyy[bin] += fpair * dely * dely;
      tpczz[bin] += fpair * delz * delz;
    }

    // Stepping bin to next bin
    if (domain->periodicity[dir1] == 1)
      bin = (bin + bin_step) % nbins1 < 0 ? (bin + bin_step) % nbins1 + nbins1
                                          : (bin + bin_step) % nbins1;
    else
      bin = bin + bin_step;
  }
}

void ComputePressureCartesian::compute_pressure_2d(double fpair, double xi, double yi, double xj,
                                                   double yj, double delx, double dely, double delz)
{
  int bin1, bin2, next_bin1, next_bin2;
  double la = 0.0, lb = 0.0, l_sum = 0.0;
  double rij[3] = {delx, dely, delz};
  double l1 = 0.0, l2, rij1, rij2;
  double SMALL = 10e-5;
  rij1 = rij[dir1];
  rij2 = rij[dir2];

  next_bin1 = (int) floor(xi / bin_width1);
  next_bin2 = (int) floor(yi / bin_width2);

  // Integrating along line
  while (lb != 1.0) {
    bin1 = next_bin1;
    bin2 = next_bin2;

    if (rij1 > 0)
      l1 = ((bin1 + 1) * bin_width1 - xi) / rij1;
    else
      l1 = (bin1 * bin_width1 - xi) / rij1;
    if (rij2 > 0)
      l2 = ((bin2 + 1) * bin_width2 - yi) / rij2;
    else
      l2 = (bin2 * bin_width2 - yi) / rij2;

    if ((l1 < l2 || l2 < lb + SMALL) && l1 <= 1.0 && l1 > lb) {
      lb = l1;
      next_bin1 = bin1 + (int) (rij1 / fabs(rij1));
    } else if (l2 <= 1.0 && l2 > lb) {
      lb = l2;
      next_bin2 = bin2 + (int) (rij2 / fabs(rij2));
    } else
      lb = 1.0;

    // Periodic boundary conditions
    if (domain->periodicity[dir1] == 1) {
      if (bin1 < 0)
        bin1 += nbins1;
      else if (bin1 >= nbins1)
        bin1 -= nbins1;
    } else if (bin1 < 0)
      bin1 = 0;
    else if (bin1 >= nbins1)
      bin1 = nbins1 - 1;

    if (domain->periodicity[dir2] == 1) {
      if (bin2 < 0)
        bin2 += nbins2;
      else if (bin2 >= nbins2)
        bin2 -= nbins2;
    } else if (bin2 < 0)
      bin2 = 0;
    else if (bin2 >= nbins2)
      bin2 = nbins2 - 1;

    if (bin1 + bin2 * nbins1 >= nbins1 * nbins2) error->all(FLERR, "Bin outside");

    tpcxx[bin1 + bin2 * nbins1] += (fpair * delx * delx * (lb - la));
    tpcyy[bin1 + bin2 * nbins1] += (fpair * dely * dely * (lb - la));
    tpczz[bin1 + bin2 * nbins1] += (fpair * delz * delz * (lb - la));

    l_sum += lb - la;
    la = lb;
  }

  if (l_sum > 1.0 + SMALL || l_sum < 1.0 - SMALL)
    error->all(FLERR, "Sum of fractional line segments does not equal 1.");
}

/* ----------------------------------------------------------------------
memory usage of data
------------------------------------------------------------------------- */

double ComputePressureCartesian::memory_usage()
{
  return (14.0 + dims + 7) * (double) (nbins1 * nbins2) * sizeof(double);
}