!
! 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 kgecchecks
use qevars
use qe_p_psi, ONLY: becp
use kgecvars
implicit none


private

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

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

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

public :: kgec_check_requirements, kgec_check_orthonormalization_nc, &
          kgec_check_orthonormalization_paw

contains

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


subroutine kgec_check_requirements
use qevars
use kgecvars, ONLY : non_local, yessym, yesinv, outdir

   !
   ! Check if the conditions for the utilization of the code are met
   !

   if( okvan .and. (.not. okpaw)     ) &
      call errore( 'kgec','Ultrasoft not implemented',              1 )
   if( okpaw .and. (.not. only_paw ) ) &
      call errore( 'kgec','Mixed norm conserving pps and PAW sets', 2 )
   if( gamma_only )                     &
      call errore( 'kgec','Gamma only code used',                   4 )
   if( noncolin )                       &
      call errore( 'kgec','Non-collinear not implemented',          5 )
   if( lsda )                           &
      call errore( 'kgec','Spin not implemented',                   6 )
   if( .not. okpaw .and. ( .not. non_local ) ) then
        write(stdout,*) 'kgec','Non-local correction not implemented'
        call errore( 'kgec','Please use non_local=.true. in KGEC if you would &
                                                         like to proceed',   7 )
   endif

   ! track both nosym and noinv by the value of nsym 
   ! (nsym=1 => identity )
   ! (nsym=2 => identity and inversion)
   ! (nsym>2 => symmetries beyond identity plus inversion)  
   if ( (nsym > 1) .and. (.not. yessym) ) then
       call errore( 'kgec','Symmetry used? (set nosym=.true. in QE or &
                                                    yessym=.true. in KGEC)', 8 )
   !endif
   !if( ( .not. noinv ) .and. ( .not. yesinv ) ) then
      call errore( 'kgec','Inversion symmetry used? (set noinv=.true. &
                                            in QE or yesinv=.true. in KGEC', 9 )
   endif

   ! 
   ! Check the kind of smearing used by pw.x
   ! 

   select case (ngauss)
   case ( -99 )
   case default
     write(stdout,*) "  WARNING: not Fermi-Dirac smearing, ngauss=", ngauss
     write(stdout,*)
   end select


   write(stdout,*) ''
   write(stdout,*) &
          '    --------------------------------------------------------------'
   write(stdout,*) &
          "               Parameters from the QE calculation "
   write(stdout,*) &
          '    --------------------------------------------------------------'
   write(stdout,'(A,2x,A20)') &
          '     System                    :', prefix
   write(stdout,'(A,2x,A20)') &
          '     Directory                 :', outdir
   write(stdout,'(A,f20.6)')  &
          '     Number of electrons       :', nelec
   write(stdout,'(A,f20.6)')  &
          '     Temperature (Ry)(degauss) :', degauss ! Ry units in QE
   write(stdout,'(A,f20.6)')  &
          '     Temperature (K) (degauss) :', degauss * RYTOEV * eV_to_kelvin
   write(stdout,'(A,f20.6)') &
          '     Temperature (eV)(degauss) :', degauss * RYTOEV
   write(stdout,'(A,i20)'  ) &
          '     Number of bands           :', nbnd
   write(stdout,'(A,i20)'  ) &
          '     Number of k-points        :', nkstot
   write(stdout,'(A,f20.6)') &
          '     Sum of k-weights          :', sum(wk) ! Should be equal to 2
   write(stdout,'(A,f20.6)') &
          '     Sum of band weights       :', sum(wg)  ! Should be equal to nelec
   write(stdout,'(A,f20.6)') &
          '     Fermi energy (eV)         :', e_f*RYTOEV
   write(stdout,'(A,f20.6)') &
          '     Volume (bohr^3)           :', omega
   write(stdout,'(A,f20.6)') &
          '     Wave function cutoff (ry) :', ecutwfc 
   if( okpaw ) then
     write(stdout,*) &
          '    Pseudo type               :                PAWs'
   else
     write(stdout,*) &
          '    Pseudo type               :     Norm-Conserving'
   endif

end subroutine kgec_check_requirements

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

   subroutine kgec_check_orthonormalization_nc

      complex(DP), allocatable :: psi_psi(:,:)

      integer :: ik                  ! counter on k-points
                                     ! always local (parallelization is hidden) 
                                         

      integer :: ib1,  ib2           ! global counters on bands
      integer :: ibb1, ibb2          ! local counters on bands
       
      complex(DP) :: caux
   
         my_rank = mp_rank( inter_bgrp_comm )

         write(stdout,*)  " "
         write(stdout,'(A)',advance='no') &
                 "     Checking ortho-normality of the Wave Functions .."

         allocate( psi_psi(1:nbnds_local, 1:nbnds_local) )

         do ik = 1, nks

#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 )

            CALL zgemm( 'C', 'N', nbnd, nbnds_local, npw, &
                        1.0_dp, evc, npwx, evc, npwx, (0.0_dp,0.0_dp), &
                        psi_psi, nbnds_local )

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

            do ibb1 = 1, nbnds_local
               ib1  = ibb1 + nbnds_per_process * my_rank
                           
               do ibb2 = 1, nbnds_local
                  ib2  = ibb2 + nbnds_per_process * my_rank
                  
                  caux = psi_psi( ibb2, ibb1 )

!! XXX ik is local, change it to global 

                  if ( (ib1/=ib2) .and. abs(caux)> 1.d-6 ) then
                     write(stdout,*) '     Non orthogonal:', ik, ib1, ib2, caux
                  endif

                  if ( (ib1==ib2) .and. abs(caux-1.0_dp)>1.0d-5 ) then
                     write(stdout,*) '     Not normalized:', ik, ib1, ib2, caux
                  endif

               enddo
            enddo

         enddo
         deallocate(psi_psi)

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

   
   end subroutine kgec_check_orthonormalization_nc
   

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


   subroutine kgec_check_orthonormalization_paw
   use kgecpaw, ONLY : psi_norm_factor
   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

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

      integer :: ib1,  ib2                    ! global counters on bands
      integer :: ibb1                         ! local counters on bands
       

IF(IONODE) THEN
         ! Check duality of PAW projectors and pseudo atomic wavefunctions
         call check_p_pswfc
ENDIF ! ionode

         my_rank = mp_rank( inter_bgrp_comm )

         write(stdout,*)  " "
         write(stdout,'(A)',advance='no') &
            "     Checking orthogonality of the Wave Functions"

         write (fname, "('non-orthogonal-wfc-k', i0,'_b', i0, '.dat')") &
                                                             my_pool_id, me_pool 

         fname=trim(tmp_dir)//fname
         open(121,file=fname, action='write' )
         write(121,'(A)') "# Non-orthogonal wavefunctions"
         write(121,'(A)') &
            "# k-point, band index, band -index, Re(<ib1|ib2>), Im(<ib1|ib2>)"
 
         ierr = 0
         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 ibb1 = 1, nbnds_local
               ib1  = ibb1 + nbnds_per_process * my_rank

            do ib2 = ib1, nbnd

               caux = cmplx( 0.0_dp, 0.0_dp )

               do ig = 1, npw
                  caux = caux + conjg( evc( ig,ib1 ) ) * evc( ig,ib2 )
               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,ib2)
               enddo
               enddo

               caux = caux * psi_norm_factor(ib1,ik) * psi_norm_factor(ib2,ik)

! XXX ik is local, change to global

               if ( ib1/=ib2 .and. abs(caux)>1.e-9 ) then
                  ierr = ierr + 1
                  write(121,'(2x,3i8,2x,1g20.7)') ik, ib2, ib1, abs(caux)
               endif

            enddo
            enddo

         enddo

         call mp_sum( ierr, inter_pool_comm )
          
         if ( ierr>0 ) then
            close(121)
            write(stdout,*) " "
            write(stdout,*) &
              "     WARNING: wave functions non orthogonal (pairs):", ierr
            write(stdout,*) "     See file non-orthogonal-wfc*.dat"
         else
            close(121,status='delete') 
            write(stdout,*) "... done"
         endif
          
         ierr=0

   end subroutine kgec_check_orthonormalization_paw
   
   

end module kgecchecks
