! Added for the calculation of epsilon_1
! 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.
! 
! 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_w
! \sigma_2(w) / w 

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

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 degenerate (degeneracies) and 
   ! non-degenerate interband (interband) contributions

   real(DP), allocatable, public :: sigma2_w(:,:,:)        
   real(DP), allocatable, public :: sigma2_w_intra(:,:,:)  
   real(DP), allocatable, public :: sigma2_W_degen(:,:,:)  
   real(DP), allocatable, public :: sigma2_w_inter(:,:,:)  

   ! Average traces
   real(DP), allocatable, public :: sigma2_w_atrace(:)     
   real(DP), allocatable, public :: sigma2_w_atrace_intra(:)     
   real(DP), allocatable, public :: sigma2_w_atrace_degen(:)     
   real(DP), allocatable, public :: sigma2_w_atrace_inter(:)     

   ! DC tensors
   real(DP), allocatable, public :: sigma2_w_dc(:,:)        
   real(DP), allocatable, public :: sigma2_w_dc_intra(:,:)  
   real(DP), allocatable, public :: sigma2_W_dc_degen(:,:)  
   real(DP), allocatable, public :: sigma2_w_dc_inter(:,:)  

   ! DC traces
   real(DP), public :: sigma2_w_dc_atrace        
   real(DP), public :: sigma2_w_dc_atrace_intra  
   real(DP), public :: sigma2_W_dc_atrace_degen  
   real(DP), public :: sigma2_w_dc_atrace_inter  


   public :: sigma2_w_sub, sigma2_w_decomp, sigma2_w_atrace_sub, &
             sigma2_w_atrace_decomp, sigma2_w_dc_sub, sigma2_w_dc_decomp, &
             sigma2_w_dc_atrace_sub, sigma2_w_dc_atrace_decomp


   ! Private variables

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

   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_w_sub

      allocate( sigma2_w(1:3,1:3,0:nw), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2:sigma2_w_sub',&
                               'allocating sigma2_w', 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

      !aux = 2 * aux                      ! factor of 2 from the reduction of this
                                          ! expression
      do iw = 0, nw

         wg2 = wgrid(iw) * 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

               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)  ) &
                                       * ( (ee*ee - wg2 -sa*sa) &
                                              /( ( eem*eem + sa*sa )           &
                                                *( 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_w(1:3,1:3,iw) = - maux(1:3,1:3) / omega

      enddo ! iw

   end subroutine sigma2_w_sub

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

   subroutine sigma2_w_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

      !aux = 2 * aux                      ! factor of 2 from the reduction of this
                                          ! expression

      ! Intra-band contribution
      allocate( sigma2_w_intra(1:3,1:3,0:nw), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2:sigma2_w_decomp',&
                              'allocating sigma2_w_intra',ABS(ierr))

      do iw = 0, nw

         eem = wgrid(iw)                       ! in Ha

         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) )&
                                       * ( -1.0_dp / ( 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_w_intra(1:3,1:3,iw) = - maux(1:3,1:3) / omega

      enddo ! iw
      ! End intraband contribution

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

      do iw = 0, nw

         wg2 = wgrid(iw) * 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

               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) )&
                                       * ( (ee*ee - wg2 -sa*sa) &
                                              /( ( eem*eem + sa*sa )           &
                                                *( 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_w_degen(1:3,1:3,iw) = - maux(1:3,1:3) / omega

      enddo ! iw
      ! End degeneracies contribution

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

      do iw = 0, nw

         wg2 = wgrid(iw) * 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

               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))&
                                       * ( (ee*ee - wg2 -sa*sa) &
                                              /( ( eem*eem + sa*sa )           &
                                                *( 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_w_inter(1:3,1:3,iw) = - maux(1:3,1:3) / omega

      enddo ! iw


      allocate( sigma2_w(1:3,1:3,0:nw), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2:sigma2_w_decomp',&
                               'allocating sigma2_w', ABS(ierr) )

      sigma2_w = sigma2_w_intra + sigma2_w_degen + sigma2_w_inter

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

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

      sigma2_w_atrace_intra(0:nw) = ( sigma2_w_intra(1,1,0:nw) + &
                    sigma2_w_intra(2,2,0:nw) + sigma2_w_intra(3,3,0:nw) )/3.0_DP


      sigma2_w_atrace_degen(0:nw) = ( sigma2_w_degen(1,1,0:nw) + &
                    sigma2_w_degen(2,2,0:nw) + sigma2_w_degen(3,3,0:nw) )/3.0_DP

      sigma2_w_atrace_inter(0:nw) = ( sigma2_w_inter(1,1,0:nw) + &
                    sigma2_w_inter(2,2,0:nw) + sigma2_w_inter(3,3,0:nw) )/3.0_DP

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

      sigma2_w_atrace = sigma2_w_atrace_intra + sigma2_w_atrace_degen &
                                              + sigma2_w_atrace_inter


   end subroutine sigma2_w_decomp


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


   subroutine sigma2_w_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
      !aux = aux * 2                      ! 2 factor coming from the reduction of this
                                          ! expresion 

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

      do iw = 0, nw
  
         wg2 = wgrid(iw) * 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

               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 )      &
                                  * (   (ee*ee- wg2 -sa*sa)      &
                                       /(  ( eem*eem + sa*sa )                 &
                                          *( 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_w_atrace(iw) = - aux2 / (3.0_DP * omega)

      enddo ! iw

   end subroutine sigma2_w_atrace_sub

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

   subroutine sigma2_w_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

      !aux = aux * 2                      ! 2 factor coming from the reduction of this
                                          ! expresion 

      ! intra-band contribution
      allocate( sigma2_w_atrace_intra(0:nw), STAT=ierr )
      if(ierr/=0) &
        call errore('kgsigma2:sigma2_w_atrace_decomp',&
                    'allocating sigma2_w_atrace_intra', abs(ierr))

      do iw = 0, nw
         aux2 = 0.0_DP

         eep = wgrid(iw)                       ! in Ha

         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 )      &
                                  * (  - 1.0_dp / ( eep*eep + 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_w_atrace_intra(iw) = - aux2 / (3.0_DP * omega)

      enddo ! iw
      ! End intra-band contribution

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

      do iw = 0, nw

         wg2 = wgrid(iw) * 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

               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 )      &
                                  * (   (ee*ee- wg2 -sa*sa)      &
                                       /(  ( eem*eem + sa*sa )                 &
                                          *( 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_w_atrace_degen(iw) = - aux2 / (3.0_DP * omega)

      enddo ! iw
      ! End degeneracies contribution

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

      do iw = 0, nw

         wg2 = wgrid(iw) * 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

               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 )         &
                                  * (   (ee*ee- wg2 -sa*sa)      &
                                       /(  ( eem*eem + sa*sa )                 &
                                          *( 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_w_atrace_inter(iw) = - aux2 / (3.0_DP * omega)

      enddo ! iw


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

      sigma2_w_atrace = sigma2_w_atrace_intra + sigma2_w_atrace_degen &
                                              + sigma2_w_atrace_inter

   end subroutine sigma2_w_atrace_decomp

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!                                  DC
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   subroutine sigma2_w_dc_sub

      allocate( sigma2_w_dc(1:3,1:3), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2:sigma2_w_dc_sub',&
                               'allocating sigma2_w_dc', 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

      !aux = aux * 2                      ! 2 factor coming from the reduction of this
                                          ! expresion 

         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

                   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)  ) &
                                       * ( (ee*ee-sa*sa) &
                                              /( ( ee*ee + sa*sa )           &
                                                *( ee*ee + 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_w_dc(1:3,1:3) = - maux(1:3,1:3) / omega

      sigma2_w_dc_atrace = ( sigma2_w_dc(1,1) + &
                             sigma2_w_dc(2,2) + &
                             sigma2_w_dc(3,3) ) / 3.0_DP

   end subroutine sigma2_w_dc_sub

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

   subroutine sigma2_w_dc_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

      !aux = aux * 2                      ! 2 factor coming from the reduction of this
                                          ! expresion 

         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

                  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 )      &
                                  * (   (ee*ee - sa*sa)                        &
                                       /(  ( ee*ee + sa*sa )                   &
                                          *( ee*ee + 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_w_dc_atrace = - aux2 / (3.0_DP * omega)

   end subroutine sigma2_w_dc_atrace_sub

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

   subroutine sigma2_w_dc_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

      !aux = aux * 2                      ! 2 factor coming from the reduction of this
                                          ! expresion 


      ! Intra-band contribution
      allocate( sigma2_w_dc_intra(1:3,1:3), STAT=ierr )
      if(ierr/=0) CALL errore('kgsigma2:sigma2_w_dc_decomp',&
                              'allocating sigma2_w_dc_intra',ABS(ierr))

         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) )&
                                       * ( -1.0_dp / ( 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_w_dc_intra(1:3,1:3) = - maux(1:3,1:3) / omega

      ! End intraband contribution


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

         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 ) 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) )&
                                       * ( (ee*ee-sa*sa) &
                                              /( ( ee*ee + sa*sa )           &
                                                *( ee*ee + 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_w_dc_degen(1:3,1:3) = - maux(1:3,1:3) / omega

      ! End degeneracies contribution

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

         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

                     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))&
                                       * ( (ee*ee - sa*sa) &
                                              /( ( ee*ee + sa*sa )           &
                                                *( ee*ee + 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_w_dc_inter(1:3,1:3) = - maux(1:3,1:3) / omega

      allocate( sigma2_w_dc(1:3,1:3), STAT=ierr )
      if (ierr/=0) CALL errore('kgsigma2:sigma2_w_dc_decomp',&
                               'allocating sigma2_w_dc', ABS(ierr) )

      sigma2_w_dc = sigma2_w_dc_intra + sigma2_w_dc_degen + sigma2_w_dc_inter


      sigma2_w_dc_atrace_intra = ( sigma2_w_dc_intra(1,1) + &
                    sigma2_w_dc_intra(2,2) + sigma2_w_dc_intra(3,3) )/3.0_DP


      sigma2_w_dc_atrace_degen = ( sigma2_w_dc_degen(1,1) + &
                    sigma2_w_dc_degen(2,2) + sigma2_w_dc_degen(3,3) )/3.0_DP

      sigma2_w_dc_atrace_inter = ( sigma2_w_dc_inter(1,1) + &
                    sigma2_w_dc_inter(2,2) + sigma2_w_dc_inter(3,3) )/3.0_DP


      sigma2_w_dc_atrace = sigma2_w_dc_atrace_intra + sigma2_w_dc_atrace_degen &
                                              + sigma2_w_dc_atrace_inter


   end subroutine sigma2_w_dc_decomp

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

   subroutine sigma2_w_dc_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

      !aux = aux * 2                      ! 2 factor coming from the reduction of this
                                          ! expresion 

      ! intra-band contribution

         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 )      &
                                  * (  - 1.0_dp / ( 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_w_dc_atrace_intra = - aux2 / (3.0_DP * omega)

      ! End intra-band contribution

      ! Degeneracies contribution

         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 ) 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 )      &
                                  * (   (ee*ee-sa*sa)      &
                                       /(  ( ee*ee + sa*sa )                 &
                                          *( ee*ee + 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_w_dc_atrace_degen = - aux2 / (3.0_DP * omega)

      ! End degeneracies contribution

      ! Inter bands contribution

         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

                  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 )         &
                                  * (   (ee*ee-sa*sa)      &
                                       /(  ( ee*ee + sa*sa )                 &
                                          *( ee*ee + 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_w_dc_atrace_inter = - aux2 / (3.0_DP * omega)


      sigma2_w_dc_atrace = sigma2_w_dc_atrace_intra + sigma2_w_dc_atrace_degen &
                                              + sigma2_w_dc_atrace_inter

   end subroutine sigma2_w_dc_atrace_decomp

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


end module kgsigma2_w
