/* ----------------------------------------------------------------------
   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: David Farrell (Northwestern University)
------------------------------------------------------------------------- */

#include "math.h"
#include "stdlib.h"
#include "string.h"
#include "run_PRD.h"
#include "universe.h"
#include "domain.h"
#include "atom.h"
#include "update.h"
#include "integrate.h"
#include "modify.h"
#include "compute.h"
#include "force.h"
#include "output.h"
#include "thermo.h"
#include "fix.h"
#include "random_park.h"
#include "finish.h"
#include "timer.h"
#include "memory.h"
#include "error.h"

using namespace LAMMPS_NS;

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

RunPRD::RunPRD(LAMMPS *lmp) : Pointers(lmp) {}

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

RunPRD::~RunPRD()
{
  MPI_Comm_free(&roots);
  delete [] world2root;
}

/* ----------------------------------------------------------------------
   perform PRD evolution and checks with inter-world communications (for checks)
 
 Syntax:
 
 run_PRD N M compute trans_dist
 
 N - total number of simulation steps to run before exiting
 M - check for a transition every M steps
 compute - compute ID that determines maximum displacement
 trans_dist - the displacement threshold criteria for a 'transition'
 
------------------------------------------------------------------------- */

void RunPRD::command(int narg, char **arg)
{
	if (universe->nworlds == 1) 
		error->all("Must have more than one processor partition to run Parallel Replica Dynamics");
	if (domain->box_exist == 0) 
		error->all("PRD command before simulation box is defined");
	if (narg != 4 ) error->universe_all("Illegal PRD command - check arguments");
	
	int nsteps = atoi(arg[0]);
	nevery = atoi(arg[1]);
	id_comp = arg[2];
	
	// some error checking on compute ID - have to access the value after integration has started
	int n = modify->find_compute(id_comp);
	if (n < 0) error->all("Could not find maximum displacement compute ID");
	class Compute *maxdispcompute;
	maxdispcompute = modify->compute[n];
	if (!maxdispcompute) error->all("Could not access compute");
	double my_maxdisp = 0.0;	// Initialize the max displacement to 0
	
	// Distance criteria for transitions
	double trans_dist = atof(arg[3]);

    // Check frequency must evenly divide total # of timesteps
	int nchecks = 0;
	if (nevery == 0) error->universe_all("Invalid frequency in run_PRD command");
		nchecks = nsteps/nevery;
	if (nchecks*nevery != nsteps) 
		error->universe_all("Non integer # of checks in run_PRD command");

	// setup for long PRD run
	update->whichflag = 0;
	update->nsteps = nsteps;
	update->beginstep = update->firststep = update->ntimestep;
	update->endstep = update->laststep = update->firststep + nsteps;

	lmp->init();

	// local storage
	me_universe = universe->me;
	MPI_Comm_rank(world,&me);
	nworlds = universe->nworlds;
	iworld = universe->iworld;

	// create MPI communicator for root proc from each world
	int color;
	if (me == 0) color = 0;
	else color = 1;
	MPI_Comm_split(universe->uworld,color,0,&roots);

	// world2root[i] = global proc that is root proc of world i
	world2root = new int[nworlds];
	if (me == 0) 
		MPI_Allgather(&me_universe,1,MPI_INT,world2root,1,MPI_INT,roots);
	MPI_Bcast(world2root,nworlds,MPI_INT,0,world);

	// signal to update displacement, access my_maxdisp
	// maxdisp[i] = maximum displacement of world i
	maxdispcompute->addstep(update->ntimestep + nevery);	// signal to update the compute at next check
	my_maxdisp = maxdispcompute->compute_scalar();
	maxdisp = new double[nworlds];
	if (me == 0) 
		MPI_Allgather(&my_maxdisp,1,MPI_DOUBLE,maxdisp,1,MPI_DOUBLE,roots);
	MPI_Bcast(maxdisp,nworlds,MPI_DOUBLE,0,world);
	
	// setup PRD runs
	MPI_Status status;
	
	if (me_universe == 0 && universe->uscreen) fprintf(universe->uscreen,"Setting up Parallel Replica Dynamics Evolve and Check run ...\n");
	if (me_universe == 0 && universe->ulogfile) fprintf(universe->ulogfile,"Setting up Parallel Replica Dynamics Evolve and Check run ...\n");
	
	update->integrate->setup();

	timer->barrier_start(TIME_LOOP);	// wait for all processes to get here

	if (me_universe == 0)
	{
		if (universe->uscreen) fprintf(universe->uscreen,"Step D1 D2 ...\n");
		if (universe->ulogfile) fprintf(universe->ulogfile,"Step D1 D2 ...\n");
		print_status();
	}

	// begin the loop over the checks
	for (int ichecks = 0; ichecks < nchecks; ichecks++)
	{
		// run for nevery timesteps before check
		update->integrate->iterate(nevery);

		// Compute the maximum displacement and communicate it
		// alert the compute to be re-computed at each check
		my_maxdisp = maxdispcompute->compute_scalar();
		maxdispcompute->addstep(update->ntimestep + nevery); // -- this doesn't seem to do what I want...
		if (me == 0) 
			MPI_Allgather(&my_maxdisp,1,MPI_DOUBLE,maxdisp,1,MPI_DOUBLE,roots);
		MPI_Bcast(maxdisp,nworlds,MPI_DOUBLE,0,world);
		
		// status report
		if (me_universe == 0) print_status();
		
	}

	timer->barrier_stop(TIME_LOOP);

	update->integrate->cleanup();
	
	Finish finish(lmp);
	finish.end(1);
	
	update->whichflag = -1;
	update->firststep = update->laststep = 0;
	update->beginstep = update->endstep = 0;

}

/* ----------------------------------------------------------------------
   proc 0 prints current status
------------------------------------------------------------------------- */

void RunPRD::print_status()
{
	if (universe->uscreen)
	{
		fprintf(universe->uscreen,"%d ",update->ntimestep);
		for (int i = 0; i < nworlds; i++) 
			fprintf(universe->uscreen,"%f ",maxdisp[i]);
		fprintf(universe->uscreen,"\n");
	}
	if (universe->ulogfile)
	{
		fprintf(universe->ulogfile,"%d ",update->ntimestep);
		for (int i = 0; i < nworlds; i++)
			fprintf(universe->ulogfile,"%f ",maxdisp[i]);
		fprintf(universe->ulogfile,"\n");
	}
}