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

using namespace std;

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

extern "C"
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);

void  CmpOverlap(function2D<double>& Olap, const function1D<double>& base)
{// Calculates overlap O(a,b)=(Pi/(a+b))^{3/2}
  for (int i=0; i<Olap.size_N(); i++)
    for (int j=0; j<Olap.size_Nd(); j++)
      Olap(i,j) = pow(M_PI/(base[i]+base[j]),3/2.);
}

void CmpHamiltn(function2D<double>& Ham, const function1D<double>& base)
{// Calculates Hamiltonian K(a,b) = 3*a*b/(a+b)*(Pi/(a+b))^{3/2}
 //                        Vext(a,b) = -4*Pi/(a+b)
  for (int i=0; i<base.size(); i++)
    for (int j=0; j<base.size(); j++){
      Ham(i,j) = 3*base[i]*base[j]*pow(M_PI/(base[i]+base[j]),3/2.)/(base[i]+base[j]);
      Ham(i,j) += -4*M_PI/(base[i]+base[j]);
    }
}

void CmpU(vector4& Uc, const function1D<double>& base)
{ // 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.);
  for (int i=0; i<base.size(); i++)
    for (int j=0; j<base.size(); j++)
      for (int k=0; k<base.size(); k++)
	for (int l=0; l<base.size(); l++){
	  double b1 = base[i]+base[k];
	  double b2 = base[j]+base[l];
	  Uc[i][j][k][l] = cp/(b1*b2*sqrt(b1+b2));
	}
}

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];
	  sumH +=  Uc[k][i][l][j]*rho[k][l];
	}
      }
      F(i,j) = Ham(i,j)+sumH+sumF;
    }
  }
}

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


int Eigensystem(int N, double& Energy, const function2D<double>& Olap, function2D<double>& F)
{// C++ wrapper function for general eigenvalue problem
  static function2D<double> tOlap(N,N);
  static function1D<double> eigenval(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, eigenval.MemPt(), work.MemPt(), &lwork, iwork.MemPt(), &liwork, &info);
  Energy = eigenval[0];
  if (info) cerr<<"Not sucessfull solving eigenvalue problem with dsygvd"<<endl;
  return info;
}

double TotEnergy(double En, const function2D<double>& rho, const function2D<double>& Ham)
{
  double sum = En;
  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;
}

double RadialDensity(double r, const function2D<double>& rho, const function1D<double>& base)
{
  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;
}

int main(int argc, char *argv[], char *env[])
{
  int N=4;// Number of base functions
  double precision = 1e-6;  // Stops when dE<precision
  int i=0;
  while (++i<argc){
    std::string str(argv[i]);
    if (str=="-precision" && i<argc-1) precision = atof(argv[++i]);
    if (str=="-h" || str=="--help"){
      std::clog<<"********* Hartree-Fock program for He **************\n";
      std::clog<<"**                                                **\n";
      std::clog<<"**      Copyright Kristjan Haule, 26.09.2005      **\n";
      std::clog<<"****************************************************\n";
      std::clog<<"\n";
      std::clog<<"He0 [-precision double]\n" ;
      std::clog<<"Options:   -precision  Stops when dE<precision ("<<precision<<")\n";
      std::clog<<"*****************************************************\n";
      return 0;
    }
  }

  // Base functions Gaussians and are taken from the Thijssen book
  function1D<double> base(N);
  base[0] = 0.298073;
  base[1] = 1.242567;
  base[2] = 5.782948;
  base[3] = 38.47497;

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

  CmpU(Uc, base);
  
  // Effective Hartree-Fock Hamiltonian
  function2D<double> F(base.size(),base.size());

  // Ground state vector
  function1D<double> C(base.size());
  //Normalized density matrix
  function2D<double> rho(base.size(),base.size());
  
  // Initial guess for the ground state
  for (int i=0; i<C.size(); i++) C[i] = 1.0;
  // Initial density matrix. Equivalent to U=0
  rho=0; 
  double En=0, Enp;
  cout.precision(16);
  do{
    Enp=En;
    CmpHeffHartree(F, Ham, Uc, rho);
    if (Eigensystem(base.size(), En, Olap, F)) return 1;
    C = F[0];
    clog<<En<<" "<<TotEnergy(En,rho,Ham)<<endl;
    CmpDensityM(rho, C, Olap);
  } while(fabs(En-Enp)>precision);

  double L=3.5;
  int M=400;
  for (int i=0; i<M; i++){
    double r = (L*i)/(M-1);
    double rhor = RadialDensity(r, rho, base);
    cout<<r<<" "<<rhor<<" "<<4*M_PI*r*r*rhor<<endl;// Radial Density
  }
  return 0;
}