/* ----------------------------------------------------------------------
   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]);
	
	// get the compute ID
	char *id_comp_full = arg[2];
	char *word = strtok(id_comp_full," \0");
	
	double my_maxdisp;
	class Compute *maxdispcompute;
	
	if (strncmp(word,"c_",2) == 0)
	{
		int n = strlen(word);
		char *id = new char[n];
		strcpy(id,&word[2]);
		char copy[9];
		strncpy(copy,id,8);
		copy[8] = '\0';
		
		n = modify->find_compute(id);
		if (n < 0) error->all("Could not find maximum displacement compute ID");
		if (modify->compute[n]->scalar_flag == 0) error->all("run_PRD compute ID does not compute scalar info");
		
		maxdispcompute = modify->compute[n];
		if (!maxdispcompute) error->all("Could not access compute");
		my_maxdisp = 0.0;	// Initialize the max displacement to 0
	}
	else error->universe_all("Illegal compute ID in run_PRD command - check arguments");
	
	
	// 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 fix, access my_maxdisp
	// maxdisp[i] = maximum displacement of world i		
	modify->clearstep_compute();
	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);
	
	// transition flag - initialize on all procs
	int trans_world = -1;
	
	// 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
		modify->clearstep_compute();
		my_maxdisp = maxdispcompute->compute_scalar();
		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 (done at every check)
		if (me_universe == 0) print_status();
		
		// Perform check against transition criteria (each master process does the check)
		if (me == 0)
		{
			// loop through the worlds to see if which ones have exceeded the threshold, take the maximum value as the one to contiue with
			double maxdispcheck = 0.0;
			for (int i=0; i < nworlds; i++)
			{
				if ((maxdisp[i] >= trans_dist) && (maxdisp[i] >= maxdispcheck))
				{
					// if transition detected, find world with highest displacement in case of multiple transitions
					maxdispcheck = maxdisp[i];
					trans_world = i;
				}
			}
		}
		// broadcast the transworld to all procs in world (now all procs in all worlds should have same value)
		MPI_Bcast(&trans_world,1,MPI_INT,0,world);	

		// if a transition has been found, hold non-transitioned states. 
		// run 1 more check phase on the transitioned state to see if the transition is stable.
		// If no transition detected, carry on as usual.
		if ((me_universe == 0) && (trans_world > -1))
		{
			if (universe->uscreen) fprintf(universe->uscreen,"Possible Transition Detected in World %d\n", trans_world);
			if (universe->ulogfile) fprintf(universe->ulogfile,"Possible Transition Detected in World %d\n", trans_world);
		}
		
		if ((trans_world > -1) && (iworld == trans_world))
		{
			// run for nevery timesteps before check
			update->integrate->iterate(nevery);
			
			// Compute the maximum displacement and communicate it
			// alert the compute to be re-computed
			modify->clearstep_compute();
			my_maxdisp = maxdispcompute->compute_scalar();
			
			if (my_maxdisp >= trans_dist)
			{
				// Stable Transition
				if (universe->uscreen && me == 0) fprintf(universe->uscreen,"Stable Transition Detected in World %d\n", trans_world);
				if (universe->ulogfile && me == 0) fprintf(universe->ulogfile,"Stable Transition Detected in World %d\n", trans_world);
				
				//MPI_Bcast(&trans_world,1,MPI_INT,world2root[trans_world],universe->uworld);
				//MPI_Bcast(&trans_world,1,MPI_INT,0,world); // have to see if needed... probably not
			}
			else
			{
				// Transition Unstable, reset the transition flag and continue
				int trans_world_temp = trans_world;
				trans_world = -1;
				MPI_Bcast(&trans_world,1,MPI_INT,0,world);
				MPI_Bcast(&trans_world,1,MPI_INT,world2root[trans_world_temp],universe->uworld);
			}
			
		}
		// Use a barrier to hold jobs - will wait for the check to be done if there was a transition, otherwise,
		// will just wait until all procs get here to begin the next check phase... easier than series of if statements
		MPI_Barrier(universe->uworld); // have to check that this stops *all* processes, not just some
		
		// if there was a stable transition on any process, all processes on all worlds break out of the loop
		if ( trans_world > -1)
		{
			if (universe->ulogfile && me_universe == 0) fprintf(universe->ulogfile,"Stable Transition Detected in World %d\n", trans_world);
			break;
		}
		
	}

	// convert simulation nsteps to the actual number of steps each world ran
	update->nsteps = (update->ntimestep - update->beginstep);
	
	// output the total number of steps run (and total time integrated) in this phase
	if (me == 0)
	{	
		int totsteps = 0;
		MPI_Allreduce(&(update->nsteps),&totsteps,1,MPI_INT,MPI_SUM,roots);
		
		if (universe->ulogfile && me_universe == 0)
		{	
			double tottime = totsteps*(update->dt);
			fprintf(universe->ulogfile,"Total Simulation Steps: %d Simulated Time: %lf\n", totsteps, tottime);
		}
	}
	
	// Output the final thermo output a restart and a dump file world with the stable transition 
	if ((trans_world > -1) && (iworld == trans_world))
	{
		for (int idump = 0; idump < output->ndump; idump++)
			output->next_dump[idump] = update->ntimestep;
		output->next_dump_any = update->ntimestep;
		if (output->restart_every) output->next_restart = update->ntimestep;
		output->next_thermo = update->ntimestep;
		
		modify->addstep_compute_all(update->ntimestep);
		output->write(update->ntimestep);
	}
	
	timer->barrier_stop(TIME_LOOP);
	
	// finish up, cleanup
	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");
	}
}