!==========================================================================
!
! Routines:
!
! (1) input_co_q()      Originally By ?         Last Modified 4/19/2009 (gsm)
!
!     input: crys, gvec,  syms, xct, flagbz types
!
!     output: kg, indexq_co, kgq, distgwfcoq types
!
!     Reads in the coarse grid q wavefunctions from file WFNq_co
!     and distributes them between processors (if xct%iwriteint=1)
!     or writes them in temporary files (if xct%iwriteint=0).
!     The k-point grid is stored in kgq.
!
!  WARNING: Since this routine is not used by a working part of the code,
!  it has not been tested since implementation of new wfn format. --DAS
!  eqp corrections, Fermi level check, scissor shift probably needed.
!==========================================================================

#include "f_defs.h"

subroutine input_co_q(kp,crys,gvec,kgq,syms,xct,flagbz,indexq_co,kg,distgwfcoq)
  
  use global_m
  use checkbz_m
  use eqpcor_m
  use fullbz_m
  use input_utils_m
  use misc_m
  use wfn_rho_vxc_io_m
  implicit none 

  type (kpoints), intent(in) :: kp
  type (crystal), intent(in) :: crys
  type (gspace), intent(in) :: gvec
  type (grid), intent(out) :: kgq
  type (symmetry), intent(in) :: syms
  type (xctinfo), intent(in) :: xct
  integer, intent(in) :: flagbz
  integer, intent(out) :: indexq_co(xct%nkpt_co)
  type (grid), intent(inout) :: kg
  type (tdistgwf), intent(out) :: distgwfcoq

  type (crystal) :: crys_co
  type (symmetry) :: syms_co
  type (wavefunction) :: wfnv,wfnc
  type (kpoints) :: kp_co
  character :: filenamev*20,filenamec*20
  character :: tmpfn*16
  integer :: iunit_v,iunit_c,ikq
  integer :: irk
  integer :: ii,jj,kk,ik,is,isp
  integer :: irks
  real(DP) :: delta,kt(3),div,qq(3)
  
  integer, allocatable :: dist(:)
  integer, allocatable :: indxk(:)
  SCALAR, allocatable :: cg(:,:), cgarray(:)

  character(len=3) :: sheader
  integer :: iflavor
  type(gspace) :: gvec_co, gvec_kpt

  logical :: skip_checkbz

  PUSH_SUB(input_co_q)

!-------------------------
! Print to stdout

  if (peinf%inode.eq.0) write(6,900)
900 format(/,1x,'Started reading coarse grid q wavefunctions from unit WFNq_co',/)
  
  ! ngmax, i.e. kp_co%ngkmax, was being read from WFN_co here
  ! seems more rational to use kp_co%ngkmax from WFNq_co -- DAS

  if (peinf%inode.eq.0) call open_file(unit=26,file='WFNq_co',form='unformatted',status='old')
  sheader = 'WFN'
  iflavor = 0
  call read_binary_header_type(26, sheader, iflavor, kp_co, gvec_co, syms_co, crys_co, warn = .false.)
  call check_trunc_kpts(xct%icutv, kp_co)

  call check_header('WFN_fi', kp, gvec, syms, crys, 'WFNq_co', kp_co, gvec_co, syms_co, crys_co, is_wfn = .true.)

  SAFE_ALLOCATE(gvec_co%components, (3, gvec_co%ng))
  call read_binary_gvectors(26, gvec_co%ng, gvec_co%ng, gvec_co%components, dont_read = .true.)
  SAFE_DEALLOCATE_P(gvec_co%components)

  if (crys_co%celvol.ne.crys%celvol) then
    call die('The crystal structure in WFNq_co is not the same as in WFN_fi.', only_root_writes = .true.)
  endif

  if (syms_co%ntran .ne. syms%ntran) then
    call die("The crystal symmetry in WFNq_co is not the same as in WFN_fi.", only_root_writes = .true.)
  endif

!-----------------------------------------------------------------------
!     Read k-points from file kpoints_co (if it exists) or from WFNq_co
!     Array indxk has the same meaning as in input

  if (xct%read_kpoints) then
    if (peinf%inode.eq.0) then
      call open_file(9,file='kpoints_co',form='formatted',status='old')
      read(9,*) kg%nr
      SAFE_ALLOCATE(kg%r, (3,kg%nr))
      do ii=1,kg%nr
        read(9,*) (kg%r(jj,ii),jj=1,3),div
        kg%r(:,ii) = kg%r(:,ii)/div
      enddo
      call close_file(9)
    endif
#ifdef MPI
    call MPI_BCAST(kg%nr,1,MPI_INTEGER,0,MPI_COMM_WORLD,mpierr)
    if (peinf%inode.ne.0) then
      SAFE_ALLOCATE(kg%r, (3,kg%nr))
    endif
    call MPI_BCAST(kg%r,3*kg%nr,MPI_REAL_DP,0,MPI_COMM_WORLD,mpierr)
#endif
    SAFE_ALLOCATE(indxk, (kg%nr))
    indxk=0
    do jj=1,kg%nr
      do ii=1,kp_co%nrk
        kt(:) = kg%r(:,jj) - kp_co%rk(:,ii)
        if (all(abs(kt(1:3)) < TOL_Small)) then
          if (indxk(jj).ne.0 .and. peinf%inode.eq.0) write(6,996) jj,indxk(jj),kg%r(:,jj)
          indxk(jj)=ii
        endif
      enddo
      if (indxk(jj).eq.0 .and. peinf%inode.eq.0) write(6,995) kg%r(:,jj)
    enddo
  else
    kgq%nr=kp_co%nrk
    SAFE_ALLOCATE(kgq%r, (3,kgq%nr))
    kgq%r(1:3,1:kgq%nr)=kp_co%rk(1:3,1:kp_co%nrk)
    SAFE_ALLOCATE(indxk, (kgq%nr))
    do ii=1,kgq%nr
      indxk(ii)=ii
    enddo
  endif
996 format(1x,'WARNING: Multiple definition of k-point',2i4,3f10.6)
995 format(1x,'WARNING: Could not find k-point',3f10.6,1x,'in WFNq_co')

!-----------------------------------------------------------------------
! Initialization of distributed wavefunctions

  if (xct%iwriteint.eq.1) then
    
    distgwfcoq%ngm=kp_co%ngkmax
    distgwfcoq%nk=kg%nr
    distgwfcoq%ns=kp_co%nspin
    distgwfcoq%nspinor=kp_co%nspinor
    distgwfcoq%nv=xct%nvb_co
    distgwfcoq%nc=xct%ncb_co
    
    SAFE_ALLOCATE(dist, (peinf%npes))
    dist=0
    jj=peinf%npes
    do ii=1,kp_co%ngkmax
      dist(jj)=dist(jj)+1
      jj=jj-1
      if (jj.eq.0) jj=peinf%npes
    enddo
    distgwfcoq%ngl=dist(peinf%inode+1)
    jj=0
    do ii=1,peinf%inode
      jj=jj+dist(ii)
    enddo
    distgwfcoq%tgl=jj
    SAFE_DEALLOCATE(dist)
    
    SAFE_ALLOCATE(distgwfcoq%ng, (distgwfcoq%nk))
    SAFE_ALLOCATE(distgwfcoq%isort, (distgwfcoq%ngl,distgwfcoq%nk))
    SAFE_ALLOCATE(distgwfcoq%zv, (distgwfcoq%ngl,distgwfcoq%nv,distgwfcoq%ns*distgwfcoq%nspinor,distgwfcoq%nk))
    SAFE_ALLOCATE(distgwfcoq%zc, (distgwfcoq%ngl,distgwfcoq%nc,distgwfcoq%ns*distgwfcoq%nspinor,distgwfcoq%nk))
    
    distgwfcoq%ng(:)=0
    distgwfcoq%isort(:,:)=0
    distgwfcoq%zv(:,:,:,:)=ZERO
    distgwfcoq%zc(:,:,:,:)=ZERO
    
  endif ! xct%iwriteint.eq.1

! DAS: Something like this is needed for eqp_co_q.dat here
!-----------------------------------------------------------------------
!     Read eqp_co_q.dat for possible interpolation

! see input_co.f90 for what to do
  
!-----------------------------------------------------------------------
!     Generate full Brillouin zone from irreducible wedge, rk -> fk

  if (flagbz.eq.1) then
    call fullbz(crys,syms,kgq,1,skip_checkbz,wigner_seitz=.true.,paranoid=.true.)
  else
    call fullbz(crys,syms,kgq,syms%ntran,skip_checkbz,wigner_seitz=.true.,paranoid=.true.)
  endif
  tmpfn='WFNq_co'
  if (.not. skip_checkbz) then
    call checkbz(kgq%nf,kgq%f,kp_co%kgrid,kp_co%shift,crys%bdot, &
      tmpfn,'k',.true.,xct%freplacebz,xct%fwritebz)
  endif
  
  if (flagbz.eq.0.and.peinf%inode.eq.0) write(6,801)
  if (flagbz.eq.1.and.peinf%inode.eq.0) write(6,802)
801 format(1x,'Using symmetries to expand the coarse grid q sampling', &
      /)
802 format(1x,'No symmetries used in the coarse grid q sampling',/)

  if (xct%nkpt_co.ne.kgq%nf) then
    if (peinf%inode.eq.0) write(0,994) xct%nkpt_co,kgq%nf
    call die("input_co_q: coarse grid mismatch")
  endif
994 format(1x,'The given number of points in the coarse grid ',i4,' does not match',/,&
      '   with the number of points in file WFNq_co ', i4,'.')

!-----------------------------------------------------------------------
!     Find correspondence with fk from WFN_co
!     indexq : correspondence between a k-point in the full BZ, kg%f,
!     and its shifted vector, kgq%f

  do ik=1,kg%nf
    ikq=0
    delta=0.1d0
    do while ((delta.gt.TOL_Small).and.(ikq.lt.kgq%nf))
      ikq=ikq+1
      qq(:) = kg%f(:,ik)-(kgq%f(:,ikq)-xct%shift(:))
      do kk=1,3
        qq(kk) = qq(kk) - anint( qq(kk) )
      enddo
      delta=sqrt((qq(1))**2+(qq(2))**2+(qq(3))**2)
    enddo
    if(delta.gt.TOL_Small) then
      if(peinf%inode.eq.0) write(0,'(a,3f10.6)') 'Could not find point equivalent to ', (kg%f(ii,ik),ii=1,3)
      call die("k-point mismatch between WFN_co and WFNq_co.", only_root_writes = .true.)
    else

! make sure that kgq%f(:,ikq) - kg%f(:,ik) = shift vector
! near the zone edge, they may differ by a lattice vector

      do jj=1,3
        ii = nint( kgq%f(jj,ikq)-kg%f(jj,ik) )
        kgq%f(jj,ikq) = kgq%f(jj,ikq) - dble(ii)
        kgq%kg0(jj,ikq) = kgq%kg0(jj,ikq) - ii
      enddo
      qq(:) = kg%f(:,ik)-(kgq%f(:,ikq)-xct%shift(:))
      delta=sqrt((qq(1))**2+(qq(2))**2+(qq(3))**2)
      if (delta.gt.TOL_Small) then
        call die("k-point mismatch between WFN_co and WFNq_co. Wrong shift.", only_root_writes = .true.)
      endif
      indexq_co(ik)=ikq
    endif
  enddo
  if (peinf%inode.eq.0) write(6,990) indexq_co(1)
990 format(1x,'Finished correspondence between shifted and unshifted coarse grids.',/,3x,'indexq_co(1) =',i4)

!-----------------------------------------------------------------------
! Read the wavefunctions and distribute or write to temp files

  if (xct%iwriteint.eq.0) then
    
    ! No other part of the code is currently trying to read INT_CWFN_CO_Q -- why not? --DAS
    write(filenamec,'(a)') 'INT_CWFN_CO_Q'
    iunit_c=123
    write(filenamev,'(a)') 'INT_VWFN_CO_Q'
    iunit_v=124
    
    if (peinf%inode.eq.0) then
      call open_file(iunit_c,file=filenamec,form='unformatted',status='replace')
      call open_file(iunit_v,file=filenamev,form='unformatted',status='replace')
    endif ! node 0
    
  endif ! xct%iwriteint.eq.0
  
  SAFE_ALLOCATE(wfnv%isort, (gvec_co%ng))
  wfnv%nband=xct%nvb_co
  wfnv%nspin=kp_co%nspin
  wfnv%nspinor=kp_co%nspinor
  wfnc%nband=xct%ncb_co
  wfnc%nspin=kp_co%nspin
  wfnc%nspinor=kp_co%nspinor

  do irk=1,kp_co%nrk
    irks = 0
    do ii=1,kg%nf
      if (kgq%indr(indexq_co(ii)) == irk) then
        irks=ii
        exit
      endif
    enddo

    SAFE_ALLOCATE(gvec_kpt%components, (3, kp_co%ngk(irk)))
    call read_binary_gvectors(26, kp_co%ngk(irk), kp_co%ngk(irk), gvec_kpt%components)

    SAFE_ALLOCATE(cg, (kp_co%ngk(irk),kp_co%nspin))
    if(irks > 0) then
      do ii = 1, kp_co%ngk(irk)
        call findvector(wfnv%isort(ii), gvec_kpt%components(:, ii), gvec)
        if (wfnv%isort(ii) == 0) call die('Could not find g-vector.')
      enddo
      
      wfnv%ng=kp_co%ngk(irk)
      wfnc%ng=kp_co%ngk(irk)
      if(peinf%inode == 0) then
        SAFE_ALLOCATE(wfnv%cg, (wfnv%ng,wfnv%nband,wfnv%nspin*wfnv%nspinor))
        SAFE_ALLOCATE(wfnc%cg, (wfnc%ng,wfnc%nband,wfnc%nspin*wfnv%nspinor))
        SAFE_ALLOCATE(cgarray, (kp_co%ngk(irk)))
      endif
    endif

! Loop over the bands
      
    do ii=1,kp_co%mnband

! Read planewave coefficients for band ii
      call read_binary_data(26, kp_co%ngk(irk), kp_co%ngk(irk), kp_co%nspin*kp_co%nspinor, cg)

      if(irks == 0) cycle
        
      if(peinf%inode == 0) then  
        do is=1, kp_co%nspin
          if (ii .gt. kp_co%ifmax(irk,is)-xct%nvb_co .and. ii .le. kp_co%ifmax(irk,is)+xct%ncb_co) then

            do isp=1, kp_co%nspinor
              do kk = 1, kp_co%ngk(irk)
                cgarray(kk)=cg(kk, is*isp)
              end do
            
              if ((ii.le.kp_co%ifmax(irk,is)).and. &
                (ii.gt.kp_co%ifmax(irk,is)-xct%nvb_co)) &
                wfnv%cg(1:wfnv%ng,kp_co%ifmax(irk,is)-ii+1,is*isp)=cgarray
            
              if ((ii.gt.kp_co%ifmax(irk,is)).and. &
                (ii.le.kp_co%ifmax(irk,is)+xct%ncb_co)) &
                wfnc%cg(1:wfnc%ng,ii-kp_co%ifmax(irk,is),is*isp)=cgarray
            enddo
            call checknorm('WFNq_co',ii,irks,kp_co%ngk(irk),is,kp_co%nspinor,cg(:,:))
          end if
        end do
      endif
        
    enddo ! ii (loop over bands)

    SAFE_DEALLOCATE(cg)
    if(peinf%inode == 0) then
      SAFE_DEALLOCATE(cgarray)
    endif
    
    if (xct%iwriteint.eq.1) then
      
#ifdef MPI
      if (peinf%inode.ne.0) then
        SAFE_ALLOCATE(wfnv%cg, (wfnv%ng,wfnv%nband,wfnv%nspin*wfnv%nspinor))
        SAFE_ALLOCATE(wfnc%cg, (wfnc%ng,wfnc%nband,wfnc%nspin*wfnc%nspinor))
      endif
      call MPI_BCAST(wfnv%cg(1,1,1),wfnv%ng*wfnv%nband*wfnv%nspin*wfnv%nspinor,MPI_SCALAR,0,MPI_COMM_WORLD,mpierr)
      call MPI_BCAST(wfnc%cg(1,1,1),wfnc%ng*wfnc%nband*wfnc%nspin*wfnc%nspinor,MPI_SCALAR,0,MPI_COMM_WORLD,mpierr)
#endif
      
      distgwfcoq%ng(irks)=wfnv%ng
      do ii=1,distgwfcoq%ngl
        if (ii+distgwfcoq%tgl.le.wfnv%ng) &
          distgwfcoq%isort(ii,irks)=wfnv%isort(ii+distgwfcoq%tgl)
      enddo
      do kk=1,distgwfcoq%ns*distgwfcoq%nspinor
        do jj=1,distgwfcoq%nv
          do ii=1,distgwfcoq%ngl
            if (ii+distgwfcoq%tgl.le.wfnv%ng) then
              distgwfcoq%zv(ii,jj,kk,irks)=wfnv%cg(ii+distgwfcoq%tgl,jj,kk)
            endif
          enddo
        enddo
      enddo
      do kk=1,distgwfcoq%ns*distgwfcoq%nspinor
        do jj=1,distgwfcoq%nc
          do ii=1,distgwfcoq%ngl
            if (ii+distgwfcoq%tgl.le.wfnv%ng) then
              distgwfcoq%zc(ii,jj,kk,irks)=wfnc%cg(ii+distgwfcoq%tgl,jj,kk)
            endif
          enddo
        enddo
      enddo
            
    endif ! xct%iwriteint.eq.1
    
    if (xct%iwriteint.eq.0) then
      
      if (peinf%inode.eq.0) then
        write(iunit_v) irks,wfnv%ng,wfnv%nband,wfnv%nspin,wfnv%nspinor
        write(iunit_v) (wfnv%isort(ii),ii=1,gvec_co%ng), &
          (((wfnv%cg(ii,jj,kk),ii=1,wfnv%ng),jj=1,wfnv%nband),kk=1,wfnv%nspin*wfnv%nspinor)
        write(iunit_c) irks,wfnc%ng,wfnc%nband,wfnc%nspin,wfnc%nspinor
        write(iunit_c) (wfnv%isort(ii),ii=1,gvec_co%ng), &
          (((wfnc%cg(ii,jj,kk),ii=1,wfnc%ng),jj=1,wfnc%nband),kk=1,wfnc%nspin*wfnc%nspinor)
      endif ! node 0
      
    endif ! xct%iwriteint.eq.0
    
    if (peinf%inode == 0 .or. xct%iwriteint == 1) then
      SAFE_DEALLOCATE_P(wfnv%cg)
      SAFE_DEALLOCATE_P(wfnc%cg)
    endif ! node 0
    
  enddo ! loop over k-points
  
  SAFE_DEALLOCATE_P(wfnv%isort)
  SAFE_DEALLOCATE(indxk)
  
  if (peinf%inode.eq.0) then
    write(6,301)
    write(6,302) kg%nr
    write(6,303) ((kg%r(ii,jj),ii=1,3),jj=1,kg%nr)
    write(6,304) kg%nf,kg%sz
301 format(/,1x,'Finished reading coarse grid q wavefunctions from unit WFNq_co',/)
302 format(6x,'nrk =',i4,/)
303 format(6x,3f10.6)
304 format(/,6x,'nfk =',i6,1x,'ksz =',f10.6,/)
    if (xct%iwriteint.eq.0) then
      call close_file(iunit_v)
      call close_file(iunit_c)
    endif ! xct%iwriteint.eq.0
    call close_file(26)
  endif ! node 0
  
  SAFE_DEALLOCATE_P(kp_co%rk)
  SAFE_DEALLOCATE_P(kp_co%ifmin)
  SAFE_DEALLOCATE_P(kp_co%ifmax)
  SAFE_DEALLOCATE_P(kp_co%el)
  
  if(xct%iwriteint == 0) then
#ifdef MPI
    call MPI_Barrier(MPI_COMM_WORLD, mpierr)
#endif
  endif

  POP_SUB(input_co_q)
  
  return
end subroutine input_co_q
