/* ----------------------------------------------------------------------
   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 "math.h"
#include "stdlib.h"
#include "string.h"
#include "fix_irradiation.h"
#include "atom.h"
#include "atom_vec.h"
#include "update.h"
#include "domain.h"
#include "region.h"
#include "comm.h"
#include "group.h"
#include "force.h"
#include "random_park.h"
#include "random_mars.h"
#include "memory.h"
#include "error.h"

using namespace LAMMPS_NS;

#define MIN(A,B) ((A) < (B)) ? (A) : (B)
#define MAX(A,B) ((A) > (B)) ? (A) : (B)

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

FixIrradiation::FixIrradiation(LAMMPS *lmp, int narg, char **arg) :
  Fix(lmp, narg, arg)
{
  if (narg != 9) error->all("Illegal fix irradiation command");

  scalar_flag = 1;
  global_freq = 1;
  extscalar = 0;

  iregion = domain->find_region(arg[3]);
  oregion = domain->find_region(arg[4]);
  
  nevery = atoi(arg[5]);
  nflux = atoi(arg[6]);
  electro_energy = atof(arg[7]);
  
  int seed = atoi(arg[8]);

  if (iregion == -1) error->all("Fix irradiation region ID1 does not exist");
  if (oregion == -1) error->all("Fix irradiation region ID2 does not exist");
  if (nevery <= 0 || nflux <= 0) error->all("Illegal fix irradiation command");
  if (electro_energy < 0) error->all("Illegal fix irradiation command");
  if (seed <= 0) error->all("Illegal fix irradiation command");

  // random number generator, same for all procs

  random = new RanPark(lmp,seed);

  // set up reneighboring

  force_reneighbor = 1;
  next_reneighbor = (update->ntimestep/nevery)*nevery + nevery;
  ndeleted = 0;

  nmax = 0;
  ilist = NULL;
  imark = NULL;
  olist = NULL;
  omark = NULL;
}

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

FixIrradiation::~FixIrradiation()
{
  delete random;
  memory->sfree(ilist);
  memory->sfree(imark);
  memory->sfree(olist);
  memory->sfree(omark);
}

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

int FixIrradiation::setmask()
{
  int mask = 0;
  mask |= PRE_EXCHANGE;
  return mask;
}

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

void FixIrradiation::init()
{
  // check that no deletable atoms are in atom->firstgroup
  // deleting such an atom would not leave firstgroup atoms first

  if (atom->firstgroup >= 0) {
    int *mask = atom->mask;
    int nlocal = atom->nlocal;
    int firstgroupbit = group->bitmask[atom->firstgroup];

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

    int flagall;
    MPI_Allreduce(&flag,&flagall,1,MPI_INT,MPI_SUM,world);

    if (flagall)
      error->all("Cannot irradiate atoms in atom_modify first group");
  }
}

/* ----------------------------------------------------------------------
   perform irradiation and particle deletion
   done before exchange, borders, reneighbor
   so that ghost atoms and neighbor lists will be correct
------------------------------------------------------------------------- */

void FixIrradiation::pre_exchange()
{
  int i,iwhichglobal,iwhichlocal;

  double **x = atom->x;
  double *mass = atom->mass;
  int *mask = atom->mask;
  int *type = atom->type;
  int nlocal = atom->nlocal;
  
  if (update->ntimestep == next_reneighbor) {

  // grow list and mark arrays if necessary

  if (atom->nlocal > nmax) {
    memory->sfree(ilist);
    memory->sfree(imark);
    memory->sfree(olist);
    memory->sfree(omark);
    nmax = atom->nmax;
    ilist = (int *) memory->smalloc(nmax*sizeof(int),"irradiation:ilist");
    imark = (int *) memory->smalloc(nmax*sizeof(int),"irradiation:imark");
    olist = (int *) memory->smalloc(nmax*sizeof(int),"irradiation:olist");
    omark = (int *) memory->smalloc(nmax*sizeof(int),"irradiation:omark");
  }

  // nall = # of atoms in irradiation region
  // nbefore = # on procs before me

  int nicount = 0;
  for (i = 0; i < nlocal; i++) {
    if (mask[i] & groupbit) {
      if (domain->regions[iregion]->match(x[i][0],x[i][1],x[i][2])) {
	     ilist[nicount++] = i;
	   }
	}
  }

  int nall,nbefore;
  MPI_Allreduce(&nicount,&nall,1,MPI_INT,MPI_SUM,world);
  MPI_Scan(&nicount,&nbefore,1,MPI_INT,MPI_SUM,world);
  nbefore -= nicount;
  
  // nirrad = total number of atoms to be irradiated
  // choose atoms randomly across all procs and mark them for irradiation
  // shrink local list of candidates as my atoms get marked for irradiation

  int nirrad = MIN(nflux,nall);
  
  for (i = 0; i < nlocal; i++) imark[i] = 0;
  
  // Irradiation
  //   
  
  for (i = 0; i < nirrad; i++) {
    iwhichglobal = static_cast<int> (nall*random->uniform());
    if (iwhichglobal < nbefore) nbefore--;
    else if (iwhichglobal < nbefore + nicount) {
      iwhichlocal = iwhichglobal - nbefore;
      imark[ilist[iwhichlocal]] = 1;
      ilist[iwhichlocal] = ilist[nicount-1];
      nicount--;
    }
    nall--;
  }

  // 
  
  double c = 2997924.58;
  double me = 9.11e-31/1.66e-27;
  
  double sita, alpha;
  double dx, dy, dr, dr2;
  double vx, vy, vz;
  double energy_transfer;
  
  double **v = atom->v;
  
  for (i = 0; i < nlocal; i++) {
  	if (imark[i]) {
  		dx = random->uniform();
  		dy = random->uniform();
  		
  		dr2 = dx*dx + dy*dy;
  		dr = sqrt(dr2);
  		
  		if (dr < 1.0) { 			
  			alpha = atan(dy/dx);
  			sita = asin(dr);
  			sita = cos(sita);
  			
  			energy_transfer = 2.0*electro_energy/force->mvv2e*(electro_energy/force->mvv2e + 2.0*me*c*c)/(mass[type[i]]*c*c);
  			energy_transfer = energy_transfer*sita*sita;
  			energy_transfer = 2.0*energy_transfer/mass[type[i]];
  			energy_transfer = sqrt(energy_transfer);
  			
  			vx = energy_transfer*dr*cos(alpha);
  			vy = energy_transfer*dr*sin(alpha);
  			vz = -energy_transfer*sita;
  			
  			if (mask[i] & groupbit) {
  				v[i][0] += vx;
  				v[i][1] += vy;
 				v[i][2] += vz;
         }
  	   }
   } 
  }
  
  next_reneighbor = update->ntimestep + nevery;
  
}
  
  // Find out atoms out of oregion for deletion
  //
  
  int nocount = 0;
  
  for (i = 0; i < nlocal; i++) {
    if (mask[i] & groupbit) {
      if (!domain->regions[oregion]->match(x[i][0],x[i][1],x[i][2])) {
	     olist[nocount++] = i;
	     omark[i] = 1;
	   }
	   else {
	     omark[i] = 0;
	   }
	 }
  }

  // nwhack = total number of atoms to delete
  	
  int nwhack;
  MPI_Allreduce(&nocount,&nwhack,1,MPI_INT,MPI_SUM,world);	
	  
  // delete my marked atoms
  // loop in reverse order to avoid copying marked atoms
  
  AtomVec *avec = atom->avec;
  
  for (i = nlocal-1; i >= 0; i--) {
    if (omark[i]) {
      avec->copy(atom->nlocal-1,i);
      atom->nlocal--;
    }
  }

  // reset global natoms
  // if global map exists, reset it now instead of waiting for comm
  // since deleting atoms messes up ghosts

  atom->natoms -= nwhack;
  if (nwhack && atom->map_style) {
    atom->nghost = 0;
    atom->map_init();
    atom->map_set();
  }

  // statistics
  
  ndeleted += nwhack;
}

/* ----------------------------------------------------------------------
   return number of deleted particles
------------------------------------------------------------------------- */

double FixIrradiation::compute_scalar()
{
  return 1.0*ndeleted;
}

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

double FixIrradiation::memory_usage()
{
  double bytes = 4*nmax * sizeof(int);
  return bytes;
}