/*
    DFT++ is a density functional package developed by the research group
    of Professor Tomas Arias

    Copyright 1996-2003 Sohrab Ismail-Beigi

    This file is part of DFT++.

    DFT++ is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    DFT++ is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with DFT++; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

    Please see the file CREDITS for a list of authors.

    For academic users, we request that publications using results obtained with
    this software reference

    "New algebraic formulation of density functional calculation," by Sohrab Ismail-Beigi
    and T.A. Arias, Computer Physics Communications 128:1-2, 1-45 (June 2000).

    and, if using the wavelet basis, further reference

    "Multiresolution analysis of electronic structure: semicardinal and wavelet bases,"
    T.A. Arias, Reviews of Modern Physics 71:1, 267-311 (January 1999).

    and 

    "Robust ab initio calculation of condensed matter: transparent convergence through
    semicardinal multiresolution analysis,'' I.P. Daykov, T.A. Arias, and
    Torkel D. Engeness, Physical Review Letters, 90:21, 216402 (May 2003).

    For your convenience, preprints of the above articles may be obtained from
    http://arXiv.org/abs/cond-mat/9909130, 9805262, and 0204411, respectively.
*/

#include "header.h"

//Sum up the total charge recursively
double totcharge(struct Grid *ingrid)
{
  int ix,iy,iz;
  double sumcharge=0.;
  int nnx,nny,nnz;

  nnx=ingrid->dim.x;
  nny=ingrid->dim.y;
  nnz=ingrid->dim.z;

  double scalefactor;

  scalefactor=ingrid->scale.x*ingrid->scale.y*ingrid->scale.z;

  for (ix=0;ix<nnx;ix++)
    for (iy=0;iy<nny;iy++)
      for (iz=0;iz<nnz;iz++)
        if ( ( (ix%2 == 1) || (iy%2 == 1) || (iz%2 == 1) ) || (!ingrid->parent))
          sumcharge+=scalefactor*ingrid->dat[iz+(nnz*(iy+nny*ix))];
  // sumcharge+=ingrid->dat[iz+(nnz*(iy+nny*ix))];
  
  dft_log("Sumcharge is: %d %20.16lf\n",ingrid->level,sumcharge);

  if (ingrid->child)
      sumcharge+=totcharge(ingrid->child);

  if (ingrid->sibling)
    sumcharge+=totcharge(ingrid->sibling);

  return(sumcharge);
}

/* Using Teter's nuclear charge distribution*/
double nuc_charge(vector3 r, double Z)
{
  const double pi = M_PI;
  const double c1=3.132576693428;
  const double c2=-2.68355838224;
  const double c3=1-c1-c2;
  
  double d,d1,d2,d3;

  double rad2=abs2(r);

  double charge;

  d=1./(4.*Z);

  d1=d/sqrt(2.); d2=d; d3=sqrt(2.)*d;
  
  charge=
    Z*c1*exp(-rad2/(d1*d1))/(pow(pi,1.5)*d1*d1*d1)+
    Z*c2*exp(-rad2/(d2*d2))/(pow(pi,1.5)*d2*d2*d2)+
    Z*c3*exp(-rad2/(d3*d3))/(pow(pi,1.5)*d3*d3*d3);

  return(charge);

}

void 
eval_nuc_charge(struct Grid *self, vector3 porg, 
           matrix3 R, int pad, double Z)
{
  // Scale of children (if any) 
  vector3 scl;

  // Origin of self (note that self->org is in the _parents_ coordinates 
  vector3 org, temp, Rtemp;

  int i,j,k;
  int i2,j2,k2;
  vector3 r;

  real A,B,C; // lengths of lattice vectors
  A = abs(R[0]); B = abs(R[1]); C = abs(R[2]);

  // Return immediately if there is nothing to do 
  // No excutables BEFORE this statement, please! 
  if (self==NULL) return;

  // set the scale
  scl.v[0]=self->scale.x; scl.v[1]=self->scale.y; scl.v[2]=self->scale.z;

  //dft_log("Level %d : scale (%lf %lf %lf)\n",self->level,scl.v[0],scl.v[1],
  //	  scl.v[2]);

  // Origin of self (note that self->org is in the _parents_ coordinates 
  temp.v[0] = self->org.x*2.0*scl.v[0]/A;
  temp.v[1] = self->org.y*2.0*scl.v[1]/B;
  temp.v[2] = self->org.z*2.0*scl.v[2]/C;
  Rtemp = R*temp;
  org = porg + Rtemp;
  // Now org is a cartesian vector.
//    org.v[0] = porg.v[0] + self->org.x *2.0*scl.v[0];
//    org.v[1] = porg.v[1] + self->org.y *2.0*scl.v[1];
//    org.v[2] = porg.v[2] + self->org.z *2.0*scl.v[2];

  //dft_log("Level %d : org (%lf %lf %lf)\n",self->level,org.v[0],org.v[1],
  //	  org.v[2]);

  //double maxval = 0.0;
  // Fill in data 
  for (i=0; i<self->dim.x; i++)
    for (j=0; j<self->dim.y; j++)
      for (k=0; k<self->dim.z; k++){
          *( self->dat+(i*self->dim.y+j)*self->dim.z+k ) = 0.0;
          for (i2=-pad; i2<=pad; i2++)
            for (j2=-pad; j2<=pad; j2++)
              for (k2=-pad; k2<=pad; k2++)
                {
		  // For a non-orthorhombic cell, we have to multiply by R,
		  // the matrix of lattice vectors.
//                    r.v[0]=org.v[0]+i*scl.v[0]+i2*top.v[0];
//                    r.v[1]=org.v[1]+j*scl.v[1]+j2*top.v[1];
//                    r.v[2]=org.v[2]+k*scl.v[2]+k2*top.v[2];
		  temp.v[0] = i*scl.v[0]/A;
		  temp.v[1] = j*scl.v[1]/B;
		  temp.v[2] = k*scl.v[2]/C;
                  Rtemp = R*temp;
		  r = org + Rtemp + i2*R[0] + j2*R[1] + k2*R[2];

                  *( self->dat+(i*self->dim.y+j)*self->dim.z+k ) 
                    += nuc_charge(r,Z);
		  
//  		  if((abs(r) <= 0.25) && (self->level < 2))
//  		    dft_log("nuc_charge value = %1.12lf\n",
//  			    *( self->dat+(i*self->dim.y+j)*self->dim.z+k )); 
//  		  if((fabs(*( self->dat+(i*self->dim.y+j)*self->dim.z+k ))
//  		     > maxval) && ( ( (i%2 == 1) || (j%2 == 1) || 
//  				      (k%2 == 1) ) || (!self->parent) ))
//  		    maxval = 
//  		      fabs(*( self->dat+(i*self->dim.y+j)*self->dim.z+k ) );
		      
                }
        }

  //dft_log("Level %d maxval = %1.12lf\n",self->level,maxval);

  // Grab data from up above to ensure consistency 
  if (self->level>0)
    for (i=0; i<self->dim.x; i+=2)
      for (j=0; j<self->dim.y; j+=2)
        for (k=0; k<self->dim.z; k+=2)
          *( self->dat+(i*self->dim.y+j)*self->dim.z+k ) = 
            *( self->parent->dat + 
               ( (self->org.x+i/2) * self->parent->dim.y + 
                 (self->org.y+j/2) ) * self->parent->dim.z +
	       (self->org.z+k/2) );

  // Call child and sibling, return immediately if none 
  if (self->child!=NULL)
    eval_nuc_charge(self->child,org,R,pad,Z);
  if (self->sibling!=NULL)
    eval_nuc_charge(self->sibling,porg,R,pad,Z);

  return;
}

// Origin must be in Cartesian coordinates!
void
get_nuc_charge(RealSpaceScalarFieldColumn &rho, vector3 origin, 
		  real charge)
{
  // This only pays attention to the real part of rho.  Zero out the
  // imaginary part before calling this routine.

  Gridspec *gridspec;
  
  // choose the appropriate <sparse|dense> gridspec
  switch(rho.embedding){
    case SPARSE:
      gridspec = rho.basis_spec->gridspec;
      break;
    case DENSE:
      gridspec = rho.basis_spec->gridspec_dense;
      break;
    default:
      die("getnuc_charge: Unknown input embedding %d\n",rho.embedding);
  }
  
  Grid *grid = mkgridnew(gridspec,NULL);

  // print the grid structure only from the master node
#ifdef DFT_MPI
  if(System::Get_procID() == 0)
    printgridstruct(stdout,grid);
#else
  printgridstruct(stdout,grid);
#endif
     
  // Evidently the eval_nuc_charge function uses the origin as a shift, so
  // negate it.
  origin *= (-1.0);

  GetRe(grid,rho);
  eval_nuc_charge(grid,origin,rho.basis_spec->lattice->latvec,1,charge);
  //dft_log("Totcharge = %1.12lf\n",totcharge(grid));
  PutRe(grid,rho);

  killgridnew(grid);
}

void
get_ionic_charge_density(RealSpaceScalarFieldColumn &rho, Ioninfo &ioninfo, 
			 Lattice &lattice)
{
  int type, ion;
  vector3 origin;
  RealSpaceScalarFieldColumn add(rho);

  rho.zero_out();
  add.zero_out();

  for(type=0; type<ioninfo.nspecies; type++){
    for(ion=0; ion<ioninfo.species[type].natoms; ion++){
	  // atomic positions are stored in lattice coordinates
	  origin = lattice.R * ioninfo.species[type].atpos[ion];
	  get_nuc_charge(add,origin,ioninfo.species[type].Z);
	  rho += add;
	  add.zero_out();
	}
  }
  // Sum the total charge on the grid.
  Gridspec *gridspec;
  
  // choose the appropriate <sparse|dense> gridspec
  switch(rho.embedding){
    case SPARSE:
      gridspec = rho.basis_spec->gridspec;
      break;
    case DENSE:
      gridspec = rho.basis_spec->gridspec_dense;
      break;
    default:
      die("getnuc_charge: Unknown input embedding %d\n",rho.embedding);
  }
  
  Grid *grid = mkgridnew(gridspec,NULL);
  //ipd: we must take care if parallel!!! - don't want to just use printf
  //dft_log("Grid Structure:\n");
  //printgridstruct(stdout,grid);

  GetRe(grid,J(rho));
  dft_log("\nTotal charge = %1.12lf\n",totcharge(grid));
  killgridnew(grid);
}

void
Vlocpot(Everything &everything)
{
  // Check the read_vloc_flag to see if we have to calculate Vlocpot, or
  // just read it in.
  if(everything.ioninfo.read_vloc_flag)
    everything.elecvars.Vlocpot.read(everything.ioninfo.vloc_filename);
  else
    {
      RealSpaceScalarFieldColumn rho(everything.elecvars.n);
      
      get_ionic_charge_density(rho,everything.ioninfo,everything.lattice);
      
      // Rho is considered by the code to be a number density, so we have to
      // explicitly add a minus sign to deal with positive charge.
      rho *= (-1.0);
      
      // HACK ALERT!!!
      // Set the value of global_abs2Yg to give convdigits = 9 in invL.
      global_abs2Yg = 1.e-10;
      
      solve_poisson(rho,everything.elecvars.Vlocpot);

//        Grid *nucpoisson=mkgridnew(rho.basis_spec->gridspec,NULL);
//        GetRe(nucpoisson,everything.elecvars.Vlocpot);
//        FILE *fpnucpoisson=fopen("nuc_poisson.dft","w");
//        writetofile(fpnucpoisson,nucpoisson);
//        killgridnew(nucpoisson);
//        fclose(fpnucpoisson);

      // This is due to differing definitions of Vlocpot between the wavelets
      // and planewave codes.
      apply_Obar(everything.elecvars.Vlocpot,everything.elecvars.Vlocpot);
    }
}