/* ----------------------------------------------------------------------
   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 author: Daniel Schwen
------------------------------------------------------------------------- */

#include "math.h"
#include "string.h"
#include "stdlib.h"
#include "compute_voronoi_atom.h"
#include "atom.h"
#include "group.h"
#include "update.h"
#include "modify.h"
#include "domain.h"
#include "memory.h"
#include "error.h"
#include "comm.h"

#include <vector>
#include "voro++.hh"

using namespace LAMMPS_NS;
using namespace voro;

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

ComputeVoronoi::ComputeVoronoi(LAMMPS *lmp, int narg, char **arg) :
  Compute(lmp, narg, arg)
{
  int sgroup;
  surface = VOROSURF_NONE;
  maxedge = 0;
  size_peratom_cols = 2;

  int iarg = 3;
  while ( iarg<narg ) {
/*    if (strcmp(arg[iarg], "surface") == 0) {
      if (iarg + 2 > narg) error->all(FLERR,"Missing group name after keyword 'surface'.");
      // group all is a special case where we just skip group testing
      if(strcmp(arg[iarg+1], "all") == 0) {
        surface = VOROSURF_ALL;
      } else {
        sgroup = group->find(arg[iarg+1]);
        if (sgroup == -1) error->all(FLERR,"Could not find compute/voronoi surface group ID");
        sgroupbit = group->bitmask[sgroup]; 
        surface = VOROSURF_GROUP;
      }
      size_peratom_cols = 3;
      iarg += 2;
    } else if (strcmp(arg[iarg], "edge_histo") == 0) {
      if (iarg + 2 > narg) error->all(FLERR,"Missing maximum edge count after keyword 'edge_histo'.");
      maxedge = atoi(arg[iarg+1]);
      iarg += 2;
    } else */
      error->all(FLERR,"Illegal compute voronoi/atom command");
  }

  peratom_flag = 1;
  nmax = 0;
  voro = NULL;

  if ( maxedge > 0 ) {
    vector_flag = 1;
    size_vector = maxedge+1;
    memory->create(edge,maxedge+1,"voronoi/atom:edge");
    memory->create(sendvector,maxedge+1,"voronoi/atom:sendvector");
    vector = edge;
  }
}

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

ComputeVoronoi::~ComputeVoronoi()
{
  memory->destroy(voro);
}

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

void ComputeVoronoi::init()
{
  int count = 0;
  for (int i = 0; i < modify->ncompute; i++)
    if (strcmp(modify->compute[i]->style,"voronoi/atom") == 0) count++;
  if (count > 1 && comm->me == 0)
    error->warning(FLERR,"More than one compute voronoi/atom command");
}

/* ----------------------------------------------------------------------
   gather compute vector data from other nodes
------------------------------------------------------------------------- */

void ComputeVoronoi::compute_peratom()
{
  int i, j;

  invoked_peratom = update->ntimestep;

  // grow per atom array if necessary

  int nlocal = atom->nlocal;
  if (nlocal > nmax) {
    memory->destroy(voro);
    nmax = atom->nmax;
    memory->create(voro,nmax,size_peratom_cols,"voronoi/atom:voro");
    array_atom = voro;
  }

  double *sublo = domain->sublo, *boxlo = domain->boxlo;
  double *subhi = domain->subhi, *boxhi = domain->boxhi;
  double *cut = comm->cutghost;
  double sublo_bound[3], subhi_bound[3], cut_bound[3];
   
  // setup bounds for voro++ domain for orthogonal and triclinic simulation boxes
  if( domain->triclinic ) {
    // triclinic box: embed parallelepiped into orthogonal voro++ domain
    double mx, my, sxy,sxz,syz;
    mx = (boxhi[0]-boxlo[0])/(subhi[0]-sublo[0]);
    my = (boxhi[1]-boxlo[1])/(subhi[1]-sublo[1]);
    sxy = domain->xy/mx;
    sxz = domain->xz/mx;
    syz = domain->yz/my;

    sublo_bound[0] = MIN(sublo[0],sublo[0]+sxy);
    sublo_bound[0] = MIN(sublo_bound[0],sublo_bound[0]+sxz);
    sublo_bound[1] = MIN(sublo[1],sublo[1]+syz);
    sublo_bound[2] = sublo[2];
    subhi_bound[0] = MAX(subhi[0],subhi[0]+sxy);
    subhi_bound[0] = MAX(subhi_bound[0],subhi_bound[0]+sxz);
    subhi_bound[1] = MAX(subhi[1],subhi[1]+syz);
    subhi_bound[2] = subhi[2];

    // cutghost is in lamda coordinates for triclinic boxes
    double *h = domain->h, cuttri[3];
    cut_bound[0] = h[0]*cut[0] + h[5]*cut[1] + h[4]*cut[2];
    cut_bound[1] = h[1]*cut[1] + h[3]*cut[2];
    cut_bound[2] = h[2]*cut[2];

  } else {
    // orthogonal box
    for( i=0; i<3; ++i ) {
      sublo_bound[i] = sublo[i];
      subhi_bound[i] = subhi[i];
      cut_bound[i] = cut[i];
    }
  }

  // n = # of voro++ spatial hash cells (with approximately cubic cells)
  int nall = nlocal + atom->nghost;
  double n[3], V;
  for( i=0; i<3; ++i ) n[i] = subhi_bound[i] - sublo_bound[i];
  V = n[0]*n[1]*n[2];
  for( i=0; i<3; ++i ) {
    n[i] = round( n[i]*pow( double(nall)/(V*8.0), 0.333333 ) );
    n[i] = n[i]==0 ? 1 : n[i];
  } 

  // initialize voro++ container
  // preallocates 8 atoms per cell
  // voro++ allocates more memory if needed
  double e = 0.01;
  container con(sublo_bound[0]-cut_bound[0]-e,subhi_bound[0]+cut_bound[0]+e,
                sublo_bound[1]-cut_bound[1]-e,subhi_bound[1]+cut_bound[1]+e,
                sublo_bound[2]-cut_bound[2]-e,subhi_bound[2]+cut_bound[2]+e,
                int(n[0]),int(n[1]),int(n[2]),false,false,false,8); 

  // pass coordinates for local and ghost atoms to voro++
  double **x = atom->x;
  for (i = 0; i < nall; i++)
    con.put(i,x[i][0],x[i][1],x[i][2]);

  // invoke voro++ and fetch results for owned atoms in group
  int *mask = atom->mask;
  std::vector<int> neigh;
  std::vector<double> narea;
  std::vector<int> norder;

  // clear edge statistics
  for (i = 0; i < maxedge; ++i) edge[i]=0;

  voronoicell_neighbor c;
  c_loop_all cl(con);
  if (cl.start()) do if (con.compute_cell(c,cl)) {
    i = cl.pid();

    // zero out surface area if surface computation was requested
    if (surface != VOROSURF_NONE) voro[i][2] = 0.0;

    if (i < nlocal && (mask[i] & groupbit)) {
      voro[i][0] = c.volume();
      c.neighbors(neigh);
      voro[i][1] = neigh.size();
      if (surface == VOROSURF_ALL) {
        voro[i][2] = c.surface_area();
      } else if (surface == VOROSURF_GROUP) {
        c.face_areas(narea);
        // loop over all faces (neighbors) and check if they are in the surface group
        for (j=0; j<voro[i][1]; ++j)  
          if (mask[neigh[j]] & sgroupbit) voro[i][2] += narea[j];
      }
      if (maxedge>0) {
        c.face_orders(norder);
        for (j=0; j<voro[i][1]; ++j)
          if (norder[j]>0) 
            if (norder[j]<=maxedge) 
              edge[norder[j]-1]++;
            else
              edge[maxedge]++;

      }
    } else if (i < nlocal) voro[i][0] = voro[i][1] = 0.0;
  } while (cl.inc());
}

/* ----------------------------------------------------------------------
   memory usage of local atom-based array
------------------------------------------------------------------------- */

double ComputeVoronoi::memory_usage()
{
  double bytes = size_peratom_cols * nmax * sizeof(double);
  return bytes;
}

void ComputeVoronoi::compute_vector()
{
  invoked_vector = update->ntimestep;
  if( invoked_peratom < invoked_vector ) compute_peratom();

  for( int i=0; i<size_vector; ++i ) sendvector[i] = edge[i];
  MPI_Allreduce(sendvector,edge,size_vector,MPI_DOUBLE,MPI_SUM,world);
}