!
! 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 kgecgradient
use qevars
use kgecvars
use kgsigma1
use kgsigma2
use parallel
implicit none

private

   integer :: ip1, ip2 ! Consecutive counters for the projectors


   integer :: iw                          ! Frequency index

   integer :: ik, ig                      ! counter on k-points and plane waves
                                          ! always local 
                                          ! (parallelization is hidden) 
                                         
   integer :: ib1,  ib2                   ! global counters on bands
   integer :: ibb1, ibb2                  ! local counters on bands

   integer :: ix1, ix2                    ! labels x, y an z

   real(DP)    :: aux, ee
   complex(DP) :: caux, psi_grad_psi_aux(3), cvaux(1:3)
   
   complex(DP), allocatable :: grad_paw_correction(:,:,:,:), psi_psi(:,:)

   real(DP) :: sumrulef, sumruleintegral

public :: kgec_gradient, psi_grad_psi, psi_psi, kgec_gradient_paw_correction, &
          sumrule_exact, sumrulef, sumrule_integral, sumruleintegral


contains

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

   subroutine kgec_gradient

   my_rank = mp_rank( inter_bgrp_comm )

   !============================================================================
   !  Gradient matrix elements over the 'solid' pseudo wave function \Psi_{n,k}
   !                                 psi_grad_psi
   !============================================================================

   write(stdout,*) ''
   write(stdout,'(A)',advance='no') &
    '     Calculating the gradient matrix elements over pseudo \Psi_{n,k} ..'


   allocate( psi_grad_psi(1:3, 1:nbnd, 1:nbnds_local, 1:nks), STAT=ierr )
   if (ierr/=0) call errore('kgec','allocating psi_grad_psi', abs(ierr) )

   do ik = 1, nks

      ! Get the G vectors for the current k-vector
#if defined(__QE512) || defined(__QE54)
          call gk_sort (xk(1,ik), ngm, g, ecutwfc / tpiba2, npw, igk, g2kin )
#elif defined(__QE6) || defined(__QE62)
          npw = ngk(ik)
#endif

      ! Read the wave function for the given k-vector
      call davcio (evc, 2*nwordwfc, iunwfc, ik, - 1)

#if defined (__KGDEBUG)
      write(*,*) "evc:", ik, iunwfc, evc
#endif

      ! write(*,*) ik,nbnd,npw,ecutwfc

      !
      ! Calculate gradient matrix elements
      !
      do ibb2=1, nbnds_local
         ib2 = ibb2 + nbnds_per_process * my_rank
         
         do ib1=1, nbnd

            psi_grad_psi_aux = cmplx( 0.0_dp, 0.0_dp )

            do  ig=1, npw
                caux= conjg( evc(ig,ib1) ) * evc(ig,ib2)
#if defined (__KGDEBUG)
                write(*,*) "caux", ib1, ig, caux
#endif

#if defined (__QE6) || defined(__QE62) || defined(__QE64)
                psi_grad_psi_aux(1:3) = psi_grad_psi_aux(1:3) &
                                   + ( xk(1:3,ik) + g(1:3,igk_k(ig,ik)) ) * caux
#elif defined (__QE512) || defined(__QE54)
                psi_grad_psi_aux(1:3) = psi_grad_psi_aux(1:3) &
                                        + ( xk(1:3,ik) + g(1:3,igk(ig)) ) * caux
#endif
            enddo
            
            call mp_sum( psi_grad_psi_aux, intra_bgrp_comm ) ! reduce on G (npw)
            
            psi_grad_psi(1:3,ib1,ibb2,ik) = &
                         cmplx( 0.0_dp, 1.0_dp ) * tpiba * psi_grad_psi_aux(1:3)

#if defined (__KGDEBUG)
           write(*,*) "psi_grad_psi:", psi_grad_psi(1:3,ib1,ibb2,ik)
#endif
         enddo
      enddo

   enddo

#if defined (__QE6) || defined(__QE62) || defined(__QE64)
   deallocate( evc, g, igk_k )
#elif defined (__QE512) || defined(__QE54)
   deallocate( evc, g, igk )
#endif

   write(stdout,'(A)') ' done. '

   
   end subroutine kgec_gradient


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


   subroutine kgec_gradient_paw_correction

   !===========================================================================
   !          Calculate PAW correction to the gradient matrix elements
   !                           ( Grad_paw_correction )
   !===========================================================================

   use kgecpaw
   use apawgm
   use qe_p_psi, ONLY: becp

     my_rank = mp_rank( inter_bgrp_comm )
   
     write(stdout,*) ''
     write(stdout,*) &
        '    Calculating the gradient matrix elements over atomic orbitals'
     write(stdout,'(A)',advance='no') &
        '     and pseudo-orbitals and PAW correction ..'

     ! < \phi | \nabbla | \phi > and < psphi|\nabla\ psphi> = DG

     call diff_awfc_grad_awfc

     ! Grad_paw_correction(1:3,ib2,ib1,ik)=
     ! \sum_{ia,l,m,l',mp'} <\Psi_{ib1}_{ik}| p_{ia,l,m} > \times
     ! (    < \phi_{ia,l,m} |\nabla| \phi_{ia,l',m'}> 
     !   -  < \hat{\phi}_{ia,l,m} |\nabla| \hat{\phi_{ia},l',m'}>) \times
     ! < p_{ia,l,m} |\Psi_{ib2}_{ik}>

     allocate( grad_paw_correction( 1:3, 1:nbnd, 1:nbnds_local,1:nks),STAT=ierr)
     if (ierr/=0) call errore('kgec:main','allocating grad_paw_correction',&
                                                                      ABS(ierr))

     grad_paw_correction( 1:3, 1:nbnd, 1:nbnds_local, 1:nks) = &
                                                         cmplx( 0.0_DP, 0.0_DP )

     do ik  = 1, nks
     
     do ibb1 = 1, nbnds_local
        ib1  = ibb1 + nbnds_per_process * my_rank
         
     do ib2 = 1, nbnd

        cvaux(1:3)=cmplx(0.0_dp, 0.0_dp)

        do ip1 = 1, n_pro_tot
!        do ip2 = 1, n_pro_tot
        do ip2 = p_index_row_start(ip1), p_index_row_end(ip1)  
            cvaux(1:3) = cvaux(1:3) &
                         + conjg( becp(ik)%k(ip2,ib2) ) &
                           * DG( 1:3, p_index(ip2), p_index(ip1) ) &
                            * becp(ik)%k(ip1,ib1)
        enddo
        enddo

        grad_paw_correction(1:3,ib2,ibb1,ik) = &
                                   cvaux(1:3) * psi_norm_factor(ib1,ik) &
                                              * psi_norm_factor(ib2,ik)

        psi_grad_psi(1:3,ib2,ibb1,ik) = ( psi_grad_psi(1:3,ib2,ibb1,ik) &
                                                         + cvaux(1:3) ) &
                                              * psi_norm_factor(ib1,ik) &
                                                * psi_norm_factor(ib2,ik)
#if defined (__KGDEBUG)
        write(*,*) ""
        write(*,*) "cvaux", cvaux(1:3)
        write(*,*) "psi_norm_factor:", psi_norm_factor(ib1,ik), psi_norm_factor(ib2,ik)
        write(*,*) "grad_paw_correction:", grad_paw_correction(1:3,ib2,ibb1,ik)
        write(*,*) "psi_grad_psi added:", psi_grad_psi(1:3,ib2,ibb1,ik)
#endif

     enddo
     enddo
     enddo

     call mp_barrier( inter_pool_comm )

     write(stdout,*) " done."

#if defined (__KGDEBUG)
     write (fname, "('grad_paw_correction-k', i0,'_b', i0, '.dat')") &
                                                             my_pool_id, me_pool
     fname=trim(tmp_dir)//fname
      open(121, file=fname,action='write')
     do ik=1, nks
     do ibb1=1, nbnds_local
        ib1 = ibb1 + nbnds_per_process * my_rank
        
     do ib2=1, nbnd
        write(121,*) ik, ib2, ib1, grad_paw_correction(1,ib2,ibb1,ik), &
                                   grad_paw_correction(2,ib2,ibb1,ik), &
                                   grad_paw_correction(3,ib2,ibb1,ik)
     enddo
     enddo
     enddo
     close(121)
#endif
   
     deallocate( grad_paw_correction )
     deallocate( psi_norm_factor )
     deallocate( becp )
     call agm_free_memory

   end subroutine kgec_gradient_paw_correction

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

   subroutine sumrule_exact
   real(DP) :: sum1, dfde

      write(stdout,*) ""
      write(stdout,*) &
        '    Calculating the exact sum (independents of frequencies):'

      sum1 = 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 ( sigma1_notintra .and. ib1 == ib2 ) cycle    ! No intra bands

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

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

           aux = 0.0_DP
           do ix1 = 1, 3
              aux = aux + conjg(psi_grad_psi(ix1,ib2,ibb1,ik)) &
                                 * psi_grad_psi(ix1,ib2,ibb1,ik)
           enddo

           sum1 = sum1 + aux * dfde

      enddo ! ib2
      enddo ! ib1
      enddo ! ik

      call mp_sum( sum1, inter_bgrp_comm ) ! reduce over bands (nbnd)
      call mp_sum( sum1, inter_pool_comm ) ! reduce over k-points (nks)

      sumrulef      = sum1 / ( 3.0_dp * nelec )

!      write(stdout,*) "      Sum rule (f)           :", sumrulef
!      write(stdout,*) ""
         
   end subroutine sumrule_exact

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

   subroutine sumrule_integral
   real(DP) :: sum1, sum2
   
   ! 
   ! Sum rule via Simpson's integration (with delta function)
   ! 
   
   sum1 = 0.0_dp
   sum2 = 0.0_dp

   select case ( AC )
   case( .true. )
      select case ( calculation )
      case( 'tensor' )

         do iw=2, nw-2, 2
            sum1 = sum1 + (sigma1(1,1,iw) + sigma1(2,2,iw) + sigma1(3,3,iw))
         enddo

         do iw=1, nw-1, 2
            sum2 = sum2 + (sigma1(1,1,iw) + sigma1(2,2,iw) + sigma1(3,3,iw))
         enddo

         sum1 = 2.0_dp * sum1 + 4.0_dp * sum2
         sum1 = sum1 + (sigma1(1,1,0) + sigma1(2,2,0) + sigma1(3,3,0))
         sum1 = (sum1 +(sigma1(1,1,nw)+sigma1(2,2,nw)+sigma1(3,3,nw)))*dw/3.0_dp

         sumruleintegral = 2.0_dp * omega * sum1 / ( 3.0_dp * pi * nelec )

      case( 'atrace' )

         do iw=2, nw-2, 2
            sum1 = sum1 + sigma1_atrace(iw)
         enddo

         do iw=1, nw-1, 2
            sum2 = sum2 + sigma1_atrace(iw)
         enddo

         sum1 = 2.0_dp * sum1 + 4.0_dp * sum2
         sum1 =   sum1 + sigma1_atrace(0)
         sum1 = ( sum1 + sigma1_atrace(nw) )*dw / 3.0_dp

         sumruleintegral = 2.0_dp * omega * sum1 / ( pi * nelec )

      end select
   end select
   
   end subroutine sumrule_integral


end module kgecgradient

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Saved subroutine with various forms of 'sum rules'
!   subroutine sumrule_exact
!   real(DP) :: sum1, sum2, sum3, sum4, focc, dfde, &
!               sumrulefintra, sumrulep, sumrulepf
!
!      write(stdout,*) ""
!      write(stdout,*) &
!        '    Calculating some sum rules (independents of frequencies):'
!
!
!      sum1 = 0.0_DP
!      sum2 = 0.0_DP
!      sum3 = 0.0_DP
!      sum4 = 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
!
!
!           if( (intra .and. ib2==ib1) .or. abs(ee)<=de ) then
!              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
!           else
!              dfde = ( wg(ib2,ik) - wg(ib1,ik) ) / ee
!           endif
!
!           aux = 0.0_DP
!           do ix1 = 1, 3
!              aux = aux + conjg(psi_grad_psi(ix1,ib2,ibb1,ik)) &
!                                 * psi_grad_psi(ix1,ib2,ibb1,ik)
!           enddo
!
!           sum2 = sum2 + aux * dfde
!
!           if( ib2 /= ib1 .and. abs(ee) > de ) then
!             if( (et(ib2,ik)-e_f) >= 0.0_dp) then
!               sum1 = sum1 + aux * ( wg(ib2,ik) - wg(ib1,ik) ) / ee
!             else
!               sum1 = sum1 - aux *                                          &
!                      (   1.0_DP/(dexp(-(et(ib2,ik)-e_f)/degauss)+1.0_DP)   &
!                        - 1.0_DP/(dexp(-(et(ib1,ik)-e_f)/degauss)+1.0_DP) ) &
!                      / ee 
!             endif
!             sum3 = sum3 - aux / ee
!             sum4 = sum4 - aux * wg(ib1,ik) / ee
!           endif
!
!
!
!      enddo ! ib2
!      enddo ! ib1
!      enddo ! ik
!
!      call mp_sum( sum1, inter_bgrp_comm ) ! reduce over bands (nbnd)
!      call mp_sum( sum1, inter_pool_comm ) ! reduce over k (nks)
!      call mp_sum( sum2, inter_bgrp_comm ) ! reduce over bands (nbnd)
!      call mp_sum( sum2, inter_pool_comm ) ! reduce over k (nks)
!      call mp_sum( sum3, inter_bgrp_comm ) ! reduce over bands (nbnd)
!      call mp_sum( sum3, inter_pool_comm ) ! reduce over k (nks)
!      call mp_sum( sum4, inter_bgrp_comm ) ! reduce over bands (nbnd)
!      call mp_sum( sum4, inter_pool_comm ) ! reduce over k (nks)
!
!      sumrulef      = sum1 / ( 3.0_dp * nelec )
!      sumrulefintra = sum2 / ( 3.0_dp * nelec )
!      sumrulep      = sum3 * 2.0_dp / 3.0_dp
!      sumrulepf     = sum4 * 2.0_dp / ( 3.0_dp * nelec )
!
!      write(stdout,*) "      Sum rule (f)           :", sumrulef
!      write(stdout,*) "      Sum rule (f with intra):", sumrulefintra
!      write(stdout,*) "      Sum rule (p)           :", sumrulep
!      write(stdout,*) "      Sum rule (pf)          :", sumrulepf
!      write(stdout,*) ""
!         
!   end subroutine sumrule_exact
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

