/*************************************************************************************

Grid physics library, www.github.com/paboyle/Grid

Source file: ./lib/qcd/hmc/GenericHmcRunner.h

Copyright (C) 2015

Author: paboyle <paboyle@ph.ed.ac.uk>
Author: Guido Cossu <guido.cossu@ed.ac.uk>

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

This program 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 General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

See the full license in the file "LICENSE" in the top level distribution
directory
*************************************************************************************/
/*  END LEGAL */
#ifndef GENERIC_HMC_RUNNER
#define GENERIC_HMC_RUNNER

namespace Grid {
namespace QCD {

// Virtual Class for HMC specific for gauge theories
// implement a specific theory by defining the BuildTheAction
template <class Implementation,
          class RepresentationsPolicy = NoHirep>
class BinaryHmcRunnerTemplate {
 public:
  INHERIT_FIELD_TYPES(Implementation);
  typedef Implementation ImplPolicy;

  enum StartType_t { ColdStart, HotStart, TepidStart, CheckpointStart };

  ActionSet<Field, RepresentationsPolicy> TheAction;
  
  // A vector of HmcObservable 
  // that can be injected from outside  
  std::vector< HmcObservable<typename Implementation::Field>* > ObservablesList;

  IntegratorParameters MDparameters;

  GridCartesian *UGrid;
  GridCartesian *FGrid;
  GridRedBlackCartesian *UrbGrid;
  GridRedBlackCartesian *FrbGrid;

  std::vector<int> SerialSeed;
  std::vector<int> ParallelSeed;

  void RNGSeeds(std::vector<int> S, std::vector<int> P){
    SerialSeed = S;
    ParallelSeed = P;
  }

  virtual void BuildTheAction(int argc, char **argv) = 0;  // necessary?

  // A couple of wrapper classes
  template <class IOCheckpointer>
  void Run(int argc, char **argv, IOCheckpointer &Checkpoint) {
    NoSmearing<Implementation> S;
    Runner(argc, argv, Checkpoint, S);
  }

  template <class IOCheckpointer, class SmearingPolicy>
  void Run(int argc, char **argv, IOCheckpointer &CP, SmearingPolicy &S) {
    Runner(argc, argv, CP, S);
  }
  //////////////////////////////

  template <class SmearingPolicy, class IOCheckpointer >
  void Runner(int argc, char **argv, IOCheckpointer &Checkpoint,
           SmearingPolicy &Smearing) {
    StartType_t StartType = HotStart;

    std::string arg;

    if (GridCmdOptionExists(argv, argv + argc, "--StartType")) {
      arg = GridCmdOptionPayload(argv, argv + argc, "--StartType");
      if (arg == "HotStart") {
        StartType = HotStart;
      } else if (arg == "ColdStart") {
        StartType = ColdStart;
      } else if (arg == "TepidStart") {
        StartType = TepidStart;
      } else if (arg == "CheckpointStart") {
        StartType = CheckpointStart;
      } else {
        std::cout << GridLogError << "Unrecognized option in --StartType\n";
        std::cout
            << GridLogError
            << "Valid [HotStart, ColdStart, TepidStart, CheckpointStart]\n";
        assert(0);
      }
    }

    int StartTraj = 0;
    if (GridCmdOptionExists(argv, argv + argc, "--StartTrajectory")) {
      arg = GridCmdOptionPayload(argv, argv + argc, "--StartTrajectory");
      std::vector<int> ivec(0);
      GridCmdOptionIntVector(arg, ivec);
      StartTraj = ivec[0];
    }

    int NumTraj = 1;
    if (GridCmdOptionExists(argv, argv + argc, "--Trajectories")) {
      arg = GridCmdOptionPayload(argv, argv + argc, "--Trajectories");
      std::vector<int> ivec(0);
      GridCmdOptionIntVector(arg, ivec);
      NumTraj = ivec[0];
    }

    int NumThermalizations = 10;
    if (GridCmdOptionExists(argv, argv + argc, "--Thermalizations")) {
      arg = GridCmdOptionPayload(argv, argv + argc, "--Thermalizations");
      std::vector<int> ivec(0);
      GridCmdOptionIntVector(arg, ivec);
      NumThermalizations = ivec[0];
    }

    GridSerialRNG sRNG;
    GridParallelRNG pRNG(UGrid);
    Field U(UGrid);

    typedef MinimumNorm2<Implementation, SmearingPolicy,
                         RepresentationsPolicy>
        IntegratorType;  // change here to change the algorithm
    IntegratorType MDynamics(UGrid, MDparameters, TheAction, Smearing);

    HMCparameters HMCpar;
    HMCpar.StartTrajectory = StartTraj;
    HMCpar.Trajectories = NumTraj;
    HMCpar.NoMetropolisUntil = NumThermalizations;

    if (StartType == HotStart) {
      // Hot start
      HMCpar.MetropolisTest = true;
      sRNG.SeedFixedIntegers(SerialSeed);
      pRNG.SeedFixedIntegers(ParallelSeed);
      Implementation::HotConfiguration(pRNG, U);
    } else if (StartType == ColdStart) {
      // Cold start
      HMCpar.MetropolisTest = true;
      sRNG.SeedFixedIntegers(SerialSeed);
      pRNG.SeedFixedIntegers(ParallelSeed);
      Implementation::ColdConfiguration(pRNG, U);
    } else if (StartType == TepidStart) {
      // Tepid start
      HMCpar.MetropolisTest = true;
      sRNG.SeedFixedIntegers(SerialSeed);
      pRNG.SeedFixedIntegers(ParallelSeed);
      Implementation::TepidConfiguration(pRNG, U);
    } else if (StartType == CheckpointStart) {
      HMCpar.MetropolisTest = true;
      // CheckpointRestart
      Checkpoint.CheckpointRestore(StartTraj, U, sRNG, pRNG);
    }

    Smearing.set_Field(U);

    HybridMonteCarlo<IntegratorType> HMC(HMCpar, MDynamics, sRNG, pRNG, U);

    for (int obs = 0; obs < ObservablesList.size(); obs++)
      HMC.AddObservable(ObservablesList[obs]); 

    // Run it
    HMC.evolve();
  }
};

// These are for gauge fields
typedef BinaryHmcRunnerTemplate<PeriodicGimplR> BinaryHmcRunner;
typedef BinaryHmcRunnerTemplate<PeriodicGimplF> BinaryHmcRunnerF;
typedef BinaryHmcRunnerTemplate<PeriodicGimplD> BinaryHmcRunnerD;

template <class RepresentationsPolicy>
using BinaryHmcRunnerTemplateHirep =
    BinaryHmcRunnerTemplate<PeriodicGimplR, RepresentationsPolicy>;



    typedef BinaryHmcRunnerTemplate<ScalarImplR, ScalarFields> ScalarBinaryHmcRunner;
}
}
#endif