!=============================================================================
!
! Module:
!
! (1) intwfn()  Originally By MLT               Last Modified 7/8/2008 (JRD)
!
!     Evaluate the transformation matrices and store them
!     in the arrays dcc,dvv.
!     The matrices are defined according to eq. (39) of Rohlfing & Louie:
!
!     dcc(ic,jc,ik,is,ivert) == d_{ic,ik}^{\tilde jc,\tilde jk)
!        ik: index of k-point in fine grid
!        ic: index of cond. band in fine grid
!        jk: index of k-point in coarse grid, fi2co_wfn(ivert,ik)=jk
!        jc: index of cond. band in coarse grid
!        is: spin index (the same in both grids)
!
!     dvv(iv,jv,ik,is,ivert) == d_{iv,ik}^{\tilde jv,\tilde jk)
!        ik: index of k-point in fine grid
!        iv: index of valence band in fine grid
!        jk: index of k-point in coarse grid, fi2co_wfn(ivert,ik)=jk
!        jv: index of valence band in coarse grid
!        is: spin index (the same in both grids)
!
!     Mixing between valence/conduction states is taken into account
!     by allowing jc to be also an occupied band, and similarly
!     for jv.
!
!     input: crys,syms,xct,gvec,kg_fi,kgq_fi types
!            indexq_fi (mapping between shifted and unshifted k-points)
!
!     output: dcc,dvv arrays
!             kco array
!             fi2co_wfn array
!
!===============================================================================

#include "f_defs.h"

module intwfn_m

  use global_m
  use fullbz_m
  use genwf_m
  use intpts_m
  use interp_m
  use mtxel_t_m
  use io_utils_m
  use misc_m

  implicit none

  private

  public :: intwfn

contains

subroutine intwfn(kp,crys,syms,xct,flag,gvec,kg_fi,kgq_fi, &
  dcc,dvv,kco,fi2co_wfn,indexq_fi,eqp,intwfnv,intwfnc,intp_coefs)
  type (kpoints), intent(inout) :: kp
  type (crystal), intent(in) :: crys
  type (symmetry), intent(in) :: syms
  type (xctinfo), intent(inout) :: xct
  type (flags), intent(in) :: flag
  type (gspace), intent(in) :: gvec
  type (grid), intent(in) :: kg_fi
  type (grid), intent(in) :: kgq_fi
  !> (xct%ncb_fi,xct%n2b_co,xct%nspin,xct%nkpt_fi,xct%npts_intp_kernel) Interp coefs.
  SCALAR, intent(out), target :: dcc(:,:,:,:,:)
  !> (xct%nvb_fi,xct%n1b_co,xct%nspin,xct%nkpt_fi,xct%npts_intp_kernel) Interp coefs.
  SCALAR, intent(out), target :: dvv(:,:,:,:,:)
  real(DP), intent(out) :: kco(:,:) !< 3,xct%nkpt_co
  integer, intent(out) :: fi2co_wfn(:,:) !< (xct%npts_intp_kernel, xct%nkpt_fi)
  integer, intent(in) :: indexq_fi(:) !< xct%nkpt_fi
  type (eqpinfo), intent(inout) :: eqp
  type (int_wavefunction), intent(inout) :: intwfnc
  type (int_wavefunction), intent(inout) :: intwfnv
  !> (xct%npts_intp_kernel, xct%nkpt_fi) Delaunay/greedy interpolation coefficients
  real(DP), intent(out) :: intp_coefs(:,:)

  integer :: closepts(4)
  real(DP) :: closeweights(4)
  type(work_genwf) :: work, workq, workco, workcoq
  type (grid) :: kg_co, kg_co_q
  type (wavefunction) :: wfnc_co,wfnc_fi
  type (wavefunction) :: wfnv_co,wfnvq_fi
  type (tdistgwf) :: distgwfco,distgwfcoq

  integer :: ik,ikt,ik_co,ik_fi,ikq_fi,vunit,eqp_dat_size
  integer :: ik_co_q, igumkq, ijk, ivert
  integer :: ic,ic_co,ic_fi,iv,iv_co,iv_fi,is
  integer :: ii,jj,kk,gumk(3),igumk
  real(DP) :: dist,qg(3),dnorm,tempshift,tempsum,dweight,worst_normc, worst_normv
  real(DP), allocatable :: normc(:,:,:,:),normv(:,:,:,:)
  integer, allocatable :: idum(:,:)
  SCALAR, allocatable :: dummy(:,:,:,:,:)
  real(DP), allocatable :: edummy(:,:,:),  dummy_coefs(:,:)
  integer, allocatable :: indexq_co(:)
  type(interp_t) :: wfn_interp
  !> (xct%nvb_fi, xct%n1b, xct%nspin, xct%nkpt_fi, xct%npts_intp_kernel)
  !> dvn points to dvv if we are restricting the transf right away,
  !! or if we never restrict it (i.e., extended_kernel==.true.). Otherwise,
  !! we allocate a buffer, calculate the unrestricted transf, and restrict+copy
  ! dvn to dvv later on.
  SCALAR, pointer :: dvn(:,:,:,:,:) 
  SCALAR, pointer :: dcn(:,:,:,:,:) !< See: dvn
  !> .true. if we calculate the unrestricted transformation, but the kernel is
  !! restricted (i.e., we`ll have to truncate the dvn to dvv later on).
  logical :: truncate_coefs
  type(progress_info) :: prog_info

  PUSH_SUB(intwfn)

  ! FHJ: Normally, xct%n1b_co depends only on xct%extended_kernel. However,
  ! we`ll temporarily overwrite this value depending on xct%unrestricted_transf
  ! till the end of the subroutine, because if xct%unrestricted_transf==.true,
  ! intwfn "thinks" it`s dealing with an extended kernel when calculating dvn/dcn.
  if (xct%unrestricted_transf) then
    xct%n1b_co = xct%nvb_co + xct%ncb_co
    xct%n2b_co = xct%n1b_co
  else
    xct%n1b_co = xct%nvb_co
    xct%n2b_co = xct%ncb_co
  endif
  truncate_coefs = xct%unrestricted_transf.and.(.not.xct%extended_kernel)
  if (truncate_coefs) then
    ! FHJ: we`ll restrict the coefficients later on, but we want to do the
    ! right expansion here => need a temp buffer.
    SAFE_ALLOCATE(dvn, (xct%nvb_fi, xct%n1b_co, xct%nspin, xct%nkpt_fi, xct%npts_intp_kernel))
    SAFE_ALLOCATE(dcn, (xct%ncb_fi, xct%n2b_co, xct%nspin, xct%nkpt_fi, xct%npts_intp_kernel))
    dvn = 0.d0
    dcn = 0.d0
  else
    ! FHJ: we either do not restrict, or restrict right from the beginning.
    dvn => dvv
    dcn => dcc
  endif
  dcc = 0.d0
  dvv = 0.d0

  SAFE_ALLOCATE(normc, (xct%nkpt_fi,xct%ncb_fi,xct%nspin,xct%npts_intp_kernel))
  SAFE_ALLOCATE(normv, (xct%nkpt_fi,xct%nvb_fi,xct%nspin,xct%npts_intp_kernel))
  SAFE_ALLOCATE(indexq_co, (xct%nkpt_co))
  intp_coefs = 0d0
  fi2co_wfn = 0
  if (flag%read_dtmat) then
    if (peinf%inode.eq.0) then

!-------------------------------
! Read transformation matrices, if they exist...
! FHJ: TODO - test this part for xct%unrestricted_transf=.true.

      if (xct%unrestricted_transf) then
        write(6,*) 'Reading dcn,dvn matrices from file'
      else
        write(6,*) 'Reading dcc,dvv matrices from file'
      endif
      
! Check if dtmat has the right parameters

      call open_file(unit=13,file='dtmat',form='unformatted',status='old')
      read(13)
      read(13) ik_co,ic_co,iv_co,ik,ic_fi,iv_fi,is

      if(ik_co /= xct%nkpt_co) then
        write(0,*) 'File has ', ik_co, '; we need ', xct%nkpt_co
        call die("dtmat does not have the correct number of k-points in coarse grid")
      endif
      if(ic_co /= xct%n2b_co) then
        write(0,*) 'File has ', ic_co, '; we need ', xct%n2b_co
        call die("dtmat does not have the correct number of conduction bands in coarse grid")
      endif
      if(iv_co /= xct%n1b_co) then
        write(0,*) 'File has ', iv_co, '; we need ', xct%n1b_co
        call die("dtmat does not have the correct number of valence bands in coarse grid")
      endif
      if(ik /= xct%nkpt_fi) then
        write(0,*) 'File has ', ik, '; we need ', xct%nkpt_fi
        call die("dtmat does not have the correct number of k-points in fine grid")
      endif
      if(ic_fi /= xct%ncb_fi) then
        write(0,*) 'File has ', ic_fi, '; we need ', xct%ncb_fi
        call die("dtmat does not have the correct number of conduction bands in fine grid")
      endif
      if(iv_fi /= xct%nvb_fi) then
        write(0,*) 'File has ', iv_fi, '; we need ', xct%nvb_fi
        call die("dtmat does not have the correct number of valence bands in fine grid")
      endif
      if(is /= xct%nspin) then
        write(0,*) 'File has ', is, '; we need ', xct%nspin
        call die("dtmat does not have the correct number of spins")
      endif
          
! kco(1:3,1:xct%nkpt_co) has the coordinates of k-points in
! the coarse grid used to generate "dtmat" (in general, this
! is equal to kg_co%fk(1:3,1:xct%nkpt_co) of a previous run)
!
! fi2co_wfn(1:xct%nkpt_fi) is a mapping array between a point in
! the fine grid and the closest ones in the coarse grid

      do jj = 1, xct%nkpt_co
        read(13) kco(1,jj),kco(2,jj),kco(3,jj)
      enddo
      do ivert = 1, xct%npts_intp_kernel
        do jj = 1, xct%nkpt_fi * xct%ncb_fi * xct%n2b_co * xct%nspin
          read(13) ik_fi,ic_fi,ik_co,ic_co,is,dcn(ic_fi,ic_co,is,ik_fi,ivert)
        enddo
      enddo
      do ivert = 1, xct%npts_intp_kernel
        do jj = 1, xct%nkpt_fi * xct%nvb_fi * xct%n1b_co * xct%nspin
          read(13) ik_fi,iv_fi,ik_co,iv_co,is,dvn(iv_fi,iv_co,is,ik_fi,ivert)
          fi2co_wfn(ivert,ik_fi) = ik_co
        enddo
      enddo

      intp_coefs(:,:) = 1d0
      if (xct%npts_intp_kernel>1) then
        read(13) intp_coefs(:,:)
      endif
      call close_file(13)
    endif

! PE # 0 distribute data

#ifdef MPI
    call MPI_BCAST(fi2co_wfn(1,1), xct%npts_intp_kernel*xct%nkpt_fi, MPI_INTEGER, 0, MPI_COMM_WORLD,mpierr)
    call MPI_BCAST(kco, 3*xct%nkpt_co, MPI_REAL_DP, 0, MPI_COMM_WORLD,mpierr)
    call MPI_BCAST(dcn(1,1,1,1,1), xct%nkpt_fi*xct%ncb_fi*xct%n2b_co*xct%nspin*xct%npts_intp_kernel, &
      MPI_SCALAR, 0, MPI_COMM_WORLD, mpierr)
    call MPI_BCAST(dvn(1,1,1,1,1), xct%nkpt_fi*xct%nvb_fi*xct%n1b_co*xct%nspin*xct%npts_intp_kernel, &
      MPI_SCALAR, 0, MPI_COMM_WORLD, mpierr)
    call MPI_BCAST(intp_coefs(1,1), xct%npts_intp_kernel*xct%nkpt_fi, MPI_REAL_DP, 0, MPI_COMM_WORLD,mpierr)
#endif

  else

!-------------------------------
! ...Otherwise, start the computation of them

! Read the wavefunctions on the coarse grid and initialize dtmat

    if (xct%qflag .eq. 0) then
      if (peinf%inode .eq. 0) then
        write(6,999)
      endif
    endif
999 format(/,1x,'We are doing a Finite Q Calculation',/)

    call timacc(41, 1)
    call input_co(kp,crys,gvec,kg_co,syms,xct,flag%bzc,distgwfco,eqp)
    call timacc(41, 2)

!---------------------------------------------------- 
! If skip_interpolation keyword is present don`t interpolate

    if (xct%skipinterp) then

      intp_coefs = 1d0
      if (peinf%inode.eq.0) then
        write(6,'(a)') ' Skipping interpolation'
        write(6,'(a)')
      endif

      do ii = 1, xct%nkpt_fi
        do jj = 1, xct%nkpt_co
          if (all(abs(kg_fi%f(1:3,ii)-kg_co%f(1:3,jj)) .lt. TOL_Small)) then
            fi2co_wfn(1,ii)=jj
            kco(1:3,jj)=kg_co%f(1:3,jj)
          endif
        enddo
      enddo

      dcn(:,:,:,:,:) = 0d0
      do ik = 1, xct%nkpt_fi
        do is = 1, xct%nspin
          do ic = 1, xct%ncb_co
            dcn(ic,ic,is,ik,1) = 1d0
          enddo
        enddo
      enddo
      dvn(:,:,:,:,:) = 0d0
      do ik = 1, xct%nkpt_fi
        do is = 1, xct%nspin
          do iv = 1, xct%nvb_co
            dvn(iv,iv,is,ik,1) = 1d0
          enddo
        enddo
      enddo

    else ! xct%skipinterp

      call timacc(42, 1)
      if (xct%delaunay_interp) then
        call interp_init(wfn_interp, crys, kg_co%f, dims=xct%idimensions, &
          periodic=.true., active_dims=xct%is_periodic)
      else
        call alloc_intpts(kg_co%nf, kg_co%f(:,:), periodic=.true.)
      endif
      call timacc(42, 2)

!---------------------------------------------------- 
! JRD: If we do finite Q we need the coarse Q shifted grid wavefunctions

      if (xct%qflag .eq. 0) then
        call timacc(41, 1)
        call input_co_q(kp,crys,gvec,kg_co_q,syms,xct,flag%bzc, &
          indexq_co,kg_co,distgwfcoq)
        call timacc(41, 2)
      else
        do ijk = 1, xct%nkpt_co
          indexq_co(ijk) = ijk
        enddo
      endif
      
      do kk=1,xct%nkpt_co
        kco(:,kk) = kg_co%f(:,kk)
      enddo
      
      if (peinf%inode.eq.0) then
        call open_file(unit=13,file='dtmat',form='unformatted',status='replace')
        write(13) xct%idimensions, xct%is_periodic(1:3), xct%npts_intp_kernel, xct%nkpt_co
        write(13) xct%nkpt_co,xct%n2b_co,xct%n1b_co, &
          xct%nkpt_fi,xct%ncb_fi,xct%nvb_fi,xct%nspin,xct%npts_intp_kernel
        do kk=1,xct%nkpt_co
          write(13) (kco(ii,kk),ii=1,3)
        enddo
        call close_file(13)
      endif

!-------------------------
! Loop over k-points in fine grid
  
      if (xct%inteqp) then
        SAFE_ALLOCATE(eqp%evshift, (xct%nvb_fi,kg_fi%nf,xct%nspin))
        SAFE_ALLOCATE(eqp%ecshift, (xct%ncb_fi,kg_fi%nf,xct%nspin))
        eqp%evshift=0D0
        eqp%ecshift=0D0
      endif

     call progress_init(prog_info, 'wavefunction interpolation', 'k-point', &
      peinf%nkpe)
      do ii=1,peinf%nkpe
        call progress_step(prog_info)
        ik_fi=peinf%ik(peinf%inode+1,ii)

! JRD: Find the points in coarse grid closest to each point in the fine grid
! We need 4 points here for 3D

        if (ik_fi .ne. 0) then
          ikq_fi = indexq_fi(ik_fi)
        else
          ! this process has no work to do
          ikq_fi = 0
        endif
          
! JRD: In principle, we should be calling the below function twice. Once for ik_fi and once for ikq_fi. If using 
! momentum operator, it makes no difference. And for velocity operator, it makes insignificant difference for typical small q`s.

        if (kg_co%nf .gt. 1 .and. ik_fi > 0) then
          call timacc(42, 1)
          if (xct%delaunay_interp) then
            call interp_eval(wfn_interp, kg_fi%f(:,ik_fi), closepts(:), closeweights(:))
          else
            call intpts_local(crys, kg_fi%f(:,ik_fi), kg_co%nf, kg_co%f(:,:), &
              xct, closepts(:), closeweights(:), periodic=.true.)
          endif
          if (xct%npts_intp_kernel>1) then
            intp_coefs(:,ik_fi) = closeweights(1:xct%npts_intp_kernel)
          else
            intp_coefs(1,ik_fi) = 1d0
          endif
          call timacc(42, 2)
        else
          closepts(:)=1
          closeweights(:)=0D0
          closeweights(1)=1D0
          if (ik_fi>0) intp_coefs(1,ik_fi) = 1d0
        endif

        do ijk = xct%idimensions + 1, 1, -1
          
          ik_co=closepts(ijk)
          dweight=closeweights(ijk)
          
          if (xct%qflag .eq. 0) then
            ik_co_q=indexq_co(ik_co)
          else
            ik_co_q=ik_co
          endif

          ! Store the umklapp G-vector, if not zero
          if (ik_fi .ne. 0) then
            qg(:)=kg_fi%f(:,ik_fi)-kg_co%f(:,ik_co)
            gumk(:)=nint(qg(:))
            call findvector(igumk, gumk, gvec)

            ! Store the umklapp G-vector, if not zero - Finite Q
            if (xct%qflag .eq. 0) then
              qg(:)=kgq_fi%f(:,ik_fi)-kg_co_q%f(:,ik_co_q)
              gumk(:)=nint(qg(:))
              call findvector(igumkq, gumk, gvec)
            else
              igumkq=igumk
            endif
          endif
      
!------------------------------
! Read needed wavefunctions: conduction bands from unshifted grid
! and valence bands from shifted grid

          if (ik_fi .ne. 0) then
            call timacc(43, 1)
            ! Need to initialize nspinor value for genwf
            intwfnc%nspinor=kp%nspinor
            intwfnv%nspinor=kp%nspinor
            call genwf(crys,gvec,kg_fi,syms,wfnc_fi,xct,ik_fi,ik_fi,work,intwfnc, is_cond = .true.)
            call genwf(crys,gvec,kgq_fi,syms,wfnvq_fi,xct,ik_fi,ikq_fi,workq,intwfnv, is_cond = .false.)
            call timacc(43, 2)
          endif
          
          if (xct%qflag .eq. 0) then
            call timacc(44, 1)
            call genwf_co(crys,gvec,kg_co,kg_co_q,syms,wfnc_co, &
              wfnv_co,xct,ik_co,ik_co_q,distgwfco,distgwfcoq,workco,workcoq)
            call timacc(44, 2)
          else
            call timacc(44, 1)
            call genwf_co(crys,gvec,kg_co,kg_co_q,syms,wfnc_co, &
              wfnv_co,xct,ik_co,ik_co_q,distgwfco,distgwfco,workco,workcoq)
            call timacc(44, 2)
          endif

          if (ik_fi .ne. 0) then
            call timacc(45, 1)
            ivert = 1
            if (xct%npts_intp_kernel>1) ivert = ijk
            call mtxel_t(gvec,wfnc_co, wfnc_fi, wfnv_co, wfnvq_fi, &
              dcn(:,:,:,ik_fi,ivert), dvn(:,:,:,ik_fi,ivert),  &
              igumk, igumkq, .not.xct%unrestricted_transf)
            call timacc(45, 2)
          endif
          
          SAFE_DEALLOCATE_P(wfnc_co%cg)
          SAFE_DEALLOCATE_P(wfnc_co%isort)
          SAFE_DEALLOCATE_P(wfnv_co%cg)
          SAFE_DEALLOCATE_P(wfnv_co%isort)
          
          if (ik_fi .ne. 0) then
            SAFE_DEALLOCATE_P(wfnc_fi%cg)
            SAFE_DEALLOCATE_P(wfnc_fi%isort)
            SAFE_DEALLOCATE_P(wfnvq_fi%cg)
            SAFE_DEALLOCATE_P(wfnvq_fi%isort)
            fi2co_wfn(ivert, ik_fi) = ik_co
          endif

! Construct Interpolated eqp shifts
          
          if (xct%inteqp .and. ik_fi .ne. 0) then
            
            do is=1,xct%nspin
              do iv_fi=1,xct%nvb_fi
                tempshift=0d0
                tempsum=0d0
                do iv_co=1,xct%nvb_co
                  tempshift = tempshift + (abs(dvn(iv_fi,iv_co,is,ik_fi,ivert))**2)*&
                    eqp%evshift_co(iv_co,kg_co%indr(ik_co),is)
                  tempsum = tempsum + (abs(dvn(iv_fi,iv_co,is,ik_fi,ivert))**2)
                enddo
                do iv_co=xct%nvb_co+1,xct%n1b_co
                  tempshift = tempshift + (abs(dvn(iv_fi,iv_co,is,ik_fi,ivert))**2)*&
                    eqp%ecshift_co(iv_co-xct%nvb_co,kg_co%indr(ik_co),is)
                  tempsum = tempsum + (abs(dvn(iv_fi,iv_co,is,ik_fi,ivert))**2)
                enddo
                if (.not.xct%renorm_transf) tempsum = 1d0
                eqp%evshift(iv_fi,ik_fi,is)= eqp%evshift(iv_fi,ik_fi,is)+tempshift*dweight/(tempsum)
              enddo
            enddo
            
            do is=1,xct%nspin
              do ic_fi=1,xct%ncb_fi
                tempshift=0d0
                tempsum=0d0
                if (xct%unrestricted_transf) then
                  do ic_co=1,xct%nvb_co
                    tempshift = tempshift + (abs(dcn(ic_fi,ic_co,is,ik_fi,ivert))**2)*&
                      eqp%evshift_co(ic_co,kg_co%indr(ik_co),is)
                    tempsum = tempsum + (abs(dcn(ic_fi,ic_co,is,ik_fi,ivert))**2)
                  enddo
                  do ic_co=xct%nvb_co+1,xct%n2b_co
                    tempshift = tempshift + (abs(dcn(ic_fi,ic_co,is,ik_fi,ivert))**2)*&
                      eqp%ecshift_co(ic_co-xct%nvb_co,kg_co%indr(ik_co),is)
                    tempsum = tempsum + (abs(dcn(ic_fi,ic_co,is,ik_fi,ivert))**2)
                  enddo
                else
                  do ic_co=1,xct%ncb_co
                    tempshift = tempshift + (abs(dcn(ic_fi,ic_co,is,ik_fi,ivert))**2)*&
                      eqp%ecshift_co(ic_co,kg_co%indr(ik_co),is)
                    tempsum = tempsum + (abs(dcn(ic_fi,ic_co,is,ik_fi,ivert))**2)
                  enddo
                endif
                if (.not.xct%renorm_transf) tempsum = 1d0
                eqp%ecshift(ic_fi,ik_fi,is)= eqp%ecshift(ic_fi,ik_fi,is)+ tempshift*dweight/(tempsum)
              enddo
            enddo
          endif
      
        enddo ! ijk
      enddo ! ii
      call progress_free(prog_info)
  
      if (xct%inteqp) then
        
        if (xct%nkpt_fi .ne. kg_fi%nf) then
          write(0,*) "WARNING: nkpt != nf", xct%nkpt_fi, kg_fi%nf
        endif
        
        call timacc(46, 1)
#ifdef MPI 
        SAFE_ALLOCATE(edummy, (xct%nvb_fi,xct%nkpt_fi,xct%nspin))
        edummy = eqp%evshift
        call MPI_Allreduce(edummy(1,1,1),eqp%evshift(1,1,1),xct%nkpt_fi*xct%nvb_fi* &
          xct%nspin,MPI_DOUBLE_PRECISION,MPI_SUM,MPI_COMM_WORLD,mpierr)
        SAFE_DEALLOCATE(edummy)
        
        SAFE_ALLOCATE(edummy, (xct%ncb_fi,xct%nkpt_fi,xct%nspin))
        edummy = eqp%ecshift
        call MPI_Allreduce(edummy(1,1,1),eqp%ecshift(1,1,1),xct%nkpt_fi*xct%ncb_fi* &
          xct%nspin,MPI_DOUBLE_PRECISION,MPI_SUM,MPI_COMM_WORLD,mpierr)
        SAFE_DEALLOCATE(edummy)
#endif
        
        if (peinf%inode .eq. 0) then
          
          ! Only write eqp_q.dat for velocity. For momentum, it is identical to valence bands of eqp.dat.
          if(flag%opr == 0) then
            call open_file(unit=21,file='eqp_q.dat',form='formatted',status='replace')
            vunit = 21
            eqp_dat_size = xct%ncb_fi*xct%nspin
          else
            vunit = 22
            eqp_dat_size = (xct%nvb_fi+xct%ncb_fi)*xct%nspin
          endif
          call open_file(unit=22,file='eqp.dat',  form='formatted',status='replace')
          do ik_fi = 1, kg_fi%nf
            if(flag%opr == 0) &
              write(21,'(3f13.9,i8)') kgq_fi%f(1,ik_fi),kgq_fi%f(2,ik_fi),kgq_fi%f(3,ik_fi),xct%nvb_fi*xct%nspin
            write(22,'(3f13.9,i8)') kg_fi%f(1,ik_fi),kg_fi%f(2,ik_fi),kg_fi%f(3,ik_fi),eqp_dat_size
            do is = 1, xct%nspin
              do iv_fi = xct%nvb_fi, 1, -1
                write(vunit,'(2i8,2f15.9)') is, xct%ifmaxq(ik_fi ,is) + 1 - iv_fi, eqp%evlda(iv_fi,ik_fi,is)*RYD, &
                  (eqp%evlda(iv_fi,ik_fi,is)+eqp%evshift(iv_fi,ik_fi,is))*RYD
              enddo
              do ic_fi = 1, xct%ncb_fi
                write(22,'(2i8,2f15.9)') is, xct%ifmax(ik_fi ,is) + ic_fi, eqp%eclda(ic_fi,ik_fi,is)*RYD, &
                  (eqp%eclda(ic_fi,ik_fi,is)+eqp%ecshift(ic_fi,ik_fi,is))*RYD
              enddo
            enddo
          enddo
          if(flag%opr == 0) call close_file(21)
          call close_file(22)
          
          call open_file(unit=23,file='bandstructure.dat',form='formatted',status='replace')
          ! Write k-points in cart. coord. instead of cryst. coord. for the 
          ! band structure plot (as it is done in espresso and wannier90) --GSM
          write(23,'(a)') '# spin    band          kx          ky          kz          E(MF)          E(QP)        Delta E'
          write(23,'(a)') '#                        (Cartesian coordinates)             (eV)           (eV)           (eV)'

          do is = 1, xct%nspin
            do iv_fi = xct%nvb_fi, 1, -1
              ! For the valence bands, the "q" structures are the ones that contained the calculated info. NOTE: kgq%fi = kg%fi
              do ik_fi = 1,kgq_fi%nf
                write(23,'(i6,i8,3f12.5,3f15.9)') is,xct%ifmaxq(ik_fi,is) + 1 - iv_fi, &
                  matmul(crys%bvec(1:3,1:3),kgq_fi%f(1:3,ik_fi))*crys%blat, &
                  eqp%evlda(iv_fi,ik_fi,is)*RYD,(eqp%evlda(iv_fi,ik_fi,is)+eqp%evshift(iv_fi,ik_fi,is))*RYD, &
                  eqp%evshift(iv_fi,ik_fi,is)*RYD
              enddo
            enddo
            do ic_fi = 1,xct%ncb_fi
              do ik_fi = 1,kg_fi%nf
                write(23,'(i6,i8,3f12.5,3f15.9)') is,xct%ifmax(ik_fi,is) + ic_fi, &
                  matmul(crys%bvec(1:3,1:3),kg_fi%f(1:3,ik_fi))*crys%blat, &
                  eqp%eclda(ic_fi,ik_fi,is)*RYD,(eqp%eclda(ic_fi,ik_fi,is)+eqp%ecshift(ic_fi,ik_fi,is))*RYD, &
                  eqp%ecshift(ic_fi,ik_fi,is)*RYD
              enddo
            enddo
          enddo
          call close_file(23)

! JRD: We now update our fine-grid eigenvalues with the interpolated ones.
  
        endif ! peinf%inode .eq. 0

        call timacc(46, 2)
        eqp%ecqp(:,:,:)=eqp%eclda(:,:,:)+eqp%ecshift(:,:,:)
        eqp%evqp(:,:,:)=eqp%evlda(:,:,:)+eqp%evshift(:,:,:)    
        
      endif ! xct%inteqp
      
      call dealloc_grid(kg_co)
      
      if (xct%qflag .eq. 0) call dealloc_grid(kg_co_q)
      
      ! typedefs initializes all of these ikolds to 0
      if (work%ikold .ne. 0) then
        SAFE_DEALLOCATE_P(work%cg)
        SAFE_DEALLOCATE_P(work%ph)
        SAFE_DEALLOCATE_P(work%ind)
        SAFE_DEALLOCATE_P(work%isort)
      endif
      if (workq%ikold .ne. 0) then
        SAFE_DEALLOCATE_P(workq%cg)
        SAFE_DEALLOCATE_P(workq%ph)
        SAFE_DEALLOCATE_P(workq%ind)
        SAFE_DEALLOCATE_P(workq%isort)
      endif
      if (workco%ikold .ne. 0) then
        SAFE_DEALLOCATE_P(workco%cg)
        SAFE_DEALLOCATE_P(workco%ph)
        SAFE_DEALLOCATE_P(workco%ind)
        SAFE_DEALLOCATE_P(workco%isort)
      endif
      if (workcoq%ikold .ne. 0) then
        SAFE_DEALLOCATE_P(workcoq%cg)
        SAFE_DEALLOCATE_P(workcoq%ph)
        SAFE_DEALLOCATE_P(workcoq%ind)
        SAFE_DEALLOCATE_P(workcoq%isort)
      endif

!--------------------------------
! Write out dcn,dvn. All PEs access dtmat
! FHJ: TODO - only one PE should do the job!
      
      do ivert = 1, xct%npts_intp_kernel
      do ii=1,peinf%npes
        if(peinf%inode+1.eq.ii) then
          call open_file(unit=13,file='dtmat',position='append',form='unformatted',status='old')
          do ikt=1,peinf%ikt(peinf%inode+1)
            ik_fi=peinf%ik(peinf%inode+1,ikt)
            do is=1,xct%nspin
              do ic_co=1,xct%n2b_co
                do ic_fi=1,xct%ncb_fi
                  write(13) ik_fi,ic_fi,fi2co_wfn(ivert,ik_fi),ic_co,is, &
                    dcn(ic_fi,ic_co,is,ik_fi,ivert)
                enddo
              enddo
            enddo
          enddo
          call close_file(13)
        endif
#ifdef MPI
        call MPI_BARRIER(MPI_COMM_WORLD,mpierr)
#endif
      enddo
      enddo
  
      do ivert = 1, xct%npts_intp_kernel
      do ii=1,peinf%npes
        if(peinf%inode+1.eq.ii) then
          call open_file(unit=13,file='dtmat',position='append',form='unformatted',status='old')
          do ikt=1,peinf%ikt(peinf%inode+1)
            ik_fi=peinf%ik(peinf%inode+1,ikt)
            do is=1,xct%nspin
              do iv_fi=1,xct%nvb_fi
                do iv_co=1,xct%n1b_co
                  write(13) ik_fi,iv_fi,fi2co_wfn(ivert,ik_fi),iv_co,is, &
                    dvn(iv_fi,iv_co,is,ik_fi,ivert)
                enddo
              enddo
            enddo
          enddo
          call close_file(13)
        endif
#ifdef MPI
        call MPI_BARRIER(MPI_COMM_WORLD,mpierr)
#endif
      enddo
      enddo
      
      if (xct%iwriteint.eq.1) then
        SAFE_DEALLOCATE_P(distgwfco%isort)
        SAFE_DEALLOCATE_P(distgwfco%zv)
        SAFE_DEALLOCATE_P(distgwfco%zc)
        if (xct%qflag .eq. 0) then
          SAFE_DEALLOCATE_P(distgwfcoq%isort)
          SAFE_DEALLOCATE_P(distgwfcoq%zv)
          SAFE_DEALLOCATE_P(distgwfcoq%zc)
        endif
      endif ! xct%iwriteint.eq.1
      

!--------------------------------
! Finished the computation of transformation matrices, dcn/dvn
! Now, PEs share the information

      call timacc(47, 1)
#ifdef MPI
      SAFE_ALLOCATE(dummy, (xct%ncb_fi,xct%n2b_co,xct%nspin,xct%nkpt_fi,xct%npts_intp_kernel))
      dummy = dcn
      call MPI_Allreduce(dummy(1,1,1,1,1),dcn(1,1,1,1,1),xct%nkpt_fi*xct%ncb_fi* &
        xct%n2b_co*xct%nspin*xct%npts_intp_kernel,MPI_SCALAR,MPI_SUM,MPI_COMM_WORLD,mpierr)
      SAFE_DEALLOCATE(dummy)
      SAFE_ALLOCATE(dummy, (xct%nvb_fi,xct%n1b_co,xct%nspin,xct%nkpt_fi,xct%npts_intp_kernel))
      dummy = dvn
      call MPI_Allreduce(dummy(1,1,1,1,1),dvn(1,1,1,1,1),xct%nkpt_fi*xct%nvb_fi* &
        xct%n1b_co*xct%nspin*xct%npts_intp_kernel,MPI_SCALAR,MPI_SUM,MPI_COMM_WORLD,mpierr)
      SAFE_DEALLOCATE(dummy)
      SAFE_ALLOCATE(idum, (xct%npts_intp_kernel, xct%nkpt_fi))
      idum = fi2co_wfn
      call MPI_Allreduce(idum(1,1),fi2co_wfn(1,1),xct%nkpt_fi*xct%npts_intp_kernel,MPI_INTEGER, &
        MPI_SUM,MPI_COMM_WORLD,mpierr)
      SAFE_DEALLOCATE(idum)
      SAFE_ALLOCATE(dummy_coefs, (xct%npts_intp_kernel, xct%nkpt_fi))
      dummy_coefs = intp_coefs
      call MPI_Allreduce(dummy_coefs(1,1), intp_coefs(1,1), xct%nkpt_fi*xct%npts_intp_kernel, MPI_REAL_DP, &
        MPI_SUM, MPI_COMM_WORLD, mpierr)
      SAFE_DEALLOCATE(dummy_coefs)
#endif
      call timacc(47, 2)

      ! I am not sure why but this barrier is necessary to avoid hanging. --DAS
#ifdef MPI
      call MPI_BARRIER(MPI_COMM_WORLD,mpierr)
#endif

      if (xct%npts_intp_kernel>1 .and. peinf%inode==0) then
        call open_file(unit=13,file='dtmat',position='append',form='unformatted',status='old')
        write(13) intp_coefs(:,:)
        call close_file(13)
      endif

#ifdef MPI
      call MPI_BARRIER(MPI_COMM_WORLD,mpierr)
#endif

      if (xct%delaunay_interp) then
        call interp_free(wfn_interp)
      else
        call dealloc_intpts()
      endif
    
    endif ! xct%skipinterp

    if (xct%iwriteint.eq.0) then
      if(peinf%inode.eq.0) then
        call open_file(126, 'INT_CWFN_CO', status='old')
        call close_file(126, delete = .true.) ! delete INT_CWFN_CO
        call open_file(127, 'INT_VWFN_CO', status='old')
        call close_file(127, delete = .true.) ! delete INT_VWFN_CO
        if (xct%qflag .eq. 0) then
          call open_file(123, 'INT_CWFN_CO_Q', status='old')
          call close_file(123, delete = .true.) ! delete INT_CWFN_CO_Q
          call open_file(124, 'INT_VWFN_CO_Q', status='old')
          call close_file(124, delete = .true.) ! delete INT_VWFN_CO_Q
        endif
      endif
    endif ! xct%iwriteint.eq.0
  endif ! flag%read_dtmat

!--------------------------------
! Check if transformation matrices are unitary and print out their
! norm (i.e., the norm of wavefunctions in the fine grid)
! The norm should be as close as possible to 1.d0. If it is not,
! important overlaps are being neglected.
  do ivert = 1, xct%npts_intp_kernel
  do is = 1, xct%nspin
    do ik_fi=1,xct%nkpt_fi
      do iv=1,xct%nvb_fi
        dnorm = 0.d0
        do iv_co=1,xct%n1b_co
          dnorm = dnorm + abs(dvn(iv,iv_co,is,ik_fi,ivert))**2
        enddo
        normv(ik_fi,iv,is,ivert) = dnorm
      enddo
      
      do ic=1,xct%ncb_fi
        dnorm = 0.d0
        do ic_co=1,xct%n2b_co
          dnorm = dnorm + abs(dcn(ic,ic_co,is,ik_fi,ivert))**2
        enddo
        normc(ik_fi,ic,is,ivert) = dnorm
      enddo
    enddo
  enddo
  enddo
      
  if (.not.flag%read_dtmat .and. peinf%inode==0) then
    call open_file(20,file='dcmat_norm.dat',form='formatted',status='replace')
    write(20,'(a,i12,a)') ' -------  Norm of dcc matrices : Spins = ', xct%nspin,'  -------'
    if (xct%nspin.eq.1) then
      write(20,'(a)') '           k-point           ik_co    c   dist    |dcc|^2      '
    else
      write(20,'(a)') '           k-point           ik_co    c   dist    |dcc|^2 (spin up)  |dcc|^2 (spin down) '
    endif
    write(20,'(21a3)') ('---',kk=1,21)
    do ik_fi=1,xct%nkpt_fi
      ik_co= fi2co_wfn(1,ik_fi)
      do kk=1,3
        qg(kk)=kg_fi%f(kk,ik_fi)-kco(kk,ik_co)
        qg(kk) = qg(kk) - anint( qg(kk) )
      enddo
      dist= DOT_PRODUCT(qg,MATMUL(crys%bdot,qg))
      do jj=1,xct%ncb_fi
        write(20,250) kg_fi%f(:,ik_fi),ik_co,jj,sqrt(dist),(normc(ik_fi,jj,is,1), is = 1, xct%nspin)
      enddo
    enddo
    call close_file(20)
    
    call open_file(20,file='dvmat_norm.dat',form='formatted',status='replace')
    write(20,'(a,i12,a)') ' -------  Norm of dvv matrices : Spins = ', xct%nspin,'  -------'
    if (xct%nspin.eq.1) then
      write(20,'(a)') '           k-point           ik_co    v   dist    |dvv|^2      '
    else
      write(20,'(a)') '           k-point           ik_co    v   dist    |dvv|^2 (spin up)  |dvv|^2 (spin down) '
    endif
    write(20,'(21a3)') ('---',kk=1,21)
    do ik_fi=1,xct%nkpt_fi
      ik_co= fi2co_wfn(1,ik_fi)
      do kk=1,3
        qg(kk)=kg_fi%f(kk,ik_fi)-kco(kk,ik_co)
        qg(kk) = qg(kk) - anint( qg(kk) )
      enddo
      dist= DOT_PRODUCT(qg,MATMUL(crys%bdot,qg))
      do jj=1,xct%nvb_fi
        write(20,250) kg_fi%f(:,ik_fi),ik_co,jj,sqrt(dist),(normv(ik_fi,jj,is,1), is = 1, xct%nspin)
      enddo
    enddo
    call close_file(20)
250   format(' ( ',f6.3,' , ',f6.3,' , ',f6.3, ' ) ',i4,1x,i4,1x,f6.3,1x,f10.6,9x,f10.6)
  endif    

!---------------------------------
! Renormalize transformation matrices: matrices dcn,dvn are
! renormalized to 1.d0 (ideally, this step should not be
! performed but the interpolation can be really bad if dcn,dvn
! do not have unitary norm).

  ! FHJ: Print worst norm before restriction/renormalization
  worst_normv = minval(normv)
  worst_normc = minval(normc)
  if (peinf%inode==0) then    
    write(6,'(1x,a)') 'Max. error in norm of transformation coefficients (1 - \sum_co |d_fi,co|^2):'
    write(6,'(1x,a,es9.3e2)') ' - For valence states: ', 1d0 - worst_normv
    write(6,'(1x,a,es9.3e2)') ' - For conduction states: ', 1d0 - worst_normc
    call warn_norm(xct, worst_normv, worst_normc)
  endif
  ! FHJ: truncate or zero out coefficients. User doesn`t care about the details...
  if (truncate_coefs) then
    if (peinf%inode==0) write(6,'(1x,a)') "Restricting transformation to dvv',dcc' subspaces."
    dvv(:,:,:,:,:) = dvn(:,:xct%nvb_co,:,:,:)
    dcc(:,:,:,:,:) = dcn(:,xct%nvb_co+1:,:,:,:)
  elseif (xct%zero_unrestricted_contrib) then
    if (peinf%inode==0) write(6,'(1x,a)') "Restricting transformation to dvv',dcc' subspaces."
    dvn(:,xct%nvb_co+1:,:,:,:) = ZERO
    dcn(:,:xct%nvb_co,:,:,:) = ZERO
  endif

  ! FHJ: And from now on, xct%n1b_co will depend only on xct%extended_kernel.
  if (.not.xct%extended_kernel) then
    xct%n1b_co = xct%nvb_co
    xct%n2b_co = xct%ncb_co
  else
    xct%n1b_co = xct%nvb_co + xct%ncb_co
    xct%n2b_co = xct%n1b_co
  endif

  ! FHJ: inteqp also renormalizes the coefs for consistency (and someone might
  ! need this in the future). And note that we go here if truncate_coefs==.true.
  ! only to recalculate the worst norm.
  if (truncate_coefs.or.xct%renorm_transf) then
    worst_normc = INF
    worst_normv = INF
    do ivert = 1, xct%npts_intp_kernel
    do is=1, xct%nspin
      do ik=1,xct%nkpt_fi
        do ic_fi=1,xct%ncb_fi
          dnorm = sum(abs(dcc(ic_fi,:,is,ik,ivert))**2)
          worst_normc = min(worst_normc, dnorm)
          if (xct%renorm_transf) then
            dnorm = sqrt(dnorm)
            dcc(ic_fi,:,is,ik,ivert) = dcc(ic_fi,:,is,ik,ivert)/dnorm
          endif
        enddo
        do iv_fi=1,xct%nvb_fi
          dnorm = sum(abs(dvv(iv_fi,:,is,ik,ivert))**2)
          worst_normv = min(worst_normv, dnorm)
          if (xct%renorm_transf) then
            dnorm = sqrt(dnorm)
            dvv(iv_fi,:,is,ik,ivert) = dvv(iv_fi,:,is,ik,ivert)/dnorm
          endif
        enddo
      enddo
    enddo
    enddo
  endif
  if (peinf%inode==0.and.(truncate_coefs.or.xct%zero_unrestricted_contrib)) then
    write(6,'(1x,a)') 'Max. error after restricting WFN transformation:'
    write(6,'(1x,a,es9.3e2)') ' - For valence states: ', 1d0 - worst_normv
    write(6,'(1x,a,es9.3e2)') ' - For conduction states: ', 1d0 - worst_normc
    call warn_norm(xct, worst_normv, worst_normc)
  endif

!-------------------------------------
! Print out information about the number of k-points in the fine grid
! interpolated from each point in the coarse grid
! ii : number of fine-grid points close to point ik_co

#ifdef VERBOSE
  if (peinf%inode==0) then
    write(6,*)
    write(6,'(1x,a)') 'K-points in coarse matrix:'
    do ik_co=1,xct%nkpt_co
      ii = 0
      do ik=1,xct%nkpt_fi
        if (any(fi2co_wfn(:,ik)==ik_co)) ii = ii + 1
      enddo
      write(6,'(i6,1x,3f10.6,1x,i6)') ik_co,kco(:,ik_co),ii
    enddo
    write(6,*)
  endif
#endif

  if (.not.associated(dvn,dvv)) then
    SAFE_DEALLOCATE_P(dvn)
  endif
  if (.not.associated(dcn,dcc)) then
    SAFE_DEALLOCATE_P(dcn)
  endif
  SAFE_DEALLOCATE(normc)
  SAFE_DEALLOCATE(normv)
  SAFE_DEALLOCATE(indexq_co)

  POP_SUB(intwfn)
  
  return
end subroutine intwfn

    subroutine warn_norm(xct, worst_normv, worst_normc)
      type(xctinfo), intent(in) :: xct
      real(DP), intent(in) :: worst_normv, worst_normc

      logical, save :: warned = .false.

      PUSH_SUB(warn_norm)

      if (.not.warned) then
        ! FHJ: for HEP calculations, change the criterion below to 5 sigma.
        if (worst_normv<95d-2 .or. worst_normc<95d-2) then
          write(0,'(/,a)') 'WARNING: there are fine/coarse transformation coefficients with 2-norm < 95%.'
          write(0,'(a)') 'To improve convergence, you might want to consider:'
          if (xct%is_absorption) then
            write(0,'(a)') '- recalculating the kernel using the "extended_kernel" flag'
          endif
          if (.not. xct%unrestricted_transf) then
            write(0,'(a)') '- using the "unrestricted_transformation" flag in the input file'
          endif
          write(0,'(a)') '- including more bands from the coarse WFN_co file (number_*_bands_coarse)'
          write(0,'(a)') '- using a coarse WFN_co file with a denser k-mesh'
          write(0,*)
          warned = .true.
        endif
      endif

      POP_SUB(warn_norm)

    end subroutine warn_norm

end module intwfn_m
