/* ----------------------------------------------------------------------
   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 "fix.h"
#include "bond_cell.h"
#include <mpi.h>
#include <cmath>
#include <cstring>
#include "atom.h"
#include "neighbor.h"
#include "comm.h"
#include "force.h"
#include "memory.h"
#include "error.h"
#include "utils.h"
#include "atom_vec.h"
#include "domain.h"


using namespace LAMMPS_NS;
using namespace FixConst;
//using namespace MathConst;

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

BondCell::BondCell(LAMMPS *lmp) : Bond(lmp),avomega(NULL),cell(NULL),nCELL(NULL), mol2cell(NULL), cell2mol(NULL)
{
  reinitflag = 1;
  cell = NULL;
  mol2cell = NULL;
  cell2mol = NULL;
}

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

BondCell::~BondCell()
{
  if (allocated && !copymode) {
    memory->destroy(nCELL);
    memory->destroy(setflag);
    memory->destroy(alpha);
    memory->destroy(avomega);
    memory->destroy(r0);
    memory->destroy(cell);
  }
}

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

void BondCell::compute(int eflag, int vflag)
{
  int i1,icell,n,type;
  tagint *molecule = atom->molecule ;
  double **omega = atom->omega;
  double **mu = atom->mu;
  int **bondlist = neighbor->bondlist;
  int nbondlist = neighbor->nbondlist;
  int nlocal = atom->nlocal;
  int newton_bond = force->newton_bond;

  // init avomega
  for (icell = 0; icell < ncell; icell ++)
       avomega[icell][0] = avomega[icell][1] = avomega[icell][2] = 0.0;

  // sum each omega
  for (n = 0; n < nbondlist; n++) {
    i1 = bondlist[n][0];
    avomega[mol2cell[molecule[i1]]][0] += mu[i1][0];
    avomega[mol2cell[molecule[i1]]][1] += mu[i1][1];
    avomega[mol2cell[molecule[i1]]][2] += mu[i1][2];
  }

  // average omega
  for (icell = 0; icell < ncell; icell++) {
    avomega[icell][0] /= nCELL[icell];
    avomega[icell][1] /= nCELL[icell];
    avomega[icell][2] /= nCELL[icell];
    //normalization
    if (domain->dimension == 2){
         avomega[icell][2] = 0.0;
    }
    double msq = sqrt(avomega[icell][0]*avomega[icell][0] + avomega[icell][1]*avomega[icell][1] + avomega[icell][2]*avomega[icell][2]); 
    if(msq > 0.0){
      avomega[icell][0] /= msq;avomega[icell][1] /= msq;avomega[icell][2] /= msq;
    }
  }
  
  for (n = 0; n < nbondlist; n++) {
    i1 = bondlist[n][0];
    type = bondlist[n][2];
    if (screen) fprintf(screen,"omega_cell-1: %f %f \n",omega[i1][0],omega[i1][1]);
    omega[i1][0] =  -alpha[type] * avomega[ mol2cell[molecule[i1]]  ][0];
    omega[i1][1] =  -alpha[type] * avomega[ mol2cell[molecule[i1]]  ][1];
    omega[i1][2] =  -alpha[type] * avomega[ mol2cell[molecule[i1]]  ][2];
    if (screen) fprintf(screen,"omega_cell-2: %f %f \n",omega[i1][0],omega[i1][1]);
  }
  //if (screen) fprintf(screen,"index of local atom %d \n",2);
}

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

void BondCell::allocate()
{
  allocated = 1;
  int n = atom->nbondtypes;

  memory->create(alpha,n+1,"bond:alpha");
  memory->create(r0,n+1,"bond:r0");
  memory->create(setflag,n+1,"bond:setflag");
  for (int i = 1; i <= n; i++) setflag[i] = 0;
}

/* ----------------------------------------------------------------------
   set coeffs for one or more types
------------------------------------------------------------------------- */

void BondCell::coeff(int narg, char **arg)
{
  if (narg != 3) error->all(FLERR,"Incorrect args for bond coefficients");
  if (!allocated) allocate();

  int ilo,ihi;
  force->bounds(FLERR,arg[0],atom->nbondtypes,ilo,ihi);

  double k_alpha = force->numeric(FLERR,arg[1]);
  double k_r0 = force->numeric(FLERR,arg[2]);

  int count = 0;
  for (int i = ilo; i <= ihi; i++) {
    alpha[i] = k_alpha;
    if (screen) fprintf(screen,"index of local atom %f \n",alpha[i]);
    r0[i] = k_r0;
    setflag[i] = 1;
    count++;
  }

  if (count == 0) error->all(FLERR,"Incorrect args for bond coefficients");


  int i,icell;

  cell = NULL;
  memory->grow(cell,atom->nmax,"bond:cell");
  atom->add_callback(0);

  MPI_Comm_rank(world,&me);
  MPI_Comm_size(world,&nprocs);

  tagint *molecule;
  molecule = atom->molecule;
  int *mask = atom->mask;
  int nlocal = atom->nlocal;

  tagint maxmol_tag = -1;
  for (i = 0; i < nlocal; i++)
    maxmol_tag = MAX(maxmol_tag,molecule[i]);

  tagint itmp;
  MPI_Allreduce(&maxmol_tag,&itmp,1,MPI_LMP_TAGINT,MPI_MAX,world);
  if (itmp+1 > MAXSMALLINT)
    error->all(FLERR,"Too many molecules for bond/cell");
  maxmol = (int) itmp;   //max molecule index
 

  //cell index= maxmol+1;

  int *ncount;
  memory->create(ncount,maxmol+1,"bondcell:ncount");
  for (i = 0; i <= maxmol; i++) ncount[i] = 0;
  for (i = 0; i < nlocal; i++)
      ncount[molecule[i]]++;
  memory->create(mol2cell,maxmol+1,"bondcell:mol2cell");
  MPI_Allreduce(ncount,mol2cell,maxmol+1,MPI_INT,MPI_SUM,world);


  //if there is an atom in each molecule
  //then ncell++
  //the final ncell is how much cells
  //if mol-index is continuous, ncell=maxmol+1
  //mol2cell preserves cell index or -1(no atoms)
  //
  ncell = 0;
  for (i = 0; i <= maxmol; i++)
     if (mol2cell[i]) mol2cell[i] = ncell++;
  else mol2cell[i] = -1;

  memory->create(cell2mol,ncell,"bondcell:cell2mol");



  //cell2mol: convert cell index to mol-ID
  //if there are atoms in each cell
  //cell2mol preserves mol index
  ncell = 0;
  for (i = 0; i <= maxmol; i++)
    if (mol2cell[i] >= 0) cell2mol[ncell++] = i;

  //cell[i] is cell index of each molecules
  for (i = 0; i < nlocal; i++) {
      cell[i] = mol2cell[molecule[i]];
  }

  memory->destroy(ncount);
  memory->create(nCELL,ncell,"CELL:ncell");
  memory->create(avomega,ncell,3,"Cell:avomega");


  //if (screen) fprintf(screen,"test segmental fault\n");
 // ncell[n] = # of atoms in Nth Cell
  // error if one or zero atoms

  int *nncount = new int[ncell];
  for (icell = 0; icell < ncell; icell++) nncount[icell] = 0;

  nlocal = atom->nlocal;

  for (i = 0; i < nlocal; i++)
    if (cell[i] >= 0) nncount[cell[i]]++;

  MPI_Allreduce(nncount,nCELL,ncell,MPI_INT,MPI_SUM,world);
  delete [] nncount;

  for (icell = 0; icell < ncell; icell++)
    if (nCELL[icell] <= 1) error->all(FLERR,"One or zero atoms in each cell");

  // print statistics

  int nsum = 0;
  for (icell = 0; icell < ncell; icell++) nsum += nCELL[icell];

  if (me == 0) {
    if (screen) fprintf(screen,"%d cells with %d atoms\n",ncell,nsum);
    if (logfile) fprintf(logfile,"%d cells with %d atoms\n",ncell,nsum);
  }
}



/* ----------------------------------------------------------------------
   return an equilbrium bond length
------------------------------------------------------------------------- */

double BondCell::equilibrium_distance(int i)
{
  return r0[i];
}

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

void BondCell::write_restart(FILE *fp)
{
  fwrite(&alpha[1],sizeof(double),atom->nbondtypes,fp);
  fwrite(&r0[1],sizeof(double),atom->nbondtypes,fp);
}

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

void BondCell::read_restart(FILE *fp)
{
  allocate();

  if (comm->me == 0) {
    utils::sfread(FLERR,&alpha[1],sizeof(double),atom->nbondtypes,fp,NULL,error);
    utils::sfread(FLERR,&r0[1],sizeof(double),atom->nbondtypes,fp,NULL,error);
  }
  MPI_Bcast(&alpha[1],atom->nbondtypes,MPI_DOUBLE,0,world);
  MPI_Bcast(&r0[1],atom->nbondtypes,MPI_DOUBLE,0,world);

  for (int i = 1; i <= atom->nbondtypes; i++) setflag[i] = 1;
}

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

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

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

double BondCell::single(int type, double rsq, int /*i*/, int /*j*/,
                        double &fforce)
{
  return 0;
}

/* ----------------------------------------------------------------------
    Return ptr to internal members upon request.
------------------------------------------------------------------------ */
void *BondCell::extract( char *str, int &dim )
{
  dim = 1;
  if (strcmp(str,"alpha")==0) return (void*) alpha;
  if (strcmp(str,"r0")==0) return (void*) r0;
  return NULL;
}