#include "samplers.h"
#include "setup_classes.h"

int main(int argc, char *argv[]){

    // we want to sample the double well potential using the OBABO sampler


    double T = 1;           // sampler hyperparameters
    double gamma = 1;
    double h = 0.01;

    int N = 50000;          // no. of iteration
    bool tavg = false;      // perform time-average?
    int t_meas = 5;         // take measurement every t_meas iterations (passed to sampler as well as print functions).

    OBABO testsampler(T, gamma, h);     // construct OBABO object defined in header "samplers.h"


    PROBLEM double_well;    /* construct object of the problem class defined by the user in header "setup_classes.h". */
    
    measurement results = testsampler.run_mpi_simulation(N, tavg, double_well, t_meas);  /* run sampler. "measurement" needs to be defined 
                                                                                            by user in header "setup_classes.h". It holds
                                                                                            information of what quantities need to be obtained
                                                                                            by the sampler and how to compute and print them.  */

    results.print_to_csv(t_meas);     // print to file, as specified by the user in the "measurement" class.

    return 0;


}

#ifndef SAMPLERS_H
#define SAMPLERS_H


#include <cmath>
#include <iostream>
#include <vector>
#include <random>
#include <string>
#include <sstream>
#include <chrono>
#include <math.h> 
#include <iomanip>
#include <algorithm>
#include <limits>
#include <numeric>
#include <iterator>
#include "setup_classes.h"
// #include <mpi.h>


/* The OBABO sampler. It requires the "PROBLEM" class and the "measurement" class in header "setup_classes.h".
   They need to be written by the user and define the sampling problem and what measurements to take. 
   Note that the member functions below require those classes to be structured in a certain way. */

class OBABO{

    private:
        const double T;
        const double gamma;
        const double h;

        measurement collect_samples(const int N, const bool tavg, PROBLEM POTCLASS, const int randomseed, const int t_meas);  // draws a single sampling trajectory

    public:
        // constructors
        OBABO(double T, double gamma, double h): T{T}, gamma{gamma}, h{h} {
        } 

        measurement run_mpi_simulation(const int N, const bool tavg, PROBLEM POTCLASS, const int t_meas);  /* sets up mpi environment and calls "collect_samples" 
                                                                                                              on each process within. Also performs averaging. */

        void print_sampler_params();    // print sampler hyperparameters.

};

#endif // SAMPLERS_H 

#include "samplers.h"



void OBABO::print_sampler_params(){
    
    std:: cout << "Sampler parameters:\n";
    std:: cout << "Temperature = " << T << ",\nFriction = " << gamma << ",\nStepsize = " << h << ".\n" << std:: endl; 

};


measurement OBABO::run_mpi_simulation(const int N, const bool tavg, PROBLEM POTCLASS, const int t_meas){
    
    int seed = 1;
    print_sampler_params();

    measurement RESULTS = OBABO::collect_samples(N, tavg, POTCLASS, seed, t_meas);

    return RESULTS;

};



measurement OBABO::collect_samples(const int N, const bool tavg, PROBLEM problem, const int randomseed, const int t_meas){

    std:: cout << "Starting OBABO simulation...\n" << std:: endl;

    // set integrator constants
    const double a = exp(-1*gamma*h);    
    const double sqrt_a = sqrt(a);
    const double sqrt_aT = sqrt((1-a)*T);
    const double h_half = 0.5*h;   

    size_t No_params = problem.parameters.size();  // number of parameters

    // COMPUTE INITIAL FORCES
    problem.compute_force();

    // COMPUTE INITIAL MEASUREMENTS
    measurement RESULTS;
    RESULTS.take_measurement(problem.parameters, problem.velocities);

    // PREPARE RNG
    std:: mt19937 twister;
    
    std:: seed_seq seq{1,20,3200,403,5*randomseed+1,12000,73667,9474+randomseed,19151-randomseed};
    std:: vector < std::uint32_t > seeds(1);
    seq.generate(seeds.begin(), seeds.end());
    twister.seed(seeds.at(0)); 

	std:: normal_distribution<> normal{0,1};
    double Rn;


	auto t1 = std:: chrono::high_resolution_clock::now();

    // MAIN LOOP -  from here, we start our journey across the "landscape"
    for ( size_t i = 1;  i <= N;  ++i ) {

        // O + B steps
        for ( size_t j = 0;  j < No_params;  ++j ) {                  
            
            Rn = normal(twister);
            problem.velocities[j] = sqrt_a * problem.velocities[j]  +  sqrt_aT * Rn  +  h_half * problem.forces[j]; 
        
        }

        // A step
        for ( size_t j = 0; j < No_params;  ++j ) {

            problem.parameters[j] += h * problem.velocities[j];
        
        }	
  
        // COMPUTE NEW FORCES
        problem.compute_force();

        // B STEP
        for ( size_t j = 0; j < No_params;  ++j ) {

            problem.velocities[j] += h_half * problem.forces[j];
        
        }   
	
        // O STEP
        for ( size_t j = 0; j < No_params;  ++j ) {
            
		    Rn = normal(twister);
            problem.velocities[j] = sqrt_a * problem.velocities[j]  +  sqrt_aT * Rn;
        
        }   								 

        // TAKE MEASUREMENT
		if( i % t_meas == 0 ) {                                                 // take measurement any t_meas steps.        
            RESULTS.take_measurement(problem.parameters, problem.velocities); 
		}
		
        if( i % int(1e5) == 0 ) std:: cout << "Iteration " << i << " done!\n";
	
	}  // END MAIN LOOP

    // FINALIZE
    auto t2 = std:: chrono:: high_resolution_clock:: now();
	auto ms_int = std:: chrono:: duration_cast < std:: chrono:: seconds > (t2 - t1);
	std:: cout << "Execution took " << ms_int.count() << " seconds!\n ";
        
    return RESULTS;
};

#ifndef SETUP_CLASSES_H
#define SETUP_CLASSES_H

#include <vector>
#include <cmath>
#include <math.h>
#include <fstream>

constexpr double PI = 3.141592653589793;    // some useful constants
constexpr double eps = 1e-12;

class PROBLEM { 
    
    // THIS PROBLEM IS THE DOUBLE WELL POTENTIAL IN 2 DIMENSIONS
    
    /* In case one wants to change the problem to something else, eg. harmonic oscillator:
       The samplers require a "compute_force" function as well as the vectors "parameters", "velocities", and "forces"
       storing the corresponding quantities. */

    public:

        // parameters, velocities, forces 
        std:: vector <double> parameters {1,0};       // CARE! THE SAMPLERS REQUIRE THIS VECTOR
        std:: vector <double> velocities {0,0};       // CARE! THE SAMPLERS REQUIRE THIS VECTOR
        std:: vector <double> forces {0,0};           // CARE! THE SAMPLERS REQUIRE THIS VECTOR

        // fill force vector
        void compute_force(){                           /* CARE! THE SAMPLERS REQUIRE THIS FUNCTION. 
                                                           It needs to be written by the user. */

                 
            double x = parameters[0];       // current parameters (positions)
            double y = parameters[1];

            double e1, e2, rho;             // some help constants
            double x_mu_diff1 = x-mu1[0]; 
            double x_mu_diff2 = x-mu2[0]; 
            double y_mu_diff1 = y-mu1[1]; 
            double y_mu_diff2 = y-mu2[1];
            
            // the exponentials 
            e1 = exp( -inv_two_det_SIG1 * ( SIG1[2]*x_mu_diff1*x_mu_diff1 - 2*SIG1[1]*x_mu_diff1*y_mu_diff1 + SIG1[0]*y_mu_diff1*y_mu_diff1 ) );
            e2 = exp( -inv_two_det_SIG2 * ( SIG2[2]*x_mu_diff2*x_mu_diff2 - 2*SIG2[1]*x_mu_diff2*y_mu_diff2 + SIG2[0]*y_mu_diff2*y_mu_diff2 ) );

            rho = pref1 * det1 * e1  +  pref2 * det2 * e2;  // the density

            //force part without noise, i.e. double well
            forces[0] = -1/rho * ( pref1 * e1 * (SIG1[2]*x_mu_diff1 - SIG1[1]*y_mu_diff1)   +   pref2 * e2 * (SIG2[2]*x_mu_diff2 - SIG2[1]*y_mu_diff2) );
            forces[1] = -1/rho * ( pref1 * e1 * (-SIG1[1]*x_mu_diff1 + SIG1[0]*y_mu_diff1)  +   pref2 * e2 * (-SIG2[1]*x_mu_diff2 + SIG2[0]*y_mu_diff2) );

            // if(sig!=0){
            //     normal_distribution<> normal{0,sig};	// add noise
            //     F.fx += normal(twister);
            //     F.fy += normal(twister);
            // }

            return;

        };


    private:

        // constants that define the potential
        const std:: vector <double> mu1 {1, 0};		                // (x,y) of the first well 
        const std:: vector <double> mu2 {-1, 0};		            // (x,y) of the second well
        const std:: vector <double> SIG1 {0.6, 0.085, 0.02};	    // elements of the two covariance matrices sig11, sig12, sig22 (note: sig12=sig21)
        const std:: vector <double> SIG2 {0.6, 0.085, 0.02};
        const double phi1 = 0.5;                                    // mixing factor (phi2 = 1-phi1)	

        // constants used by the force computation 
        const double det1 = SIG1[0]*SIG1[2] - SIG1[1]*SIG1[1];      // determinants of cov. matrices
        const double det2 = SIG2[0]*SIG2[2] - SIG2[1]*SIG2[1];      
        const double inv_two_det_SIG1 = 1 / (2*det1) ;              // used in the exponents in the force
        const double inv_two_det_SIG2 = 1 / (2*det2) ;
        const double pref1 = phi1 / ( 2*PI*pow( det1, 1.5) );  	    // prefactors in the force
        const double pref2 = (1-phi1) / ( 2*PI*pow( det2, 1.5) );


};



class measurement{

    /* The measurement class that defines what quantities are collected by the samplers, how they are computed, and
       printed to a file. This class needs to be modified by the user.
       In this example, we only save the first parameter (corresponding to the x-coordinate
       when used in the double well problem above), as well as the kinetic energy. */

    public:

        void take_measurement(std:: vector <double> parameters, std:: vector <double> velocities){  /* CARE!!! The samplers need this function.
                                                                                                       It needs to be written by the user. */
            
            x.push_back(parameters[0]);
            kin_energy.push_back( 0.5* (velocities[0]*velocities[0] + velocities[1]*velocities[1]) );
        
        }

        void print_to_csv(const int t_meas){    /* print results to file. this routine would be called in the main-file. 
                                                   It needs to be written by the user. "t_meas" gives the number of sampler iterations 
                                                   between two measurements (it is passed here to get the correct iteration count in the
                                                   output file.)*/

            std:: ofstream file{"TEST.csv"};
            std:: cout << "Writing to file...\n";

            for ( size_t i = 0; i<x.size(); ++i )
            {
                file << i*t_meas << " " << x.at(i) << " " << kin_energy.at(i) <<  "\n";  
            }
            file.close();

        }

        std:: vector <double> x; 
        std:: vector <double> kin_energy;



};


#endif // SETUP_CLASSES_H