module energy_bands

    use constants, only : dp
    use pw_types, only : wave_function

    implicit none

    complex(DP), target, allocatable, dimension(:,:) :: evc
    real(DP),            allocatable, dimension(:)   :: et
    type(wave_function), allocatable, dimension(:)   :: psi
    integer,             parameter                   :: unit_wvfct=2
    integer,             parameter                   :: unit_wvfct_aux=3

    integer                       :: unit_psi_in  = unit_wvfct
    integer                       :: unit_psi_out = unit_wvfct_aux

    logical, save                 :: first = .true.

contains

  subroutine c_bands(kvect,nbnd,isolve)

    use constants,              only: ry_to_ev
    use lapack_wrap,            only: cdiagh
    use pw_types,               only: k_point, construct, destroy, &
                                      average_field
    use scf,                    only: Vscf, V0
    use hamiltonian,            only: fill_h
    use iterative,              only: cegter

    integer, intent(in)       :: nbnd
    complex(DP), allocatable  :: hc(:,:)
    integer                   :: i, npw, nks, swap, isolve

    integer                  :: iter, nbndx
    real(DP)                 :: iterations, avg_iter
    type(k_point), dimension(:)   :: kvect
    real(DP)                 :: x
    logical, parameter       :: reorder = .true.

    rewind(unit_psi_in) 
    rewind(unit_psi_out) 
    
    nks = size(kvect(:))
    for_each_k_point: do i=1, nks
       npw = kvect(i)%npw
       write(*,fmt="(' k = (',f6.3,',',f6.3,',',f6.3,'), npw = ',i5)")  &
            kvect(i)%k, kvect(i)%npw
       
       if(first.or.isolve.eq.0) then
          call fetch_psi(kvect(i), nbands=nbnd)
          allocate (hc(npw,npw))
          call fill_h(hc,kvect(i))
          call cdiagh(hc,et,evc)
          deallocate(hc)
       else
          V0 = average_field(Vscf)
          call fetch_psi(kvect(i), unit=unit_psi_in)
          call cegter ( psi, et, &
               maxter=20, maxdim = 4*nbnd, threshold = 1.0d-9 )
       end if

       write(*,fmt="(8f9.5)") et*Ry_to_eV
       
       call bury_psi(unit_psi_out,kvect(i))
       
    end do for_each_k_point
    
    first        = .false.              
    swap         = unit_psi_in
    unit_psi_in  = unit_psi_out
    unit_psi_out = swap
    
  end subroutine c_bands

  subroutine sum_bands(kvect,nbnd,nelec)

    use constants, only: zero, one, degspin, eps
    use pw_types,  only: field, available, direct_space, assignment(=), &
                         average_field, weight, construct, destroy, axpy, &
                         scal, set_field_state, field_square, k_point
    use reciprocal_lattice, only: g, ngm
    use structure,          only: omega
    use scf,                only: eband, rho
    use utilities, only : error

    integer, intent(in) :: nbnd, nelec

    integer             :: i, n, m, npw, nks
    type(k_point), dimension(:)   :: kvect
    type(field)         :: psi_field
    real(DP)            :: total_charge, echarge

    write (*,'(a)') ' calculating the charge density'

    if (mod(nelec,2)/=0) call error ('sum_bands', &
         'number of electrons must be even in a band isulator',1)
    if (nelec/2 > nbnd) call error ('sum_bands', &
         'number of bands is too small to contain all the electrons', nelec/2)

    rewind(unit_psi_in)
    
    if(.not.available(rho)) then
       call construct(rho,DIRECT_SPACE)
    else
       call set_field_state(rho,DIRECT_SPACE)
       rho = ZERO
    end if

    call construct(psi_field,DIRECT_SPACE)

    eband = ZERO
    
    nks = size(kvect(:))
    do i=1, nks

       call fetch_psi(kvect(i),unit=unit_psi_in)

       do n=1, nelec/2
          psi_field = psi(n)
          call field_square(psi_field) ! psi_field = |psi(n)|^2 (real space)
          ! rho <- rho + 2*w(k) * |psi(n)|^2
          call axpy(DEGSPIN * weight(kvect(i)), psi_field,rho)
          ! eband <- eband + 2*w(k) * e(n)
          eband = eband + et(n) * DEGSPIN * weight(kvect(i))
       end do

       call kill_psi
    end do

    call destroy(psi_field)

    call scal(ONE/omega, rho)

    total_charge = omega*average_field(rho)

    write(6,'(a,f12.8)') '   total charge =',total_charge
    if  (abs(total_charge-nelec).gt.eps) call error('sum_bands', &
         'total charge is not equal to the number of electrons',1)

    return
  end subroutine sum_bands
  
  subroutine fetch_psi(kvect, unit, nbands)

    use constants, only : zero
    use pw_types, only : k_point, construct
    use utilities, only : error

    type(k_point), intent(in)     :: kvect
    integer, intent(in), optional :: nbands, unit
    real(DP)                      :: k(3)
    integer                       :: npw, nbnd

    if(present(unit)) then
       read(unit) k(:), nbnd, npw
       if ( any( k(:) /= kvect%k(:) ) ) &
            call error("fetch_psi","k-vector mismatch",1)
       if ( npw /= kvect%npw ) &
            call error("fetch_psi","size mismatch",2)
       if( present(nbands) ) then
          if( nbnd /= nbands ) call error("fetch_psi","size mismatch",3)
       end if
    else
       if(.not.present(nbands)) &
            call error("fetch_psi","too few arguments",4)
          npw  = kvect%npw
          nbnd = nbands
    end if
       
    allocate (evc(npw,nbnd),et(nbnd),psi(nbnd))
    et(:) = ZERO

    if(present(unit)) read(unit) evc, et

    call construct(psi(1:nbnd),kvect,evc(:,1:nbnd))

    return
  end subroutine fetch_psi

  subroutine bury_psi(unit, kvect)

    use pw_types, only : k_point
    type(k_point), intent(in) :: kvect
    integer, intent(in)       :: unit

    write(unit) kvect%k, size(evc(1,:)), size(evc(:,1))
    write(unit) evc, et
     
    call kill_psi
    return
  end subroutine bury_psi

  subroutine kill_psi
    use pw_types, only : destroy

    if(allocated(psi)) then
       call destroy(psi)
       deallocate(psi)
    end if
    deallocate (evc,et)
    return
  end subroutine kill_psi

end module energy_bands
