/****************************************************************************
 *
 * DFT++:  density functional package developed by
 *         the research group of Prof. Tomas Arias, MIT.
 *
 * Principal author: Sohrab Ismail-Beigi
 *
 * Modifications for MPI version: Kenneth P Esler,
 *                                Sohrab Ismail-Beigi, and
 *                                Tairan Wang.
 *
 * Modifications for LSD version: Jason A Cline,
 *                                Gabor Csanyi,
 *                                Tairan Wang, and
 *                                Evan Reed
 *
 * Modifications for lattice/Pulay forces: Gabor Csanyi and
 *                                         Sohrab Ismail-Beigi
 *
 * K-point symmetries:  Nicolaj Moll
 *
 * Ionic dynamics and relaxation:  Tairan Wang and
 *                                 Evan Reed
 * 
 * Copyright (C) 1996-2000 Sohrab Ismail-Beigi
 *
 * Permission granted to the academic community to 
 * freely distribute and use this software provided that this notice is
 * maintained intact, and all publications resulting from its use
 * cite the following reference:
 *
 * Sohrab Ismail-Beigi and T. A. Arias, "New Algebraic Formulation of
 * Density Functional Calculation", Computer Physics Communications,
 * volume 128, issue 1-2, pages 1-45, June 2000.
 *
 ****************************************************************************/

/* 
 *   Sohrab Ismail-Beigi             May 1997
 *
 * A set of routines to calculate the empty bands via CG minimization
 * of an objective function (sum of empty state eigenvalues) with
 * constraints of orthonormality of the empty states with each other
 * and orthogonality to the fillled states.
 *
 */

/* $Id: calcempties.cpp_defunct,v 1.14.2.2 2002/02/25 18:42:15 ivan Exp $ */

#include <stdio.h>
#include <math.h>
#include <time.h>

#include "header.h"

/*
 * Does A_empty*Y where A_empty = Y - C_filled*C_filled^*O*Y,
 * i.e., project out the filled states from Y
 */
ColumnBundle
A_empty(const ColumnBundle &C_filled, const ColumnBundle &Y)
{
  ColumnBundle result(Y);
  result -= C_filled*(C_filled^(O(Y)));
  return result;
}

/*
 * U = (A*Y)^O(A*Y)   and   C = A*Y*U^(-1/2),
 * i.e. orthonormalized the Y into C for empty bands.
 */
void
calc_UC_empty(Elecinfo &empty_info,Elecvars &filled_vars,
              Elecvars &empty_vars)
{
  int q, nstates = empty_info.nstates;

  for (q=0; q < nstates; q++)
    {
      BlochState &s = empty_vars.states[q];
      ColumnBundle AY(s.Y);
      AY = A_empty(filled_vars.states[q].C,s.Y);
      s.U = AY^(O(AY));
      s.C = AY*Uminusonehalf(s.U,s.W,s.mu);
    }
}

/*
 * Returns the empty band energy function which is just
 *        E = sum_k { trace(C[k]^Hsp*C[k]) }
 */
real
empty_energy(Elecinfo &empty_info, Elecvars &filled_vars,
             Elecvars &empty_vars, Ioninfo &ioninfo,Lattice &lattice)
{
  int q;
  real Ener;
  
  Ener = 0.0;
  for (q=0; q < empty_info.nstates; q++){
    BlochState &s=empty_vars.states [q];
    Ener += REAL( trace(s.C^Hsp(s.C,s)) );
  }

  return Ener;
}

/*
 * Returns the empty band energy function which is just
 *        E = sum_k { trace(C[k]^Hsp*C[k]) }
 * as well as calculating its derivate versus Y[k] (placed into grad[])
 * and computing the subspace Hamiltonians and diagonalizing them.
 */
static real
empty_energy_grad_Hsub(Elecinfo &empty_info,Elecvars &filled_vars,
		       Elecvars &empty_vars,
		       Ioninfo &ioninfo,
		       Lattice &lattice,
		       ColumnBundle **grad)
{
  int q;
  real Ener;

  Ener = 0.0;
  for (q=0; q < empty_info.nstates; q++)
  {
      BlochState &s=empty_vars.states[q];
      ColumnBundle HspC(s.C);

      HspC = Hsp(s.C,s);
      s.Hsub = s.C^HspC;
      s.Hsub.hermetian = 1;
      diagonalize_herm(s.Hsub_eigs,s.Hsub_evecs,s.Hsub,s.Hsub.nr);

      Ener += REAL(trace(s.Hsub));
      
      *grad[q] = Pbar(filled_vars.states[q].C,HspC);
      *grad[q] -= O(s.C*(s.C^HspC));
      HspC = (*grad[q])*s.Umhalf;
      *grad[q] = HspC;
    }

  return Ener;
}

// --ipd: i updated this file up to this point
//        will try to reuse the CG minimizer that we have for the filled

/*
 * Does a linmin along dir for the empty bands:
 * the configuration we start with in empty_vars has energy E0
 * and gradient grad.  Stepsize is the initial stepsize to use; the final
 * stepsize actually used in returned.
 */
static real
do_linmin_empties(ColumnBundle *dir,
                  Basis *basis,
                  Ioninfo &ioninfo,
                  Lattice &lattice,
                  Elecinfo &empty_info,Elecvars &filled_vars,
                  Elecvars &empty_vars,
                  real E0,
                  ColumnBundle *grad,
                  real stepsize)
{
  real gamma,dderiv,curvature,E;
  ColumnBundle *Y = empty_vars.Y;
  int nstates = empty_info.nstates;

  /* Directional derivative */
  dderiv = 2.0*REAL(dot(nstates,grad,dir));
  for (;;)
    {
      /* Shift Y by stepsize */
             scale_accumulate(nstates,stepsize,dir,Y);

      /* Calculate energy of shifted position */
      calc_UC_empty(empty_info,filled_vars,empty_vars);
      E = empty_energy(empty_info,filled_vars,empty_vars,ioninfo,lattice);

      /* Shift back */
      scale_accumulate(nstates,-stepsize,dir,Y);

      /* Do a parabolic fit to the E0, E, and dderiv
       * to get curvature for quadratic fit and proposed minimum (gamma) */
      curvature = 2.0*(E-E0-stepsize*dderiv)/(stepsize*stepsize);
      gamma = -dderiv/curvature;
      dft_log("dderiv = %8.1le  curvature = %8.1le\n",
		dderiv,curvature);
      dft_log("stepsize = %8.1le    gamma = %8.1le\n",
		stepsize,gamma);
      dft_log_flush();

      /* If curvature is wrong way, take a bigger step */
      if (curvature < 0.0)
	{ stepsize *= 4.0; continue; }
      /* If the proposed minimum is much larger than the stepsize,
       * increase stepsize */
      else if (fabs(gamma/stepsize) > 10.0)
	{ stepsize *= 10.0; continue; }
      /* If much smaller, decrease */
      else if (fabs(gamma/stepsize) <  0.1)
	{ stepsize *= 0.1; continue; }
      /* Otherwise, it was a good linmin so stop! */
      else
	break;
    }

  /* Move to the bottom of the parabola, given by gamma */
  scale_accumulate(nstates,gamma,dir,Y);

  /* Return the stepsize */
  return stepsize;
}


/*
 * Do niter CG steps on the empty bands
 */
static real
minimize_empties_cg(int niter,
		    real stepsize,
		    Basis *basis,
		    Ioninfo &ioninfo,
		    Lattice &lattice,
		    Elecinfo &empty_info,Elecvars &filled_vars,
		    Elecvars &empty_vars)
{
  int q,iter;
  time_t timenow;
  real linminfact=0.0,cosgg=0.0,alpha,abs2_grad_now,abs2_grad_old,E;
  int nstates = empty_info.nstates;
  int nbands = empty_info.nbands;

  /* gradients and search direction */
  ColumnBundle *grad_now,*grad_old,*dir;

  /* Allocate space for the gradients and search directions */
  grad_now = alloc_ColumnBundle_array(nstates,nbands,basis);
  grad_old = alloc_ColumnBundle_array(nstates,nbands,basis);
  dir = alloc_ColumnBundle_array(nstates,nbands,basis);

  dft_log("\n----- minimize_empties_cg() -----\n");
  dft_log("Starting %d iterations of conjugate gradients with\n",
	    niter);
  dft_log("initial stepsize = %lg.\n\n",stepsize);
  dft_log_flush();
  
  for (iter=0; iter < niter; iter++)
    {
      /* Calculate energy and gradient of current config */
      calc_UC_empty(empty_info,filled_vars,empty_vars);
      E = empty_energy_grad_Hsub(empty_info,filled_vars,
				 empty_vars,ioninfo,lattice,grad_now);

      /* Square length of gradients */
      abs2_grad_now = abs2(nstates,grad_now);
      abs2_grad_old = abs2(nstates,grad_old);

      /* Some linimin statistics from last iteration */
      if (iter > 0)
	{
	  linminfact = 
	    REAL(dot(nstates,grad_now,dir))/REAL(dot(nstates,grad_old,dir));
	  cosgg = 4.0*REAL(dot(nstates,grad_old,grad_now))/
	    sqrt( 16.0*abs2_grad_now*abs2_grad_old );
	  dft_log("\nlinmin = %8.1le   cosgg = %8.1le\n\n",
		    linminfact,cosgg);
	  dft_log_flush();
	}

      /* Time stamp and print out current energies and subspace Hamiltonian */
      timenow = time(0);

      dft_log("------------------------------------------------------\n");
      dft_log("Iteration %d   %s\n",iter,ctime(&timenow));
      dft_log("E_empty = %23.13le\n",E);
      dft_log("\n");
      dft_log("nkpts=%d:  Hsub eigenvalues are:\n",nstates);
      for (int band=0; band < empty_info.nbands; band++)
	{
	  for (q=0; q < nstates; q++)
	    dft_log("%le  ",empty_vars.Hsub_eigs[q][band]);
	  dft_log("\n");
	}
      dft_log("\n");
      dft_log_flush();


      /* If this is the last iteration, don't linmin, but
       * quit the iteration loop and exit the routine */
      if (iter==niter-1)
	break;

      /* Calculate search direction */
      alpha = 0.0;
      /* If we've done a "good" linmin, use CG: */
      /* i.e. calculate alpha */
      if (iter>0 && fabs(linminfact) < 0.05 && fabs(cosgg) < 0.1)
	alpha = abs2_grad_now/abs2_grad_old;
      dft_log("|grad| = %le\n",2.0*sqrt(abs2_grad_now));
      dft_log("alpha = %8.1le\n",alpha);
      dft_log_flush();

      /* Calculate current search direction:
       * d_now = 2.0*grad_now + alpha*d_old */
      for (q=0; q < nstates; q++)
	{
	  dir[q] *= alpha;
	  dir[q] += grad_now[q];
	  dir[q] += grad_now[q];
	}

      /* Do a linmin along dir */
      stepsize = do_linmin_empties(dir,basis,
				   ioninfo,lattice,empty_info,filled_vars,
				   empty_vars,
				   E,grad_now,stepsize);
      /* copy old gradient */
      for (q=0; q < nstates; q++)
	{
	  grad_old[q] = grad_now[q];
	}
    }

  /* Free up memory used by all the gradients and diretions */
  free_ColumnBundle_array(nstates,grad_now);
  free_ColumnBundle_array(nstates,grad_old);
  free_ColumnBundle_array(nstates,dir);

  /* return stepsize used */
  return stepsize;

}
		    

/*
 * This function either reads C_empty_filename or uses random numbers
 * for initial empty band C variables.  It allocates space and variables
 * for the empty bands, and then does niter_cg CG cycles on the
 * empty bands.  It then frees up the memory used, writes the final empty C
 * values to the file 'C.empty', returns.
 */
void
calc_empties(int niter_cg,real stepsize,
	     int nbands_empty,
	     Basis *basis,
	     Elecinfo &filled_info,Elecvars &filled_vars,
	     Ioninfo &ioninfo, Lattice &lattice,
	     char *init_C_empty_action,char *C_empty_filename)
{
  Everything empty_everything;
  Elecinfo &empty_info = empty_everything.elecinfo;
  Elecvars &empty_vars = empty_everything.elecvars;
  int q,nstates;

  /* Copy requisite info from filled states */
  empty_info = filled_info;
  empty_info.nbands = nbands_empty;
  empty_info.nelectrons = 0;
  nstates = filled_info.nstates;

  dft_log("\n----- calc_empties() -----\n");
  dft_log_flush();

  /* Allocate space for the empty variables; make their Vscloc point
   * to the right place */
  empty_everything.setup_elecvars();
  for (q=0; q < nstates; q++)
    {
      empty_vars.Y[q].Vscloc = empty_vars.C[q].Vscloc = 
	&(filled_vars.Vscloc);
      empty_vars.Y[q].Vscloc_up = empty_vars.C[q].Vscloc_up = 
	&(filled_vars.Vscloc_up);
      empty_vars.Y[q].Vscloc_dn = empty_vars.C[q].Vscloc_dn = 
	&(filled_vars.Vscloc_dn);
    }

  /* Either read the initial empty states from disk... */
  if (strcmp(init_C_empty_action,"read") == 0)
    {
      dft_log("\n-----> Reading C_empty from '%s'\n\n",
		C_empty_filename);
      dft_log_flush();
      read_ColumnBundle_array(C_empty_filename,nstates,empty_vars.Y);
    }
  /* or randomize the empty states. */
  else
    {
      dft_log("\n---->  Randomizing empty bands.\n\n");
      dft_log_flush();
      System::seed_with_time();
      randomize_ColumnBundle_array(nstates,empty_vars.Y);
    }
  /* Do CG on empty states */
  minimize_empties_cg(niter_cg,stepsize,basis,ioninfo,lattice,
		      empty_info,filled_vars,empty_vars);

  /* Change C_empty to be the digonal basis of Hsub */
  dft_log("\nSwitching to diagonal basis of Hsub:  C_empty <- HsubC_emtpy\n\n");
  dft_log_flush();

  for (q=0; q < empty_info.nstates; q++)
    {
      do_ColumnBundle_matrix_mult(empty_vars.C[q],
				   empty_vars.Hsub_evecs[q],
				   empty_vars.Y[q],0);
      empty_vars.C[q] = empty_vars.Y[q];
    }
  
  /* Write out empty state variables */
  dft_log("\n\nWriting out C.empty.\n\n");
  dft_log_flush();
  write_ColumnBundle_array("C.empty",nstates,empty_vars.C);
  
  /* Make Vscloc point back to the one in empty_vars */
  for (q=0; q < nstates; q++)
    {
      empty_vars.Y[q].Vscloc = empty_vars.C[q].Vscloc =
	&(empty_vars.Vscloc); 
      empty_vars.Y[q].Vscloc_up = empty_vars.C[q].Vscloc_up =
	&(empty_vars.Vscloc_up); 
      empty_vars.Y[q].Vscloc_dn = empty_vars.C[q].Vscloc_dn =
	&(empty_vars.Vscloc_dn); 
    }

  /* free up used memory */
}