#include <mpi.h>
#include <iostream>
#include <math.h>
#include "lib.h"
#include <fstream>

using namespace std;

//Step length for numerical derivation
const double h = .0001;
const double hh = h*h;
int NumberMCGlobal;
int Z_global;
double delta_t_global;
long idum_global;

void dfpmin(double p[], int n, double gtol, int *iter, double *fret,
	    double(*func)(double []), void (*dfunc)(double [], double []));

//Gaussian random deviate
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;
  }
}

//Distance from nucleus and electron i
double r_i(double** R, int i) {
  
  double result = 0;
  for (int j = 0; j < 3; j++)
    result += R[i][j]*R[i][j];
  
  return sqrt(result);
}

//Distance between r_i and r_j
double r_ij(double** R, int i, int j) {

  double result = 0;
  for (int k = 0; k < 3; k++)
    result += (R[i][k]-R[j][k]) * ((R[i][k]-R[j][k]));

  return sqrt(result);
}

//Trial wave function
double Psi_trial(double** R, double alpha, double beta) {

  double r12 = r_ij(R,0,1);

  return exp(-alpha*(r_i(R,0) + r_i(R,1))) * exp(r12 / (2*(1+beta*r12)));
}

//check singularity at r=0
bool Singularity(double** R, int Z) {
  
  for (int i = 0; i < Z; i++)
    if (r_i(R,i) < 1e-10)
      return true;

  for (int i = 0; i < Z - 1; i++)
    for (int j = i+1; j < Z; j++)
      if (r_ij(R,i,j) < 1e-10)
	return true;

  return false;
}

//Local energy
double E_local(double** R, double alpha, double beta, int Z) {

  double psi0 = Psi_trial(R,alpha,beta);

  double** R_plus = (double**) matrix(Z,3,sizeof(double));
  double** R_minus = (double**) matrix(Z,3,sizeof(double));

  for (int i = 0; i < Z; i++) {
    for (int j = 0; j < 3; j++) {
      R_plus[i][j] = R[i][j];
      R_minus[i][j] = R[i][j];
    }
  }

  //Kinetic energy, numerical derivative
  double kinetic = 0;
  for (int i = 0; i < Z; i++)
    for (int j = 0; j < 3; j++) {
      R_plus[i][j] += h;
      R_minus[i][j] -= h;
      kinetic -= Psi_trial(R_plus,alpha,beta) + Psi_trial(R_minus,alpha,beta) - 2*psi0;
      R_plus[i][j] = R[i][j];
      R_minus[i][j] = R[i][j];
    }
  kinetic *= .5 / (hh*psi0);

  //Potential energy from electron-nucleus
  double potential = 0;
  for (int i = 0; i < Z; i++)
    potential -= Z / r_i(R,i);

  //Electron-electron repulsion
  for (int i = 0; i < Z - 1; i++)
    for (int j = i+1; j < Z; j++)
      potential += 1 / r_ij(R,i,j);

  //Free memory
  free_matrix((void**) R_plus);
  free_matrix((void**) R_minus);

  return potential + kinetic;
}

//Quantum force, numerical derivation
void calcQF(double** R, double** F, double alpha, double beta, int Z) {

  double** R_plus = (double**) matrix(Z,3,sizeof(double));
  double** R_minus = (double**) matrix(Z,3,sizeof(double));

  for (int i = 0; i < Z; i++) {
    for (int j = 0; j < 3; j++) {
      R_plus[i][j] = R[i][j];
      R_minus[i][j] = R[i][j];
    }
  }
  
  double psi0 = Psi_trial(R,alpha,beta);

  for (int i = 0; i < Z; i++) {
    for (int j = 0; j < 3; j++) {
      R_plus[i][j] += h;
      R_minus[i][j] -= h;
      F[i][j] = (Psi_trial(R_plus,alpha,beta) - Psi_trial(R_minus,alpha,beta)) / (h*psi0);
      R_plus[i][j] = R[i][j];
      R_minus[i][j] = R[i][j];
    }
  }

  //Free memory
  free_matrix((void**) R_plus);
  free_matrix((void**) R_minus);
}

//MC calculation of <H> with importance sampling
void kjoerMC(double& sum, double& ESquared, int NumberMC, 
	     double delta_t, long& idum, double alpha, int& N, double beta, int Z, double& sum1alpha, double& sum2alpha,
	     double& sum1beta, double& sum2beta, bool calcderiv) {

  //Initialize positions
  double** R = (double**) matrix(Z,3,sizeof(double));
  for (int i = 0; i < Z; i++)
    for (int j = 0; j < 3; j++)
      R[i][j] = gaussian_deviate(&idum);

  //Initialize quantum force
  double** F = (double**) matrix(Z,3,sizeof(double));
  calcQF(R,F,alpha,beta,Z);

  double EL; //Local energy
  double sqrtdt = sqrt(delta_t);
  double D = .5;

  //For Metropolis-Hastings algo:
  double** R_new = (double**) matrix(Z,3,sizeof(double));
  double** F_new = (double**) matrix(Z,3,sizeof(double));

  double P_new; //Probability in trial position
  double P; //Probability in present positions
  double greensratio; // Ratio between Green's functions
  
  double alphaderiv,betaderiv,psi0;

  //MC looping starts here
  for (int k = 0; k < NumberMC; k++) {

    //Loop over electrons
    for (int i = 0; i < Z; i++) {

      //Move electron i
      for (int j = 0; j < 3; j++)
	R_new[i][j] = R[i][j] + D*delta_t*F[i][j] + gaussian_deviate(&idum)*sqrtdt;
 
      //og la de andre elektronene bli hvor de er
      for (int ii = 0; ii < Z; ii++)
	for (int j = 0; j < 3; j++)
	  if (ii != i)
	    R_new[ii][j] = R[ii][j];

      //Quantum force in new position
      calcQF(R_new,F_new,alpha,beta,Z);

      //Ratio between Green's functions
      greensratio = 0;
      for (int ii = 0; ii < Z; ii++)
	for (int j = 0; j < 3; j++)
	  greensratio += .5*(F_new[ii][j]+F[ii][j]) * (.5*D*delta_t*(F[ii][j]-F_new[ii][j]) + R[ii][j] - R_new[ii][j]);
      greensratio = exp(greensratio);
	
      //Metropolis-Hastings test
      P = Psi_trial(R,alpha,beta);
      P = P * P;
      P_new = Psi_trial(R_new,alpha,beta);
      P_new = P_new * P_new;

      if (ran2(&idum) < greensratio*P_new/P) {
	//Accept move
	for (int ii = 0; ii < Z; ii++) {
	  for (int j = 0; j < 3; j++) {
	    R[ii][j] = R_new[ii][j];
	    F[ii][j] = F_new[ii][j];
	  }
	}
      }

    } //End loop over electron to be moved
    
    if (!Singularity(R,Z) && k > NumberMC/20) {

      EL = E_local(R, alpha, beta, Z);
      
      sum += EL;
      ESquared += EL * EL;
      
      if (calcderiv) {
	psi0 = Psi_trial(R,alpha,beta);
	
	alphaderiv = (Psi_trial(R,alpha+h,beta) - Psi_trial(R,alpha-h,beta))/(2*h*psi0);
	sum1alpha += alphaderiv*EL;
	sum2alpha += alphaderiv;
	
	betaderiv = (Psi_trial(R,alpha,beta+h) - Psi_trial(R,alpha,beta-h))/(2*h*psi0);
	sum1beta += betaderiv*EL;
	sum2beta += betaderiv;
      }

      N++;
    }
  } //End loop over MC cycles

  //Free memory
  free_matrix((void**) R);
  free_matrix((void**) R_new);
  free_matrix((void**) F);
  free_matrix((void**) F_new);  
}

double energy(double* param) {

  double sum = 0;
  double ESquared = 0;
  int N = 0;
  double sum1alpha = 0;
  double sum2alpha = 0;
  double sum1beta = 0;
  double sum2beta = 0;

  kjoerMC(sum,ESquared,NumberMCGlobal,delta_t_global,idum_global,param[1],N,param[2],
	  Z_global,sum1alpha,sum2alpha,sum1beta,sum2beta,false);
  
  cout << "alpha = " << param[1] << "   beta = " << param[2] << "    energy = " << sum/N << endl;

  return sum / N;
}

void energy_deriv(double* param, double* values) {
  
  double sum = 0;
  double ESquared = 0;
  int N = 0;
  double sum1alpha = 0;
  double sum2alpha = 0;
  double sum1beta = 0;
  double sum2beta = 0;

  kjoerMC(sum,ESquared,NumberMCGlobal,delta_t_global,idum_global,param[1],N,param[2],
	  Z_global,sum1alpha,sum2alpha,sum1beta,sum2beta,true);
  
  double E = sum / N;

  values[1] = 2 * ( (sum1alpha/N) - (sum2alpha/N)*E );
  values[2] = 2 * ( (sum1beta/N) - (sum2beta/N)*E );

  cout << "alpha = " << param[1] << "   beta =  " << param[2] << 
    "    alphaderiv = " << values[1] << "    betaderiv = " << values[2] << endl;
}


int main() {

  //MPI initialization
  int numprocs;
  int MyRank;
  MPI_Init(NULL,NULL);
  MPI_Comm_size(MPI_COMM_WORLD, &numprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &MyRank);

  int NumberMC = 100000; //Number of Monte Carlo cycles
  NumberMCGlobal = NumberMC;
  int Z = 2; //Nuclear charge
  Z_global = Z;
  double alpha = 1.7; //Variational paramater
  double beta = 0.5; //Variational paramater
  long idum = - time(NULL) +  MyRank*1000;
  idum_global = idum;
  double delta_t = .05;
  delta_t_global = delta_t;

  double* param = new double[3];
  double* values = new double[3];
  param[1] = alpha;
  param[2] = beta;

  double gtol = .001;
  int iter;
  double fret;

  dfpmin(param, 2, gtol, &iter, &fret, &energy, &energy_deriv);

  cout << "alpha = " << param[1] << endl;
  cout << "beta = " << param[2] << endl;
  cout << "energy = " << fret << endl;
  cout << "iter = " << iter << endl << endl;
    
  //Stop MPI
  MPI_Finalize();

  return 0;
}