// Variational Monte Carlo for atoms with importance sampling, no slater det
// Test case for helium
#include "mpi.h"
#include <cmath>
#include <iostream>
#include <fstream>
#include <iomanip>

using namespace  std;
// output file as global variable
ofstream ofile;  
// the step length and its squared inverse for the second derivative 
#define h 0.001
#define h2 1000000
#define alpha 1.6875
// declaration of functions 

// The Mc sampling for the variational Monte Carlo 
void  mc_sampling(int, int, double *, double *);

// The variational wave function
double  wave_function(double **, double);

// The variational wave function for the move of one particle at the time
void  wave_function_onemove(double **, double **, double *, double, int);

// The local energy 
double  local_energy(double **, double, double);

// The quantum force
void  quantum_force(double **, double **, double, double);

//  allocate space for a matrix
void  **matrix(int, int, int);

//  free space for  a matrix
void free_matrix(void **);

// ran2 for uniform deviates, initialize with negative seed.
double ran2(long *);

// function for gaussian random numbers
double gaussian_deviate(long *);


//  Here we define global variables  used in various functions
//  These can be changed by reading from file the different parameters
int dimension = 3; // three-dimensional system
int charge = 2;  //  we fix the charge to be that of the helium atom
int my_rank, numprocs;  //  these are the parameters used by MPI  to define which node and how many
double timestep = 0.05;  //  we fix the time step  for the gaussian deviate
int number_particles  = 2;  //  we fix also the number of electrons to be 2


// Begin of main program   

int main(int argc, char* argv[])
{
  char *outfilename;
  int i;
  double total_number_cycles;
  double *cumulative_e, *cumulative_e2;
  double *total_cumulative_e, *total_cumulative_e2;
  double  time_start, time_end, total_time;
  int number_cycles = 100000000; //  default number of cycles
  int max_variations = 10;    //  default number of variations
  double beta, variance, energy, error;

  //  MPI initializations
  MPI_Init (&argc, &argv);
  MPI_Comm_size (MPI_COMM_WORLD, &numprocs);
  MPI_Comm_rank (MPI_COMM_WORLD, &my_rank);
  time_start = MPI_Wtime();

  if (my_rank == 0 && argc <= 1) {
    cout << "Bad Usage: " << argv[0] << 
      " read also output file on same line" << endl;
    exit(1);
  }
  if (my_rank == 0 && argc > 1) {
    outfilename=argv[1];
    ofile.open(outfilename); 
  }
  total_cumulative_e = new double[max_variations+1];
  total_cumulative_e2 = new double[max_variations+1];
  cumulative_e = new double[max_variations+1];
  cumulative_e2 = new double[max_variations+1];

  //  initialize the arrays  by zeroing them
  for( i=1; i <= max_variations; i++){
    cumulative_e[i] = cumulative_e2[i]  = total_cumulative_e[i] = total_cumulative_e2[i]  = 0.0;
  }
  
  // broadcast the total number of  variations
  MPI_Bcast (&max_variations, 1, MPI_INT, 0, MPI_COMM_WORLD);
  MPI_Bcast (&number_cycles, 1, MPI_INT, 0, MPI_COMM_WORLD);

  total_number_cycles = number_cycles*numprocs; 
  //  Do the mc sampling  and accumulate data with MPI_Reduce
  mc_sampling(max_variations, number_cycles, cumulative_e, cumulative_e2);
  //  Collect data in total averages
  for( i=1; i <= max_variations; i++){
    MPI_Reduce(&cumulative_e[i], &total_cumulative_e[i], 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
    MPI_Reduce(&cumulative_e2[i], &total_cumulative_e2[i], 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
  }
  time_end = MPI_Wtime();
  total_time = time_end-time_start;
  // Print out results  
  if ( my_rank == 0) {
    cout << "Time = " <<  total_time  << " on number of processors: "  << numprocs  << endl;
    beta = 0.1;
    for( i=1; i <= max_variations; i++){
      beta += 0.1;  
      energy = total_cumulative_e[i]/numprocs;
      variance = total_cumulative_e2[i]/numprocs-energy*energy;
      error=sqrt(variance/(total_number_cycles-1.0));
      ofile << setiosflags(ios::showpoint | ios::uppercase);
      ofile << setw(15) << setprecision(8) << beta;
      ofile << setw(15) << setprecision(8) << energy;
      ofile << setw(15) << setprecision(8) << variance;
      ofile << setw(15) << setprecision(8) << error << endl;
    }
    ofile.close();  // close output file
  }
  delete [] total_cumulative_e; delete [] total_cumulative_e2; 
  delete [] cumulative_e; delete [] cumulative_e2; 
  // End MPI
  MPI_Finalize ();  
  return 0;
}  //  end of main function


// Monte Carlo sampling with the Metropolis algorithm  

void mc_sampling(int max_variations, int number_cycles, double *cumulative_e, double *cumulative_e2)
{
  int cycles, variate, i, j, k;
  long idum;
  double wfnew, wfold, beta, energy, energy2, delta_e;
  double **r_old, **r_new, **qforce_new, **qforce_old;
  double greensfunction, D; 
  // diffusion constant from Schroedinger equation
  D = 0.5; 
  // every node has its own seed for the random numbers
  idum = -1-my_rank;
  // allocate matrices which contain the position of the particles  
  r_old = (double **) matrix( number_particles, dimension, sizeof(double));
  r_new = (double **) matrix( number_particles, dimension, sizeof(double));
  qforce_old = (double **) matrix( number_particles, dimension, sizeof(double));
  qforce_new = (double **) matrix( number_particles, dimension, sizeof(double));
  for (i = 0; i < number_particles; i++) { 
    for ( j=0; j < dimension; j++) {
      r_old[i][j] = r_new[i][j] = qforce_new[i][j] = qforce_old[i][j] = 0 ;
    }
  }
  beta = 0.1;
  // loop over variational parameters  
  for (variate=1; variate <= max_variations; variate++){
    // initialisations of variational parameters and energies 
    beta += 0.1;  
    energy = energy2 = delta_e = 0.0;
    //  initial trial position, note calling with beta 
    for (i = 0; i < number_particles; i++) { 
      for ( j=0; j < dimension; j++) {
	r_old[i][j] = gaussian_deviate(&idum)*sqrt(timestep);
      }
    }
    wfold = wave_function(r_old, beta);
    quantum_force(r_old, qforce_old, beta, wfold);
    // loop over monte carlo cycles 
    for (cycles = 1; cycles <= number_cycles; cycles++){ 
      // new position 
      for (i = 0; i < number_particles; i++) { 
	for ( j=0; j < dimension; j++) {
	  // gaussian deviate to compute new positions using a given timestep
	  r_new[i][j] = r_old[i][j] + gaussian_deviate(&idum)*sqrt(timestep)+qforce_old[i][j]*timestep*D;
	}  
        //  for the other particles we need to set the position to the old position since
        //  we move only one particle at the time
        for (k = 0; k < number_particles; k++) {
	  if ( k != i) {
	    for ( j=0; j < dimension; j++) {
	      r_new[k][j] = r_old[k][j];
	    }
	  } 
        }
	//        wave_function_onemove(r_new, qforce_new, &wfnew, beta); 
        wfnew = wave_function(r_new, beta); 
        quantum_force(r_new, qforce_new, beta, wfnew);
	//  we compute the log of the ratio of the greens functions to be used in the 
	//  Metropolis-Hastings algorithm
        greensfunction = 0.0;            
	for ( j=0; j < dimension; j++) {
	  greensfunction += 0.5*(qforce_old[i][j]+qforce_new[i][j])*
	    (D*timestep*0.5*(qforce_old[i][j]-qforce_new[i][j])-r_new[i][j]+r_old[i][j]);
        }
        greensfunction = exp(greensfunction);
 	// The Metropolis test is performed by moving one particle at the time
	if(ran2(&idum) <= greensfunction*wfnew*wfnew/wfold/wfold ) { 
	  for ( j=0; j < dimension; j++) {
	    r_old[i][j] = r_new[i][j];
	    qforce_old[i][j] = qforce_new[i][j];
	  }
	  wfold = wfnew;
	}
      }  //  end of loop over particles
      // compute local energy  
      delta_e = local_energy(r_old, beta, wfold);
      // update energies
      energy += delta_e;
      energy2 += delta_e*delta_e;
    }   // end of loop over MC trials   
    // update the energy average and its squared 
    cumulative_e[variate] = energy/number_cycles;
    cumulative_e2[variate] = energy2/number_cycles;
  }    // end of loop over variational  steps 
  free_matrix((void **) r_old); // free memory
  free_matrix((void **) qforce_old); // free memory
  free_matrix((void **) r_new); // free memory
  free_matrix((void **) qforce_new); // free memory
}   // end mc_sampling function  


// Function to compute the squared wave function and the quantum force

double  wave_function(double **r, double beta)
{
  int i, j, k;
  double argument, wf, r_single_particle, r_ij;
  
  argument = wf = 0;
  for (i = 0; i < number_particles; i++) { 
    r_single_particle = 0;
    for (j = 0; j < dimension; j++) { 
      r_single_particle  += r[i][j]*r[i][j];
    }
    argument += sqrt(r_single_particle);
  }
  wf = exp(-argument*alpha) ;
  // contribution from electron-electron potential  
  for (i = 0; i < number_particles-1; i++) { 
    for (j = i+1; j < number_particles; j++) {
      r_ij = 0;  
      for (k = 0; k < dimension; k++) { 
	r_ij += (r[i][k]-r[j][k])*(r[i][k]-r[j][k]);
      }
      r_ij=sqrt(r_ij);
      wf *= exp(r_ij/(2*(1+beta*r_ij)));
    }
  }
  return wf;
}

// Function to calculate the local energy with num derivative

double  local_energy(double **r, double beta, double wfold)
{
  int i, j , k;
  double e_local, wfminus, wfplus, e_kinetic, e_potential, r_12, 
    r_single_particle;
  double **r_plus, **r_minus;
  
  // allocate matrices which contain the position of the particles  
  // the function matrix is defined in the progam library 
  r_plus = (double **) matrix( number_particles, dimension, sizeof(double));
  r_minus = (double **) matrix( number_particles, dimension, sizeof(double));
  for (i = 0; i < number_particles; i++) { 
    for ( j=0; j < dimension; j++) {
      r_plus[i][j] = r_minus[i][j] = r[i][j];
    }
  }
  // compute the kinetic energy  
  e_kinetic = 0;
  for (i = 0; i < number_particles; i++) {
    for (j = 0; j < dimension; j++) { 
      r_plus[i][j] = r[i][j]+h;
      r_minus[i][j] = r[i][j]-h;
      wfminus = wave_function(r_minus, beta); 
      wfplus  = wave_function(r_plus, beta); 
      e_kinetic -= (wfminus+wfplus-2*wfold);
      r_plus[i][j] = r[i][j];
      r_minus[i][j] = r[i][j];
    }
  }
  // include electron mass and hbar squared and divide by wave function 
  e_kinetic = 0.5*h2*e_kinetic/wfold;
  // compute the potential energy 
  e_potential = 0;
  // contribution from electron-proton potential  
  for (i = 0; i < number_particles; i++) { 
    r_single_particle = 0;
    for (j = 0; j < dimension; j++) { 
      r_single_particle += r[i][j]*r[i][j];
    }
    e_potential -= charge/sqrt(r_single_particle);
  }
  // contribution from electron-electron potential  
  for (i = 0; i < number_particles-1; i++) { 
    for (j = i+1; j < number_particles; j++) {
      r_12 = 0;  
      for (k = 0; k < dimension; k++) { 
	r_12 += (r[i][k]-r[j][k])*(r[i][k]-r[j][k]);
      }
      e_potential += 1/sqrt(r_12);          
    }
  }
  free_matrix((void **) r_plus); // free memory
  free_matrix((void **) r_minus);
  e_local = e_potential+e_kinetic;
  return e_local;
}

// Function to calculate the quantum force with numerical derivative

void  quantum_force(double **r, double **qforce, double beta, double wf)
{
  int i, j;
  double wfminus, wfplus; 
  double **r_plus, **r_minus;

  r_plus = (double **) matrix( number_particles, dimension, sizeof(double));
  r_minus = (double **) matrix( number_particles, dimension, sizeof(double));
  for (i = 0; i < number_particles; i++) { 
    for ( j=0; j < dimension; j++) {
      r_plus[i][j] = r_minus[i][j] = r[i][j];
    }
  }
  // compute the first derivative
  for (i = 0; i < number_particles; i++) {
    for (j = 0; j < dimension; j++) { 
      r_plus[i][j] = r[i][j]+h;
      r_minus[i][j] = r[i][j]-h;
      wfminus = wave_function(r_minus, beta); 
      wfplus  = wave_function(r_plus, beta); 
      qforce[i][j] = (wfplus-wfminus)*h*2.0/wf;
      r_plus[i][j] = r[i][j];
      r_minus[i][j] = r[i][j];
    }
  }
  free_matrix((void **) r_plus); // free memory
  free_matrix((void **) r_minus);

} // end of quantum_force function



/*
 * The function                             
 *      void  **matrix()                    
 * reserves dynamic memory for a two-dimensional matrix 
 * using the C++ command new . No initialization of the elements. 
 * Input data:                      
 *  int row      - number of  rows          
 *  int col      - number of columns        
 *  int num_bytes- number of bytes for each 
 *                 element                  
 * Returns a void  **pointer to the reserved memory location.                                
 */

void **matrix(int row, int col, int num_bytes)
{
  int      i, num;
  char     **pointer, *ptr;

  pointer = new(nothrow) char* [row];
  if(!pointer) {
    cout << "Exception handling: Memory allocation failed";
    cout << " for "<< row << "row addresses !" << endl;
    return NULL;
  }
  i = (row * col * num_bytes)/sizeof(char);
  pointer[0] = new(nothrow) char [i];
  if(!pointer[0]) {
    cout << "Exception handling: Memory allocation failed";
    cout << " for address to " << i << " characters !" << endl;
    return NULL;
  }
  ptr = pointer[0];
  num = col * num_bytes;
  for(i = 0; i < row; i++, ptr += num )   {
    pointer[i] = ptr; 
  }

  return  (void **)pointer;

} // end: function void **matrix()

  /*
   * The function                         
   *      void free_matrix()              
   * releases the memory reserved by the function matrix() 
   *for the two-dimensional matrix[][] 
   * Input data:                          
   *  void far **matr - pointer to the matrix
   */

void free_matrix(void **matr)
{

  delete [] (char *) matr[0];
  delete [] matr;

}  // End:  function free_matrix() 


/*
** The function 
**         ran2()
** is a long periode (> 2 x 10^18) random number generator of 
** L'Ecuyer and Bays-Durham shuffle and added safeguards.
** Call with idum a negative integer to initialize; thereafter,
** do not alter idum between sucessive deviates in a
** sequence. RNMX should approximate the largest floating point value
** that is less than 1.
** The function returns a uniform deviate between 0.0 and 1.0
** (exclusive of end-point values).
*/

#define IM1 2147483563
#define IM2 2147483399
#define AM (1.0/IM1)
#define IMM1 (IM1-1)
#define IA1 40014
#define IA2 40692
#define IQ1 53668
#define IQ2 52774
#define IR1 12211
#define IR2 3791
#define NTAB 32
#define NDIV (1+IMM1/NTAB)
#define EPS 1.2e-7
#define RNMX (1.0-EPS)

double ran2(long *idum)
{
  int            j;
  long           k;
  static long    idum2 = 123456789;
  static long    iy=0;
  static long    iv[NTAB];
  double         temp;

  if(*idum <= 0) {
    if(-(*idum) < 1) *idum = 1;
    else             *idum = -(*idum);
    idum2 = (*idum);
    for(j = NTAB + 7; j >= 0; j--) {
      k     = (*idum)/IQ1;
      *idum = IA1*(*idum - k*IQ1) - k*IR1;
      if(*idum < 0) *idum +=  IM1;
      if(j < NTAB)  iv[j]  = *idum;
    }
    iy=iv[0];
  }
  k     = (*idum)/IQ1;
  *idum = IA1*(*idum - k*IQ1) - k*IR1;
  if(*idum < 0) *idum += IM1;
  k     = idum2/IQ2;
  idum2 = IA2*(idum2 - k*IQ2) - k*IR2;
  if(idum2 < 0) idum2 += IM2;
  j     = iy/NDIV;
  iy    = iv[j] - idum2;
  iv[j] = *idum;
  if(iy < 1) iy += IMM1;
  if((temp = AM*iy) > RNMX) return RNMX;
  else return temp;
}
#undef IM1
#undef IM2
#undef AM
#undef IMM1
#undef IA1
#undef IA2
#undef IQ1
#undef IQ2
#undef IR1
#undef IR2
#undef NTAB
#undef NDIV
#undef EPS
#undef RNMX

// End: function ran2()



// random numbers with gaussian distribution
double gaussian_deviate(long * idum)
{
  static int iset = 0;
  static double gset;
  double fac, rsq, v1, v2;

  if ( idum < 0) iset =0;
  if (iset == 0) {
    do {
      v1 = 2.*ran2(idum) -1.0;
      v2 = 2.*ran2(idum) -1.0;
      rsq = v1*v1+v2*v2;
    } while (rsq >= 1.0 || rsq == 0.);
    fac = sqrt(-2.*log(rsq)/rsq);
    gset = v1*fac;
    iset = 1;
    return v2*fac;
  } else {
    iset =0;
    return gset;
  }
} // end function for gaussian deviates