/* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
   Copyright (c) 2016-2023 The plumed team
   (see the PEOPLE file at the root of the distribution for a list of names)

   See http://www.plumed.org for more information.

   This file is part of plumed, version 2.

   plumed is free software: you can redistribute it and/or modify
   it under the terms of the GNU Lesser General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   plumed is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public License
   along with plumed.  If not, see <http://www.gnu.org/licenses/>.
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
#include "CLTool.h"
#include "core/CLToolRegister.h"
#include "core/PlumedMain.h"
#include "tools/Vector.h"
#include "tools/Random.h"
#include "tools/Communicator.h"
#include <string>
#include <cstdio>
#include <vector>
#include <memory>

//+PLUMEDOC TOOLS pesmd
/*
Pesmd allows one to do (biased) Langevin dynamics on a two-dimensional potential energy surface.

The energy landscape that you are moving about on is specified using in a plumed input file.
For example the following example input uses [MATHEVAL](MATHEVAL.md) to define a two dimensional potential.

```plumed
d1: DISTANCE ATOMS=1,2 COMPONENTS
ff: MATHEVAL ARG=d1.x,d1.y PERIODIC=NO FUNC=x*x+y*y
bb: BIASVALUE ARG=ff
```

Atom 1 is placed at the origin.  The x and y components on our surface are the
positions of the particle on our two dimensional energy landscape.  By calculating the
vector connecting atom 1 (the origin) to atom 2 (the position of our particle) we are thus
getting the position of the atom on the energy landscape.  This is then inserted into the function
that is calculated on the second line.  The value of this function is then used as a bias.

We can also specify a potential on a grid and look at the dynamics on this function using pesmd.  To do
so you would use the DISTANCE command as in the input above together with an [EXTERNAL](EXTERNAL.md) action
that takes d1.x and d1.y as arguments.  In this way you can use pesmd to do a dynamics on a free energy surface calculated using metadynamics
and sum_hills.  Notice also, that once we have defined our potential we can use all the biasing functions
within plumed in addition in order to do a biased dynamics on the potential energy landscape of interest.

You run a Langevin simulation using pesmd with the following command:

```plumed
plumed pesmd < input
```

The input to pesmd is specified in an input file that contains one directive per line like the one shown below.

```plumed
#TOOL=pesmd
temperature 0.722
tstep 0.005
friction 1
dimension 2
nstep 50
ipos 0.0 0.0
```

This file instructs pesmd to do 50 steps of Langevin dynamics on a 2D potential energy surface
at a temperature of 0.722

If you run the following command a description of all the directives that can be used in the input file will be output.

```plumed
plumed pesmd --help
```

*/
//+ENDPLUMEDOC

namespace PLMD {
namespace cltools {

class PesMD  : public PLMD::CLTool {
  std::string description() const override {
    return "Langevin dynamics on PLUMED energy landscape";
  }
public:
  static void registerKeywords( Keywords& keys ) {
    CLTool::registerKeywords( keys );
    keys.add("compulsory","nstep","The number of steps of dynamics you want to run");
    keys.add("compulsory","temperature","NVE","the temperature at which you wish to run the simulation in LJ units");
    keys.add("compulsory","friction","off","The friction (in LJ units) for the Langevin thermostat that is used to keep the temperature constant");
    keys.add("compulsory","tstep","0.005","the integration timestep in LJ units");
    keys.add("compulsory","dimension","the dimension of your energy landscape");
    keys.add("compulsory","plumed","plumed.dat","the name of the plumed input file containing the potential");
    keys.add("compulsory","ipos","0.0","the initial position of the system");
    keys.add("compulsory","idum","0","The random number seed");
    keys.addFlag("periodic",false,"are your input coordinates periodic");
    keys.add("optional","min","minimum value the coordinates can take for a periodic domain");
    keys.add("optional","max","maximum value the coordinates can take for a periodic domain");
  }

  explicit PesMD( const CLToolOptions& co ) :
    CLTool(co) {
    inputdata=inputType::ifile;
  }

private:

  void read_input(double& temperature,
                  double& tstep,
                  double& friction,
                  int& dim,
                  std::string& plumedin,
                  std::vector<double>& ipos,
                  int&    nstep,
                  bool&   lperiod,
                  std::vector<double>& periods,
                  int&    idum) {
    // Read everything from input file
    std::string tempstr;
    parse("temperature",tempstr);
    if( tempstr!="NVE" ) {
      Tools::convert(tempstr,temperature);
    }
    parse("tstep",tstep);
    std::string frictionstr;
    parse("friction",frictionstr);
    if( tempstr!="NVE" ) {
      if(frictionstr=="off") {
        error("pecify friction for thermostat");
      }
      Tools::convert(frictionstr,friction);
    }
    parse("plumed",plumedin);
    parse("dimension",dim);
    parse("nstep",nstep);
    parse("idum",idum);
    ipos.resize( dim );
    parseVector("ipos",ipos);

    parseFlag("periodic",lperiod);
    if( lperiod ) {
      if( dim>3 ) {
        error("can only do three dimensional periodic functions");
      }
      std::vector<double> min( dim );
      parseVector("min",min);
      std::vector<double> max( dim );
      parseVector("max",max);
      periods.resize( dim );
      for(int i=0; i<dim; ++i) {
        if( max[i]<min[i] ) {
          error("invalid periods specified max is less than min");
        }
        periods[i]=max[i]-min[i];
      }
    }
  }


public:

  int main( FILE* in, FILE* out, PLMD::Communicator& pc) override {
    std::string plumedin;
    std::vector<double> ipos;
    double temp, tstep, friction;
    bool lperiod;
    int dim, nsteps, seed;
    std::vector<double> periods;
    int plumedWantsToStop;
    Random random;

    read_input( temp, tstep, friction, dim, plumedin, ipos, nsteps, lperiod, periods, seed );
    // Setup random number generator
    random.setSeed(seed);

    // Setup box if we have periodic domain
    std::vector<double> box(9, 0.0);
    if( lperiod && dim==1 ) {
      box[0]=box[4]=box[8]=periods[0];
    } else if( lperiod && dim==2 ) {
      box[0]=periods[0];
      box[4]=box[8]=periods[1];
    } else if( lperiod && dim==3 ) {
      box[0]=periods[0];
      box[4]=periods[1];
      box[8]=periods[2];
    } else if( lperiod ) {
      error("invalid dimension for periodic potential must be 1, 2 or 3");
    }

    // Create plumed object and initialize
    auto plumed=Tools::make_unique<PLMD::PlumedMain>();
    int s=sizeof(double);
    plumed->cmd("setRealPrecision",&s);
    if(Communicator::initialized()) {
      plumed->cmd("setMPIComm",&pc.Get_comm());
    }
    int natoms=( std::floor(dim/3) +  2 );
    plumed->cmd("setNatoms",&natoms);
    plumed->cmd("setNoVirial");
    plumed->cmd("setMDEngine","pesmd");
    plumed->cmd("setTimestep",&tstep);
    plumed->cmd("setPlumedDat",plumedin.c_str());
    plumed->cmd("init");

    // Now create some fake atoms
    int nat = std::floor( dim/3 ) + 1;
    std::vector<double> masses( 1+nat, 1 );
    std::vector<Vector> velocities( nat ), positions( nat+1 ), forces( nat+1 );
    // Will set these properly eventually
    // Atom zero is fixed at origin
    positions[0].zero();
    // the for loop initializes two variable:
    // k is used only here, so its scope is limited
    for(int i=0,k=0; i<nat; ++i) {
      for(unsigned j=0; j<3; ++j) {
        if( k<dim ) {
          positions[1+i][j]=ipos[k];
        } else {
          positions[1+i][j]=0;
        }
        ++k;
      }
    }
    // And initialize the velocities
    for(int i=0; i<nat; ++i)
      for(int j=0; j<3; ++j) {
        velocities[i][j]=random.Gaussian() * std::sqrt( temp );
      }
    // And calculate the kinetic energy
    double tke=0;
    for(int i=0; i<nat; ++i) {
      for(int j=0; j<3; ++j) {
        if( 3*i+j>dim-1 ) {
          break;
        }
        tke += 0.5*velocities[i][j]*velocities[i][j];
      }
    }

    // Now call plumed to get initial forces
    {
      int istep=0;
      double zero=0;
      plumed->cmd("setStep",&istep);
      plumed->cmd("setMasses",&masses[0]);
      Tools::set_to_zero(forces);
      plumed->cmd("setForces",&forces[0][0]);
      plumed->cmd("setEnergy",&zero);
    }
    if( lperiod ) {
      plumed->cmd("setBox",&box[0]);
    }
    plumed->cmd("setPositions",&positions[0][0]);
    plumed->cmd("calc");

    double therm_eng=0;
    FILE* fp=fopen("stats.out","w+");

    for(int istep=0; istep<nsteps; ++istep) {

      if( istep%20==0 && pc.Get_rank()==0 ) {
        std::printf("Doing step %i\n",istep);
      }

      // Langevin thermostat
      double lscale=std::exp(-0.5*tstep/friction);
      double lrand=std::sqrt((1.-lscale*lscale)*temp);
      for(int j=0; j<nat; ++j) {
        for(int k=0; k<3; ++k) {
          if( 3*j+k>dim-1 ) {
            break;
          }
          therm_eng=therm_eng+0.5*velocities[j][k]*velocities[j][k];
          velocities[j][k]=lscale*velocities[j][k]+lrand*random.Gaussian();
          therm_eng=therm_eng-0.5*velocities[j][k]*velocities[0][k];
        }
      }

      // First step of velocity verlet
      for(int j=0; j<nat; ++j) {
        for(int k=0; k<3; ++k) {
          if( 3*j+k>dim-1 ) {
            break;
          }
          velocities[j][k] = velocities[j][k] + 0.5*tstep*forces[1+j][k];
          positions[1+j][k] = positions[1+j][k] + tstep*velocities[j][k];
        }
      }

      int istepplusone=istep+1;
      plumedWantsToStop=0;
      plumed->cmd("setStep",&istepplusone);
      plumed->cmd("setMasses",&masses[0]);
      Tools::set_to_zero(forces);
      plumed->cmd("setForces",&forces[0][0]);
      double fenergy=0.0;
      plumed->cmd("setEnergy",&fenergy);
      plumed->cmd("setPositions",&positions[0][0]);
      plumed->cmd("setStopFlag",&plumedWantsToStop);
      plumed->cmd("calc");
      // if(istep%2000==0) plumed->cmd("writeCheckPointFile");
      if(plumedWantsToStop) {
        nsteps=istep;
      }

      // Second step of velocity verlet
      for(int j=0; j<nat; ++j) {
        for(int k=0; k<3; ++k) {
          if( 3*j+k>dim-1 ) {
            break;
          }
          velocities[j][k] = velocities[j][k] + 0.5*tstep*forces[1+j][k];
        }
      }

      // Langevin thermostat
      lscale=std::exp(-0.5*tstep/friction);
      lrand=std::sqrt((1.-lscale*lscale)*temp);
      for(int j=0; j<nat; ++j) {
        for(int k=0; k<3; ++k) {
          if( 3*j+k>dim-1) {
            break;
          }
          therm_eng=therm_eng+0.5*velocities[j][k]*velocities[j][k];
          velocities[j][k]=lscale*velocities[j][k]+lrand*random.Gaussian();
          therm_eng=therm_eng-0.5*velocities[j][k]*velocities[j][k];
        }
      }
      // Calculate total kinetic energy
      tke=0;
      for(int i=0; i<nat; ++i) {
        for(int j=0; j<3; ++j) {
          if( 3*i+j>dim-1 ) {
            break;
          }
          tke += 0.5*velocities[i][j]*velocities[i][j];
        }
      }

      // Print everything
      // conserved = potential+1.5*ttt+therm_eng;
      if( pc.Get_rank()==0 ) {
        std::fprintf(fp,"%i %f %f %f \n", istep, istep*tstep, tke, therm_eng );
      }
    }

    fclose(fp);

    return 0;
  }
};

PLUMED_REGISTER_CLTOOL(PesMD,"pesmd")

}
}
