// clang-format off
/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   https://www.lammps.org/, Sandia National Laboratories
   LAMMPS development team: developers@lammps.org

   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.
------------------------------------------------------------------------- */

/* ----------------------------------------------------------------------
   Contributing authors: Timothy Sirk (ARL), Pieter in't Veld (BASF)

This pair style srp command calculates a segmental repulsive force
between bonds. This is useful for preventing the crossing of bonds if
soft non-bonded potentials are used, such as DPD polymer chains.

See the doc page for pair_style srp command for usage instructions.

There is an example script for this package in examples/PACKAGES/srp.

Please contact Timothy Sirk for questions (tim.sirk@us.army.mil).
------------------------------------------------------------------------- */

#include "pair_srp.h"

#include "atom.h"
#include "citeme.h"
#include "comm.h"
#include "domain.h"
#include "error.h"
#include "fix_srp.h"
#include "force.h"
#include "memory.h"
#include "modify.h"
#include "neigh_list.h"
#include "neighbor.h"
#include "output.h"
#include "thermo.h"

#include <cmath>
#include <cstring>

using namespace LAMMPS_NS;

static constexpr double SMALL = 1.0e-10;
static constexpr double BIG = 1e10;
#define ONETWOBIT 0x40000000

static const char cite_srp[] =
  "pair srp command: doi:10.1063/1.3698476\n\n"
  "@Article{Sirk2012\n"
  " author = {T. W. Sirk and Y. R. Sliozberg and J. K. Brennan and M. Lisal and J. W. Andzelm},\n"
  " title = {An Enhanced Entangled Polymer Model for Dissipative Particle Dynamics},\n"
  " journal = {J.~Chem.\\ Phys.},\n"
  " year =    2012,\n"
  " volume =  136,\n"
  " pages =   {134903}\n"
  "}\n\n";

static int srp_instance = 0;

/* ----------------------------------------------------------------------
 set size of pair comms in constructor
 ---------------------------------------------------------------------- */

PairSRP::PairSRP(LAMMPS *lmp) : Pair(lmp), fix_id(nullptr)
{
  writedata = 1;
  single_enable = 0;

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

  nextra = 1;
  segment = nullptr;

  // create fix SRP instance here with unique fix id
  // similar to granular pair styles with history,
  //   this should be early enough that FixSRP::pre_exchange()
  //   will be invoked before other fixes that migrate atoms
  //   this is checked for in FixSRP

  f_srp = dynamic_cast<FixSRP *>(modify->add_fix(fmt::format("{:02d}_FIX_SRP all SRP", srp_instance)));
  ++srp_instance;
}

/* ----------------------------------------------------------------------
 allocate all arrays
 ------------------------------------------------------------------------- */

void PairSRP::allocate()
{
    allocated = 1;
    // particles of bptype inserted by fix srp
    // bptype is the highest numbered atom type
    int n = bptype;
    memory->create(cutsq, n + 1, n + 1, "pair:cutsq");
    memory->create(cut, n + 1, n + 1, "pair:cut");
    memory->create(a0, n + 1, n + 1, "pair:a0");

    // setflag for atom types
    memory->create(setflag,n+1,n+1,"pair:setflag");
    for (int i = 1; i <= n; i++)
        for (int j = i; j <= n; j++)
            setflag[i][j] = 0;

    maxcount = 0;
}

/* ----------------------------------------------------------------------
 free
 ------------------------------------------------------------------------- */

PairSRP::~PairSRP()
{
  if (allocated) {
    memory->destroy(setflag);
    memory->destroy(cutsq);
    memory->destroy(cut);
    memory->destroy(a0);
    memory->destroy(segment);
  }

  // check nfix in case all fixes have already been deleted
  if (modify->nfix && modify->get_fix_by_id(f_srp->id)!=nullptr) modify->delete_fix(f_srp->id);
}

/* ----------------------------------------------------------------------
 compute bond-bond repulsions
 ------------------------------------------------------------------------- */

void PairSRP::compute(int eflag, int vflag)
{
    // setup energy and virial
    ev_init(eflag, vflag);

    double **x = atom->x;
    double **f = atom->f;
    int nlocal = atom->nlocal;
    int nall = nlocal + atom->nghost;
    int i0, i1, j0, j1;
    int i,j,ii,jj,inum,jnum;
    double dijsq, dij;

    int *ilist,*jlist,*numneigh,**firstneigh;
    inum = list->inum;
    ilist = list->ilist;
    numneigh = list->numneigh;
    firstneigh = list->firstneigh;

    double dx,dy,dz,ti,tj;
    double wd, lever0, lever1, evdwl, fpair;
    double fxlever0, fylever0, fzlever0, fxlever1, fylever1, fzlever1;
    double fx, fy, fz;
    evdwl = 0.0;

    // mapping global to local for atoms inside bond particles
    // exclude 1-2 neighs if requested
    if (neighbor->ago == 0) {
      remapBonds(nall);
      if (exclude) onetwoexclude(ilist, inum, jlist, numneigh, firstneigh);
    }

  // this pair style only used with hybrid
  // due to exclusions
  // each atom i is type bptype
  // each neigh j is type bptype

  // using midpoint distance option
  if (midpoint) {

    for (ii = 0; ii < inum; ii++) {

      i = ilist[ii];
      jnum = numneigh[i];
      // two atoms inside bond particle
      i0 = segment[i][0];
      j0 = segment[i][1];

      for (jj = 0; jj < jnum; jj++) {

        jlist = firstneigh[i];
        j = jlist[jj];

        // enforce 1-2 exclusions
        if ((sbmask(j) & exclude))
          continue;

        j &= NEIGHMASK;
        //retrieve atoms from bond particle
        i1 = segment[j][0];
        j1 = segment[j][1];

        // midpt dist bond 0 and 1
        dx = 0.5*(x[i0][0] - x[i1][0] + x[j0][0] - x[j1][0]);
        dy = 0.5*(x[i0][1] - x[i1][1] + x[j0][1] - x[j1][1]);
        dz = 0.5*(x[i0][2] - x[i1][2] + x[j0][2] - x[j1][2]);
        dijsq = dx*dx + dy*dy + dz*dz;

        if (dijsq < cutsq[bptype][bptype]) {
        dij = sqrt(dijsq);

        if (dij < SMALL)
          continue;     // dij can be 0.0 with soft potentials

        wd = 1.0 - dij / cut[bptype][bptype];
        fpair = 0.5 * a0[bptype][bptype] * wd / dij; // 0.5 factor for lever rule

        // force for bond 0, beads 0,1
        //force between bonds
        fx = fpair * dx;
        fy = fpair * dy;
        fz = fpair * dz;

        f[i0][0] += fx; //keep force sign for bond 0
        f[i0][1] += fy;
        f[i0][2] += fz;

        f[j0][0] += fx;
        f[j0][1] += fy;
        f[j0][2] += fz;

        f[i1][0] -= fx; //flip force sign for bond 1
        f[i1][1] -= fy;
        f[i1][2] -= fz;

        f[j1][0] -= fx;
        f[j1][1] -= fy;
        f[j1][2] -= fz;

        // ************************************************* //

        if (eflag) {
          evdwl = 0.5 * a0[bptype][bptype] * cut[bptype][bptype] * wd * wd;
        }

        if (evflag) {
          ev_tally(i0,i1,nlocal,1,0.5*evdwl,0.0,fpair,dx,dy,dz);
          ev_tally(j0,j1,nlocal,1,0.5*evdwl,0.0,fpair,dx,dy,dz);
        }

        if (vflag_fdotr) virial_fdotr_compute();

        }
      }
   }
 } else {
  // using min distance option

    for (ii = 0; ii < inum; ii++) {

      i = ilist[ii];
      jnum = numneigh[i];
      i0 = segment[i][0];
      j0 = segment[i][1];

      for (jj = 0; jj < jnum; jj++) {

        jlist = firstneigh[i];
        j = jlist[jj];

        // enforce 1-2 exclusions
        if ((sbmask(j) & exclude))
          continue;

        j &= NEIGHMASK;

        i1 = segment[j][0];
        j1 = segment[j][1];

        getMinDist(x, dx, dy, dz, ti, tj, i0, j0, i1, j1);
        dijsq = dx*dx + dy*dy + dz*dz;

        if (dijsq < cutsq[bptype][bptype]) {

        dij = sqrt(dijsq);

        if (dij < SMALL)
      continue;     // dij can be 0.0 with soft potentials

        wd = 1.0 - dij / cut[bptype][bptype];
        fpair = a0[bptype][bptype] * wd / dij;

        // force for bond 0, beads 0,1
        lever0 = 0.5 + ti; // assign force according to lever rule
        lever1 = 0.5 + tj; // assign force according to lever rule
        //force between bonds
        fx = fpair * dx;
        fy = fpair * dy;
        fz = fpair * dz;

        //decompose onto atoms
        fxlever0 = fx * lever0;
        fylever0 = fy * lever0;
        fzlever0 = fz * lever0;
        fxlever1 = fx * lever1;
        fylever1 = fy * lever1;
        fzlever1 = fz * lever1;

        f[i0][0] += fxlever0; //keep force sign for bond 0
        f[i0][1] += fylever0;
        f[i0][2] += fzlever0;

        f[j0][0] += (fx - fxlever0);
        f[j0][1] += (fy - fylever0);
        f[j0][2] += (fz - fzlever0);

        f[i1][0] -= fxlever1; //flip force sign for bond 1
        f[i1][1] -= fylever1;
        f[i1][2] -= fzlever1;

        f[j1][0] -= (fx - fxlever1);
        f[j1][1] -= (fy - fylever1);
        f[j1][2] -= (fz - fzlever1);

        // ************************************************* //

        if (eflag) {
          evdwl = 0.5 * a0[bptype][bptype] * cut[bptype][bptype] * wd * wd;
        }

        if (evflag) {
          ev_tally(i0,i1,nlocal,1,0.5*evdwl,0.0,0.5*fpair,dx,dy,dz);
          ev_tally(j0,j1,nlocal,1,0.5*evdwl,0.0,0.5*fpair,dx,dy,dz);
        }

       if (vflag_fdotr) virial_fdotr_compute();

      }
    }
  }
 }
}

/* ----------------------------------------------------------------------
 global settings
 ------------------------------------------------------------------------- */

void PairSRP::settings(int narg, char **arg)
{
  if (narg < 3 || narg > 7)
    error->all(FLERR,"Illegal pair_style command");

  if (atom->tag_enable == 0)
    error->all(FLERR,"Pair_style srp requires atom IDs");

  cut_global = utils::numeric(FLERR, arg[0], false, lmp);
  // wildcard
  if (strcmp(arg[1],"*") == 0) {
    btype = 0;
  } else {
    btype_str = arg[1];
  }

  // settings
  midpoint = false;
  min = false;

  if (strcmp(arg[2],"min") == 0) min = true;
  else if (strcmp(arg[2],"mid") == 0) midpoint = true;
  else
    error->all(FLERR,"Illegal pair_style command");

  int iarg = 3;
  // default exclude 1-2
  // scaling for 1-2, etc not supported
  exclude = 1;

  // use last atom type by default for bond particles
  bptype = atom->ntypes;

  while (iarg < narg) {
    if (strcmp(arg[iarg],"exclude") == 0) {
      if (iarg+2 > narg) error->all(FLERR,"Illegal pair srp command");
      exclude = utils::logical(FLERR, arg[iarg+1], false, lmp);
      if (min && !exclude) error->all(FLERR,"Illegal exclude option in pair srp command");
      iarg += 2;
    } else if (strcmp(arg[iarg],"bptype") == 0) {
      if (iarg+2 > narg) error->all(FLERR,"Illegal pair srp command");
      bptype_str = arg[iarg+1];
      iarg += 2;
    } else error->all(FLERR,"Illegal pair srp command");
  }
}

/* ----------------------------------------------------------------------
 set coeffs
 ------------------------------------------------------------------------- */

void PairSRP::coeff(int narg, char **arg)
{
  if (narg < 3 || narg > 4)
    error->all(FLERR,"PairSRP: Incorrect args for pair coeff");
  if (!allocated) allocate();

  if (btype_str.size() > 0)
    btype = utils::expand_type_int(FLERR, btype_str, Atom::BOND, lmp);
  if ((btype > atom->nbondtypes) || (btype <= 0))
    error->all(FLERR,"Invalid bond type {} for pair style srp", btype);

  if (bptype_str.size() > 0)
    bptype = utils::expand_type_int(FLERR, bptype_str, Atom::ATOM, lmp);
  if ((bptype < 1) || (bptype > atom->ntypes))
    error->all(FLERR,"Invalid bond particle type {} for pair style srp", bptype);

  // reset cutoffs if explicitly set
  if (allocated) {
    int i,j;
    for (i = 1; i <= bptype; i++)
      for (j = i; j <= bptype; j++)
        if (setflag[i][j]) cut[i][j] = cut_global;
  }

  // set ij bond-bond cutoffs
  int ilo, ihi, jlo, jhi;
  utils::bounds(FLERR,arg[0], 1, bptype, ilo, ihi, error);
  utils::bounds(FLERR,arg[1], 1, bptype, jlo, jhi, error);

  double a0_one = utils::numeric(FLERR,arg[2],false,lmp);
  double cut_one = cut_global;
  if (narg == 4)  cut_one = utils::numeric(FLERR,arg[3],false,lmp);

  int count = 0;
  for (int i = ilo; i <= ihi; i++) {
    for (int j = MAX(jlo,i); j <= jhi; j++) {
      a0[i][j] = a0_one;
      cut[i][j] = cut_one;
      cutsq[i][j] = cut_one * cut_one;
      setflag[i][j] = 1;
      count++;
    }
  }

  if (count == 0) error->warning(FLERR,"PairSRP: No pair coefficients were set");
}

/* ----------------------------------------------------------------------
 init specific to this pair style
 ------------------------------------------------------------------------- */

void PairSRP::init_style()
{
  if (!force->newton_pair)
    error->all(FLERR,"Pair srp requires newton pair on");

  // verify that fix SRP is still defined and has not been changed.

  if (strcmp(f_srp->style,"SRP") != 0)
    error->all(FLERR,"Fix SRP has been changed unexpectedly");

  if (comm->me == 0)
    utils::logmesg(lmp,"Using type {} for bond particles\n",bptype);

  // set bond and bond particle types in fix srp
  // bonds of this type will be represented by bond particles
  // if bond type is 0, then all bonds have bond particles
  // btype = bond type

  char c0[20];
  char* arg0[2];
  sprintf(c0, "%d", btype);
  arg0[0] = (char *) "btype";
  arg0[1] = c0;
  f_srp->modify_params(2, arg0);

  // bptype = bond particle type
  sprintf(c0, "%d", bptype);
  arg0[0] = (char *) "bptype";
  arg0[1] = c0;
  f_srp->modify_params(2, arg0);

  // bond particles do not contribute to energy or virial
  // bond particles do not belong to group all
  // but thermo normalization is by nall
  // therefore should turn off normalization
  char *arg1[2];
  arg1[0] = (char *) "norm";
  arg1[1] = (char *) "no";
  output->thermo->modify_params(2, arg1);
  if (comm->me == 0) error->message(FLERR,"Thermo normalization turned off by pair srp");

  neighbor->add_request(this);
}

/* ----------------------------------------------------------------------
 init for one type pair i,j and corresponding j,i
 ------------------------------------------------------------------------- */

double PairSRP::init_one(int i, int j)
{
 if (setflag[i][j] == 0) error->all(FLERR,"PairSRP: All pair coeffs are not set");

  cut[j][i] = cut[i][j];
  a0[j][i] = a0[i][j];

  return cut[i][j];
}

/* ----------------------------------------------------------------------
 find min distance for bonds i0/j0 and i1/j1
 ------------------------------------------------------------------------- */
inline void PairSRP::getMinDist(double** &x, double &dx, double &dy, double &dz, double &ti, double &tj, int &i0, int &j0, int &i1, int &j1)
{
    // move these outside the loop
    double diffx0, diffy0, diffz0, diffx1, diffy1, diffz1, dPx, dPy, dPz, RiRi, RiRj, RjRj;
    double denom, termx0, termy0, termz0, num0, termx1, termy1, termz1, num1;

    // compute midpt dist from 1st atom, 1st bond
    diffx0 = x[j0][0] - x[i0][0]; // x,y,z from bond 0
    diffy0 = x[j0][1] - x[i0][1];
    diffz0 = x[j0][2] - x[i0][2];

    // compute midpt dist from 1st atom, 2nd bond
    diffx1 = x[j1][0] - x[i1][0];
    diffy1 = x[j1][1] - x[i1][1];
    diffz1 = x[j1][2] - x[i1][2];

    // midpoint distance
    dPx = 0.5*(diffx0-diffx1) + x[i0][0]-x[i1][0];
    dPy = 0.5*(diffy0-diffy1) + x[i0][1]-x[i1][1];
    dPz = 0.5*(diffz0-diffz1) + x[i0][2]-x[i1][2];

    // Ri^2 Rj^2
    RiRi = diffx0*diffx0 + diffy0*diffy0 + diffz0*diffz0;
    RiRj = diffx0*diffx1 + diffy0*diffy1 + diffz0*diffz1;
    RjRj = diffx1*diffx1 + diffy1*diffy1 + diffz1*diffz1;
    denom = RiRj*RiRj - RiRi*RjRj;

    // handle case of parallel lines
    // reduce to midpt distance
    if (fabs(denom) < SMALL) {
        if (denom < 0) denom = -BIG;
        else denom = BIG;
    }

    // calc ti
    termx0 = RiRj*diffx1 - RjRj*diffx0;
    termy0 = RiRj*diffy1 - RjRj*diffy0;
    termz0 = RiRj*diffz1 - RjRj*diffz0;
    num0 = dPx*termx0 + dPy*termy0 + dPz*termz0;
    ti = num0 / denom;
    if (ti > 0.5) ti = 0.5;
    if (ti < -0.5) ti = -0.5;

    // calc tj
    termx1 = RiRj*diffx0 - RiRi*diffx1;
    termy1 = RiRj*diffy0 - RiRi*diffy1;
    termz1 = RiRj*diffz0 - RiRi*diffz1;
    num1 = dPx*termx1 + dPy*termy1 + dPz*termz1;
    tj = -num1/ denom;
    if (tj > 0.5)  tj = 0.5;
    if (tj < -0.5) tj = -0.5;

    // min dist
    dx = dPx - ti*diffx0 + tj*diffx1;
    dy = dPy - ti*diffy0 + tj*diffy1;
    dz = dPz - ti*diffz0 + tj*diffz1;
}

/* --------------------------------------------------------
map global id of atoms in stored by each bond particle
 ------------------------------------------------------- */
inline void PairSRP::remapBonds(int &nall)
{
  if (nall > maxcount) {
    memory->grow(segment, nall, 2, "pair:segment");
    maxcount = nall;
  }

  // loop over all bond particles
  // each bond paricle holds two bond atoms
  // map global ids of bond atoms to local ids
  // might not be able to map both bond atoms of j, if j is outside neighcut
  // these are not on neighlist, so are not used
  int tmp;
  srp = f_srp->array_atom;

    for (int i = 0; i < nall; i++) {
      if (atom->type[i] == bptype) {
        // tmp is local id
        // tmp == -1 is ok
        tmp = atom->map((int)srp[i][0]);
        segment[i][0] = domain->closest_image(i,tmp);
        // repeat with other id
        tmp = atom->map((int)srp[i][1]);
        segment[i][1] = domain->closest_image(i,tmp);
      }
    }
}

/* --------------------------------------------------------
add exclusions for 1-2 neighs, if requested
more complex exclusions or scaling probably not needed
 ------------------------------------------------------- */
inline void PairSRP::onetwoexclude(int* &ilist, int &inum, int* &jlist, int* &numneigh, int** &firstneigh)
{
    int i0, i1, j0, j1;
    int i,j,ii,jj,jnum;

    // encode neighs with exclusions
    // only need 1-2 info for normal uses of srp
    // add 1-3, etc later if ever needed

    for (ii = 0; ii < inum; ii++) {

      i = ilist[ii];
      jnum = numneigh[i];
      // two atoms inside bond particle
      i0 = segment[i][0];
      j0 = segment[i][1];

      for (jj = 0; jj < jnum; jj++) {

        jlist = firstneigh[i];
        j = jlist[jj];
        j &= NEIGHMASK;
        //two atoms inside bond particle
        i1 = segment[j][0];
        j1 = segment[j][1];

        // check for a 1-2 neigh
        if (i0 == i1 || i0 == j1 || i1 == j0 || j0 == j1) {
          j |= ONETWOBIT;
          jlist[jj] = j;
        }
      }
    }
}

/* ----------------------------------------------------------------------
proc 0 writes to data file
------------------------------------------------------------------------- */

void PairSRP::write_data(FILE *fp)
{
  for (int i = 1; i <= atom->ntypes; i++)
    fprintf(fp,"%d %g\n",i,a0[i][i]);
}

/* ----------------------------------------------------------------------
proc 0 writes all pairs to data file
------------------------------------------------------------------------- */

void PairSRP::write_data_all(FILE *fp)
{
  for (int i = 1; i <= atom->ntypes; i++)
    for (int j = i; j <= atom->ntypes; j++)
      fprintf(fp,"%d %d %g %g\n",i,j,a0[i][j],cut[i][j]);
}

/* ----------------------------------------------------------------------
   proc 0 writes to restart file
------------------------------------------------------------------------- */

void PairSRP::write_restart(FILE *fp)
{
  write_restart_settings(fp);

  int i,j;
  for (i = 1; i <= atom->ntypes; i++)
    for (j = i; j <= atom->ntypes; j++) {
      fwrite(&setflag[i][j],sizeof(int),1,fp);
      if (setflag[i][j]) {
        fwrite(&a0[i][j],sizeof(double),1,fp);
        fwrite(&cut[i][j],sizeof(double),1,fp);
      }
    }
}

/* ----------------------------------------------------------------------
   proc 0 reads from restart file, bcasts
------------------------------------------------------------------------- */

void PairSRP::read_restart(FILE *fp)
{
  read_restart_settings(fp);
  allocate();

  int i,j;
  int me = comm->me;
  for (i = 1; i <= atom->ntypes; i++)
    for (j = i; j <= atom->ntypes; j++) {
      if (me == 0) utils::sfread(FLERR,&setflag[i][j],sizeof(int),1,fp,nullptr,error);
      MPI_Bcast(&setflag[i][j],1,MPI_INT,0,world);
      if (setflag[i][j]) {
        if (me == 0) {
          utils::sfread(FLERR,&a0[i][j],sizeof(double),1,fp,nullptr,error);
          utils::sfread(FLERR,&cut[i][j],sizeof(double),1,fp,nullptr,error);
        }
        MPI_Bcast(&a0[i][j],1,MPI_DOUBLE,0,world);
        MPI_Bcast(&cut[i][j],1,MPI_DOUBLE,0,world);
      }
    }
}
/* ----------------------------------------------------------------------
   proc 0 writes to restart file
------------------------------------------------------------------------- */

void PairSRP::write_restart_settings(FILE *fp)
{
  fwrite(&cut_global,sizeof(double),1,fp);
  fwrite(&bptype,sizeof(int),1,fp);
  fwrite(&btype,sizeof(int),1,fp);
  fwrite(&min,sizeof(int),1,fp);
  fwrite(&midpoint,sizeof(int),1,fp);
  fwrite(&exclude,sizeof(int),1,fp);
}

/* ----------------------------------------------------------------------
   proc 0 reads from restart file, bcasts
------------------------------------------------------------------------- */

void PairSRP::read_restart_settings(FILE *fp)
{
  if (comm->me == 0) {
    utils::sfread(FLERR,&cut_global,sizeof(double),1,fp,nullptr,error);
    utils::sfread(FLERR,&bptype,sizeof(int),1,fp,nullptr,error);
    utils::sfread(FLERR,&btype,sizeof(int),1,fp,nullptr,error);
    utils::sfread(FLERR,&min,sizeof(int),1,fp,nullptr,error);
    utils::sfread(FLERR,&midpoint,sizeof(int),1,fp,nullptr,error);
    utils::sfread(FLERR,&exclude,sizeof(int),1,fp,nullptr,error);
  }
  MPI_Bcast(&cut_global,1,MPI_DOUBLE,0,world);
}
