#include "math.h"
#include "mpi.h"
#include "string.h"
#include "stdlib.h"
#include "compute_nnn.h"
#include "update.h"
#include "atom.h"
#include "atom_vec.h"
#include "force.h"
#include "pair.h"
#include "comm.h"
#include "neighbor.h"
#include "neigh_list.h"
#include "neigh_request.h"
#include "random_mars.h"
#include "memory.h"
#include "error.h"


using namespace LAMMPS_NS;

#define BIG 1.0e20

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

ComputeNNN::ComputeNNN(LAMMPS *lmp, int narg, char **arg) :
  Compute(lmp, narg, arg)
{
	
  MPI_Comm_rank(world,&me);
	
  local_flag=1;           // 0/1 if compute_local() function exists
  size_local_cols =2;
  create_attribute = 1;
	
  double cutoff = force->numeric(FLERR,arg[3]);	
  iatomtype = force->inumeric(FLERR,arg[4]);
  imaxbond = force->inumeric(FLERR,arg[5]);
  jatomtype = force->inumeric(FLERR,arg[6]);
  jmaxbond = force->inumeric(FLERR,arg[7]);
  btype = force->inumeric(FLERR,arg[8]);
	
  cutsq = cutoff*cutoff;
  bondcount = NULL;
 
  memory->grow(bondcount,nmax,"nnn:bondcount");
  countflag = 0;
  
  nmax = 0;
  partner = finalpartner = NULL;
  distsq = NULL;

  // zero out stats
  createcount = 0;
  createcounttotal = 0;
  
  array=NULL;
 }

ComputeNNN::~ComputeNNN()
{
  memory->destroy(bondcount);
  memory->destroy(partner);
  memory->destroy(finalpartner);
  memory->destroy(distsq);
  
  memory->destroy(array);
}


void ComputeNNN::init()
{
  // check cutoff for iatomtype,jatomtype

  if (force->pair == NULL || cutsq > force->pair->cutsq[iatomtype][jatomtype])
    error->all(FLERR,"Fix bond/create cutoff is longer than pairwise cutoff");

  int irequest = neighbor->request(this,instance_me);
  neighbor->requests[irequest]->pair = 0;
  neighbor->requests[irequest]->compute = 1;
  neighbor->requests[irequest]->occasional = 1;
}

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

void ComputeNNN::setup()
{
  int i,j,m;

  // compute initial bondcount if this is first run
  // can't do this earlier, in constructor or init, b/c need ghost info

  if (countflag) return;
  countflag = 1;

  // count bonds stored with each bond I own
  // if newton bond is not set, just increment count on atom I
  // if newton bond is set, also increment count on atom J even if ghost
  // bondcount is long enough to tally ghost atom counts

  int *num_bond = atom->num_bond;
  int **bond_type = atom->bond_type;
  tagint **bond_atom = atom->bond_atom;
  int nlocal = atom->nlocal;
  int nghost = atom->nghost;
  int nall = nlocal + nghost;
  int newton_bond = force->newton_bond;

  for (i = 0; i < nall; i++) bondcount[i] = 0;

  for (i = 0; i < nlocal; i++) 
    for (j = 0; j < num_bond[i]; j++) {
      if (bond_type[i][j] == btype) {
        bondcount[i]++;
        if (newton_bond) {
          m = atom->map(bond_atom[i][j]);
          if (m < 0) 
            error->one(FLERR,"Fix bond/create needs ghost atoms from further away");
          bondcount[m]++;
        }
      }
    }
  	
  // if newton_bond is set, need to sum bondcount
 
  commflag = 1;
  if (newton_bond) comm->reverse_comm_compute(this);
}

void ComputeNNN::compute_local()
{
  invoked_array = update->ntimestep;
  reallocate(1);
	
  int i,j,k,m,n,ii,jj,inum,jnum,itype,jtype,n1,n2,n3,possible;
  double xtmp,ytmp,ztmp,delx,dely,delz,rsq;
  int *ilist,*jlist,*numneigh,**firstneigh;
  tagint *slist;
 
  int nlocal = atom->nlocal;
  int nghost = atom->nghost;
  int nall = nlocal + nghost;
 
  // if (update->ntimestep % nevery) return;

  comm->forward_comm();

  // forward comm of bondcount, so ghosts have it

  commflag = 1;
  comm->forward_comm_compute(this);
  
  // resize bond partner list and initialize it
  // needs to be atom->nmax in length

  if (atom->nmax > nmax) {
    memory->destroy(partner);
    memory->destroy(finalpartner);
    memory->destroy(distsq);
    nmax = atom->nmax;
    memory->create(partner,nmax,"NNN:partner");
    memory->create(finalpartner,nmax,"NNN:finalpartner");
    memory->create(distsq,nmax,"NNN:distsq");
  }
  
  for (i = 0; i < nmax; i++) {
    partner[i] = 0;
    finalpartner[i] = 0;
    distsq[i] = BIG;
	
  }
  
  // loop over neighbors of my atoms
  // each atom sets one closest eligible partner atom ID to bond with

  double **x = atom->x;
  tagint *tag = atom->tag;
  int **nspecial = atom->nspecial;
  tagint **special = atom->special;
  int *mask = atom->mask;
  int *type = atom->type;
  
  neighbor->build_one(list,1);
  inum = list->inum;
  ilist = list->ilist;
  numneigh = list->numneigh;
  firstneigh = list->firstneigh;
  
  int num_iatoms = 0;
	
  for (int iii = 0; iii < inum; iii++) {
	int i=ilist[iii];
	  if (type[i]==iatomtype) {
		num_iatoms++;
	  }
  }
  
  size_local_rows = num_iatoms;
  
  for (ii = 0; ii < inum; ii++) {
    i = ilist[ii];
    if (!(mask[i] & groupbit)) continue;
    itype = type[i];
    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;
      if (!(mask[j] & groupbit)) continue;
      jtype = type[j];

      possible = 0;
      if (itype == iatomtype && jtype == jatomtype) {
        if ((imaxbond == 0 || bondcount[i] < imaxbond) &&
            (jmaxbond == 0 || bondcount[j] < jmaxbond))
          possible = 1;
		  
      } else if (itype == jatomtype && jtype == iatomtype) {
        if ((jmaxbond == 0 || bondcount[i] < jmaxbond) &&
            (imaxbond == 0 || bondcount[j] < imaxbond))
          possible = 1;
      }
      if (!possible) continue;

      // do not allow a duplicate bond to be created
      // check 1-2 neighbors of atom I

      for (k = 0; k < nspecial[i][0]; k++)
        if (special[i][k] == tag[j]) possible = 0;
      if (!possible) continue;

      delx = xtmp - x[j][0];
      dely = ytmp - x[j][1];
      delz = ztmp - x[j][2];
      rsq = delx*delx + dely*dely + delz*delz;
	  if (rsq >= cutsq) continue;

      if (rsq < distsq[i]) {
		partner[i] = tag[j];
        distsq[i] = rsq;
      }
      if (rsq < distsq[j]) {
		partner[j] = tag[i];
		distsq[j] = rsq;
      }
    }
  }

  // reverse comm of distsq and partner
  // not needed if newton_pair off since I,J pair was seen by both procs

   commflag = 2;
  if (force->newton_pair) comm->reverse_comm_compute(this);
  
  //write atom ID (itype) and distance to nearest non-bonded neighbor of jtype
  
  int on_my_count = 0;
  
  for (int iii = 0; iii < inum; iii++) {
	int i=ilist[iii];
	  if (type[i]==iatomtype) {
		array[on_my_count][0]=tag[i];
		array[on_my_count][1]=distsq[i];
		on_my_count++;
	  }
  }
}

void ComputeNNN::reallocate(int n)
{
  // grow vector or array and indices array

 // while (nmax < n) nmax += DELTA;
    memory->destroy(array);
    memory->create(array,atom->nlocal,2,"NNN:array");
    array_local = array;
 }