/* ----------------------------------------------------------------------
   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 "mpi.h"
#include "string.h"
#include "compute_heatflux3.h"
#include "atom.h"
#include "force.h"
#include "modify.h"
#include "comm.h"
#include "memory.h"
#include "group.h"
#include "error.h"
#include "update.h"
#include "neighbor.h"
#include "math.h"
#include "pair.h"
#include "universe.h"

#include <iostream>
#include <iomanip>
//outf "myfile";

using namespace LAMMPS_NS;
using namespace std;

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

ComputeHeatFlux3::ComputeHeatFlux3(LAMMPS *lmp, int narg, char **arg) : 
  Compute(lmp, narg, arg)
{
  if (narg != 3) error->all("Illegal compute heatflux command");

  scalar_flag = 1;
  vector_flag = 1;
  size_vector = 8;   // ? size of vector returned by compute
  
  comm_reverse = 1;
  neigh_half_once = 1;

  extensive = 0;
  energy1 = NULL;
  force1 = NULL;
  force2 = NULL;
  nmax =0;
  vector = new double[8];   // ? size of vector returned by compute
}

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

ComputeHeatFlux3::~ComputeHeatFlux3()
{ 
  memory->sfree(energy1);
  memory->sfree(force1);
  memory->sfree(force2);
  delete [] vector;
}

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

void ComputeHeatFlux3::init()
{

}

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

void ComputeHeatFlux3::compute_vector()
{
  int *neighs;
  int i,j,k,m,n,itype,jtype,numneigh;
  double **v = atom->v; 
  double **x = atom->x; 

  if (atom->nmax > nmax) {
    memory->sfree(energy1);
    memory->sfree(force1);
    memory->sfree(force2);
    nmax = atom->nmax;
    energy1 = (double *)
      memory->smalloc(nmax*sizeof(double),"compute/epair/atom:energy");
    scalar_atom = energy1;
    force1 = (double *)
      memory->smalloc(nmax*sizeof(double),"compute/epair/atom:energy");
    scalar_atom = force1;
    force2 = (double *)
      memory->smalloc(nmax*sizeof(double),"compute/epair/atom:energy");
    scalar_atom = force2;


  }


  if (force->newton_pair) n = atom->nlocal + atom->nghost;
  else n = atom->nlocal;
  for (i = 0; i < n; i++) {
    energy1[i] = 0.0;
    force1[i] = 0.0;
  }
  
  if (!neighbor->half_every) neighbor->build_half();
  
  double *special_coul = force->special_coul;
  double *special_lj = force->special_lj;
  double **cutsq = force->pair->cutsq;
  double *mass = atom->mass;
  double *rmass = atom->rmass;
  int *type = atom->type;
  int *mask = atom->mask;
  
  int nlocal = atom->nlocal;
  double mvv2e = force->mvv2e;
  int nall = atom->nlocal + atom->nghost;
  double xtmp,ytmp,ztmp,delx,dely,delz,rsq;
  double f,factor_coul,factor_lj,e; 
  
  Pair::One one;
  double ke;

  double t[8];

  for (i=0; i<8;i++) 
    t[i] = 0.0;

  for (i = 0; i < nlocal; i++) {
    xtmp = x[i][0];
    ytmp = x[i][1];
    ztmp = x[i][2];
    itype = type[i];
    neighs = neighbor->firstneigh[i];
    numneigh = neighbor->numneigh[i];
    for (k = 0; k < numneigh; k++) {
      j = neighs[k];
      
      if (j < nall) factor_coul = factor_lj = 1.0;
      else {
	factor_coul = special_coul[j/nall];
	factor_lj = special_lj[j/nall];
	j %= nall;
      }
      
      delx = xtmp - x[j][0];
      dely = ytmp - x[j][1];
      delz = ztmp - x[j][2];
      rsq = delx*delx + dely*dely + delz*delz;
      jtype = type[j];
      if (rsq < cutsq[itype][jtype]) {
	// Neighbor lists are constructed such that each pair is only visited once
	force->pair->single(i,j,itype,jtype,rsq,factor_coul,factor_lj,1,one);	
	e = one.fforce;
        // Allocate the total interaction energy to each of i and j
	// Later need to divide by 2 to compensate for this double-counting
	energy1[i]+=one.eng_coul + one.eng_vdwl;
	energy1[j]+=one.eng_coul + one.eng_vdwl;
	force1[i]+=e*delx;
	force1[j]+=-e*delx;

      }		
    }
  }
 
  
  // If j>nlocal arises in the loop above, force is allocated to ghost atoms
  // Now transfer this from the ghosts to their corresponding "source atoms"
  if (force->newton_pair)
    comm->reverse_comm_compute(this);

  // For the sake of clarity only print out the first few ghosts
  std::cout << "\n";
  for (i = 0; i < 12; i++) {
    std::cout << "atom " << i << setiosflags(ios::fixed) << setprecision(10);
    std::cout << "   x= " << x[i][0] << "   y= " << x[i][1] << "   z= "  << x[i][2];
    std::cout << "   vx= " << v[i][0] << "   vy= " << v[i][1] <<"   vz= " << v[i][2] << "\n";
  }


  for (i = 0; i < nlocal; i++)
        { 
	  // Compensation for the double-allocation mentioned above
      energy1[i] *= 0.500000;
      if (mass)
	{
	  if (mask[i] & groupbit)
	    ke = 0.5 * mvv2e * mass[type[i]] *
	      (v[i][0]*v[i][0] + v[i][1]*v[i][1] + v[i][2]*v[i][2]);
	}
      else
	{
	  if (mask[i] & groupbit)
	    ke = 0.5 * mvv2e * rmass[i] *
	      (v[i][0]*v[i][0] + v[i][1]*v[i][1] + v[i][2]*v[i][2]);
	}
      t[1]+=ke;
      t[2]+=energy1[i]; 
      } 
  MPI_Allreduce(t,vector,8,MPI_DOUBLE,MPI_SUM,world);
}


int ComputeHeatFlux3::pack_reverse_comm(int n, int first, double *buf)
{
  int i,m,last;
  
  m = 0;
  last = first + n;
  for (i = first; i < last; i++) buf[m++] = energy1[i];
  return 1;
}

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

void ComputeHeatFlux3::unpack_reverse_comm(int n, int *list, double *buf)
{
  int i,j,m;
  m = 0;
  for (i = 0; i < n; i++) {
    j = list[i];
    energy1[j] += buf[m++];
  }
}