!==========================================================================
!
! Routines:
!
! (1) genwf_co()        Originally By MLT       Last Modified 4/19/2009 (gsm)
!
!     input: crys, gvec, kg, kgq,  syms, xct, distgwfco, distgwfcoq types
!            ik    label of k-point in FBZ
!            ikq   label of k-point in FBZ Q Shifted
!
!     output: wfnc   conduction wavefunctions at point k
!             wfnv   valence wavefunctions at point k
!
!==========================================================================

#include "f_defs.h"

subroutine genwf_co(crys,gvec,kg,kgq,syms,wfnc,wfnv,xct, &
  ik,ikq,distgwfco,distgwfcoq,workco,workcoq)

  use global_m
  use gmap_m
  use input_utils_m
  use misc_m
  use sort_m
  implicit none

  type (crystal), intent(in) :: crys
  type (gspace), intent(in) :: gvec
  type (grid), intent(in) :: kg
  type (grid), intent(in) :: kgq
  type (symmetry), intent(in) :: syms
  type (wavefunction), intent(out) :: wfnc
  type (wavefunction), intent(out) :: wfnv
  type (xctinfo), intent(in) :: xct
  integer, intent(in) :: ik
  integer, intent(in) :: ikq
  type (tdistgwf), intent(in) :: distgwfco
  type (tdistgwf), intent(in) :: distgwfcoq
  type (work_genwf), intent(inout) :: workco
  type (work_genwf), intent(inout) :: workcoq

  character :: filenamev*20,filenamec*20,errmsg*100
  integer :: irk,iunit_v,iunit_c,kval
  integer :: ig,hh,ii,jj,kk,eof,ipe

  integer, allocatable :: distint(:),distkval(:)
  integer, allocatable :: distik(:),distikold(:)
  integer, allocatable :: isorti(:),distisort(:)
  real(DP) :: qk(3)
  real(DP), allocatable :: ekin(:)
  SCALAR, allocatable :: distz(:,:,:)
  PUSH_SUB(genwf_co)

!-----------------------------------------------------------------------
! Initialize temp file names

  if (xct%iwriteint.eq.0) then
    if (xct%qflag .eq. 0) then
      write(filenamev,'(a)') 'INT_VWFN_CO_Q'
    else
      write(filenamev,'(a)') 'INT_VWFN_CO'
    endif
    write(filenamec,'(a)') 'INT_CWFN_CO'
    if (xct%qflag .eq. 0) then
      iunit_v=124
    else
      iunit_v=127
    endif
    iunit_c=126
  endif

!-----------------------------------------------------------------------
! Deal with the (possibly shifted) valence wavefunctions

  if (xct%qflag .eq. 0) then
    kval=kgq%indr(ikq)
  else
    kval=kg%indr(ikq)
  endif
    
  if (ikq.ne.workcoq%ikold) then
    
    if (xct%iwriteint.eq.1) then
      workcoq%ng=distgwfcoq%ng(kval)
      workcoq%nb=distgwfcoq%nv
      workcoq%ns=distgwfcoq%ns
      workcoq%nspinor=distgwfcoq%nspinor
    endif
    if (xct%iwriteint.eq.0) then
      call open_file(iunit_v,file=filenamev,form='unformatted',status='old')
      eof=0
      read(iunit_v) irk,workcoq%ng,workcoq%nb,workcoq%ns,workcoq%nspinor
      do while ((irk.ne.kval).and.(eof.eq.0))
        read(iunit_v)
        read(iunit_v,iostat=eof) irk,workcoq%ng,workcoq%nb,workcoq%ns,workcoq%nspinor
      enddo
      if (eof.ne.0) then
        write(errmsg,999) ikq, kval, trim(filenamev)
        call die(errmsg)
      endif
999   format('Could not find the valence wavefunctions for k-point',i4,&
        '(equivalent to k-point ',i4,' in the IBZ) in file ',a)
    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
    SAFE_ALLOCATE(workcoq%cg, (workcoq%ng,workcoq%nb,workcoq%ns*workcoq%nspinor))
    SAFE_ALLOCATE(workcoq%ind, (workcoq%ng))
    SAFE_ALLOCATE(workcoq%ph, (workcoq%ng))
    SAFE_ALLOCATE(workcoq%isort, (gvec%ng))
    
  endif ! ikq.ne.workcoq%ikold
  
  wfnv%ng=workcoq%ng
  wfnv%nband=workcoq%nb
  wfnv%nspin=workcoq%ns
  wfnv%nspinor=workcoq%nspinor
  if (workcoq%ns.ne.xct%nspin) then
    write(0,998) xct%nspin,workcoq%ns,trim(filenamev)
    write(errmsg,*) ' Spin mismatch in workcoq ', workcoq%ikold, ikq
    call die(errmsg)
  endif
998 format(1x,'The given number of spins',i2,1x,'does not match with', &
      1x,'the number of spins',i2,1x,'in file',1x,a,'.')
  
  SAFE_ALLOCATE(wfnv%cg, (wfnv%ng,wfnv%nband,wfnv%nspin*wfnv%nspinor))
  SAFE_ALLOCATE(wfnv%isort, (gvec%ng))
  
  if (xct%iwriteint.eq.1) then
    
! Share kval, ikq, ikoldq among processors

    SAFE_ALLOCATE(distint, (peinf%npes))
    SAFE_ALLOCATE(distkval, (peinf%npes))
    SAFE_ALLOCATE(distik, (peinf%npes))
    SAFE_ALLOCATE(distikold, (peinf%npes))
    
    distint(:)=0
    distint(peinf%inode+1)=kval
#ifdef MPI
    call MPI_Allreduce(distint,distkval,peinf%npes,MPI_INTEGER,MPI_SUM,MPI_COMM_WORLD,mpierr)
#else
    distkval(:)=distint(:)
#endif

    distint(:)=0
    distint(peinf%inode+1)=ikq
#ifdef MPI
    call MPI_Allreduce(distint,distik,peinf%npes,MPI_INTEGER,MPI_SUM,MPI_COMM_WORLD,mpierr)
#else
    distik(:)=distint(:)
#endif
    
    distint(:)=0
    distint(peinf%inode+1)=workcoq%ikold
#ifdef MPI
    call MPI_Allreduce(distint,distikold,peinf%npes,MPI_INTEGER,MPI_SUM,MPI_COMM_WORLD,mpierr)
#else
    distikold(:)=distint(:)
#endif

    SAFE_DEALLOCATE(distint)

! Gather distributed wavefunctions for rk-kpoint

    SAFE_ALLOCATE(distisort, (distgwfcoq%ngm))
    SAFE_ALLOCATE(distz, (distgwfcoq%ngm,distgwfcoq%nv,distgwfcoq%ns*distgwfcoq%nspinor))
    
    do ipe=1,peinf%npes
      if (distik(ipe).ne.distikold(ipe)) then
        distisort(:)=0
        distz(:,:,:)=ZERO
        do ii=1,distgwfcoq%ngl
          distisort(ii+distgwfcoq%tgl)=distgwfcoq%isort(ii,distkval(ipe))
        enddo
        do kk=1,distgwfcoq%ns*distgwfcoq%nspinor
          do jj=1,distgwfcoq%nv
            do ii=1,distgwfcoq%ngl
              distz(ii+distgwfcoq%tgl,jj,kk)=distgwfcoq%zv(ii,jj,kk,distkval(ipe))
            enddo
          enddo
        enddo
#ifdef MPI
        ii=distgwfcoq%ng(distkval(ipe))
        call MPI_Reduce(distisort,workcoq%isort,ii,MPI_INTEGER,MPI_SUM,ipe-1,MPI_COMM_WORLD,mpierr)
        do kk=1,workcoq%ns*workcoq%nspinor
          do jj=1,workcoq%nb
            call MPI_Reduce(distz(:,jj,kk),workcoq%cg(:,jj,kk),ii,MPI_SCALAR,MPI_SUM,ipe-1,MPI_COMM_WORLD,mpierr)
          enddo
        enddo
#else
        do ii=1,workcoq%ng
          workcoq%isort(ii)=distisort(ii)
        enddo
        do kk=1,workcoq%ns*workcoq%nspinor
          do jj=1,workcoq%nb
            do ii=1,workcoq%ng
              workcoq%cg(ii,jj,kk)=distz(ii,jj,kk)
            enddo
          enddo
        enddo
#endif
      endif ! distik.ne.distikold
    enddo ! ipe
    
    SAFE_DEALLOCATE(distisort)
    SAFE_DEALLOCATE(distz)
    
    SAFE_DEALLOCATE(distkval)
    SAFE_DEALLOCATE(distik)
    SAFE_DEALLOCATE(distikold)
    
  endif ! xct%iwriteint.eq.1

! Read wavefunctions for rk-kpoint from temp file
  if (xct%iwriteint.eq.0 .and. ikq.ne.workcoq%ikold) then
    read(iunit_v) (workcoq%isort(ii),ii=1,gvec%ng),(((workcoq%cg(ii,jj,kk), &
      ii=1,wfnv%ng),jj=1,wfnv%nband),kk=1,wfnv%nspin*wfnv%nspinor)
    call close_file(iunit_v)
  endif ! xct%iwriteint.eq.0
  
  if (ikq.ne.workcoq%ikold) then

! Compute inverse index array of Fourier components around rk-kpoint

    SAFE_ALLOCATE(isorti, (gvec%ng))
    isorti(:)=0
    do ii=1,wfnv%ng
      isorti(workcoq%isort(ii))=ii
    enddo
    
! Compute index array of Fourier components around fk-kpoint
    
    SAFE_ALLOCATE(ekin, (gvec%ng))
    if (xct%qflag .eq. 0) then
      qk(1:3)=kgq%f(1:3,ikq)
    else
      qk(1:3)=kg%f(1:3,ikq)
    endif
    call kinetic_energies(gvec, crys%bdot, ekin, qvec = qk)
    call sortrx(gvec%ng, ekin, workcoq%isort, gvec = gvec%components)
    SAFE_DEALLOCATE(ekin)

! Find ind and ph relating wavefunctions in fk to rk-kpoint

    workcoq%ind(:)=0
    workcoq%ph(:)=ZERO
    if (xct%qflag .eq. 0) then
      call gmap(gvec,syms,wfnv%ng,kgq%itran(ikq), &
        kgq%kg0(:,ikq),workcoq%isort,isorti,workcoq%ind,workcoq%ph,.true.)
    else
      call gmap(gvec,syms,wfnv%ng,kg%itran(ikq), &
        kg%kg0(:,ikq),workcoq%isort,isorti,workcoq%ind,workcoq%ph,.true.)
    endif
    SAFE_DEALLOCATE(isorti)

! Compute and renormalize valence wavefunctions

    do kk=1,wfnv%nspin
      do jj=1,wfnv%nband
        do hh=1,wfnv%nspinor
          do ii=1,wfnv%ng
            if (workcoq%ind(ii) .gt. 0) then
              wfnv%cg(ii,jj,kk*hh)=workcoq%ph(ii)*workcoq%cg(workcoq%ind(ii),jj,kk*hh)
            endif
          enddo
        enddo
        call checknorm('wfnv%cg',jj,ik,wfnv%ng,kk,wfnv%nspinor,wfnv%cg(1:wfnv%ng,jj,:))
      enddo
    enddo

    workcoq%cg(:,:,:)=wfnv%cg(:,:,:)
    
  endif ! ikq.ne.workcoq%ikold


  wfnv%cg(:,:,:)=workcoq%cg(:,:,:)
  wfnv%isort(:)=workcoq%isort(:)
  workcoq%ikold=ikq

!-----------------------------------------------------------------------
! Deal with the conduction wavefunctions

  kval=kg%indr(ik)
    
  if (ik.ne.workco%ikold) then
    
    if (xct%iwriteint.eq.1) then
      workco%ng=distgwfco%ng(kval)
      workco%nb=distgwfco%nc
      workco%ns=distgwfco%ns
      workco%nspinor=distgwfco%nspinor
    endif
    if (xct%iwriteint.eq.0) then
      call open_file(iunit_c,file=filenamec,form='unformatted',status='old')
      eof=0
      read(iunit_c) irk,workco%ng,workco%nb,workco%ns,workco%nspinor
      do while ((irk.ne.kval).and.(eof.eq.0))
        read(iunit_c)
        read(iunit_c,iostat=eof) irk,workco%ng,workco%nb,workco%ns,workco%nspinor
      enddo
      if (eof.ne.0) then
        write(errmsg,999) ik, kval, trim(filenamec)
        call die(errmsg)
      endif
    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
    SAFE_ALLOCATE(workco%cg, (workco%ng,workco%nb,workco%ns*workco%nspinor))
    SAFE_ALLOCATE(workco%ind, (workco%ng))
    SAFE_ALLOCATE(workco%ph, (workco%ng))
    SAFE_ALLOCATE(workco%isort, (gvec%ng))
    
  endif ! ik.ne.workco%ikold
  
  wfnc%ng=workco%ng
  wfnc%nband=workco%nb
  wfnc%nspin=workco%ns
  wfnc%nspinor=workco%nspinor
  if (workco%ns.ne.xct%nspin) then
    write(errmsg,998) xct%nspin,workco%ns,trim(filenamec)
    call die(errmsg)
  endif
  
  SAFE_ALLOCATE(wfnc%cg, (wfnc%ng,wfnc%nband,wfnc%nspin*wfnc%nspinor))
  SAFE_ALLOCATE(wfnc%isort, (gvec%ng))
  
  if (xct%iwriteint.eq.1) then

! Share kval, ik, workco%ikold among processors
    
    SAFE_ALLOCATE(distint, (peinf%npes))
    SAFE_ALLOCATE(distkval, (peinf%npes))
    SAFE_ALLOCATE(distik, (peinf%npes))
    SAFE_ALLOCATE(distikold, (peinf%npes))
    
    distint(:)=0
    distint(peinf%inode+1)=kval
#ifdef MPI
    call MPI_Allreduce(distint,distkval,peinf%npes,MPI_INTEGER,MPI_SUM,MPI_COMM_WORLD,mpierr)
#else
    distkval(:)=distint(:)
#endif
    distint(:)=0
    distint(peinf%inode+1)=ik
#ifdef MPI
    call MPI_Allreduce(distint,distik,peinf%npes,MPI_INTEGER,MPI_SUM,MPI_COMM_WORLD,mpierr)
#else
    distik(:)=distint(:)
#endif
    distint(:)=0
    distint(peinf%inode+1)=workco%ikold
#ifdef MPI
    call MPI_Allreduce(distint,distikold,peinf%npes,MPI_INTEGER,MPI_SUM,MPI_COMM_WORLD,mpierr)
#else
    distikold(:)=distint(:)
#endif
    
    SAFE_DEALLOCATE(distint)

! Gather distributed wavefunctions for rk-kpoint

    SAFE_ALLOCATE(distisort, (distgwfcoq%ngm))
    SAFE_ALLOCATE(distz, (distgwfcoq%ngm,distgwfcoq%nc,distgwfcoq%ns*distgwfcoq%nspinor))

    do ipe=1,peinf%npes
      if (distik(ipe).ne.distikold(ipe)) then
        distisort(:)=0
        distz(:,:,:)=ZERO
        do ii=1,distgwfco%ngl
          distisort(ii+distgwfco%tgl)=distgwfco%isort(ii,distkval(ipe))
        enddo
        do kk=1,distgwfco%ns*distgwfco%nspinor
          do jj=1,distgwfco%nc
            do ii=1,distgwfco%ngl
              distz(ii+distgwfco%tgl,jj,kk)=distgwfco%zc(ii,jj,kk,distkval(ipe))
            enddo
          enddo
        enddo
#ifdef MPI
        ii=distgwfco%ng(distkval(ipe))
        call MPI_Reduce(distisort,workco%isort,ii,MPI_INTEGER,MPI_SUM,ipe-1,MPI_COMM_WORLD,mpierr)
        do kk=1,workco%ns*workco%nspinor
          do jj=1,workco%nb
            call MPI_Reduce(distz(:,jj,kk),workco%cg(:,jj,kk),ii,MPI_SCALAR,MPI_SUM,ipe-1,MPI_COMM_WORLD,mpierr)
          enddo
        enddo
#else
        do ii=1,workco%ng
          workco%isort(ii)=distisort(ii)
        enddo
        do kk=1,workco%ns*workco%nspinor
          do jj=1,workco%nb
            do ii=1,workco%ng
              workco%cg(ii,jj,kk)=distz(ii,jj,kk)
            enddo
          enddo
        enddo
#endif
      endif ! distik.ne.distikold
    enddo ! ipe
    
    SAFE_DEALLOCATE(distisort)
    SAFE_DEALLOCATE(distz)
    
    SAFE_DEALLOCATE(distkval)
    SAFE_DEALLOCATE(distik)
    SAFE_DEALLOCATE(distikold)
    
  endif ! xct%iwriteint.eq.1
  
! Read wavefunctions for rk-kpoint from temp file
  if (xct%iwriteint.eq.0 .and. ik.ne.workco%ikold) then
    read(iunit_c) (workco%isort(ii),ii=1,gvec%ng),(((workco%cg(ii,jj,kk), &
      ii=1,wfnc%ng),jj=1,wfnc%nband),kk=1,wfnc%nspin*wfnc%nspinor)
    call close_file(iunit_c)
  endif ! xct%iwriteint.eq.0
  
  if (ik.ne.workco%ikold) then

! JRD: Below is now necessary because kg might be different from kgq

! Compute inverse index array of Fourier components around rk-kpoint

    SAFE_ALLOCATE(isorti, (gvec%ng))
    isorti(:)=0
    do ii=1,wfnc%ng
      isorti(workco%isort(ii))=ii
    enddo

! Compute index array of Fourier components around fk-kpoint

    SAFE_ALLOCATE(ekin, (gvec%ng))
    call kinetic_energies(gvec, crys%bdot, ekin, qvec = kg%f(:, ik))
    call sortrx(gvec%ng, ekin, workco%isort, gvec = gvec%components)
    SAFE_DEALLOCATE(ekin)
    
! Find ind and ph relating wavefunctions in fk to rk-kpoint

    workco%ind(:)=0
    workco%ph(:)=ZERO
    call gmap(gvec,syms,wfnc%ng,kg%itran(ik), &
      kg%kg0(:,ik),workco%isort,isorti,workco%ind,workco%ph,.true.)
    SAFE_DEALLOCATE(isorti)

! Compute and renormalize conduction wavefunctions

    do kk=1,wfnc%nspin
      do jj=1,wfnc%nband
        do hh=1,wfnc%nspinor
          do ii=1,wfnc%ng
            if (workco%ind(ii) .gt. 0) then
              wfnc%cg(ii,jj,kk*hh)=workco%ph(ii)*workco%cg(workco%ind(ii),jj,kk*hh)
            endif
          enddo
        enddo
        call checknorm('wfnc%cg',jj,ik,wfnc%ng,kk,wfnc%nspinor,wfnc%cg(1:wfnc%ng,jj,:))
      enddo
    enddo

    workco%cg(:,:,:)=wfnc%cg(:,:,:)
    
  endif ! ik.ne.workco%ikold
  

  wfnc%cg(:,:,:)=workco%cg(:,:,:)
  wfnc%isort(:)=workco%isort(:)
  workco%ikold=ik

  POP_SUB(genwf_co)
  
  return
end subroutine genwf_co
