#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "green.h"
#include <complex>
#include "global.h"

/*******************************************************************************
 * The class of Green is designed to evaluate the LDOS via the Green's Function
 * method. The meaning of input/output parameters are as follows:
 *
 *   ntm     (input, value)  total number of atoms in system
 *   sdim    (input, value)  dimension of the system; usually 3
 *   niter   (input, value)  maximum iterations during Lanczos diagonalization
 *   min     (input, value)  minimum value for the angular frequency
 *   max     (input, value)  maximum value for the angular frequency
 *   ndos    (input, value)  total number of points in LDOS
 *   eps     (input, value)  epson that govens the width of delta-function
 *   Hessian (input, pointer of pointer) mass-weighted force constant matrix, of
 *                           dimension [natom*sysdim][natm*sysdim]; it is actually
 *                           the dynamical matrix at gamma point
 *   itm     (input, value)  index of the atom to evaluate local phonon DOS, from 0
 *   lpdos   (output, array) double array of size (ndos, sdim)
 *******************************************************************************
 * References:
 *  1. Z. Tang and N. R. Aluru, Phys. Rev. B 74, 235441 (2006).
 *  2. C. Hudon, R. Meyer, and L.J. Lewis, Phys. Rev. B 76, 045409 (2007).
 *  3. L.T. Kong and L.J. Lewis, Phys. Rev. B 77, 165422 (2008).
 *
 * NOTE: The real-space Green's function method is not expected to work accurately
 *       for small systems, say, system (unit cell) less than 500 atoms.
 *******************************************************************************/

/*------------------------------------------------------------------------------
 * Constructor is used as the main driver
 *----------------------------------------------------------------------------*/
Green::Green(const int ntm, const int sdim, const int niter, const double min, const double max,
             const int ndos, const double eps, double **Hessian, const int itm, double **lpdos)
{
   const double tpi = 8.*atan(1.);
   natom = ntm; sysdim = sdim; nit = niter; epson = eps;
   wmin = min*tpi; wmax = max*tpi; nw = ndos + (ndos+1)%2;
   H = Hessian; iatom = itm;
   ldos = lpdos;
 
   memory = new Memory();
   if (natom < 1 || iatom < 0 || iatom >= natom){
     printf("\nError: Wrong number of total atoms or wrong index of interested atom!\n");
     return;
   }
   ndim = natom * sysdim;
 
   if (nit < 1){printf("\nError: Wrong input of maximum iterations!\n"); return;}
   if (nit > ndim){printf("\nError: # Lanczos iterations is not expected to exceed the degree of freedom!\n"); return;}
   if (nw  < 1){printf("\nError: Wrong input of points in LDOS!\n"); return;}
 
   // initialize variables and allocate local memories
   dw = (wmax - wmin)/double(nw-1);
   memory->create(alpha, sysdim,nit,  "Green_Green:alpha");
   memory->create(beta,  sysdim,nit+1,"Green_Green:beta");
   //memory->create(ldos,  nw,sysdim, "Green_Green:ldos");
 
   // use Lanczos algorithm to diagonalize the Hessian
   Lanczos();
 
   // Get the inverser of the treated hessian by continued fractional method
   Recursion();
 
   return;
}

/*------------------------------------------------------------------------------
 * Deconstructor is used to free memory
 *----------------------------------------------------------------------------*/
Green::~Green()
{
   H = NULL;
   ldos = NULL;
   memory->destroy(alpha);
   memory->destroy(beta);
 
   delete memory;
 
   return;
}
      
/*------------------------------------------------------------------------------
 * Private method to diagonalize a matrix by the Lanczos algorithm
 *----------------------------------------------------------------------------*/
void Green::Lanczos()
{
   double *vp, *v, *w, *ptr;
 
   vp = new double [ndim];
   v  = new double [ndim];
   w  = new double [ndim];
   
   int ipos = iatom*sysdim;
 
   // Loop over dimension
   for (int idim = 0; idim < sysdim; ++idim){
      beta[idim][0] = 0.;
      for (int i = 0; i < ndim; ++i) vp[i] = v[i] = 0.;
      v[ipos+idim] = 1.;
 
      // Loop on fraction levels
      for (int i = 0; i < nit; ++i){
         double sum_a = 0.;
         for (int j = 0; j < ndim; ++j){
            double sumHv = 0.;
            for (int k = 0; k < ndim; ++k) sumHv += H[j][k]*v[k];
            w[j] = sumHv - beta[idim][i]*vp[j];
            sum_a += w[j]*v[j];
         }
         alpha[idim][i] = sum_a;
 
         for (int k = 0; k < ndim; ++k) w[k] -= alpha[idim][i]*v[k];
 
         double gamma = 0.;
         for (int k = 0; k < ndim; ++k) gamma += w[k]*v[k];
         for (int k = 0; k < ndim; ++k) w[k] -= gamma*v[k];
 
         double sum_b = 0.;
         for (int k = 0; k < ndim; ++k) sum_b += w[k]*w[k];
         beta[idim][i+1] = sqrt(sum_b);
 
         ptr = vp; vp = v; v = ptr;
         double tmp = 1./beta[idim][i+1];    
         for (int k = 0; k < ndim; ++k) v[k] = w[k] * tmp;
      }
   }
 
   ptr = NULL;
   delete []vp;
   delete []v;
   delete []w;
 
   return;
}

/*------------------------------------------------------------------------------
 * Private method to compute the LDOS via the recusive method for system with
 * many atoms
 *----------------------------------------------------------------------------*/
void Green::Recursion()
{
   // local variables
   double *alpha_inf, *beta_inf, *xmin, *xmax;
   alpha_inf = new double [sysdim];
   beta_inf  = new double [sysdim];
   xmin  = new double [sysdim];
   xmax  = new double [sysdim];
 
   int nave = nit/4;
   for (int idim = 0; idim < sysdim; ++idim){
      alpha_inf[idim] = beta_inf[idim] = 0.;
 
      for (int i = nit-nave; i < nit; ++i){
         alpha_inf[idim] += alpha[idim][i];
         beta_inf[idim] += beta[idim][i+1];
      }
 
      alpha_inf[idim] /= double(nave);
      beta_inf[idim]  /= double(nave);
 
      xmin[idim] = alpha_inf[idim] - 2.*beta_inf[idim];
      xmax[idim] = alpha_inf[idim] + 2.*beta_inf[idim];
   }
 
   std::complex<double> Z, z_m_a, r_x, rec_x, rec_x_inv;
   double sr, si;
 
   double w = wmin;
   for (int i = 0; i < nw; ++i){
      double a = w*w, ax, bx;
      Z = std::complex<double>(w*w, epson);
 
      for (int idim = 0; idim < sysdim; ++idim){
         double two_b = 2.*beta_inf[idim]*beta_inf[idim];
         double rtwob = 1./two_b;
 
         z_m_a = Z - alpha_inf[idim]*alpha_inf[idim];
 
         if ( a < xmin[idim] ){
            r_x = sqrt(-2.*two_b + z_m_a);
            ax = std::real(r_x) * rtwob;
            bx = std::imag(r_x) * rtwob;
         } else if (a > xmax[idim]) {
            r_x = sqrt(-2.*two_b + z_m_a);
            ax = -std::real(r_x) * rtwob;
            bx = -std::imag(r_x) * rtwob;
         } else {
            r_x = sqrt(2.*two_b - z_m_a);
            ax =  std::imag(r_x) * rtwob;
            bx = -std::real(r_x) * rtwob;
         }
 
         sr = (a - alpha_inf[idim])*rtwob + ax;
         si = epson * rtwob + bx;
         rec_x = std::complex<double> (sr, si);
 
         for (int j = 0; j < nit; ++j){
            rec_x_inv = Z - alpha[idim][nit-j-1] - beta[idim][nit-j]*beta[idim][nit-j]*rec_x;
            rec_x = 1./rec_x_inv;
         }
         ldos[i][idim] = std::imag(rec_x)*w;
      }
      w += dw;
   }
   delete []alpha_inf;
   delete []beta_inf;
   delete []xmin;
   delete []xmax;
 
   return;
}

/*------------------------------------------------------------------------------
 * Private method to compute the LDOS via the recusive method for system with
 * a few atoms (less than NMAX)
 *----------------------------------------------------------------------------*/
void Green::recursion()
{
   // local variables
   std::complex<double> Z, rec_x, rec_x_inv;
   std::complex<double> cunit = std::complex<double>(0.,1.);
 
   double w = wmin;
 
   for (int i = 0; i < nw; ++i){
      Z = std::complex<double>(w*w, epson);
 
      for (int idim = 0; idim < sysdim; ++idim){
         rec_x = std::complex<double>(0.,0.);
 
         for (int j = 0; j < nit; ++j){
            rec_x_inv = Z - alpha[idim][nit-j-1] - beta[idim][nit-j]*beta[idim][nit-j]*rec_x;
            rec_x = 1./rec_x_inv;
         }
         ldos[i][idim] = std::imag(rec_x)*w;
      }
      w += dw;
   }
   return;
}

/*------------------------------------------------------------------------------*/