/*** COPYRIGHT AND LICENSE STATEMENT *** 
 *
 *  This file is part of the MLMCLangevin code.
 *  
 *  (c) Copyright Eike Mueller, Grigoris Katsiolides and Tony Shardlow, University of Bath
 *  
 *  Main Developer: Eike Mueller
 *  
 *  Contributors:
 *    Grigoris Katsiolides
 *    Tony Shardlow
 *  
 *  MLMCLangevin 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.
 *  
 *  MLMCLangevin 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 MLMCLangevin (see files COPYING and COPYING.LESSER). If not,
 *  see <http://www.gnu.org/licenses/>.
 *
 *** COPYRIGHT AND LICENSE STATEMENT ***/

#include <iostream>
#include <fstream>
#include <cstring>
#include <cstdio>
#include <time.h>
#include <vector>

// Configuration file
#include "config.h"

// Parameters
#include "parameters.hh"

// Langevin equation
#include "langevin_ho.hh"
#include "langevin_qo.hh"
#include "langevin_hdm.hh"
#include "langevin_idm.hh"

// Distributions for random numbers
#include "distribution_normal.hh"
#include "distribution_twopoint.hh"
#include "distribution_threepoint.hh"
#include "distribution_fourpoint.hh"

// Time stepping methods
#include "timestep_euler.hh"
#include "timestep_geometric_langevin.hh"
#include "timestep_stoermer_verlet.hh"
#include "timestep_bivariate_ou.hh"
#include "timestep_ou.hh"
#include "timestep_baoab.hh"

// Initial distribution
#include "statedistribution_fixed.hh"

// Monte Carlo integrators
#include "montecarlo.hh"
#include "montecarlo_standard.hh"
#include "montecarlo_distributionbias.hh"
#include "montecarlo_multilevel.hh"

// Quantity of interest
#include "quantityofinterest.hh"

// Trajectory logger
#include "trajectorylogger.hh"

// Estimators
#include "estimator.hh"

// Timer
#include "timer.hh"

/* ************************************* *
 * General parameters for main program
 * ************************************* */
class ParametersGeneral : public Parameters {
public:
  // Constructor
  ParametersGeneral() :
    Parameters("general"),
    do_standardmc_(true),
    do_distbiasmc_(true),
    do_multilevelmc_(true) {
    addKey("do_standardmc",Bool);
    addKey("do_distbiasmc",Bool);
    addKey("do_multilevelmc",Bool);
  }

  // Read from file
  int readFile(const std::string filename) {
    int readSuccess = Parameters::readFile(filename);
    if (!readSuccess) {
      do_standardmc_ = contents["do_standardmc"]->getBool();
      do_distbiasmc_ = contents["do_distbiasmc"]->getBool();
      do_multilevelmc_ = contents["do_multilevelmc"]->getBool();
    }
    return readSuccess;
  }

  // Return data
  bool do_standardmc() const { return do_standardmc_; }
  bool do_distbiasmc() const { return do_distbiasmc_; }
  bool do_multilevelmc() const { return do_multilevelmc_; }

  // Class data
private:
  bool do_standardmc_;
  bool do_distbiasmc_;
  bool do_multilevelmc_;
};

/* ************************************* *
 * M A I N
 * ************************************* */
int main(int argc, char * argv[]) {

  // Read command line parameters and parameter files
  if (argc != 2) {
    std::cout << "Usage: " << argv[0] << " <parameterfile>" << std::endl;
    return 0;
  }
  std::string parameterFilename = argv[1];

  ParametersGeneral parameters_general;
  ParametersMonteCarlo parameters_montecarlo;
  ParametersDistribution parameters_distribution;

  std::cout << std::endl;
  std::cout << " +-------------------+" << std::endl;
  std::cout << " !   Multilevel MC   !" << std::endl;
  std::cout << " +-------------------+" << std::endl;
  std::cout << std::endl;
  std::cout << " ------- Parameters ------- " << std::endl;
  std::cout << std::endl;

  if (parameters_general.readFile(parameterFilename)) return 1;
  if (parameters_distribution.readFile(parameterFilename)) return 1;
  if (parameters_montecarlo.readFile(parameterFilename)) return 1;

  ParametersMultilevel parameters_multilevel(parameters_montecarlo.L());
  if (parameters_multilevel.readFile(parameterFilename)) return 1;

  std::cout << parameters_distribution << std::endl;
  std::cout << parameters_montecarlo << std::endl;
  std::cout << parameters_multilevel << std::endl;

  // *** Langevin equation ***

  std::cout << "  Langevin equation = ";
#ifdef LANGEVIN_HO
  std::cout << "Harmonic Oscillator " << std::endl;
  ParametersHO parameters_ho;
  if (parameters_ho.readFile(parameterFilename)) return 1;
  std::cout << parameters_ho << std::endl;
  typedef LangevinHO Langevin;
  Langevin langevin(parameters_ho,parameters_distribution);
#endif

#ifdef LANGEVIN_QO
  std::cout << "Quartic Oscillator " << std::endl;
  ParametersQO parameters_qo;
  if (parameters_qo.readFile(parameterFilename)) return 1;
  std::cout << parameters_qo << std::endl;
  typedef LangevinQO Langevin;
  Langevin langevin(parameters_qo,parameters_distribution);
#endif

#ifdef LANGEVIN_HDM
  std::cout << "Homogeneous Dispersion Model " << std::endl;
  ParametersHDM parameters_hdm;
  if (parameters_hdm.readFile(parameterFilename)) return 1;
  std::cout << parameters_hdm << std::endl;
  typedef LangevinHDM Langevin;
  Langevin langevin(parameters_hdm,parameters_distribution);
#endif

#ifdef LANGEVIN_IDM
  std::cout << "Inhomogeneous Dispersion Model " << std::endl;
  ParametersIDM parameters_idm;
  if (parameters_idm.readFile(parameterFilename)) return 1;
  std::cout << parameters_idm << std::endl;
  typedef LangevinIDM Langevin;
  Langevin langevin(parameters_idm,parameters_distribution);
#endif
  std::cout << std::endl;

  // *** Time stepping method ***

  std::cout << "  Time stepping method = ";
#ifdef TIMESTEP_EULER_MARUYAMA
  std::cout << "Euler Maruyama" << std::endl;
  typedef TimestepEuler<Langevin,false> Timestep;
#endif

#ifdef TIMESTEP_SYMPLECTIC_EULER
  std::cout << "Symplectic Euler" << std::endl;
  typedef TimestepEuler<Langevin,true> Timestep;
#endif

#ifdef TIMESTEP_GEOMETRIC_LANGEVIN
  std::cout << "Geometric Langevin" << std::endl;
  typedef TimestepGeometricLangevin<Langevin> Timestep;
#endif

#ifdef TIMESTEP_STOERMER_VERLET
  std::cout << "   Stoermer Verlet" << std::endl;
  typedef TimestepStoermerVerlet<Langevin> Timestep;
#endif

#ifdef TIMESTEP_BIVARIATE_OU
  std::cout << "Bivariate Ornstein Uhlenbeck" << std::endl;
  typedef TimestepBivariateOU<Langevin> Timestep;
#endif

#ifdef TIMESTEP_OU
  std::cout << "Ornstein Uhlenbeck" << std::endl;
  typedef TimestepOU<Langevin> Timestep;
#endif

#ifdef TIMESTEP_BAOAB
  std::cout << "Symmetric Langevin Velocity-Verlet (BAOAB)" << std::endl;
  typedef TimestepBAOAB<Langevin> Timestep;
#endif
  std::cout << std::endl;

  Timestep timestep(langevin);

  // *** Quantity of Interest ***

  std::cout << "  Quantity of interest = ";
#ifdef QOI_GAUSSIAN
  ParametersQoIGaussian parameters_qoigaussian;
  std::cout << "Gaussian function [position]" << std::endl;
  if (parameters_qoigaussian.readFile(parameterFilename)) return 1;
  std::cout << parameters_qoigaussian << std::endl;
  typedef QoIGaussian QoI;
  QoI qoi(parameters_qoigaussian);
#endif // QOI_GAUSSIAN

#ifdef QOI_GAUSSIAN_VELOCITY
  ParametersQoIGaussian parameters_qoigaussian;
  std::cout << "Gaussian function [velocity]" << std::endl;
  if (parameters_qoigaussian.readFile(parameterFilename)) return 1;
  std::cout << parameters_qoigaussian << std::endl;
  typedef QoIGaussianVelocity QoI;
  QoI qoi(parameters_qoigaussian);
#endif // QOI_GAUSSIAN_VELOCITY

#ifdef QOI_INDICATOR
  ParametersQoIIndicator parameters_qoiindicator;
  std::cout << "Indicator function" << std::endl;
  if (parameters_qoiindicator.readFile(parameterFilename)) return 1;
  std::cout << parameters_qoiindicator << std::endl;
  typedef QoIIndicator QoI;
  QoI qoi(parameters_qoiindicator);
#endif // QOI_INDICATOR

#ifdef QOI_TWONORM
  std::cout << "Two Norm function" << std::endl;
  typedef QoITwoNorm QoI;
  QoI qoi;
#endif // QOI_TOWNORM

#ifdef QOI_POSITION
  std::cout << "Position" << std::endl;
  typedef QoIPosition QoI;
  QoI qoi;
#endif // QOI_POSITION

#ifdef QOI_SMOOTHED_DENSITY
  ParametersQoISmoothedDensity parameters_qoismootheddensity;
  std::cout << "Smoothed Density function" << std::endl;
  if (parameters_qoismootheddensity.readFile(parameterFilename)) return 1;
  std::cout << parameters_qoismootheddensity << std::endl;
  typedef QoISmoothedDensity QoI;
  QoI qoi(parameters_qoismootheddensity);
#endif // QOI_SMOOTHED_DENSITY
  std::cout << std::endl;

  // *** Random number Distribution ***

  std::cout << "  Random number distribution = ";
#ifdef DISTRIBUTION_NORMAL
  std::cout << "Normal distribution " << std::endl;
  typedef DistributionNormal Distribution;
#endif // DISTRIBUTION_NORMAL

#ifdef DISTRIBUTION_TWOPOINT
  std::cout << "Two point distribution " << std::endl;
  typedef DistributionTwoPoint Distribution;
#endif // DISTRIBUTION_TWOPOINT

#ifdef DISTRIBUTION_THREEPOINT
  std::cout << "Three point distribution " << std::endl;
  typedef DistributionThreePoint Distribution;
#endif // DISTRIBUTION_THREEPOINT

#ifdef DISTRIBUTION_FOURPOINT
  std::cout << "Four point distribution " << std::endl;
  typedef DistributionFourPoint Distribution;
#endif // DISTRIBUTION_FOURPOINT

  std::cout << std::endl;

  Distribution distribution(parameters_distribution);

  // *** Initial state distribution ***

  std::cout << "  Initial State distribution = ";
#ifdef INITIALDISTRIBUTION_GIBBS
  std::cout << "Gibbs distribution " << std::endl;
  typedef Langevin InitialDistribution;
  InitialDistribution &initialdistribution = langevin;
#endif // INITIALDISTRIBUTION_GIBBS

#ifdef INITIALDISTRIBUTION_FIXED
  std::cout << "Fixed point " << std::endl;
  ParametersInitialDistributionFixed parameters_initialdistributionfixed;
  if (parameters_initialdistributionfixed.readFile(parameterFilename)) return 1;
  std::cout << parameters_initialdistributionfixed << std::endl;
  typedef StateDistributionFixed InitialDistribution;
  InitialDistribution initialdistribution(parameters_initialdistributionfixed,
                                          parameters_distribution);
#endif // INITIALDISTRIBUTION_FIXED

#ifdef RECORD_TRAJECTORIES
  typedef TrajectoryLoggerFile TrajectoryLogger;
  TrajectoryLogger trajectory_logger("trajectories.dat",
                                     MAX_TRAJECTORIES);
#else
  typedef TrajectoryLoggerDummy TrajectoryLogger;
  TrajectoryLogger trajectory_logger;
#endif // RECORD_TRAJECTORIES

  std::cout << std::endl;

  std::cout << "  Timer = " << Timer::type;
  std::cout << ", resolution = ";
  std::cout << std::scientific << Timer::resolution()*1000.;
  std::cout << " ms " << std::endl;

  std::cout << std::endl;

  std::cout << " ------- Results ---------- " << std::endl;
  // STANDARD MC
  if (parameters_general.do_standardmc()) {
    std::cout << std::endl;
    std::cout << "*** STANDARD MC ***" << std::endl;
    MonteCarloStandard<Distribution,
                       Timestep,
                       QoI,
                       InitialDistribution,
                       TrajectoryLogger>
      standardmc(distribution,
                 timestep,
                 qoi,
                 initialdistribution,
                 parameters_montecarlo,
                 trajectory_logger);
    standardmc.run();
    standardmc.show_results();
    unsigned int N = standardmc.Npart();
    unsigned int M = parameters_montecarlo.M();
    std::cout << std::endl;
  }
  // DISTRIBUTION BIAS MC
  if (parameters_general.do_distbiasmc()) {
    std::cout << std::endl;
    std::cout << "*** DISTRIBUTIONBIAS MC ***" << std::endl;
    MonteCarloDistributionBias<Distribution,
                               Timestep,
                               QoI,
                               InitialDistribution>
      distributionbiasmc(distribution,
                         timestep,
                         qoi,
                         initialdistribution,
                         parameters_montecarlo);
    distributionbiasmc.run();
    distributionbiasmc.show_results();
    unsigned int N = distributionbiasmc.Npart();
    unsigned int M = parameters_montecarlo.M();
    std::cout << std::endl;
  }
  // MULTILEVEL MC
  if (parameters_general.do_multilevelmc()) {
    std::cout << std::endl;
    std::cout << "*** MULTILEVEL MC ***" << std::endl;
    MonteCarloMultilevel<Distribution,
                         Timestep,
                         QoI,
                         InitialDistribution>
      multilevelmc(distribution,
                   timestep,
                   qoi,
                   initialdistribution,
                   parameters_montecarlo,
                   parameters_multilevel);

    multilevelmc.run();
    multilevelmc.show_results();
    std::cout << std::endl;
    multilevelmc.saveTimings("timing.dat");

    multilevelmc.save_results("results.dat");
    // save data to MATLAB
    // parameterFilename should have no dots
    // i.e., "p.in" can cause problems later
    std::string stem(parameterFilename);
    multilevelmc.save_matlab(stem);
  }
}