#include "f_defs.h"

!-----------------------------------------------------------------------
subroutine input_kernel(crys,gvec,kg,kp,syms,xct,flagbz, &
  intwfnv,intwfnc)
!-----------------------------------------------------------------------
!
!     Read parameters from file WFN_co
!     Initialize k-points sampling, kg type
!     Initialize G-space, gvec
!
!     input: xct type
!
!     output: crys,gvec,syms,kg types
!             INT_VWFN_* and INT_CWFN_* files
!
  use global_m
  use checkbz_m
  use fullbz_m
  use input_utils_m
  use misc_m
  use sort_m
  use wfn_rho_vxc_io_m
  use read_rho_vxc_m
  implicit none

  type (crystal), intent(out) :: crys
  type (gspace), intent(out) :: gvec
  type (grid), intent(out) :: kg
  type (kpoints), intent(out) :: kp
  type (symmetry), intent(out) :: syms
  type (xctinfo), intent(inout) :: xct
  integer, intent(in) :: flagbz
  type (int_wavefunction), intent(out) :: intwfnc,intwfnv


  type (wavefunction) :: wfnc,wfnv
  character :: tmpfn*16
  integer :: iwritev,iwritec,iwritek
  integer :: ickmem,irk
  integer :: ii,jj,kk,is
  integer :: irks,ivband,icband

  real(DP) :: diffvol,vcell,raux,kt(3),div,tol
  
  real(DP), allocatable :: ek_tmp(:)
  integer, allocatable :: index(:),indxk(:),k_tmp(:,:)
  integer, allocatable :: isrti(:)
  SCALAR, allocatable :: cg(:,:), cgarray(:)

  character(len=3) :: sheader
  integer :: iflavor
  type(gspace) :: gvec_kpt
  logical :: skip_checkbz

  PUSH_SUB(input_kernel)

#ifdef VERBOSE
  if (peinf%inode .eq. 0) write(6,*) ' '
#endif

  call logit('input_kernel:  reading WFN_co')

  if (peinf%inode == 0) call open_file(25,file='WFN_co',form='unformatted',status='old')

  sheader = 'WFN'
  iflavor = 0
  call read_binary_header_type(25, sheader, iflavor, kp, gvec, syms, crys, dont_warn_kgrid=xct%patched_sampling_co)
  call check_trunc_kpts(xct%icutv, kp)

  if(any(kp%shift > TOL_Zero) .and. peinf%inode == 0) then
    write(0,'(a)') "WARNING: WFN_co has a shift. This is not recommended."
  endif

  call logit('input_kernel:  reading gvec info')
  SAFE_ALLOCATE(gvec%components, (3, gvec%ng))
  call read_binary_gvectors(25, gvec%ng, gvec%ng, gvec%components)

  call get_volume(vcell,crys%bdot)
  diffvol=abs(crys%celvol-vcell)
  if (diffvol.gt.TOL_Small) then
    call die('volume mismatch', only_root_writes = .true.)
  endif

  ! there is no fine grid to set the Fermi level in kernel
  call find_efermi(xct%rfermi, xct%efermi, xct%efermi_input, kp, kp%mnband, 1, &
    "coarse grid", should_search = .true., should_update = .true., write7 = .false.)

  call assess_degeneracies(kp, kp%el(kp%mnband, :, :), kp%mnband - 1, xct%efermi, TOL_Degeneracy)

  xct%nspin=kp%nspin

  if(any(kp%ifmax(:,:) == 0)) &
    call die("BSE codes cannot handle a system where some k-points have no occupied bands.", only_root_writes = .true.)

  kp%nvband=minval(kp%ifmax(:,:)-kp%ifmin(:,:))+1
  kp%ncband=kp%mnband-maxval(kp%ifmax(:,:))

!----------------------------------------------------------------
! (gsm) check whether the requested number of bands
!       is available in the wavefunction file

  if(xct%nvb_co.gt.kp%nvband) then
    call die("The requested number of valence bands is not available in WFN_co.", &
             only_root_writes=.true.)
  endif
  if(xct%ncb_co.gt.kp%ncband) then
    call die("The requested number of conduction bands is not available in WFN_co.", &
             only_root_writes=.true.)
  endif

! DAS: degenerate subspace check

  if (peinf%inode.eq.0) then
    if(xct%ncb_co.eq.kp%ncband) then
      call die("You must provide one more conduction band in WFN_co in order to assess degeneracy.", &
               only_root_writes=.true.)
    endif
    do jj = 1, kp%nspin
      do ii = 1, kp%nrk
        if(kp%ifmax(ii, jj) - xct%nvb_co > 0) then
          ! no need to compare against band 0 if all valence are included
          if(abs(kp%el(kp%ifmax(ii, jj) - xct%nvb_co + 1, ii, jj) &
            - kp%el(kp%ifmax(ii, jj) - xct%nvb_co, ii, jj)) .lt. TOL_Degeneracy) then
            if(xct%degeneracy_check_override) then
              write(0,'(a)') &
                "WARNING: Selected number of valence bands breaks degenerate subspace in WFN_co. " // &
                "Run degeneracy_check.x for allowable numbers."
              write(0,*)
            else
              write(0,'(a)') &
                "Run degeneracy_check.x for allowable numbers, or use keyword " // &
                "degeneracy_check_override to run anyway (at your peril!)."
              call die("Selected number of valence bands breaks degenerate subspace in WFN_co.", &
                       only_root_writes=.true.)
            endif
          endif
        endif
        if(abs(kp%el(kp%ifmax(ii, jj) + xct%ncb_co, ii, jj) &
          - kp%el(kp%ifmax(ii, jj) + xct%ncb_co + 1, ii, jj)) .lt. TOL_Degeneracy) then 
          if(xct%degeneracy_check_override) then
            write(0,'(a)') &
              "WARNING: Selected number of conduction bands breaks degenerate subspace in WFN_co. " // &
              "Run degeneracy_check.x for allowable numbers."
            write(0,*)
          else
            write(0,'(a)') &
              "Run degeneracy_check.x for allowable numbers, or use keyword " // &
              "degeneracy_check_override to run anyway (at your peril!)."
            call die("Selected number of conduction bands breaks degenerate subspace in WFN_co.",&
                     only_root_writes=.true.)
          endif
        endif
      enddo
    enddo
  endif

!-----------------------------------------------------------------------
!     Read the k-point sampling from kpoints (if it exists) or from
!     WFN_co

  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

!----------------------------------------------------------------
!     indxk : stores the correspondence between k-points kg%r and kp%rk
!     (it is used to select the set of wavefunctions to be stored)
!     tol : tolerance in the coordinates of k-points

    tol = 1.d-4
    SAFE_ALLOCATE(indxk, (kg%nr))
    indxk=0
    do jj=1,kg%nr
      do ii=1,kp%nrk
        kt(:) = kg%r(:,jj) - kp%rk(:,ii)
        if ((abs(kt(1)).lt.tol).and.(abs(kt(2)).lt.tol) &
          .and.(abs(kt(3)).lt.tol)) then
          if (indxk(jj).ne.0) write(0,*) 'WARNING: multiple definition of k-point',jj,indxk(jj),kg%r(:,jj)
          indxk(jj)=ii
        endif
      enddo
      if (indxk(jj).eq.0) write(0,*) 'WARNING: could not find vector ',kg%r(:,jj),' in WFN_co'
!
!     no need to stop here; if indxk.eq.0, the job will stop in genwf
!
    enddo
  else
    kg%nr=kp%nrk
    SAFE_ALLOCATE(kg%r, (3,kg%nr))
    kg%r(1:3,1:kg%nr)=kp%rk(1:3,1:kp%nrk)
    SAFE_ALLOCATE(indxk, (kg%nr))
    do ii=1,kg%nr
      indxk(ii) = ii
    enddo
  endif

!-----------------------------------------------------------------------
!     Order g-vectors with respect to their kinetic energy
!
!     xct%ecute = energy cutoff used in dielectric matrix
!     xct%ecutg = energy cutoff used in wavefunctions and interaction
!                   kernel, see Rohlfing & Louie, PRB 62(8),p. 4938
!                   (must be slightly longer than xct%ecute because
!                   of umklapp vectors)
!
  call logit('input_kernel:  reordering gvecs')
  raux=0.d0
  do ii=1,3
    if (crys%bdot(ii,ii).gt.raux) raux=crys%bdot(ii,ii)
  enddo
  if (abs(xct%ecutg).lt.TOL_Zero) xct%ecutg = xct%ecute + raux
  
  SAFE_ALLOCATE(index, (gvec%ng))
  SAFE_ALLOCATE(gvec%ekin, (gvec%ng))
  call kinetic_energies(gvec, crys%bdot, gvec%ekin)
  call sortrx(gvec%ng, gvec%ekin, index, gvec = gvec%components)


  xct%ng   = gcutoff(gvec%ng, gvec%ekin, index, xct%ecutg)
  if (xct%theory.eq.0) &
  xct%neps = gcutoff(gvec%ng, gvec%ekin, index, xct%ecute) ! cuteness energy??

  if(peinf%inode == 0) then
    if (xct%theory.eq.0) &
    write(6,'(a,f12.6,a,i8,a)') 'Screened coulomb cutoff = ', xct%ecute, ' Ry :', xct%neps, ' G-vectors'
    write(6,'(a,f12.6,a,i8,a)') 'Bare coulomb cutoff     = ', xct%ecutg, ' Ry :',   xct%ng, ' G-vectors'
  endif
  if(xct%ecute > xct%ecutg .and. xct%theory.eq.0) then
    call die("The screened_coulomb_cutoff cannot be greater than the bare_coulomb_cutoff.", only_root_writes = .true.)
  endif

  SAFE_ALLOCATE(ek_tmp, (gvec%ng))
  ek_tmp = gvec%ekin
  SAFE_ALLOCATE(k_tmp, (3,gvec%ng))
  k_tmp = gvec%components
  do ii=1,gvec%ng
    gvec%ekin(ii) = ek_tmp(index(ii))
    gvec%components(:,ii) = k_tmp(:,index(ii))
  enddo

  call gvec_index(gvec)

  ! If we are not doing just purely TDHF then
  ! read the charge density/fxc
  if ((xct%theory == 1) .and. ((1.0d0 - xct%coulomb_mod%long_range_frac_fock > TOL_SMALL) .or. &
     (1.0d0 - xct%coulomb_mod%short_range_frac_fock > TOL_SMALL))) then
    xct%coul_mod_flag=.true.
    SAFE_ALLOCATE(isrti, (gvec%ng))
    do ii=1,gvec%ng
      isrti(index(ii)) = ii
    enddo
    call read_rho(xct%wpg, gvec, kp, syms, crys, isrti, index, 'WFN_co')
    SAFE_DEALLOCATE(isrti)
  endif

  SAFE_DEALLOCATE(index)
  SAFE_DEALLOCATE(ek_tmp)
  SAFE_DEALLOCATE(k_tmp)

  
!-----------------------------------------------------------------------
!     Generate full brillouin zone from irreducible wedge, rk -> fk
!
!     If flagbz.eq.1, only Identity will be used as
!     symmetry operation. In this case, kg%r (irreducible BZ) and kg%f
!     (full BZ) will be identical.
!
  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 sampling',/)
802 format(1x,'No symmetries used in the coarse grid sampling',/)
!
  call timacc(7,1)
  if (flagbz.eq.1) then
    call fullbz(crys,syms,kg,1,skip_checkbz,wigner_seitz=.true.,paranoid=.true.)
  else
    call fullbz(crys,syms,kg,syms%ntran,skip_checkbz,wigner_seitz=.true.,paranoid=.true.)
  endif
  tmpfn='WFN_co'
  if (.not. skip_checkbz) then
    call checkbz(kg%nf,kg%f,kp%kgrid,kp%shift,crys%bdot, &
      tmpfn,'k',.true.,xct%freplacebz,xct%fwritebz)
  endif
  call timacc(7,2)
  xct%nkpt_co=kg%nf

#ifdef VERBOSE
  if (peinf%inode.eq.0) then
    write(6,*) 'xct%nkpt_co = ',xct%nkpt_co
    write(6,'(a5,3a9,2a7,5x,3a5)') &
      'i',' ','k',' ','indr','itran','kg0'
    do ii=1,kg%nf
      write(6,'(i5,3f9.4,2i7,5x,3i5)') &
        ii,kg%f(:,ii),kg%indr(ii),kg%itran(ii),kg%kg0(:,ii)
    enddo
  endif
#endif

!------------------------------------------------
! Distribute vcks-quadruplets among the PEs

  call logit('input_kernel:  calling distrib_kernel')
  call distrib_kernel(xct,kp%ngkmax,kg,kg,gvec)

#ifdef VERBOSE
  if (peinf%inode .eq. 0) then
    write(6,*)
    write(6,*) "Allocating isort", peinf%iownwfv(peinf%inode+1), peinf%iownwfc(peinf%inode+1), gvec%ng
  endif
#endif
  SAFE_ALLOCATE(intwfnv%cg, (kp%ngkmax,peinf%iownwfv(peinf%inode+1),kp%nspin))
  SAFE_ALLOCATE(intwfnc%cg, (kp%ngkmax,peinf%iownwfc(peinf%inode+1),kp%nspin))
  SAFE_ALLOCATE(intwfnv%isort, (gvec%ng,peinf%iownwfk(peinf%inode+1)))
  SAFE_ALLOCATE(intwfnc%isort, (gvec%ng,peinf%iownwfk(peinf%inode+1)))
  SAFE_ALLOCATE(intwfnv%ng, (peinf%iownwfk(peinf%inode+1)))
  SAFE_ALLOCATE(intwfnc%ng, (peinf%iownwfk(peinf%inode+1)))
  intwfnv%nspin=kp%nspin
  intwfnc%nspin=kp%nspin

!------------------------------------------------
! Begin Loop

  SAFE_ALLOCATE(wfnv%isort, (gvec%ng))
  SAFE_ALLOCATE(wfnc%isort, (gvec%ng))
  wfnv%nspin=kp%nspin
  wfnc%nspin=kp%nspin

  do irk=1,kp%nrk
    irks = 0
    do ii=1,kg%nr
      if (irk == indxk(ii)) then
        irks=ii
        exit
      endif
    enddo

    SAFE_ALLOCATE(gvec_kpt%components, (3, kp%ngk(irk)))
    call read_binary_gvectors(25, kp%ngk(irk), kp%ngk(irk), gvec_kpt%components)

    SAFE_ALLOCATE(cg, (kp%ngk(irk),kp%nspin))
    if(irks > 0) then
      do ii = 1, kp%ngk(irk)
        call findvector(wfnv%isort(ii), gvec_kpt%components(:, ii), gvec)
        if(wfnv%isort(ii) == 0) call die('could not find gvec', only_root_writes=.true.)
      enddo

      wfnv%ng=kp%ngk(irk)
      wfnc%ng=kp%ngk(irk)
      SAFE_ALLOCATE(cgarray, (kp%ngk(irk)))
      wfnc%isort(1:gvec%ng)=wfnv%isort(1:gvec%ng)
      ickmem=0
    endif

!        write(6,*) peinf%inode, 'loop wfnup', irk

! Loop Over Bands

    do ii=1,kp%mnband
          
      call read_binary_data(25, kp%ngk(irk), kp%ngk(irk), kp%nspin, cg)

! If we do not need this band, skip it...
      if(irks == 0) cycle
!
      do is=1, kp%nspin
        if (ii .ge. (kp%ifmax(irk,is)-xct%nvb_co+1) .and. ii .le. (kp%ifmax(irk,is)+xct%ncb_co)) then
! Check normalization of this band
          if (peinf%inode.eq.0) then
            do kk=1, kp%ngk(irk)
              cgarray(kk)=cg(kk, is)
            end do
            call checknorm('WFN_co',ii,irks,kp%ngk(irk),kp%nspin,kp%nspinor,cg(:,:))
          end if
!          write(6,*) peinf%inode, 'Read', irk, ii
              
!-----------------------
! If ii is one of the selected valence band...
          if((ii.le.kp%ifmax(irk,is)).and. &
            (ii.ge.(kp%ifmax(irk,is)-xct%nvb_co+1))) then
                
!            write(6,*) peinf%inode, 'in val',irk,ii
                
                
#ifdef MPI
            call MPI_BCAST(cgarray,kp%ngk(irk),MPI_SCALAR,0, &
              MPI_COMM_WORLD,mpierr)
#endif
                
            ivband=kp%ifmax(irk,is)-ii+1
            iwritev=peinf%ipev(peinf%inode+1,ivband,irk)
            iwritek=peinf%ipek(peinf%inode+1,irk)
!            write(6,*) peinf%inode, 'bcast val',irk,ii,ivband,
!     >       iwritev
            if(iwritev.ne.0) then
              intwfnv%cg(1:wfnv%ng,iwritev,is)=cgarray(1:kp%ngk(irk))
            endif
            if(iwritek.ne.0) then
              intwfnv%isort(1:gvec%ng,iwritek)=wfnv%isort(1:gvec%ng)
              intwfnv%ng(iwritek)=kp%ngk(irk)
            endif

!            write(*,'(a, 4i4)') 'valence', peinf%inode+1, ivband, irk, iwritev
!            write(*, '(a, 6i5)') "valence", ivband, ii, is, peinf%inode, irk, iwritev
!            write(6,*) peinf%inode, 'wrote val',irk,ii

          endif !ii is one of the selected valence band

!-----------------------
! If ii is one of the selected conduction band...
              
          if((ii.ge.(kp%ifmax(irk,is)+1)) &
            .and.(ii.le.(kp%ifmax(irk,is)+xct%ncb_co))) then
                
#ifdef MPI
            call MPI_BCAST(cgarray,kp%ngk(irk),MPI_SCALAR,0, &
              MPI_COMM_WORLD,mpierr)
#endif
                
            icband=ii-kp%ifmax(irk,is)
            iwritec=peinf%ipec(peinf%inode+1,icband,irk)
            iwritek=peinf%ipek(peinf%inode+1,irk)
!            write(6,*) 'icband',ii,icband,iwritec
            if(iwritec.ne.0) then
              intwfnc%cg(1:wfnc%ng,iwritec,is)=cgarray(1:kp%ngk(irk))
            endif
            if(iwritek.ne.0) then
              intwfnc%isort(1:gvec%ng,iwritek)=wfnc%isort(1:gvec%ng)
              intwfnc%ng(iwritek)=kp%ngk(irk)
            endif

!            write(*,'(a, 4i4)') 'conduction', peinf%inode+1, icband, irk, iwritec
!            write(*, '(a, 6i5)') "conduction", icband, ii, is, peinf%inode, irk, iwritec

          endif ! ii is one of the selected conduction bands
        end if
      end do
    enddo ! ii (loop on bands)

    SAFE_DEALLOCATE(cg)
    SAFE_DEALLOCATE(cgarray)
    
  enddo !end loop over k-points

  SAFE_DEALLOCATE_P(wfnv%isort)
  SAFE_DEALLOCATE_P(wfnc%isort)
  SAFE_DEALLOCATE(indxk)
      
! Write out info about xtal

  if(peinf%inode.eq.0) then
    write(6,'(1x,a)') 'Crystal wavefunctions read from unit WFN_co'
    write(6,'(/,1x,a,i0)') 'nrk = ', kg%nr
    write(6,'(/,1x,a,i0,2x,a,f10.6)') 'nfk = ', kg%nf, 'ksz = ', kg%sz
#ifdef VERBOSE
    write(6,*)
    write(6,'(1x,3f10.6)') kg%r(1:3,1:kg%nr)
    write(6,*)
#endif
    call close_file(25)
  endif ! node 0

  if(xct%iwriteint == 0) then
#ifdef MPI
    call MPI_Barrier(MPI_COMM_WORLD, mpierr)
#endif
  endif

  POP_SUB(input_kernel)

  return
end subroutine input_kernel
