/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   www.cs.sandia.gov/~sjplimp/lammps.html
   Steve Plimpton, sjplimp@sandia.gov, Sandia National Laboratories

   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 "string.h"
#include "fix_dipole.h"
#include "mpi.h"
#include "atom.h"
#include "update.h"
#include "domain.h"
#include "output.h"
#include "group.h"
#include "modify.h"
#include "memory.h"
#include "error.h"
#include "math.h"
using namespace LAMMPS_NS;
/* ---------------------------------------------------------------------- */

FixDipole::FixDipole(LAMMPS *lmp, int narg, char **arg) : Fix(lmp, narg, arg)
{
  if (narg != 5) error->all("Illegal fix dipole command");
  nevery = atoi(arg[3]);
  if (nevery <= 0) error->all("Illegal fix dipole command");
  first = 1;
  restart_peratom = 1;

  MPI_Comm_rank(world,&me);
  if (me == 0) {
    fp = fopen(arg[4],"w");
    if (fp == NULL) {
      char str[128];
      sprintf(str,"Cannot open fix dipole file %s",arg[4]);
      error->one(str);
    }
  }

  if (me == 0) {
    fprintf(fp,"Dipole moment for group %s\n",
	    group->names[igroup]);
    fprintf(fp,"x y z total\n");
  }

  // perform initial allocation of atom-based array
  // register with atom class

  xoriginal = NULL;
  grow_arrays(atom->nmax);
  atom->add_callback(0);
  atom->add_callback(1);

  // xoriginal = initial unwrapped positions of atoms
  
  double **x = atom->x;
  int *mask = atom->mask;
  int *image = atom->image;
  int nlocal = atom->nlocal;
  
  double xprd = domain->xprd;
  double yprd = domain->yprd;
  double zprd = domain->zprd;
  int xbox,ybox,zbox;

  for (int i = 0; i < nlocal; i++) {
    if (mask[i] & groupbit) {
      //xbox = image[i]%1000 - 500;
      //ybox = (image[i]/1000) % 1000 - 500;
      //zbox = image[i]/1000000 - 500;
       xbox = (image[i] & 1023) - 512;
       ybox = (image[i] >> 10 & 1023) - 512;
       zbox = (image[i] >> 20) - 512;
      xoriginal[i][0] = x[i][0] + xbox*xprd;
      xoriginal[i][1] = x[i][1] + ybox*yprd;
      xoriginal[i][2] = x[i][2] + zbox*zprd;
    } else xoriginal[i][0] = xoriginal[i][1] = xoriginal[i][2] = 0.0;
  }

  // ndipole = # of atoms in group

  ndipole = 0;
  for (int i = 0; i < nlocal; i++)
    if (mask[i] & groupbit) ndipole++;

  int ndipole_all;
  MPI_Allreduce(&ndipole,&ndipole_all,1,MPI_INT,MPI_SUM,world);
  ndipole = ndipole_all;

  if (ndipole == 0) error->all("No atoms to compute dipole for");
}

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

FixDipole::~FixDipole()
{
  if (me == 0) fclose(fp);

  // if atom class still exists:
  //   unregister this fix so atom class doesn't invoke it any more
 
  if (atom) atom->delete_callback(id,0);
  if (atom) atom->delete_callback(id,1);

  // delete locally stored array

  memory->destroy_2d_double_array(xoriginal);
}

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

int FixDipole::setmask()
{
  int mask = 0;
  mask |= END_OF_STEP;
  return mask;
}

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

void FixDipole::init()
{
  // warn if more than one dipole fix

  int count = 0;
  for (int i = 0; i < modify->nfix; i++)
    if (strcmp(modify->fix[i]->style,"dipole") == 0) count++;
  if (count > 1 && me == 0) error->warning("More than one dipole fix");
}

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

void FixDipole::setup()
{
  if (first) end_of_step();
  first = 0;
}

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

void FixDipole::end_of_step()
{
//Caculation of dipole moment based on: Cerius 2.0 documentation  
  double **x = atom->x;
  int *mask = atom->mask;
  int *image = atom->image;
  int nlocal = atom->nlocal;
  double *q = atom->q;
  double qmin = 100.0;
  double qp[3],qn[3];
  double qpsum,qnsum; 
  double qpr[3],qnr[3]; 
  double xprd = domain->xprd;
  double yprd = domain->yprd;
  double zprd = domain->zprd;
  int xbox,ybox,zbox;
  double dipole[8];
  dipole[0] = dipole[1] = dipole[2] = dipole[3] = dipole[4] = dipole[5] = dipole[6] = dipole[7] = 0.0;
  qpr[0] = qpr[1] = qpr[2] = 0.0;
  qnr[0] = qnr[1] = qnr[2] = 0.0;
  qpsum = qnsum = 0.0;

  for (int i = 0; i < nlocal; i++)
{    
if (mask[i] & groupbit) {
    //  xbox = image[i]%1000 - 500;
    //  ybox = (image[i]/1000) % 1000 - 500;
    //  zbox = image[i]/1000000 - 500;
       xbox = (image[i] & 1023) - 512;
       ybox = (image[i] >> 10 & 1023) - 512;
       zbox = (image[i] >> 20) - 512;
    if ((q[i]) >= 0)
      {      
        qpsum += (q[i]);
        qpr[0] += (x[i][0]+xbox*xprd)*(q[i]);
        qpr[1] += (x[i][1]+ybox*yprd)*(q[i]);
        qpr[2] += (x[i][2]+zbox*zprd)*(q[i]);
      }
    else
     {
        qnsum += (q[i]);
        qnr[0] += (x[i][0]+xbox*xprd)*(q[i]);
        qnr[1] += (x[i][1]+ybox*yprd)*(q[i]);
        qnr[2] += (x[i][2]+zbox*zprd)*(q[i]);
     } 
    }
}
double tmp;
MPI_Allreduce(&qpsum,&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qpsum = tmp;
MPI_Allreduce(&qnsum,&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qnsum = tmp;
MPI_Allreduce(&qpr[0],&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qpr[0] = tmp;
MPI_Allreduce(&qpr[1],&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qpr[1] = tmp;
MPI_Allreduce(&qpr[2],&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qpr[2] = tmp;
MPI_Allreduce(&qnr[0],&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qnr[0] = tmp;
MPI_Allreduce(&qnr[1],&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qnr[1] = tmp;
MPI_Allreduce(&qnr[2],&tmp,1,MPI_DOUBLE,MPI_SUM,world);
qnr[2] = tmp;
     qmin = qpsum;
     if (qpsum > fabs(qnsum))
      {      
        qmin = fabs(qnsum);
      }  
  qp[0] = qpr[0]/qpsum;
  qp[1] = qpr[1]/qpsum;
  qp[2] = qpr[2]/qpsum;
  qn[0] = qnr[0]/qnsum;
  qn[1] = qnr[1]/qnsum;
  qn[2] = qnr[2]/qnsum;
  dipole[0]=4.802*qmin*(qp[0]-qn[0]); 
  dipole[1]=4.802*qmin*(qp[1]-qn[1]); 
  dipole[2]=4.802*qmin*(qp[2]-qn[2]); 
  dipole[3]=sqrt(pow(dipole[0],2)+pow(dipole[1],2)+pow(dipole[2],2)); 
  dipole[4]=4.802*(qpr[0] + qnr[0]); 
  dipole[5]=4.802*(qpr[1] + qnr[1]); 
  dipole[6]=4.802*(qpr[2] + qnr[2]); 
  dipole[7]=sqrt(pow(dipole[4],2)+pow(dipole[5],2)+pow(dipole[6],2)); 

/*  if (me == 0) fprintf(fp,"%g %g %g %g %g  %g %g %g %g %g \n",
		       dipole[0],dipole[1],dipole[2],dipole[3],qpsum,qnsum,dipole[4],dipole[5],dipole[6],dipole[7]);*/
  if (me == 0) fprintf(fp,"%g %g %g %g \n",
		       dipole[0],dipole[1],dipole[2],dipole[3]);
}

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

int FixDipole::memory_usage()
{
  int bytes = atom->nmax*3 * sizeof(double);
  return bytes;
}

/* ----------------------------------------------------------------------
   allocate atom-based array
------------------------------------------------------------------------- */

void FixDipole::grow_arrays(int nmax)
{
  xoriginal =
    memory->grow_2d_double_array(xoriginal,nmax,3,"fix_dipole:xoriginal");
}

/* ----------------------------------------------------------------------
   copy values within local atom-based array
------------------------------------------------------------------------- */

void FixDipole::copy_arrays(int i, int j)
{
  xoriginal[j][0] = xoriginal[i][0];
  xoriginal[j][1] = xoriginal[i][1];
  xoriginal[j][2] = xoriginal[i][2];
}

/* ----------------------------------------------------------------------
   pack values in local atom-based array for exchange with another proc
------------------------------------------------------------------------- */

int FixDipole::pack_exchange(int i, double *buf)
{
  buf[0] = xoriginal[i][0];
  buf[1] = xoriginal[i][1];
  buf[2] = xoriginal[i][2];
  return 3;
}

/* ----------------------------------------------------------------------
   unpack values in local atom-based array from exchange with another proc
------------------------------------------------------------------------- */

int FixDipole::unpack_exchange(int nlocal, double *buf)
{
  xoriginal[nlocal][0] = buf[0];
  xoriginal[nlocal][1] = buf[1];
  xoriginal[nlocal][2] = buf[2];
  return 3;
}

/* ----------------------------------------------------------------------
   pack values in local atom-based arrays for restart file
------------------------------------------------------------------------- */

int FixDipole::pack_restart(int i, double *buf)
{
  int m = 0;
  buf[m++] = 4;
  buf[m++] = xoriginal[i][0];
  buf[m++] = xoriginal[i][1];
  buf[m++] = xoriginal[i][2];
  return m;
}

/* ----------------------------------------------------------------------
   unpack values from atom->extra array to restart the fix
------------------------------------------------------------------------- */

void FixDipole::unpack_restart(int nlocal, int nth)
{
  double **extra = atom->extra;

  // skip to Nth set of extra values

  int m = 0;
  for (int i = 0; i < nth; i++) m += static_cast<int> (extra[nlocal][m]);
  m++;

  xoriginal[nlocal][0] = extra[nlocal][m++];
  xoriginal[nlocal][1] = extra[nlocal][m++];
  xoriginal[nlocal][2] = extra[nlocal][m++];
}

/* ----------------------------------------------------------------------
   maxsize of any atom's restart data
------------------------------------------------------------------------- */

int FixDipole::maxsize_restart()
{
  return 4;
}

/* ----------------------------------------------------------------------
   size of atom nlocal's restart data
------------------------------------------------------------------------- */

int FixDipole::size_restart(int nlocal)
{
  return 4;
}