//#include <mpi.h>
#include <cmath>
#include <iostream>
#include <fstream>
#include <iomanip>
#include "vectormatrixclass.h"
//#include "lib.h"

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;


static double sqrarg;
#define SQR(a) ((sqrarg=(a)) == 0.0 ? 0.0 : sqrarg*sqrarg)


static double maxarg1,maxarg2;
#define FMAX(a,b) (maxarg1=(a),maxarg2=(b),(maxarg1) > (maxarg2) ?\
        (maxarg1) : (maxarg2))


#define ALF 1.0e-4
#define TOLX 1.0e-7

void lnsrch(int n, Vector &xold, double fold, Vector &g, Vector &p, Vector &x,
	    double *f, double stpmax, int *check, double (*func)(Vector &p))
{
  int i;
  double a,alam,alam2,alamin,b,disc,f2,fold2,rhs1,rhs2,slope,sum,temp,
    test,tmplam;

  *check=0;
  for (sum=0.0,i=0;i<n;i++) sum += p(i)*p(i);
  sum=sqrt(sum);
  if (sum > stpmax)
    for (i=0;i<n;i++) p(i) *= stpmax/sum;
  for (slope=0.0,i=0;i<n;i++)
    slope += g(i)*p(i);
  test=0.0;
  for (i=0;i<n;i++) {
    temp=fabs(p(i))/FMAX(fabs(xold(i)),1.0);
    if (temp > test) test=temp;
  }
  alamin=TOLX/test;
  alam=1.0;
  for (;;) {
    for (i=0;i<n;i++) x(i)=xold(i)+alam*p(i);
    *f=(*func)(x);
    if (alam < alamin) {
      for (i=0;i<n;i++) x(i)=xold(i);
      *check=1;
      return;
    } else if (*f <= fold+ALF*alam*slope) return;
    else {
      if (alam == 1.0)
	tmplam = -slope/(2.0*(*f-fold-slope));
      else {
	rhs1 = *f-fold-alam*slope;
	rhs2=f2-fold2-alam2*slope;
	a=(rhs1/(alam*alam)-rhs2/(alam2*alam2))/(alam-alam2);
	b=(-alam2*rhs1/(alam*alam)+alam*rhs2/(alam2*alam2))/(alam-alam2);
	if (a == 0.0) tmplam = -slope/(2.0*b);
	else {
	  disc=b*b-3.0*a*slope;
	  if (disc<0.0) cout << "Roundoff problem in lnsrch." << endl;
	  else tmplam=(-b+sqrt(disc))/(3.0*a);
	}
	if (tmplam>0.5*alam)
	  tmplam=0.5*alam;
      }
    }
    alam2=alam;
    f2 = *f;
    fold2=fold;
    alam=FMAX(tmplam,0.1*alam);
  }
}
#undef ALF
#undef TOLX


#define ITMAX 200
#define EPS 3.0e-8
#define TOLX (4*EPS)
#define STPMX 100.0


void dfpmin(Vector &p, int n, double gtol, int *iter, double *fret,
	    double(*func)(Vector &p), void (*dfunc)(Vector &p, Vector &g))
{

  int check,i,its,j;
  double den,fac,fad,fae,fp,stpmax,sum=0.0,sumdg,sumxi,temp,test;
  Vector dg(n), g(n), hdg(n), pnew(n), xi(n);
  Matrix hessian(n,n);

  fp=(*func)(p);
  (*dfunc)(p,g);
  for (i = 0;i < n;i++) {
    for (j = 0; j< n;j++) hessian(i,j)=0.0;
    hessian(i,i)=1.0;
    xi(i) = -g(i);
    sum += p(i)*p(i);
  }
  stpmax=STPMX*FMAX(sqrt(sum),(double)n);
  for (its=1;its<=ITMAX;its++) {
    *iter=its;
    lnsrch(n,p,fp,g,xi,pnew,fret,stpmax,&check,func);
    fp = *fret;
    for (i = 0; i< n;i++) {
      xi(i)=pnew(i)-p(i);
      p(i)=pnew(i);
    }
    test=0.0;
    for (i = 0;i< n;i++) {
      temp=fabs(xi(i))/FMAX(fabs(p(i)),1.0);
      if (temp > test) test=temp;
    }
    if (test < TOLX) {
      return;
    }
    for (i=0;i<n;i++) dg(i)=g(i);
    (*dfunc)(p,g);
    test=0.0;
    den=FMAX(*fret,1.0);
    for (i=0;i<n;i++) {
      temp=fabs(g(i))*FMAX(fabs(p(i)),1.0)/den;
      if (temp > test) test=temp;
    }
    if (test < gtol) {
      return;
    }
    for (i=0;i<n;i++) dg(i)=g(i)-dg(i);
    for (i=0;i<n;i++) {
      hdg(i)=0.0;
      for (j=0;j<n;j++) hdg(i) += hessian(i,j)*dg(j);
    }
    fac=fae=sumdg=sumxi=0.0;
    for (i=0;i<n;i++) {
      fac += dg(i)*xi(i);
      fae += dg(i)*hdg(i);
      sumdg += SQR(dg(i));
      sumxi += SQR(xi(i));
    }
    if (fac*fac > EPS*sumdg*sumxi) {
      fac=1.0/fac;
      fad=1.0/fae;
      for (i=0;i<n;i++) dg(i)=fac*xi(i)-fad*hdg(i);
      for (i=0;i<n;i++) {
	for (j=0;j<n;j++) {
	  hessian(i,j) += fac*xi(i)*xi(j)
	    -fad*hdg(i)*hdg(j)+fae*dg(i)*dg(j);
	}
      }
    }
    for (i=0;i<n;i++) {
      xi(i)=0.0;
      for (j=0;j<n;j++) xi(i) -= hessian(i,j)*g(j);
    }
  }
  cout << "too many iterations in dfpmin" << endl;
}
#undef ITMAX
#undef EPS
#undef TOLX
#undef STPMX

     /*
     ** The function
     **           ran1()
     ** is an "Minimal" random number generator of Park and Miller
     ** (see Numerical recipe page 280) with 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 IA 16807
#define IM 2147483647
#define AM (1.0/IM)
#define IQ 127773
#define IR 2836
#define NTAB 32
#define NDIV (1+(IM-1)/NTAB)
#define EPS 1.2e-7
#define RNMX (1.0-EPS)

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

   if (*idum <= 0 || !iy) {
      if (-(*idum) < 1) *idum=1;
      else *idum = -(*idum);
      for(j = NTAB + 7; j >= 0; j--) {
         k     = (*idum)/IQ;
         *idum = IA*(*idum - k*IQ) - IR*k;
         if(*idum < 0) *idum += IM;
         if(j < NTAB) iv[j] = *idum;
      }
      iy = iv[0];
   }
   k     = (*idum)/IQ;
   *idum = IA*(*idum - k*IQ) - IR*k;
   if(*idum < 0) *idum += IM;
   j     = iy/NDIV;
   iy    = iv[j];
   iv[j] = *idum;
   if((temp=AM*iy) > RNMX) return RNMX;
   else return temp;
}
#undef IA
#undef IM
#undef AM
#undef IQ
#undef IR
#undef NTAB
#undef NDIV
#undef EPS
#undef RNMX

// End: function ran1()





//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.*ran1(idum) -1.0;
      v2 = 2.*ran1(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(Matrix &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(Matrix &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(Matrix &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(Matrix &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(Matrix &R, double alpha, double beta, int Z) {

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

  Matrix R_plus(Z,3), R_minus(Z,3);
  R_plus = R; R_minus = R;

  //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);

  return potential + kinetic;
}

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

  Matrix R_plus(Z,3), R_minus(Z,3); 
  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);
    }
  }
}

//MC calculation of <H> with importance sampling
void RunMC(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
  Matrix R(Z,3), F(Z,3), R_new(Z,3), F_new(Z,3);
  for (int i = 0; i < Z; i++)
    for (int j = 0; j < 3; j++)
      R(i,j) = gaussian_deviate(&idum);

  //Initialize quantum force
  calcQF(R,F,alpha,beta,Z);

  double EL; //Local energy
  double sqrtdt = sqrt(delta_t);
  double D = .5;
  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 (ran1(&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
}

double energy(Vector &param) {

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

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

  return sum / N;
}

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

  RunMC(sum,ESquared,NumberMCGlobal,delta_t_global,idum_global,param(0),N,param(1),
	  Z_global,sum1alpha,sum2alpha,sum1beta,sum2beta,true);
  
  double E = sum / N;

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

  cout << "alpha = " << param(0) << "   beta =  " << param(1) << 
    "    alphaderiv = " << values(0) << "    betaderiv = " << values(1) << 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.5; //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;

  Vector param(2); param(0) = alpha; param(1) = beta;

  double gtol = .001;
  int iter;
  double fret;
  dfpmin(param, 2, gtol, &iter, &fret, &energy, &energy_deriv);

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

  return 0;
}