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

/* ----------------------------------------------------------------------
   Contributing authors: Mike Salerno (NRL) added single methods
------------------------------------------------------------------------- */

#include <stdlib.h>
#include <string.h>
#include "create_bonds.h"
#include "atom.h"
#include "domain.h"
#include "force.h"
#include "neighbor.h"
#include "neigh_request.h"
#include "neigh_list.h"
#include "comm.h"
#include "group.h"
#include "special.h"
#include "error.h"

using namespace LAMMPS_NS;

enum{MANY,SBOND,SANGLE,SDIHEDRAL};

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

CreateBonds::CreateBonds(LAMMPS *lmp) : Pointers(lmp) {}

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

void CreateBonds::command(int narg, char **arg)
{
  if (domain->box_exist == 0)
    error->all(FLERR,"Create_bonds command before simulation box is defined");
  if (atom->tag_enable == 0)
    error->all(FLERR,"Cannot use create_bonds unless atoms have IDs");
  if (atom->molecular != 1)
    error->all(FLERR,"Cannot use create_bonds with non-molecular system");

  if (narg < 4) error->all(FLERR,"Illegal create_bonds command");

  // parse args

  int style;

  int iarg;
  if (strcmp(arg[0],"many") == 0) {
    style = MANY;
    if (narg != 6) error->all(FLERR,"Illegal create_bonds command");
    igroup = group->find(arg[1]);
    if (igroup == -1) error->all(FLERR,"Cannot find create_bonds group ID");
    group1bit = group->bitmask[igroup];
    igroup = group->find(arg[2]);
    if (igroup == -1) error->all(FLERR,"Cannot find create_bonds group ID");
    group2bit = group->bitmask[igroup];
    btype = force->inumeric(FLERR,arg[3]);
    rmin = force->numeric(FLERR,arg[4]);
    rmax = force->numeric(FLERR,arg[5]);
    if (rmin > rmax) error->all(FLERR,"Illegal create_bonds command");
    iarg = 6;
  } else if (strcmp(arg[0],"single/bond") == 0) {
    style = SBOND;
    if (narg < 4) error->all(FLERR,"Illegal create_bonds command");
    btype = force->inumeric(FLERR,arg[1]);
    batom1 = force->tnumeric(FLERR,arg[2]);
    batom2 = force->tnumeric(FLERR,arg[3]);
    iarg = 4;
  } else if (strcmp(arg[0],"single/angle") == 0) {
    style = SANGLE;
    if (narg < 5) error->all(FLERR,"Illegal create_bonds command");
    atype = force->inumeric(FLERR,arg[1]);
    aatom1 = force->tnumeric(FLERR,arg[2]);
    aatom2 = force->tnumeric(FLERR,arg[3]);
    aatom3 = force->tnumeric(FLERR,arg[4]);
    iarg = 5;
  } else if (strcmp(arg[0],"single/dihedral") == 0) {
    style = SDIHEDRAL;
    if (narg < 6) error->all(FLERR,"Illegal create_bonds command");
    dtype = force->inumeric(FLERR,arg[1]);
    datom1 = force->tnumeric(FLERR,arg[2]);
    datom2 = force->tnumeric(FLERR,arg[3]);
    datom3 = force->tnumeric(FLERR,arg[4]);
    datom4 = force->tnumeric(FLERR,arg[5]);
    iarg = 6;
  } else error->all(FLERR,"Illegal create_bonds command");

  // optional args

  int specialflag = 1;

  while (iarg < narg) {
    if (strcmp(arg[iarg],"special") == 0) {
      if (iarg+2 > narg) error->all(FLERR,"Illegal create_bonds command");
      if (strcmp(arg[iarg+1],"yes") == 0) specialflag = 1;
      else if (strcmp(arg[iarg+1],"no") == 0) specialflag = 0;
      else error->all(FLERR,"Illegal create_bonds command");
      iarg += 2;
    } else error->all(FLERR,"Illegal create_bonds command");
  }

  // error checks

  if (style == MANY) {
    if (btype <= 0 || btype > atom->nbondtypes)
      error->all(FLERR,"Invalid bond type in create_bonds command");
    if (specialflag == 0)
      error->all(FLERR,"Cannot use special no with create_bonds many");
  } else if (style == SBOND) {
    if (btype <= 0 || btype > atom->nbondtypes)
      error->all(FLERR,"Invalid bond type in create_bonds command");
  } else if (style == SANGLE) {
    if (atype <= 0 || atype > atom->nangletypes)
      error->all(FLERR,"Invalid angle type in create_bonds command");
  } else if (style == SDIHEDRAL) {
    if (dtype <= 0 || dtype > atom->ndihedraltypes)
      error->all(FLERR,"Invalid dihedral type in create_bonds command");
  }

  // invoke creation method

  if (style == MANY) many();
  else if (style == SBOND) single_bond();
  else if (style == SANGLE) single_angle();
  else if (style == SDIHEDRAL) single_dihedral();

  // trigger special list build

  if (specialflag) {
    Special special(lmp);
    special.build();
  }
}

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

void CreateBonds::many()
{
  double rminsq = rmin*rmin;
  double rmaxsq = rmax*rmax;

  // store state before bond creation

  bigint nbonds_previous = atom->nbonds;

  // request a full neighbor list for use by this command

  int irequest = neighbor->request(this);
  neighbor->requests[irequest]->pair = 0;
  neighbor->requests[irequest]->command = 1;
  neighbor->requests[irequest]->half = 0;
  neighbor->requests[irequest]->full = 1;
  neighbor->requests[irequest]->occasional = 1;
  neighbor->requests[irequest]->command_style = "create_bonds";

  // init entire system since comm->borders and neighbor->build is done
  // comm::init needs neighbor::init needs pair::init needs kspace::init, etc

  lmp->init();

  // error check on cutoff
  // if no pair style, neighbor list will be empty

  if (force->pair == NULL)
    error->all(FLERR,"Create_bonds requires a pair style be defined");
  if (rmax > neighbor->cutneighmax)
    error->all(FLERR,"Create_bonds max distance > neighbor cutoff");
  if (rmax > neighbor->cutneighmin && comm->me == 0)
    error->warning(FLERR,"Create_bonds max distance > minimum neighbor cutoff");

  // require special_bonds 1-2 weights = 0.0 and KSpace = NULL
  // so that already bonded atom pairs do not appear in neighbor list
  // otherwise with newton_bond = 1,
  //   would be hard to check if I-J bond already existed
  // note that with KSpace, pair with weight = 0 could still be in neigh list

  if (force->special_lj[1] != 0.0 || force->special_coul[1] != 0.0)
    error->all(FLERR,"Create_bonds command requires "
               "special_bonds 1-2 weights be 0.0");
  if (force->kspace)
    error->all(FLERR,"Create_bonds command requires "
               "no kspace_style be defined");

  // setup domain, communication and neighboring
  // acquire ghosts and build standard neighbor lists

  if (domain->triclinic) domain->x2lamda(atom->nlocal);
  domain->pbc();
  domain->reset_box();
  comm->setup();
  if (neighbor->style) neighbor->setup_bins();
  comm->exchange();
  comm->borders();
  if (domain->triclinic) domain->lamda2x(atom->nlocal+atom->nghost);
  neighbor->build();

  // build neighbor list this command needs based on earlier request

  NeighList *list = neighbor->lists[irequest];
  neighbor->build_one(list);

  // loop over all neighs of each atom
  // compute distance between two atoms consistently on both procs
  // add bond if group and distance criteria are met
  // check that bond list does not overflow

  tagint *tag = atom->tag;
  int *mask = atom->mask;
  double **x = atom->x;
  int *num_bond = atom->num_bond;
  int **bond_type = atom->bond_type;
  tagint **bond_atom = atom->bond_atom;
  double newton_bond = force->newton_bond;
  int nlocal = atom->nlocal;

  int i,j,ii,jj,inum,jnum,flag;
  double xtmp,ytmp,ztmp,delx,dely,delz,rsq;
  int *ilist,*jlist,*numneigh,**firstneigh;

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

  for (ii = 0; ii < inum; ii++) {
    i = ilist[ii];
    xtmp = x[i][0];
    ytmp = x[i][1];
    ztmp = x[i][2];
    jlist = firstneigh[i];
    jnum = numneigh[i];

    for (jj = 0; jj < jnum; jj++) {
      j = jlist[jj];
      j &= NEIGHMASK;

      // only consider bond creation if I,J distance between 2 cutoffs
      // compute rsq identically on both I,J loop iterations
      // if I,J tags equal, do not bond atom to itself

      if (tag[i] < tag[j]) {
        delx = xtmp - x[j][0];
        dely = ytmp - x[j][1];
        delz = ztmp - x[j][2];
      } else if (tag[i] > tag[j]) {
        delx = x[j][0] - xtmp;
        dely = x[j][1] - ytmp;
        delz = x[j][2] - ztmp;
      } else continue;
      rsq = delx*delx + dely*dely + delz*delz;
      if (rsq < rminsq || rsq > rmaxsq) continue;

      // only consider bond creation if igroup and jgroup match I,J atoms

      flag = 0;
      if ((mask[i] & group1bit) && (mask[j] & group2bit)) flag = 1;
      if ((mask[i] & group2bit) && (mask[j] & group1bit)) flag = 1;
      if (!flag) continue;

      // create bond, check for overflow
      // on I,J loop iterations, store with 1 or 2 atoms based on newton_bond

      if (!newton_bond || tag[i] < tag[j]) {
        if (num_bond[i] == atom->bond_per_atom)
          error->one(FLERR,
                     "New bond exceeded bonds per atom in create_bonds");
        bond_type[i][num_bond[i]] = btype;
        bond_atom[i][num_bond[i]] = tag[j];
        num_bond[i]++;
      }
    }
  }

  // recount bonds

  bigint nbonds = 0;
  for (int i = 0; i < nlocal; i++) nbonds += num_bond[i];

  MPI_Allreduce(&nbonds,&atom->nbonds,1,MPI_LMP_BIGINT,MPI_SUM,world);
  if (!force->newton_bond) atom->nbonds /= 2;

  // print new bond count

  bigint nadd_bonds = atom->nbonds - nbonds_previous;

  if (comm->me == 0) {
    if (screen) {
      fprintf(screen,"Added " BIGINT_FORMAT
              " bonds, new total = " BIGINT_FORMAT "\n",
              nadd_bonds,atom->nbonds);
    }

    if (logfile) {
      fprintf(logfile,"Added " BIGINT_FORMAT
              " bonds, new total = " BIGINT_FORMAT "\n",
              nadd_bonds,atom->nbonds);
    }
  }
  // trigger clearing the list of available neighbor list requests
  // and a full rebuild of them during the next run setup.
  // otherwise the request from this command may linger around.
  neighbor->init();
}

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

void CreateBonds::single_bond()
{
  int m;

  // check that 2 atoms exist

  int count = 0;
  if (atom->map(batom1) >= 0) count++;
  if (atom->map(batom2) >= 0) count++;

  int allcount;
  MPI_Allreduce(&count,&allcount,1,MPI_INT,MPI_SUM,world);
  if (allcount != 2) 
    error->all(FLERR,"Create_bonds single/bond atoms do not exist");

  // create bond once or 2x if newton_bond set

  int *num_bond = atom->num_bond;
  int **bond_type = atom->bond_type;
  tagint **bond_atom = atom->bond_atom;

  if ((m = atom->map(batom1)) >= 0) {
    if (num_bond[m] == atom->bond_per_atom)
      error->one(FLERR,"New bond exceeded bonds per atom in create_bonds");
    bond_type[m][num_bond[m]] = btype;
    bond_atom[m][num_bond[m]] = batom2;
    num_bond[m]++;
  }
  ++atom->nbonds;

  if (force->newton_bond) return;

  if ((m = atom->map(batom2)) >= 0) {
    if (num_bond[m] == atom->bond_per_atom)
      error->one(FLERR,"New bond exceeded bonds per atom in create_bonds");
    bond_type[m][num_bond[m]] = btype;
    bond_atom[m][num_bond[m]] = batom1;
    num_bond[m]++;
  }
}

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

void CreateBonds::single_angle()
{
  int m;

  // check that 3 atoms exist

  int count = 0;
  if (atom->map(aatom1) >= 0) count++;
  if (atom->map(aatom2) >= 0) count++;
  if (atom->map(aatom3) >= 0) count++;

  int allcount;
  MPI_Allreduce(&count,&allcount,1,MPI_INT,MPI_SUM,world);
  if (allcount != 3) 
    error->all(FLERR,"Create_bonds single/angle atoms do not exist");

  // create angle once or 3x if newton_bond set

  int *num_angle = atom->num_angle;
  int **angle_type = atom->angle_type;
  tagint **angle_atom1 = atom->angle_atom1;
  tagint **angle_atom2 = atom->angle_atom2;
  tagint **angle_atom3 = atom->angle_atom3;

  if ((m = atom->map(aatom2)) >= 0) {
    if (num_angle[m] == atom->angle_per_atom)
      error->one(FLERR,"New angle exceeded angles per atom in create_bonds");
    angle_type[m][num_angle[m]] = atype;
    angle_atom1[m][num_angle[m]] = aatom1;
    angle_atom2[m][num_angle[m]] = aatom2;
    angle_atom3[m][num_angle[m]] = aatom3;
    num_angle[m]++;
  }
  ++atom->nangles;

  if (force->newton_bond) return;

  if ((m = atom->map(aatom1)) >= 0) {
    if (num_angle[m] == atom->angle_per_atom)
      error->one(FLERR,"New angle exceeded angles per atom in create_bonds");
    angle_type[m][num_angle[m]] = atype;
    angle_atom1[m][num_angle[m]] = aatom1;
    angle_atom2[m][num_angle[m]] = aatom2;
    angle_atom3[m][num_angle[m]] = aatom3;
    num_angle[m]++;
  }

  if ((m = atom->map(aatom3)) >= 0) {
    if (num_angle[m] == atom->angle_per_atom)
      error->one(FLERR,"New angle exceeded angles per atom in create_bonds");
    angle_type[m][num_angle[m]] = atype;
    angle_atom1[m][num_angle[m]] = aatom1;
    angle_atom2[m][num_angle[m]] = aatom2;
    angle_atom3[m][num_angle[m]] = aatom3;
    num_angle[m]++;
  }
}

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

void CreateBonds::single_dihedral()
{
  int m;

  // check that 4 atoms exist

  int count = 0;
  if (atom->map(datom1) >= 0) count++;
  if (atom->map(datom2) >= 0) count++;
  if (atom->map(datom3) >= 0) count++;
  if (atom->map(datom4) >= 0) count++;

  int allcount;
  MPI_Allreduce(&count,&allcount,1,MPI_INT,MPI_SUM,world);
  if (allcount != 4) 
    error->all(FLERR,"Create_bonds single/dihedral atoms do not exist");

  // create bond once or 4x if newton_bond set

  int *num_dihedral = atom->num_dihedral;
  int **dihedral_type = atom->dihedral_type;
  tagint **dihedral_atom1 = atom->dihedral_atom1;
  tagint **dihedral_atom2 = atom->dihedral_atom2;
  tagint **dihedral_atom3 = atom->dihedral_atom3;
  tagint **dihedral_atom4 = atom->dihedral_atom4;

  if ((m = atom->map(datom2)) >= 0) {
    if (num_dihedral[m] == atom->dihedral_per_atom)
      error->one(FLERR,
                 "New dihedral exceeded dihedrals per atom in create_bonds");
    dihedral_type[m][num_dihedral[m]] = dtype;
    dihedral_atom1[m][num_dihedral[m]] = datom1;
    dihedral_atom2[m][num_dihedral[m]] = datom2;
    dihedral_atom3[m][num_dihedral[m]] = datom3;
    dihedral_atom4[m][num_dihedral[m]] = datom4;
    num_dihedral[m]++;
  }
  ++atom->ndihedrals;

  if (force->newton_bond) return;

  if ((m = atom->map(datom1)) >= 0) {
    if (num_dihedral[m] == atom->dihedral_per_atom)
      error->one(FLERR,
                 "New dihedral exceeded dihedrals per atom in create_bonds");
    dihedral_type[m][num_dihedral[m]] = dtype;
    dihedral_atom1[m][num_dihedral[m]] = datom1;
    dihedral_atom2[m][num_dihedral[m]] = datom2;
    dihedral_atom3[m][num_dihedral[m]] = datom3;
    dihedral_atom4[m][num_dihedral[m]] = datom4;
    num_dihedral[m]++;
  }

  if ((m = atom->map(datom3)) >= 0) {
    if (num_dihedral[m] == atom->dihedral_per_atom)
      error->one(FLERR,
                 "New dihedral exceeded dihedrals per atom in create_bonds");
    dihedral_type[m][num_dihedral[m]] = dtype;
    dihedral_atom1[m][num_dihedral[m]] = datom1;
    dihedral_atom2[m][num_dihedral[m]] = datom2;
    dihedral_atom3[m][num_dihedral[m]] = datom3;
    dihedral_atom4[m][num_dihedral[m]] = datom4;
    num_dihedral[m]++;
  }

  if ((m = atom->map(datom4)) >= 0) {
    if (num_dihedral[m] == atom->dihedral_per_atom)
      error->one(FLERR,
                 "New dihedral exceeded dihedrals per atom in create_bonds");
    dihedral_type[m][num_dihedral[m]] = dtype;
    dihedral_atom1[m][num_dihedral[m]] = datom1;
    dihedral_atom2[m][num_dihedral[m]] = datom2;
    dihedral_atom3[m][num_dihedral[m]] = datom3;
    dihedral_atom4[m][num_dihedral[m]] = datom4;
    num_dihedral[m]++;
  }
}