#include <iostream>
#include <iomanip>
#include <vector>
#include "function.h"

using namespace std;

typedef vector<vector<vector<vector<double> > > > vector4;

extern "C" {
  double erf(double x);
  void dsygvd_(const int* ITYPE, const char* JOBZ, const char* UPLO, const int* N,
	       double* A, const int* LDA, double* B, const int* LDB, double* W, double* WORK, const int* LWORK,
	       int* IWORK, const int* LIWORK, int* INFO);
};


inline double F0(double x)
{// This function is needed for overlap calculation of molecules in GTO base
 //  This is call to external erf function
  if (fabs(x)<1e-12) return 1.0;
  double sx = sqrt(x);
  return sqrt(M_PI)*erf(sx)/(2*sx);
}


void CmpHeffHartree(function2D<double>& F, const function2D<double>& Ham, const vector4& Uc, const function2D<double>& rho)
{// Computes effective Hartree hamiltonian
  for (int i=0; i<F.size_N(); i++){
    for (int j=0; j<F.size_Nd(); j++){
      double sumH=0, sumF=0;
      for (int k=0; k<rho.size_N(); k++){
	for (int l=0; l<rho.size_Nd(); l++){
 	  sumH +=  Uc[k][i][l][j]*2*rho[k][l];
 	  sumF -=  Uc[k][i][j][l]*1*rho[k][l];
	}
      }
      F(i,j) = Ham(i,j)+sumH+sumF;
    }
  }
}

void CmpDensityM(function2D<double>& rho, const function2D<double>& C, int Nel, const function2D<double>& olap)
{// Computes density matrix
  rho=0;
  for (int a=0; a<Nel; a++){
    for (int i=0; i<C.size_Nd(); i++)
      for (int j=0; j<C.size_Nd(); j++)
	rho[i][j] += 0.5*C[a/2][i]*C[a/2][j];
  }
}


int Eigensystem(int N, function1D<double>& Energy, const function2D<double>& Olap, function2D<double>& F)
{// C++ wrapper function for general eigenvalue problem
  static function2D<double> tOlap(N,N);
  int lwork = 1 + 6*N + 2*N*N + 10;
  static function1D<double> work(lwork);
  int liwork = 3 + 5*N + 10;
  static function1D<int> iwork(liwork);
  int itype=1;
  int info=0;
  tOlap = Olap;
  int lF = F.size_Nd();
  int lO = Olap.size_Nd();
  dsygvd_(&itype, "V", "U", &N,	F.MemPt(), &lF, tOlap.MemPt(), &lO, Energy.MemPt(), work.MemPt(), &lwork, iwork.MemPt(), &liwork, &info);
  if (info) cerr<<"Not sucessfull solving eigenvalue problem with dsygvd"<<endl;
  return info;
}

double TotEnergy(const function1D<double>& Ei, int Nel, const function2D<double>& rho, const function2D<double>& Ham)
{// Calculates Total Energy within HF approximation
  double sum = 0;
  for (int i=0; i<Nel; i++) sum += 0.5*Ei[i/2];
  for (int i=0; i<rho.size_N();  i++)
    for (int j=0; j<rho.size_Nd(); j++)
      sum += rho[i][j]*Ham[i][j];
  return sum;
}

class He{
  int Z;
  function1D<double> base;
public:
  He(double=0) : base(4), Z(2)
  {
    base[0] = 0.298073;
    base[1] = 1.242567;
    base[2] = 5.782948;
    base[3] = 38.47497;
  };
  int size(){return base.size();}
  void CmpOverlap(function2D<double>& Olap);
  void CmpHamiltn(function2D<double>& Ham, const function2D<double>& Olap);
  void CmpU(vector4& Uc, const function2D<double>& Olap);
  double RadialDensity(double r, const function2D<double>& rho);
  void SetDistance(double){};
  int Nel(){return Z;}
  double NucleiRepulsion(){ return 0.0;}
};

void  He::CmpOverlap(function2D<double>& Olap)
{// Calculates overlap O(a,b)=(Pi/(a+b))^{3/2}
  for (int i=0; i<base.size(); i++)
    for (int j=0; j<base.size(); j++)
      Olap(i,j) = pow(M_PI/(base[i]+base[j]),3/2.);
}
void He::CmpHamiltn(function2D<double>& Ham, const function2D<double>& Olap)
{// Calculates Hamiltonian K(a,b) = 3*a*b/(a+b)*(Pi/(a+b))^{3/2}
 //                        Vext(a,b) = -2*Z*Pi/(a+b)
  double prf = 2*Z/sqrt(M_PI);
  for (int i=0; i<base.size(); i++)
    for (int j=0; j<base.size(); j++){
      Ham(i,j) = 3*base[i]*base[j]/(base[i]+base[j])*Olap(i,j);
      Ham(i,j) -= prf*sqrt(base[i]+base[j])*Olap(i,j);
    }
}
void He::CmpU(vector4& Uc, const function2D<double>& Olap)
{ // Coulomb matrix is U_{abcd}=2*pi^(5/2)/((a+c)(b+d)Sqrt(a+b+c+d))
  //  double cp = 2*pow(M_PI,5/2.);
  double prf = 2/sqrt(M_PI);
  for (int i=0; i<base.size(); i++)
    for (int j=0; j<base.size(); j++)
      for (int k=0; k<base.size(); k++){
	double prf_Olap_ik = prf*Olap(i,k);
	double b1 = base[i]+base[k];
	for (int l=0; l<base.size(); l++){
	  double b2 = base[j]+base[l];
	  Uc[i][j][k][l] = prf_Olap_ik*Olap(j,l)*sqrt(b1*b2/(b1+b2));
	}
      }
}
double He::RadialDensity(double r, const function2D<double>& rho)
{
  double sum=0;
  for (int a=0; a<base.size(); a++)
    for (int b=0; b<base.size(); b++)
      sum += rho[a][b]*exp(-(base[a]+base[b])*r*r);
  return sum;
}

class H2{
  int N;
  int Z1, Z2;
  function1D<double> base;
  double Dst, Dst_2, Dst05;
public:
  H2(double Dst_=1.) : N(4), Z1(1), Z2(1), base(2*N), Dst(Dst_), Dst_2(Dst*Dst), Dst05(0.5*Dst)
  {
    base[0] = 13.00773;
    base[1] = 1.962079;
    base[2] = 0.444529;
    base[3] = 0.1219492;
    base[4] = base[0];
    base[5] = base[1];
    base[6] = base[2];
    base[7] = base[3];
  };
  int size(){return base.size();}
  void CmpOverlap(function2D<double>& Olap);
  void CmpHamiltn(function2D<double>& Ham, const function2D<double>& Olap);
  void CmpU(vector4& Uc, const function2D<double>& Olap);
  double NucleiRepulsion(){return Z1*Z2/Dst;}
  int Nel(){return Z1+Z2;}
  double RadialDensity(double r, double z, const function2D<double>& rho);
  void SetDistance(double Dst_)
  { Dst = Dst_; Dst_2 = Dst*Dst; Dst05=0.5*Dst; }
};

void  H2::CmpOverlap(function2D<double>& Olap)
{// Calculates overlap O(a,b)=(Pi/(a+b))^{3/2}
  for (int i=0; i<base.size(); i++){
    for (int j=0; j<base.size(); j++){
      int ii = i/N, jj = j/N; // which molecule (0 or 1)
      Olap(i,j)     = pow(M_PI/(base[i]+base[j]),3/2.);
      if (ii!=jj) Olap(i,j) *= exp(-Dst_2*base[i]*base[j]/(base[i]+base[j]));
    }
  }
}
void H2::CmpHamiltn(function2D<double>& Ham, const function2D<double>& Olap)
{// Calculates Hamiltonian K(a,b) = 3*a*b/(a+b)*(Pi/(a+b))^{3/2}
 //                        Vext(a,b) = -2*Z*Pi/(a+b)
  double prf[2] = {2*Z1/sqrt(M_PI),2*Z2/sqrt(M_PI)};
  for (int i=0; i<base.size(); i++){
    int ii = i/N; // which molecule (0 or 1)
    for (int j=0; j<base.size(); j++){
      int jj = j/N; // which molecule (0 or 1)
      double bij = base[i]+base[j];
      double R_pq = Dst05*((2*ii-1)*base[i]+(2*jj-1)*base[j])/bij; // center of mass
      double a_ij = base[i]*base[j]/bij;
      Ham(i,j) = 3*a_ij*Olap(i,j);                 // Basic kinetic part
      if (ii!=jj) Ham(i,j) *= (1-2./3.*a_ij*Dst_2); // If they sit on different atoms
      //      Ham(i,j) -= sqrt(bij)*Olap(i,j)*(prf[0]*F0(bij*sqr(R_pq+Dst05))+prf[1]*F0(bij*sqr(R_pq-Dst05))); // Nucleous contribution
      double F1 = F0(bij*sqr(R_pq+Dst05));
      double F2 = F0(bij*sqr(R_pq-Dst05));
      double dH = -sqrt(bij)*Olap(i,j)*(prf[0]*F1+prf[1]*F2);
      Ham(i,j) += dH;
    }
  }
}
void H2::CmpU(vector4& Uc, const function2D<double>& Olap)
{ // Coulomb matrix is U_{abcd}=2*pi^(5/2)/((a+c)(b+d)Sqrt(a+b+c+d))
  //  double cp = 2*pow(M_PI,5/2.);
  double prf = 2/sqrt(M_PI);
  for (int i=0; i<base.size(); i++){
    int ii = i/N; // which molecule (0 or 1)
    for (int j=0; j<base.size(); j++){
      int jj = j/N; // which molecule (0 or 1)
      for (int k=0; k<base.size(); k++){
	int kk = k/N; // which molecule (0 or 1)
	double R_ik = Dst05*((2*ii-1)*base[i]+(2*kk-1)*base[k])/(base[i]+base[k]);
	double prf_Olap_ik = prf*Olap(i,k);
	double b_ik = base[i]+base[k];
	for (int l=0; l<base.size(); l++){
	  int ll = l/N; // which molecule (0 or 1)
	  double R_jl = Dst05*((2*jj-1)*base[j]+(2*ll-1)*base[l])/(base[j]+base[l]);
	  double b_jl = base[j]+base[l];
	  double a_ijkl = b_ik*b_jl/(b_ik+b_jl);
	  Uc[i][j][k][l] = prf_Olap_ik*Olap(j,l)*sqrt(a_ijkl)*F0(a_ijkl*sqr(R_ik-R_jl));
	}
      }
    }
  }
}
double H2::RadialDensity(double r, double z, const function2D<double>& rho)
{
  double sum=0;
  for (int i=0; i<base.size(); i++){
    int ii = i/N; // which molecule (0 or 1)
    double R_ii = Dst05*(2*ii-1);
    for (int j=0; j<base.size(); j++){
      int jj = j/N; // which molecule (0 or 1)
      double R_jj = Dst05*(2*jj-1);
      double r2 = sqr(r);
      sum += rho[i][j]*exp(-base[i]*(r2+sqr(z-R_ii))-base[j]*(r2+sqr(z-R_jj)));
    }
  }
  return sum;
}


int main(int argc, char *argv[], char *env[])
{
  double precision = 1e-6;  // Stops when dE<precision
  double Dst = 1.388;// Distance between molecules
  int i=0;
  while (++i<argc){
    std::string str(argv[i]);
    if (str=="-precision" && i<argc-1) precision = atof(argv[++i]);
    if (str=="-R" && i<argc-1) Dst = atof(argv[++i]);
    if (str=="-h" || str=="--help"){
      std::clog<<"** Hartree-Fock program for He atom and H2 molecule **\n";
      std::clog<<"**                                                  **\n";
      std::clog<<"**      Copyright Kristjan Haule, 26.09.2005        **\n";
      std::clog<<"******************************************************\n";
      std::clog<<"\n";
      std::clog<<"hartreeF [-precision double]\n" ;
      std::clog<<"Options:   -precision  Stops when dE<precision ("<<precision<<")\n";
      std::clog<<"           -R          Distance between nucleous ("<<Dst<<")\n";
      std::clog<<"*****************************************************\n";
      return 0;
    }
  }


  H2 mol(Dst);

  // Hamiltonian and overlap storage
  function2D<double> Olap(mol.size(),mol.size()), Ham(mol.size(), mol.size());
  // Coulomb matrix storage
  vector4 Uc;
  
  // Resizing Coulomb matrix
  Uc.resize(mol.size());
  for (int i=0; i<mol.size(); i++){
    Uc[i].resize(mol.size());
    for (int j=0; j<mol.size(); j++) {
      Uc[i][j].resize(mol.size());
      for (int k=0; k<mol.size(); k++) Uc[i][j][k].resize(mol.size());
    }
  }

  // Ground state vector
  function1D<double> C(mol.size());
  // Initial guess for tmol ground state
  for (int i=0; i<C.size(); i++) C[i] = 1.0;
  //Normalized density matrix
  function2D<double> rho(mol.size(),mol.size());
  // Initial density matrix. Equivalent to U=0
  rho=0; 
  // Effective Hartree-Fock Hamiltonian
  function2D<double> F(mol.size(),mol.size());
  // Ionization energies
  function1D<double> Ei(mol.size());
  
  mol.CmpOverlap(Olap);
  mol.CmpHamiltn(Ham, Olap);
  mol.CmpU(Uc, Olap);

  double TEn=0, pTEn;
  cout.precision(16);
  do{
    pTEn=TEn;
    CmpHeffHartree(F, Ham, Uc, rho);
    if (Eigensystem(mol.size(), Ei, Olap, F)) return 1;
    CmpDensityM(rho, F, mol.Nel(), Olap);
    TEn = TotEnergy(Ei,mol.Nel(), rho,Ham) + mol.NucleiRepulsion();
    clog<<Ei[0]<<" "<<TEn<<endl;
  } while(fabs(pTEn-TEn)>precision);
  cout<<"Total energy of H2 molecule at distance R="<<Dst<<" is "<<TEn<<endl;
  
  return 0;
}