!
! Copyright (C) 2017- Orbital-free DFT group at University of Florida
!
! This program 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 
! any later version. See the file LICENSE in the root directory of the 
! present distribution, or http://www.gnu.org/copyleft/gpl.txt ,
! or contact developers via e-mail: calderin@lcalderin.net, or paper mail:
!
! Quantum Theory Project
! University of Florida
! P.O. Box 118435
! Gainesville, FL 32611
!
! This program 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.
! 
! Authors: Lazaro Calderin and Valentin Karasiev
!-----------------------------------------------------------------------
! 


module kgsigma2
! \sigma_2

use kinds, ONLY    : DP
use qevars
use kgecvars, ONLY : deltarep, deltawidth, nw, wmin, wmax, wgrid, dw, sa, sg, &
                     psi_grad_psi, de, sigma2_notintra,                       &
                     nbnds_local, nbnds_per_process, nbnds_per_group,         &
                     my_rank, my_rank_old, my_group
use parallel

implicit none
 
   private

   ! Internal units : atomic Ha
   ! Note on multipliying by the prefactor 2 / omega for sigma2 in a.u.
   ! - The 2 went into the wg ( the combined fermi-dirac and k-weighting factor)
   
   ! Gradient matrix elements needed as input
   ! complex(DP), allocatable    :: psi_grad_psi(:,:,:,:)

   ! \sigma_2 and its decomposition in intra bands, degenerate and inter bands 
   ! contributions

   real(DP), allocatable, public :: sigma2(:,:,:)
   real(DP), allocatable, public :: sigma2_intra(:,:,:)        
   real(DP), allocatable, public :: sigma2_inter(:,:,:)  
   real(DP), allocatable, public :: sigma2_degen(:,:,:)  

   ! Average traces
   real(DP), allocatable, public :: sigma2_atrace(:)
   real(DP), allocatable, public :: sigma2_atrace_intra(:)     
   real(DP), allocatable, public :: sigma2_atrace_degen(:)     
   real(DP), allocatable, public :: sigma2_atrace_inter(:)     

   public :: sigma2_sub, sigma2_decomp, sigma2_atrace_sub, &
             sigma2_atrace_decomp


   ! Private variables

   real(DP) :: aux, aux2, ee, eem, eep, maux(3,3)

   integer :: ios , ierr ! Input and memory allocation error status variables

   integer :: iw                          ! Frequency index
   integer :: ik, ib1, ib2, ibb1          ! counters on k-points and bands
   integer :: ig                          ! counter on plane waves
   integer :: ix1, ix2                    ! labels x, y an z

   real(DP) :: focc, dfde


   contains

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   subroutine frequency_grid()
   !
   ! Frequency grid
   !

      dw   = ( wmax - wmin ) / dble(nw)  ! Frecuency step

      allocate( wgrid(0:nw), STAT=ierr )
      if (ierr/=0) call errore('kgsigma2','allocating wgrid', abs(ierr))
      
      ! Frequency grid in eV
      do iw = 0, nw
         wgrid(iw) = wmin + iw * dw
      enddo


      write(stdout,*) '    Frequency grid:'
      write(stdout,*) "       Frequency step (eV)    :", dw
      write(stdout,*) "       Frequency steps        :", nw
      write(stdout,'(a, f10.4,a,f10.4 )') &
                 "        Frequency range (eV)   :", wgrid(0),  ' to ',wgrid(nw)

     ! Change units to Ha
     wgrid = wgrid / AUTOEV
     dw    = dw    / AUTOEV

   end subroutine frequency_grid


   !===========================================================================
   !                 Full conductivity tensor \sigma_2
   !===========================================================================
   !

   subroutine sigma2_sub

      allocate( sigma2(1:3,1:3,0:nw), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2','allocating sigma2', ABS(ierr) )
 
      sa   = 0.5_dp * deltawidth / AUTOEV ! Lorentzian half-width in Ha
      aux  = 0.5_dp                 ! 1/2 from ( (e-w)/( (e-w)^2 + sa^2/4 ) - 
                                    !            (e+w)/( (e+w)^2 + sa^2/4 ) ) /2
      do iw = 0, nw

         do ix1 = 1, 3
         do ix2 = 1, 3
            maux(ix2,ix1) = 0.0_DP

            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group

               do ib2 = 1, nbnd

                   if ( sigma2_notintra .and. ib2 == ib1 ) cycle   

                   ee  = ( et(ib1,ik) - et(ib2,ik) ) * 0.5_DP ! in Ha
                   eem = ee - wgrid(iw)                       ! in Ha
                   eep = ee + wgrid(iw)                       ! in Ha

                   if ( ib1 == ib2 .or. abs(ee) <=  de ) then
                      focc = 1.0_DP/(dexp(( et(ib2,ik) - e_f )/degauss)+1.0_DP )
                      dfde = -wk(ik) * focc * ( focc - 1.0_DP ) / degauss*2.0_DP
                   endif

                   if ( abs(ee) > de ) &
                      dfde = ( wg(ib2,ik) - wg(ib1,ik) )/ ee

                      maux(ix2,ix1) = maux(ix2,ix1)                            &
                                 + dfde                                        &
                                  *real( conjg(psi_grad_psi(ix1,ib2,ibb1,ik))  &
                                          *    psi_grad_psi(ix2,ib2,ibb1,ik)  )&
                                          * (    eem/( eem*eem + sa*sa )       &
                                               - eep/( eep*eep + sa*sa ) )

               enddo ! ib2
               enddo ! ib1

            enddo ! ik

      enddo ! ix1
      enddo ! ix2

      call mp_sum( maux, inter_bgrp_comm ) ! reduce on bands (nbnd)
      call mp_sum( maux, inter_pool_comm ) ! reduce on k (nks)

      sigma2(1:3,1:3,iw) = - aux * maux(1:3,1:3) / omega

      enddo ! iw

      allocate( sigma2_atrace(0:nw), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2','allocating sigma1_atrace', ABS(ierr))

    
      sigma2_atrace(0:nw) = ( sigma2(1,1,0:nw) + & 
                              sigma2(2,2,0:nw) + &
                              sigma2(3,3,0:nw) ) / 3.0_DP


   end subroutine sigma2_sub

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   subroutine sigma2_decomp

      sa   = 0.5_dp * deltawidth / AUTOEV ! Lorentzian half-width in Ha
      aux  = 0.5_dp                 ! 1/2 from ( (e-w)/( (e-w)^2 + sa^2/4 ) - 
                                    !            (e+w)/( (e+w)^2 + sa^2/4 ) ) /2

      ! Intra bands contribution

      allocate( sigma2_intra(1:3,1:3,0:nw), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2','allocating sigma2_intra',ABS(ierr))
 
      do iw = 0, nw
         eem = wgrid(iw)

         do ix1 = 1, 3
         do ix2 = 1, 3
            maux(ix2,ix1) = 0.0_DP

            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group

                  focc =  1.0_DP/( dexp((et(ib1,ik)-e_f)/degauss) + 1.0_DP )
                  dfde = -wk(ik) * focc * ( focc - 1.0_DP ) / degauss * 2.0_DP

                  maux(ix2,ix1) = maux(ix2,ix1)                                &
                                  + dfde                                       &
                                    *real( conjg(psi_grad_psi(ix1,ib1,ibb1,ik))&
                                            *   psi_grad_psi(ix2,ib1,ibb1,ik) )&
                                     * ( eem/( eem*eem + sa*sa ) )

               enddo ! ib1

            enddo ! ik

         enddo ! ix1
         enddo ! ix2

         call mp_sum( maux, inter_bgrp_comm ) ! reduce on bands (nbnd)
         call mp_sum( maux, inter_pool_comm ) ! reduce on k (nks)

         sigma2_intra(1:3,1:3,iw) = - maux(1:3,1:3) / omega

      enddo ! iw
      ! End intra bands contribution


      ! Degeneracies contribution
      allocate( sigma2_degen(1:3,1:3,0:nw), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2','allocating sigma2_degen',ABS(ierr))

      do iw = 0, nw

         do ix1 = 1, 3
         do ix2 = 1, 3
            maux(ix2,ix1) = 0.0_DP

            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group

               do ib2 = 1, nbnd

                  if ( ib1 == ib2 ) cycle

                  ee  = ( et(ib1,ik) - et(ib2,ik) ) * 0.5_DP ! in Ha
                  eem = ee - wgrid(iw)                       ! in Ha
                  eep = ee + wgrid(iw)                       ! in Ha

                  if( abs(ee) <=  de ) then
                     focc =  1.0_DP/( dexp((et(ib2,ik)-e_f)/degauss) + 1.0_DP )
                     dfde = -wk(ik) * focc * ( focc - 1.0_DP ) / degauss * 2.0_DP

                     maux(ix2,ix1) = maux(ix2,ix1)                             &
                                  + dfde                                       &
                                    *real( conjg(psi_grad_psi(ix1,ib2,ibb1,ik))&
                                            *   psi_grad_psi(ix2,ib2,ibb1,ik) )&
                                     * ( eem/( eem*eem + sa*sa ) - &
                                         eep/( eep*eep + sa*sa ) )
                   endif

               enddo ! ib2
               enddo ! ib1

            enddo ! ik

         enddo ! ix1
         enddo    ! ix2

         call mp_sum( maux, inter_bgrp_comm ) ! reduce on bands (nbnd)
         call mp_sum( maux, inter_pool_comm ) ! reduce on k (nks)

         sigma2_degen(1:3,1:3,iw) = - aux * maux(1:3,1:3) / omega

      enddo ! iw
      ! End degeneracies contribution

      ! Inter bands contribution
      allocate( sigma2_inter(1:3,1:3,0:nw), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2','allocating sigma2_inter',ABS(ierr))

      do iw = 0, nw

         do ix1 = 1, 3
         do ix2 = 1, 3
            maux(ix2,ix1) = 0.0_DP

            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group

               do ib2 = 1, nbnd

                  if ( ib1 == ib2 ) cycle 

                  ee  = ( et(ib1,ik) - et(ib2,ik) ) * 0.5_DP ! in Ha
                  if ( abs(ee) <= de ) cycle

                  eem = ee - wgrid(iw)                       ! in Ha
                  eep = ee + wgrid(iw)                       ! in Ha

                  dfde = ( wg(ib2,ik) - wg(ib1,ik) )/ee

                  maux(ix2,ix1) = maux(ix2,ix1)                                &
                                  + dfde                                       &
                                    *real( conjg(psi_grad_psi(ix1,ib2,ibb1,ik))&
                                           *     psi_grad_psi(ix2,ib2,ibb1,ik))&
                                             * (   eem/( eem*eem + sa*sa )     &
                                                 - eep/( eep*eep + sa*sa ) )

               enddo ! ib2
               enddo    ! ib1

            enddo ! ik

         enddo ! ix1
         enddo ! ix2

         call mp_sum( maux, inter_bgrp_comm ) ! reduce on bands (nbnd)
         call mp_sum( maux, inter_pool_comm ) ! reduce on k (nks)

         sigma2_inter(1:3,1:3,iw) = - aux * maux(1:3,1:3) / omega

      enddo ! iw


      allocate( sigma2(1:3,1:3,0:nw), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2','allocating sigma2', ABS(ierr) )
      sigma2 = sigma2_degen + sigma2_inter

      allocate( sigma2_atrace_degen(0:nw), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2','allocating sigma2_atrace_degen', &
                                                                      ABS(ierr))
      allocate( sigma2_atrace_inter(0:nw), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2','allocating sigma2_atrace_inter', &
                                                                      ABS(ierr))

      sigma2_atrace_degen(0:nw) = ( sigma2_degen(1,1,0:nw) + &
                    sigma2_degen(2,2,0:nw) + sigma2_degen(3,3,0:nw) )/3.0_DP
      sigma2_atrace_inter(0:nw) = ( sigma2_inter(1,1,0:nw) + &
                    sigma2_inter(2,2,0:nw) + sigma2_inter(3,3,0:nw) )/3.0_DP

      allocate( sigma2_atrace(0:nw), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2','allocating sigma2_atrace', &
                                                                      ABS(ierr))

      sigma2_atrace = sigma2_atrace_degen + sigma2_atrace_inter


   end subroutine sigma2_decomp


   !============================================================================
   !             Average trace of \sigma_2 = tr(sigma2)/3
   !============================================================================


   subroutine sigma2_atrace_sub

      sa   = 0.5_dp * deltawidth / AUTOEV ! Lorentzian half-width in Ha
      aux  = 0.5_dp                 ! 1/2 from ( (e-w)/( (e-w)^2 + sa^2/4 ) - 
                                    !            (e+w)/( (e+w)^2 + sa^2/4 ) ) /2

      allocate( sigma2_atrace(0:nw), STAT=ierr )
      if(ierr/=0) call errore('kgcond','allocating sigma2_atrace', abs(ierr))

      do iw = 0, nw

         aux2 = 0.0_DP

         do ix1 = 1, 3

            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group

               do ib2 = 1, nbnd
              
                  if ( sigma2_notintra .and. ib2 == ib1 ) cycle  

                  ee  = ( et(ib1,ik) - et(ib2,ik) ) * 0.5_DP ! in Ha
                  eem = ee - wgrid(iw)                       ! in Ha
                  eep = ee + wgrid(iw)                       ! in Ha

                  if ( ib1 == ib2 .or. abs(ee) <=  de ) then
                     focc = 1.0_DP / (dexp(( et(ib2,ik) - e_f )/degauss)+1.0_DP)
                     dfde =-wk(ik) * focc * ( focc - 1.0_DP ) / degauss * 2.0_DP
                  endif

                  if ( abs(ee) > de ) &
                     dfde =  ( wg(ib2,ik) - wg(ib1,ik) )/ ee

                     aux2 = aux2                                               &
                            + dfde                                             &
                                * conjg( psi_grad_psi( ix1,ib2,ibb1,ik ) )     &
                                 *        psi_grad_psi( ix1,ib2,ibb1,ik )      &
                                  * (   eem/( eem*eem + sa*sa )                &
                                      - eep/( eep*eep + sa*sa ) )

              enddo ! ib2
              enddo ! ib1

           enddo ! ik

        enddo    ! ix1

        call mp_sum( aux2, inter_bgrp_comm ) ! reduce on bands (nbnd)
        call mp_sum( aux2, inter_pool_comm ) ! reduce on k (nks)

        sigma2_atrace(iw) = - aux * aux2 / (3.0_DP * omega)

      enddo ! iw

   end subroutine sigma2_atrace_sub

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   subroutine sigma2_atrace_decomp

      sa   = 0.5_dp * deltawidth / AUTOEV ! Lorentzian half-width in Ha
      aux  = 0.5_dp                 ! 1/2 from ( (e-w)/( (e-w)^2 + sa^2/4 ) - 
                                    !            (e+w)/( (e+w)^2 + sa^2/4 ) ) /2

      ! Intra bands contribution
      allocate( sigma2_atrace_intra(0:nw), STAT=ierr )
      if(ierr/=0) &
        call errore('kgsigma2','allocating sigma2_atrace_intra', abs(ierr))

      do iw = 0, nw
         eem = wgrid(iw)
         aux2 = 0.0_DP

         do ix1 = 1, 3
            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group


                   focc =1.0_DP / ( dexp(( et(ib1,ik) - e_f )/degauss)+1.0_DP)
                   dfde = -wk(ik) * focc * ( focc - 1.0_DP ) / degauss * 2.0_DP

                   aux2 = aux2                                                 &
                            +   dfde                                           &
                                * conjg( psi_grad_psi( ix1,ib1,ibb1,ik ) )     &
                                  *       psi_grad_psi( ix1,ib1,ibb1,ik )      &
                                   * ( eem/( eem*eem + sa*sa )  )               

               enddo ! ib1
            enddo ! ik
         enddo ! ix1

         call mp_sum( aux2, inter_bgrp_comm ) ! reduce on bands (nbnd)
         call mp_sum( aux2, inter_pool_comm ) ! reduce on k (nks)

         sigma2_atrace_intra(iw) = - aux2 / (3.0_DP * omega)

      enddo ! iw
      ! End intra bands contribution

      ! Degeneracies contribution
      allocate( sigma2_atrace_degen(0:nw), STAT=ierr )
      if(ierr/=0) call errore('kgsigma2','allocating sigma2_atrace', &
                                                                      abs(ierr))

      do iw = 0, nw
         aux2 = 0.0_DP

         do ix1 = 1, 3
            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group

               do ib2 = 1, nbnd

                  if ( ib1 == ib2 ) cycle

                  ee  = ( et(ib1,ik) - et(ib2,ik) ) * 0.5_DP ! in Ha
                  eem = ee - wgrid(iw)                       ! in Ha
                  eep = ee + wgrid(iw)                       ! in Ha

                  if ( abs(ee) <= de ) then
                     focc =1.0_DP / ( dexp(( et(ib2,ik) - e_f )/degauss)+1.0_DP)
                     dfde = -wk(ik) * focc * ( focc - 1.0_DP ) / degauss * 2.0_DP

                     aux2 = aux2                                               &
                            +   dfde                                           &
                                * conjg( psi_grad_psi( ix1,ib2,ibb1,ik ) )     &
                                  *       psi_grad_psi( ix1,ib2,ibb1,ik )      &
                                   * ( eem/( eem*eem + sa*sa )                 &
                                       -eep/( eep*eep + sa*sa ) )

                 endif

               enddo ! ib2
               enddo ! ib1
            enddo ! ik
         enddo ! ix1

         call mp_sum( aux2, inter_bgrp_comm ) ! reduce on bands (nbnd)
         call mp_sum( aux2, inter_pool_comm ) ! reduce on k (nks)

         sigma2_atrace_degen(iw) = - aux * aux2 / (3.0_DP * omega)

      enddo ! iw
      ! End degeneracies contribution

      ! Inter bands contribution
      allocate( sigma2_atrace_inter(0:nw), STAT=ierr )
      if (ierr/=0) call errore('kgsigma2','allocating sigma2_atrace_inter', &
                                                                      abs(ierr))

      do iw = 0, nw
         aux2 = 0.0_DP

         do ix1 = 1, 3
            do ik = 1, nks

               do ibb1 = 1, nbnds_local
                  ib1 = ibb1 + nbnds_per_process * my_rank_old &
                             + nbnds_per_group   * my_group

               do ib2 = 1, nbnd

                  if ( ib1 == ib2 ) cycle

                  ee  = ( et(ib1,ik) - et(ib2,ik) ) * 0.5_DP ! in Ha
                  if ( abs(ee) <=  de ) cycle

                  eem = ee - wgrid(iw)                       ! in Ha
                  eep = ee + wgrid(iw)                       ! in Ha

                  dfde =  ( wg(ib2,ik) - wg(ib1,ik) )/ ee

                  aux2 = aux2                                                  &
                         + dfde                                                &
                             * conjg( psi_grad_psi( ix1,ib2,ibb1,ik ) )        &
                                *      psi_grad_psi( ix1,ib2,ibb1,ik )         &
                                  * ( eem/( eem*eem + sa*sa )                  &
                                      - eep/( eep*eep + sa*sa ) )

               enddo ! ib2
               enddo ! ib1
            enddo ! ik
         enddo    ! ix1

         call mp_sum( aux2, inter_bgrp_comm ) ! reduce on bands (nbnd)
         call mp_sum( aux2, inter_pool_comm ) ! reduce on k (nks)

         sigma2_atrace_inter(iw) = - aux * aux2 / (3.0_DP * omega)

      enddo ! iw


      allocate( sigma2_atrace(0:nw), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2','allocating sigma2_atrace', &
                                                                     ABS(ierr) )

      sigma2_atrace = sigma2_atrace_degen + sigma2_atrace_inter

   end subroutine sigma2_atrace_decomp

end module kgsigma2
