!
! 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 kgecpaw
use kinds, ONLY : DP
implicit none

      real(DP),    allocatable :: psi_norm_factor(:,:)

contains

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

   subroutine kgec_paw_renorm
   use qevars
   use qe_p_psi, ONLY: p_psi, becp
   use kgecvars
   use apawgm, ONLY : awfc_grad_awfc,         &
                      n_pro_tot, check_p_pswfc, diff_awfc_grad_awfc,   &
                      psi_norm_correction_paw_at, psi_norm_correction_at,  &
                      p_index, p_index_row_start, p_index_row_end, agm_free_memory

   implicit none


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

      integer :: ik, ig                ! counter on k-points and plane waves
                                       ! always local (parallelization is hidden) 
                                         

      integer :: ib1                    ! global counters on bands

      complex(DP) :: caux

      my_rank = mp_rank( inter_bgrp_comm )

      write(stdout,*) " "
      write(stdout,'(A)',advance='no')  &
         "     Calculating the re-normalization factors of | \Psi_{nk} > ..."

      ! If PAW, renormalize the wave function psi
      ! The norm of |psi> is given by
      ! <psi|psi> =  <pssi|pspsi> 
      ! + sum_anlm sum_an'lm <pspsi|p_anlm> 
      ! ( <Ranl|Ran'l> - <pseudo R_anl| pseudo R_an'l> ) <p_an'lm|pspsi>

      ! Calculating the 
      ! ( < R_anl | R_an'l >  - < pseudo R_anl | pseudo R_an'l > ) term
       call psi_norm_correction_paw_at()

       ! Get <p|\Psi> ( becp%k )
       call p_psi

       ! Calculating the renormalization factors 1/<psi|psi>^(1/2)
       allocate( psi_norm_factor( 1:nbnd, 1:nks ) )

       do ik=1, nks
          ! Get npw 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) || defined(__QE64)
          npw = ngk(ik)
#endif
          ! Read the wave function for the current k-vector
          call davcio ( evc , 2*nwordwfc, iunwfc, ik, - 1 )

          do ib1=1, nbnd

             caux = cmplx( 0.0_dp, 0.0_dp )

             do ig=1, npw
                caux = caux + conjg( evc(ig,ib1) ) * evc(ig,ib1)
             enddo

             call mp_sum( caux, intra_bgrp_comm ) ! reduce on G (npw)

             do ip1=1, n_pro_tot
!             do ip2=1, n_pro_tot
             do ip2=p_index_row_start(ip1), p_index_row_end(ip1)
                caux = caux + conjg( becp(ik)%k(ip1,ib1) ) &
                              * psi_norm_correction_at( p_index(ip1), p_index(ip2) ) &
                                * becp(ik)%k(ip2,ib1)
             enddo
             enddo

             psi_norm_factor( ib1, ik) = 1.0_dp / dsqrt(abs(caux))
!             psi_norm_factor( ib1, ik ) = 1.0_dp


          enddo
       enddo

       write(stdout,*) ' done.'

   end subroutine kgec_paw_renorm


end module kgecpaw
