/*
    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"

/* 
 * Implementing the member functions of PW_Basis and PW_BasisSpec
 */


///////////// PW_BasisSpec  ///////////////////////

PW_BasisSpec::PW_BasisSpec()
{
  Ecut = 0;
  size_realgrid =Nx = Ny = Nz = NxNyNz = 0;
  kpt_fold[0] = kpt_fold[1] = kpt_fold[2] = 1;
  auto_fftbox = 0;
  basis_flag = 0;
}

void PW_BasisSpec::setup(Everything &e)
{
  lattice = &(e.lattice);
}

///////////// PW_Basis  ///////////////////////////

PW_Basis::PW_Basis()
{
  basis_spec = NULL;
  qnum = NULL;

  nbasis = 0;

  Gxmin = Gxmax = Gymin = Gymax = Gzmin = Gzmax = 0;
  Gx = Gy = Gz = NULL;
  index = NULL;
  kplusG2 = NULL;
  kplusGx = NULL; 
  kplusGy = NULL;
  kplusGz = NULL;
  invLkernel = NULL;
}

// destructor
PW_Basis::~PW_Basis()
{
  if(Gx != NULL)
    myfree(Gx);
  if(Gy != NULL)
    myfree(Gy);
  if(Gz != NULL)
    myfree(Gz);
  if(index != NULL)
    myfree(index);
}

void PW_Basis::operator=(PW_Basis &in)
{
  int i;
  
  basis_spec=in.basis_spec;
  qnum = in.qnum;
  Gxmin = in.Gxmin;
  Gxmax = in.Gxmax;
  Gymin = in.Gymin;
  Gymax = in.Gymax;
  Gzmin = in.Gzmin;
  Gzmax = in.Gzmax;
  nbasis = in.nbasis;
  kplusG2 = NULL;
  kplusGx = NULL; 
  kplusGy = NULL;
  kplusGz = NULL;
  invLkernel = NULL;

  if(nbasis > 0){
    Gx = (int *) mymalloc(sizeof(int)*nbasis, "Gx", "PW_Basis=()");
    Gy = (int *) mymalloc(sizeof(int)*nbasis, "Gy", "PW_Basis=()");
    Gz = (int *) mymalloc(sizeof(int)*nbasis, "Gz", "PW_Basis=()");

    index = (int *) mymalloc(sizeof(int)*nbasis, "index", "PW_Basis=()");
  
    // separate loops for cache
    for(i=0; i<nbasis; i++){
      Gx[i] = in.Gx[i];
    }
    for(i=0; i<nbasis; i++){
      Gy[i] = in.Gy[i];
    }
    for(i=0; i<nbasis; i++){
      Gz[i] = in.Gz[i];
    }
    for(i=0; i<nbasis; i++){
      index[i] = in.index[i];
    }
    
  } else{
    Gx = Gy = Gx = NULL;
    index = NULL;
  }

}

/*
 * First stage Basis initialization.
 * Set up crystal structure parameters.
 */

void PW_Basis::setup_template(PW_BasisSpec &bs)
{
  int i, j;
  matrix3 invGGT,box;
  vector3 e,f;
  int ibox[3],fftbox[3];


  dft_log("\n----- PW_Basis::setup_grid() -----\n");
  dft_log("Setting up the real space grid\n\n");

  basis_spec = &bs;


  /* We want to know for what vectors v lying on the constant energy surface
   * Ecut = 0.5*(v,GGT*v) the x, y, or z component/projection of v is maximal.
   * This is easy:  the gradient at that point v must lie in the x, y, or z
   * direction!  I.e. if e is the unit direction we are interested in,
   * GGT*v = mu*e or v = mu*inv(GGT)*e, where mu is chosen to be
   * mu = sqrt(2*Ecut/(e,inv(GGT)*e) to ensure v is on the const. energy
   * surface. We solve this for e being x,y,z unit vectors and put the
   * resulting vectors into the rows of box. */
  invGGT = basis_spec->lattice->invGGT;
  for (i=0; i < 3; i++) {
    for (j=0; j < 3; j++)
      e.v[j] = 0.0;
    e.v[i] = 1.0;
    f = invGGT*e;
    f = sqrt(2.0*basis_spec->Ecut/(e*f))*f;
    for (j=0; j < 3; j++)
      box.m[j][i] = f.v[j];
  }

  dft_log("\nEnergy cutoff Ecut = %lg Hartrees\n",basis_spec->Ecut);
  dft_log("On the surface Ecut = 0.5*|G|^2, the vector extremizing\n");
  dft_log("(G,e_x) = "); box[0].print(dft_global_log,"%lg ");
  dft_log("(G,e_y) = "); box[1].print(dft_global_log,"%lg ");
  dft_log("(G,e_z) = "); box[2].print(dft_global_log,"%lg ");

  /* Truncate the values on the diagonal of box to integers.  This will
   * be the size of a box along x/y/z which will contain all G-vectors
   * of energy < Ecut. */
  for (i=0; i < 3; i++)
    ibox[i] = (int)fabs(box.m[i][i]);
  Gxmax = ibox[0]; Gxmin = -ibox[0];
  Gymax = ibox[1]; Gymin = -ibox[1];
  Gzmax = ibox[2]; Gzmin = -ibox[2];

  dft_log("\nSize of box containing G-vectors = ");
  dft_log("[-%d,%d] by [-%d,%d] by [-%d,%d]\n",
	    ibox[0],ibox[0],ibox[1],ibox[1],ibox[2],ibox[2]);
  dft_log("Size of box containing density = ");
  dft_log("[-%d,%d] by [-%d,%d] by [-%d,%d]\n\n",
	    2*ibox[0],2*ibox[0],2*ibox[1],2*ibox[1],2*ibox[2],2*ibox[2]);

  /* Find the minimal FFT box size the factors into the primes (2,3,5,7).
   * The minimum value for the size of the fftbox is 2*2*ibox+1 because
   * we square the wave-functions and so the FFT box G-vectors must
   * at least range over -2*ibox to 2*ibox inclusive.
   * However, various other routines require the FFT box size to be an
   * even integer, so we start the fftbox-size at 4*ibox+2.
   * The loop tries to factorize the fftbox-size into (2,3,5,7)...if that
   * isn't doable, it increases the fftbox-size by 2 and tries again. */
  for (i=0; i < 3; i++) {
    int b,n2,n3,n5,n7,done_factoring;

    fftbox[i] = 4*ibox[i]+2  -2;
    /* increase fftbox[i] by 2 and try to factor it into (2,3,5,7) */
    do {
	  fftbox[i] += 2;
	  b = fftbox[i];
	  n2 = n3 = n5 = n7 = done_factoring = 0;
	  while (!done_factoring) {
        if (b%2==0) { n2++; b /= 2; continue; }
        if (b%3==0) { n3++; b /= 3; continue; }
        if (b%5==0) { n5++; b /= 5; continue; }
        if (b%7==0) { n7++; b /= 7; continue; }
        done_factoring = 1;
      }
	}
    while (b != 1); /*  b==1 means fftbox[i] is (2,3,5,7) factorizable */
    dft_log("fftbox[%d] = %d =(2^%d)*(3^%d)*(5^%d)*(7^%d)\n",
            i,fftbox[i],n2,n3,n5,n7);
  }

  /* If we're given FFT box sizes to use, then use them! */
  if (basis_spec->auto_fftbox == FALSE) {
    fftbox[0] = basis_spec->Nx;
    fftbox[1] = basis_spec->Ny;
    fftbox[2] = basis_spec->Nz;
    dft_log("==============================================\n");
    dft_log("Overiding with specified sizes: %d by %d by %d\n",
            basis_spec->Nx,basis_spec->Ny,basis_spec->Nz);
    dft_log("==============================================\n");
  }

  for (i=0; i < 3; i++)
    if (fftbox[i] < 4*ibox[i]+2)
      die("\nBasis::setup_grid():  fftbox[%d] is too small.\nIt should be AT LEAST %d = 4*%d+2\n\n",i,4*ibox[i]+2,ibox[i]);

  dft_log("Using fftbox = %d by %d by %d\n\n",
	    fftbox[0],fftbox[1],fftbox[2]);

  basis_spec->Nx = fftbox[0];
  basis_spec->Ny = fftbox[1];
  basis_spec->Nz = fftbox[2];
  basis_spec->NxNyNz = basis_spec->Nx * basis_spec->Ny * basis_spec->Nz;

  basis_spec->size_realgrid = basis_spec->NxNyNz;
  
  // Set up the FFT routines
  setupFFT3D(basis_spec->Nx,basis_spec->Ny,basis_spec->Nz);
}


/*
 * Second part of the initialization:
 * setting up the G-vectors and index-arrays.
 */
void PW_Basis::setup(const vector3 &kvec)
{
  /* Figure out the G-vectors within the cutoff energy.
   * First count up how many there are, then put them in the basis struct.
   * In the continuum limit, there should be (4*pi/3)*(2*Ecut)^(3/2)/det(G)
   * points:  divide the volume of the maximum energy sphere in G-space by
   * the volume of the primitive cell in G-space. */
  real G2;
  real G2max = 0;

  int i, j, k, n, ind, ibox[3];
  int Nx = basis_spec->Nx;
  int Ny = basis_spec->Ny;
  int Nz = basis_spec->Nz;
  vector3 f;

  ibox[0] = (Nx - 2)/4 + 1;
  ibox[1] = (Ny - 2)/4 + 1;
  ibox[2] = (Nz - 2)/4 + 1;

  // figure out nbasis
  matrix3 GGT = basis_spec->lattice->GGT;
  nbasis = 0;
  for (i = -ibox[0]; i <= ibox[0]; i++)
    for (j = -ibox[1]; j <= ibox[1]; j++)
      for (k = -ibox[2]; k <= ibox[2]; k++) {
        f.v[0] = i;
        f.v[1] = j;
        f.v[2] = k;
        f += kvec;
        G2 = f*(GGT*f);
        if (0.5*G2 <= basis_spec->Ecut)
          nbasis++;
      }

  dft_log("nbasis = %d for k = [%6.3f %6.3f %6.3f]\n",
          nbasis, kvec.v[0],kvec.v[1],kvec.v[2]);
  
  // allocate memory for the G vectors and index
  Gx = (int *)mymalloc(sizeof(int)*nbasis,"Gx[]","Basis::setup_Gvectors()");
  Gy = (int *)mymalloc(sizeof(int)*nbasis,"Gy[]","Basis::setup_Gvectors()");
  Gz = (int *)mymalloc(sizeof(int)*nbasis,"Gz[]","Basis::setup_Gvectors()");
  index = (int *)mymalloc(sizeof(int)*nbasis,
                          "index[]","Basis::setup_Gvectors()");
  n = 0;
  for (i = -ibox[0]; i <= ibox[0]; i++)
    for (j = -ibox[1]; j <= ibox[1]; j++)
      for (k = -ibox[2]; k <= ibox[2]; k++) {
        f.v[0] = i;
        f.v[1] = j;
        f.v[2] = k;
        f += kvec;
        G2 = f*(GGT*f);
        if (0.5*G2 <= basis_spec->Ecut) {
	      if (G2 > G2max)
            G2max = G2;
	      Gx[n] = i;
	      Gy[n] = j;
	      Gz[n] = k;
	      ind = 0;
	      if (k >= 0) ind += k;
	      else        ind += k+ Nz;
	      if (j >= 0) ind += Nz*j;
	      else        ind += Nz*(j+Ny);
	      if (i >= 0) ind += Nz*Ny*i;
	      else        ind += Nz*Ny*(Nx+i);
	      index[n] = ind;
	      dft_log(DFT_ANAL_LOG,
                  "G = [ %3d %3d %3d ]  index = %7d  G2 = %lg\n",
                  Gx[n], Gy[n], Gz[n], index[n], G2);
	      n++;
	    }
      }

  // Spit out some diagnostics
  dft_log("\n0.5*|G|^2 = %lg for highest energy plane-wave\n",
          0.5*G2max);
  dft_log("\n");


  dft_log(DFT_ANAL_LOG, "all done! :-)\n");
}