/* ----------------------------------------------------------------------
   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.
------------------------------------------------------------------------- */

/* ---------------------------------------------------------------------- 
fix setvelbyforce 
 written by Marco Kalweit
 Cranfield, 23/01/2006
------------------------------------------------------------------------- */

#include "string.h"
#include "stdlib.h"
#include "fix_setvelbyforce.h"
#include "atom.h"
#include "force.h"
#include "update.h"
#include "respa.h"
#include "error.h"
#include "mpi.h"
#include "stdio.h"
#include "math.h"

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

FixSetVelbyForce::FixSetVelbyForce(int narg, char **arg) : Fix(narg, arg)
{
    int me;
    MPI_Comm_rank(world,&me);

    if (narg < 7) error->all("Illegal fix setvelbyforce command");
    
    flagx = flagy = flagz = 1;
    memset(&print, 0, sizeof print);

    if (strcmp(arg[3],"NULL") == 0) flagx = 0;
    else vx = atof(arg[3]);
    if (strcmp(arg[4],"NULL") == 0) flagy = 0;
    else vy = atof(arg[4]);
    if (strcmp(arg[5],"NULL") == 0) flagz = 0;
    else vz = atof(arg[5]);
    
    a = atof(arg[6]);
    
    if(me == 0)
    if(narg > 7)
    {
	if(narg < 9) error->all("Illegal fix setvelbyforce command");
	print.pFileOut = fopen(arg[7], "a+");
	if(!print.pFileOut) error->all("Unable to open/create file for writing");
	print.ntsout = atoi(arg[8]);

	//print header
	if(force->dimension == 3)
	    fprintf(print.pFileOut, "%10s %20s %20s %20s %20s %20s %20s %20s\n", 
			"time step", "time", "v_x", "v_y", "v_z", "f_total_x", "f_total_y", "f_total_z");
	else
	    fprintf(print.pFileOut, "%10s %20s %20s %20s %20s %20s\n", 
		"time step", "time", "v_x", "v_y", "f_total_x", "f_total_y");
    }
}

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

FixSetVelbyForce::~FixSetVelbyForce()
{
    if(print.pFileOut) fclose(print.pFileOut);
}

/* ---------------------------------------------------------------------- */
int FixSetVelbyForce::setmask()
{
  int mask = 0;
  mask |= POST_FORCE;
  return mask;
}


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

void FixSetVelbyForce::setup()
{
  if (strcmp(update->integrate_style,"verlet") == 0) post_force(1);
}

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

void FixSetVelbyForce::post_force(int vflag)
{
    int me;
    int *type = atom->type;
    double *mass = atom->mass;
    double **v = atom->v;
    double **f = atom->f;
    int *mask = atom->mask;
    int nlocal = atom->nlocal;
    double massl; // local mass of the group
    double masst; // global mass of the group
    double mcl[3]; // current local momentum vector
    double mct[3]; // current global velocity vector
    double vct[3]; // current global velocity vector
    double dv[3]; // velocity difference
    double da[3]; // acceleration
    double ftl[3]; // total force local
    double ftg[3]; // total force global
    
    MPI_Comm_rank(world,&me);
    ftl[0] = ftl[1] = ftl[2] = 0.0;
    ftg[0] = ftg[1] = ftg[2] = 0.0;

    // calc the current local velocity vector
    massl = 0.0;
    mcl[0] = mcl[1] = mcl[2] = 0.0;
    for(int i = 0; i < nlocal; i++)
    if(mask[i] & groupbit)
    {  
	double massi = mass[type[i]];
	massl += massi;
	mcl[0] += massi*v[i][0];
	mcl[1] += massi*v[i][1];
	mcl[2] += massi*v[i][2];
    }
    
    // sum up and redistribute the total momentum of the group
    masst = 0.0;
    mct[0] = mct[1] = mct[2] = 0.0;
    MPI_Allreduce(&massl,&masst,1,MPI_DOUBLE,MPI_SUM,world);
    MPI_Allreduce(mcl,mct,3,MPI_DOUBLE,MPI_SUM,world);
    
    // calc the total velocity
    vct[0] = mct[0] / masst;
    vct[1] = mct[1] / masst;
    vct[2] = mct[2] / masst;

    // calc the dv
    dv[0] = (flagx) ?vx - vct[0] :0.0;
    dv[1] = (flagy) ?vx - vct[1] :0.0;
    dv[2] = (flagz) ?vx - vct[2] :0.0;

    // normalisation factor for dv
    double fnorm = sqrt(dv[0]*dv[0] + dv[1]*dv[1] + dv[2]*dv[2]);
    
    // prevent to overshoot due to the discrete time step
    double ac = (a*update->dt > fnorm)? fnorm/update->dt : a;
     
    if(fnorm > 0.0)
    {
	da[0] = (dv[0] / fnorm)*ac;
	da[1] = (dv[1] / fnorm)*ac;
	da[2] = (dv[2] / fnorm)*ac;

	// now apply the force
	for(int i = 0; i < nlocal; i++)
	if(mask[i] & groupbit)
	{
	    double massi = mass[type[i]];
	    double dma2f = massi/force->ftm2v;
	    if (flagx){ f[i][0] += dma2f*da[0]; ftl[0] += dma2f*da[0];}
	    if (flagy){ f[i][1] += dma2f*da[1]; ftl[1] += dma2f*da[1];}
	    if (flagz){ f[i][2] += dma2f*da[2]; ftl[2] += dma2f*da[2];}
	}
    }
    
    MPI_Reduce(ftl,ftg,3,MPI_DOUBLE,MPI_SUM,0,world);

    if(print.pFileOut)
    {
	print.nSum++;
	print.v[0] += vct[0];
	print.v[1] += vct[1];
	print.v[2] += vct[2];
	print.f[0] += ftg[0];
	print.f[1] += ftg[1];
	print.f[2] += ftg[2];

	if(update->ntimestep % print.ntsout == 0)
	{
	    print.v[0] /= (double)print.nSum;
	    print.v[1] /= (double)print.nSum;
	    print.v[2] /= (double)print.nSum;
	    print.f[0] /= (double)print.nSum;
	    print.f[1] /= (double)print.nSum;
	    print.f[2] /= (double)print.nSum;
	    print.nSum = 0;

	    if(force->dimension == 3)
		fprintf(print.pFileOut, "%10d %20g %20g %20g %20g %20g %20g %20g\n", 
			update->ntimestep, update->ntimestep*update->dt,
			print.v[0], print.v[1], print.v[2], print.f[0], print.f[1], print.f[2]);
	    else
		fprintf(print.pFileOut, "%10d %20g %20g %20g %20g %20g\n", 
			update->ntimestep, update->ntimestep*update->dt,
			print.v[0], print.v[1], print.f[0], print.f[1]);
	}
    }
}