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

/*
 * Code by Gabor Csanyi,  June 2001
 * 
 * spin-orbit matrix elements
 *
 */


#include "header.h"


/*
 * applies the 'del' operator to a columnbundle. for planewaves, del = k + G
 *
 * the parameter `a' specifies which component of del is calculated, valid values are
 * 'x', 'y', 'z'
 */

void apply_del(const column_bundle &Y,column_bundle &DY, char a)
{
  int i,j,nbasis;
  real kx,ky,kz, kplusG;
  Basis *basis;

  if (Y.basis == 0)
    die("apply_del() called with Y.basis == 0\n");
  if (Y.col_length != Y.basis->nbasis)
    die("apply_del() called with Y.col_length != nbasis\n");
  if ( (Y.col_length!=DY.col_length) || (Y.tot_ncols!=DY.tot_ncols) )
    die("apply_del() called with different sizes for Y and DY\n");

  /* Let's go! */
  nbasis = Y.col_length;
  basis = Y.basis;
  kx = Y.qnum->kvec.v[0];
  ky = Y.qnum->kvec.v[1];
  kz = Y.qnum->kvec.v[2];
  for (j=0; j < nbasis; j++)
    {
      switch (a){
      case 'x':
	kplusG = kx + basis->Gx[j];
	break;
      case 'y':
	kplusG = ky + basis->Gy[j];
	break;
      case 'z':
	kplusG = kz + basis->Gz[j];
	break;
      default:
	die("Invalid component in apply_del: %c\n", a);
      }
      for (i=0; i < Y.my_ncols; i++)
        DY.col[i].c[j] = kplusG*Y.col[i].c[j];

    }

}


/*
 * applies the 'del' operator to a vector. for planewaves, del = G
 *
 * del v = I(G*(J(v)))
 * 
 * the parameter `a' specifies which component of del is calculated, valid values are
 * 'x', 'y', 'z'
 */

void apply_del(const vector &v, vector &dv, char a)
{

  Basis *basis;
  Lattice *lattice;

  if (v.basis == 0)
    die("apply_del() called with v.basis == 0\n");
  if ( v.n!= dv.n )
    die("apply_del() called with different sizes for v and dv\n");


  basis = v.basis;
  lattice = basis->lattice;
  int Nx = basis->Nx;
  int Ny = basis->Ny;
  int Nz = basis->Nz;
  int Nx2 = Nx/2;
  int Ny2 = Ny/2;
  int Nz2 = Nz/2;
  int NyNz = Ny * Nz;
  int NxNyNz;

  // transform the vector into g-space
  apply_J(v, dv);
  

  // get the G info from the lattice
  real g0, g1, g2;
  g0 = lattice->G.m[0][a-'x'];
  g1 = lattice->G.m[1][a-'x'];
  g2 = lattice->G.m[2][a-'x'];
  int index, i, j, k;

  for (i=-Nx2; i < Nx2; i++)
    for (j=-Ny2; j < Ny2; j++)
      for (k=-Nz2; k < Nz2; k++){
	index = 0;
	if (k < 0) index += k+Nz;        else index += k;
	if (j < 0) index += Nz*(j+Ny);   else index += Nz*j;
	if (i < 0) index += NyNz*(i+Nx); else index += NyNz*i;
	dv.c[index] *= (i * g0 + j * g1 + k * g2);
      }

  // transform back to real space
  dv = I(dv);

}




/*
 *  this is real space code, achieves del(vector), kept here for debugging purposes
 *
 *
 */



vector
gradx(vector v){
  int i,j,k;
  vector dxv(v.n, v.basis);
  int Nx = dxv.basis->Nx;
  int Ny = dxv.basis->Ny;
  int Nz = dxv.basis->Nz;

  for(i = 0; i < Nx; i++)
    for(j = 0; j < Ny; j++)
      for(k = 0; k < Nz; k++)
	if(i == 0)
	  dxv.c[(i*Ny+j)*Nz+k] = (v.c[((i+1)*Ny+j)*Nz+k]-v.c[((i-0)*Ny+j)*Nz+k]);
	else if(i == Nx-1)
	  dxv.c[(i*Ny+j)*Nz+k] = (v.c[((i+0)*Ny+j)*Nz+k]-v.c[((i-1)*Ny+j)*Nz+k]);
	else
	  dxv.c[(i*Ny+j)*Nz+k] = (v.c[((i+1)*Ny+j)*Nz+k]-v.c[((i-1)*Ny+j)*Nz+k])/2.0; 

  return dxv;
}

vector
grady(vector v){
  int i,j,k;
  vector dyv(v.n, v.basis);
  int Nx = dyv.basis->Nx;
  int Ny = dyv.basis->Ny;
  int Nz = dyv.basis->Nz;

  for(i = 0; i < Nx; i++)
    for(j = 0; j < Ny; j++)
      for(k = 0; k < Nz; k++)
	if(j == 0)
	  dyv.c[(i*Ny+j)*Nz+k] = (v.c[(i*Ny+(j+1))*Nz+k]-v.c[(i*Ny+(j-0))*Nz+k]);
	else if(j == Ny-1)
	  dyv.c[(i*Ny+j)*Nz+k] = (v.c[(i*Ny+(j+0))*Nz+k]-v.c[(i*Ny+(j-1))*Nz+k]);
	else
	  dyv.c[(i*Ny+j)*Nz+k] = (v.c[(i*Ny+(j+1))*Nz+k]-v.c[(i*Ny+(j-1))*Nz+k])/2.0; 

  return dyv;
}

vector
gradz(vector v){
  int i,j,k;
  vector dzv(v.n, v.basis);
  int Nx = dzv.basis->Nx;
  int Ny = dzv.basis->Ny;
  int Nz = dzv.basis->Nz;

  for(i = 0; i < Nx; i++)
    for(j = 0; j < Ny; j++)
      for(k = 0; k < Nz; k++)
	if(k == 0)
	  dzv.c[(i*Ny+j)*Nz+k] = (v.c[(i*Ny+j)*Nz+k+1]-v.c[(i*Ny+j)*Nz+k-0]);
	else if(k == Nz-1)
	  dzv.c[(i*Ny+j)*Nz+k] = (v.c[(i*Ny+j)*Nz+k+0]-v.c[(i*Ny+j)*Nz+k-1]);
	else
	  dzv.c[(i*Ny+j)*Nz+k] = (v.c[(i*Ny+j)*Nz+k+1]-v.c[(i*Ny+j)*Nz+k-1])/2.0; 

  return dzv;
}

/*
 * this routine takes a vector and a columnbundle, whose column length
 * is the same as the vector length, and performs the operation that
 * is equivalent to making a diagonal matrix with the vector on the diagonal,
 * and leftmultiplying the column bundle by this matrix:
 * 
 * Diag(v) * Y
 *
 */

column_bundle 
diag_mult(vector &v, column_bundle &Y)
{
  if(v.n != Y.col_length)
    die("In diag_mult(), v.n != Y.col_length\n");

  column_bundle vY(Y);

  scalar vn;

  for(int n = 0; n < v.n; n++){
    vn = v.c[n];
    for(int col = 0; col < Y.tot_ncols; col++)
      vY.col[col].c[n] *= vn;
  }

  return vY;
} 


/*
 * this is the main spin-orbit routine, to be called from the
 * outside
 *
 * it only works with LSDA, and dumps 4 matrices, corresponding
 * to the 4 spin subblocks of Hso(s,k,s',k')
 *
 * Hso is diagonal in k
 *
 */

void
calc_spinorbit(Everything &e)
{
  Elecinfo &einfo = e.elecinfo;
  Elecvars &evars = e.elecvars;
  Ioninfo &ioninfo = e.ioninfo;
  Basis *basis = e.basis;
  Lattice &lattice = e.lattice;


  // We can only do a spin-orbit calculation if spin is explicitly included in
  // the calculation
  if(einfo.spintype != ZSPIN)
    die("Spin-orbit calculation is only implemented for spintype==ZSPIN.\nCannot continue.\n");

  // We will need the total self-consistent potential.
  calc_Vscloc(evars,einfo);

  column_bundle *C = evars.C;  

  vector Vsc_up = *(C[0].Vscloc_up);
  vector Vsc_dn = *(C[0].Vscloc_dn);

  vector dVsc(Vsc_up.basis->NxNyNz, Vsc_up.basis);

  // we are computing the matrix element <Psi_k_s | Hso | Psi_k_s'>
  // the matrix is diagonal in k

  matrix Hso;
  Hso.init(einfo.nbands, einfo.nbands);

  char filename[255];
  for (int q_up = 0; q_up < einfo.nstates; q_up++){
      Hso.zero_out();
      
      /*
       * for each spin-up state, find the corresponding spin-down state, and
       * then compute the 4 quadrants of the Hso matrix
       *
       */

      // when q corresponds to a spin-down state, just skip
      if(einfo.qnums[q_up].spin == -1) continue;

      int q_dn;
      for(q_dn = 0; q_dn < einfo.nstates; q_dn++)
	if((einfo.qnums[q_dn].spin == -1) && (einfo.qnums[q_dn].kvec == einfo.qnums[q_up].kvec))
	  break;
      
      complex i(0,1);

      column_bundle DC(C[q_up].tot_ncols,C[q_up].basis->nbasis);
      copy_innards_column_bundle(&C[q_up],&DC);
      column_bundle IDC(C[q_up].tot_ncols,C[q_up].basis->NxNyNz);
      copy_innards_column_bundle(&C[q_up],&IDC);

      // up up
      Hso.zero_out();

      apply_del(Vsc_up, dVsc, 'x');
      apply_del(C[q_up], DC, 'y');
      IDC = I(DC);
      Hso += C[q_up]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_up, dVsc, 'y');
      apply_del(C[q_up], DC, 'x');
      IDC = I(DC);
      Hso -= C[q_up]^(Idag(diag_mult(dVsc,IDC))); 

      sprintf(filename, "Hso_%d.uu", q_up);
      Hso.write(filename);

      // up dn
      Hso.zero_out();

      apply_del(Vsc_dn, dVsc, 'y');
      apply_del(C[q_dn], DC, 'z');
      IDC = I(DC);
      Hso += C[q_up]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_dn, dVsc, 'z');
      apply_del(C[q_dn], DC, 'y');
      IDC = I(DC);
      Hso -= C[q_up]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_dn, dVsc, 'x');
      apply_del(C[q_dn], DC, 'z');
      IDC = I(DC);
      Hso += i*C[q_up]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_dn, dVsc, 'z');
      apply_del(C[q_dn], DC, 'x');
      IDC = I(DC);
      Hso -= i*C[q_up]^(Idag(diag_mult(dVsc,IDC))); 

      sprintf(filename, "Hso_%d.ud", q_up);
      Hso.write(filename);

      // dn up
      Hso.zero_out();

      apply_del(Vsc_up, dVsc, 'y');
      apply_del(C[q_up], DC, 'z');
      IDC = I(DC);
      Hso += C[q_dn]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_up, dVsc, 'z');
      apply_del(C[q_up], DC, 'y');
      IDC = I(DC);
      Hso -= C[q_dn]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_up, dVsc, 'x');
      apply_del(C[q_up], DC, 'z');
      IDC = I(DC);
      Hso -= i*C[q_dn]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_up, dVsc, 'z');
      apply_del(C[q_up], DC, 'x');
      IDC = I(DC);
      Hso += i*C[q_dn]^(Idag(diag_mult(dVsc,IDC))); 

      sprintf(filename, "Hso_%d.du", q_up);
      Hso.write(filename);

      // dn dn
      Hso.zero_out();

      apply_del(Vsc_dn, dVsc, 'x');
      apply_del(C[q_dn], DC, 'y');
      IDC = I(DC);
      Hso -= C[q_dn]^(Idag(diag_mult(dVsc,IDC))); 
      apply_del(Vsc_dn, dVsc, 'y');
      apply_del(C[q_dn], DC, 'x');
      IDC = I(DC);
      Hso += C[q_dn]^(Idag(diag_mult(dVsc,IDC))); 

      sprintf(filename, "Hso_%d.dd", q_up);
      Hso.write(filename);

  }
}


/*
 * same as above, but done column-by-column, which is nasty
 * and inefficient
 *
 */



void
calc_spinorbit_col_by_col(Everything &e)
{
  Elecinfo &einfo = e.elecinfo;
  Elecvars &evars = e.elecvars;
  Ioninfo &ioninfo = e.ioninfo;
  Basis *basis = e.basis;
  Lattice &lattice = e.lattice;


  // We can only do a spin-orbit calculation if spin is explicitly included in
  // the calculation
  if(einfo.spintype != ZSPIN)
    die("Spin-orbit calculation is only implemented for spintype==ZSPIN.\nCannot continue.\n");

  // We will need the total self-consistent potential.
  calc_Vscloc(evars,einfo);

  column_bundle *C = evars.C;  

  vector Vsc_up = *(C[0].Vscloc_up);
  vector Vsc_dn = *(C[0].Vscloc_dn);

  // derivatives of Vsc
  vector dxVsc_up = gradx(Vsc_up);
  vector dyVsc_up = grady(Vsc_up);
  vector dzVsc_up = gradz(Vsc_up);
  vector dxVsc_dn = gradx(Vsc_dn);
  vector dyVsc_dn = grady(Vsc_dn);
  vector dzVsc_dn = gradz(Vsc_dn);

  
  // store the real space version of a single column
  int NxNyNz = C[0].basis->NxNyNz;
  column_bundle Ccol(1,NxNyNz,"local");
  column_bundle ICcol(1,NxNyNz,"local");
  column_bundle DxCcol(1,NxNyNz,"local");
  column_bundle IDxCcol(1,NxNyNz,"local");
  column_bundle DyCcol(1,NxNyNz,"local");
  column_bundle IDyCcol(1,NxNyNz,"local");
  column_bundle DzCcol(1,NxNyNz,"local");
  column_bundle IDzCcol(1,NxNyNz,"local");
 
  // we are computing the matrix element <Psi_q1 | Hso | Psi_q2>

  matrix Hso;
  Hso.init(einfo.nbands, einfo.nbands);

  for (int q1 = 0; q1 < einfo.nstates; q1++){
    for (int q2 = 0; q2 < einfo.nstates; q2++){

      // derivatives of the wavefunctions
      column_bundle DxC(C[q2].tot_ncols,C[q2].basis->nbasis);
      copy_innards_column_bundle(&C[q2],&DxC);
      apply_del(C[q2],DxC, 'x'); 
      column_bundle DyC(C[q2].tot_ncols,C[q2].basis->nbasis);
      copy_innards_column_bundle(&C[q2],&DyC);
      apply_del(C[q2],DyC, 'y'); 
      column_bundle DzC(C[q2].tot_ncols,C[q2].basis->nbasis);
      copy_innards_column_bundle(&C[q2],&DzC);
      apply_del(C[q2],DzC, 'z'); 

      Hso.zero_out();

      for (int i1=0; i1 < C[q1].tot_ncols; i1++){
	for (int i2=0; i2 < DxC.tot_ncols; i2++){
	  /* Local vars */
	  int j;
	  
	  /* Copy all the descriptive junk inside of X */
	  copy_innards_column_bundle(&C[q1],&Ccol);
	  copy_innards_column_bundle(&C[q1],&ICcol);
	  copy_innards_column_bundle(&C[q2],&DxCcol);
	  copy_innards_column_bundle(&C[q2],&IDxCcol);
	  copy_innards_column_bundle(&C[q2],&DyCcol);
	  copy_innards_column_bundle(&C[q2],&IDyCcol);
	  copy_innards_column_bundle(&C[q2],&DzCcol);
	  copy_innards_column_bundle(&C[q2],&IDzCcol);
	  
	  /* Copy in i1'th column of C into Ccol */
	  for (j=0; j < C[q1].col_length; j++)
	    Ccol.col[0].c[j] = C[q1].col[i1].c[j];
	  /* Apply I (i.e. FFT3D) */
	  apply_I(Ccol,ICcol);
	  
	  /* Copy in i2'th column of DC into DCcol */
	  for (j=0; j < DxC.col_length; j++)
	    DxCcol.col[0].c[j] = DxC.col[i2].c[j];
	  /* Apply I (i.e. FFT3D) */
	  apply_I(DxCcol,IDxCcol);
	  /* Copy in i2'th column of DC into DCcol */
	  for (j=0; j < DyC.col_length; j++)
	    DyCcol.col[0].c[j] = DyC.col[i2].c[j];
	  /* Apply I (i.e. FFT3D) */
	  apply_I(DyCcol,IDyCcol);
	  /* Copy in i2'th column of DC into DCcol */
	  for (j=0; j < DzC.col_length; j++)
	    DzCcol.col[0].c[j] = DzC.col[i2].c[j];
	  /* Apply I (i.e. FFT3D) */
	  apply_I(DzCcol,IDzCcol);


	  scalar s;
	  s.x = s.y = 0;
	  complex ii(0,1);
	  if(einfo.qnums[q1].spin == 1 && einfo.qnums[q2].spin == 1){
	    for (j = 0; j < ICcol.col_length; j++)
	      s += conj(ICcol.col[0].c[j]) * dxVsc_up.c[j] * IDyCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s -= conj(ICcol.col[0].c[j]) * dyVsc_up.c[j] * IDxCcol.col[0].c[j];
	  }
	  else if(einfo.qnums[q1].spin == 1 && einfo.qnums[q2].spin == -1){
	    for (j = 0; j < ICcol.col_length; j++)
	      s += conj(ICcol.col[0].c[j]) * dyVsc_dn.c[j] * IDzCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s -= conj(ICcol.col[0].c[j]) * dzVsc_dn.c[j] * IDyCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s += ii*conj(ICcol.col[0].c[j]) * dxVsc_dn.c[j] * IDzCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s -= ii*conj(ICcol.col[0].c[j]) * dzVsc_dn.c[j] * IDxCcol.col[0].c[j];
	  }
	  else if(einfo.qnums[q1].spin == -1 && einfo.qnums[q2].spin == 1){
	    for (j = 0; j < ICcol.col_length; j++)
	      s += conj(ICcol.col[0].c[j]) * dyVsc_up.c[j] * IDzCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s -= conj(ICcol.col[0].c[j]) * dzVsc_up.c[j] * IDyCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s -= ii*conj(ICcol.col[0].c[j]) * dxVsc_up.c[j] * IDzCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s += ii*conj(ICcol.col[0].c[j]) * dzVsc_up.c[j] * IDxCcol.col[0].c[j];
	  }
	  if(einfo.qnums[q1].spin == -1 && einfo.qnums[q2].spin == -1){
	    for (j = 0; j < ICcol.col_length; j++)
	      s -= conj(ICcol.col[0].c[j]) * dxVsc_dn.c[j] * IDyCcol.col[0].c[j];
	    for (j = 0; j < ICcol.col_length; j++)
	      s += conj(ICcol.col[0].c[j]) * dyVsc_dn.c[j] * IDxCcol.col[0].c[j];
	  }

	  Hso(i1, i2) = s;
	}
      }
      char filename[255];
      sprintf(filename, "Hso_%d.%d", q1, q2);
      Hso.write(filename);
    }
  }
}