module iterative
  implicit none
contains
  
  subroutine cegter ( psi, e, maxter, maxdim, threshold )

    !     iterative solution of the eigenvalue problem:
    !
    !     ( h - e ) * psi = 0
    !
    !     where h is a complex hermitean matrix, e is a real scalar,
    !     and psi a complex vector

    use constants, only : dp, zero, one
    use pw_types, only: wave_function, clonate, axpy, scal, construct, &
                        destroy, operator(.dot.), assignment(=), psi_link
    use lapack_wrap, only: cdiagh
    use hamiltonian, only: h_psi, g_psi
    use utilities, only : error

    type(wave_function), intent(inout), dimension(:) :: psi 
    real(DP),            intent(out),   dimension(:) :: e
    integer,             intent(in),        optional :: maxter, maxdim
    real(DP),            intent(in),        optional :: threshold
    type(wave_function),     pointer,   dimension(:) :: phi, hphi ! temp
    integer,             allocatable,   dimension(:) :: pivot     ! temp

    logical,     allocatable  :: conv(:) ! flag: .true. when e(n) is convgd
    complex(DP), allocatable  :: hp(:,:) ! projection of the hamiltonian
                                         ! onto the current basis set
    complex(DP), allocatable  :: vc(:,:) ! eigenvectors of hp (column-wise)
    real(DP),    allocatable  :: ew(:)   ! temp array for eigenvalues

    integer                   :: n,m, np, neigv, iter, itermax, nbase, &
                                 pbase, basemax, notcnv, info
    real(DP)                  :: thresh
    type(wave_function), pointer :: swap

    ! set up input variables

    neigv = size(e)
    if(present(maxter)) then
       itermax = maxter
    else
       itermax = 50
    end if
    
    if(present(maxdim)) then
       basemax = maxdim
    else
       basemax = 4*neigv
    end if
    
    if(present(threshold)) then
       thresh = threshold
    else
       thresh = 1.0d-7
    end if
    
    ! prepare work arrays

    allocate ( phi(basemax),        &
               hphi(basemax),       &
               hp(basemax,basemax), & 
               vc(basemax,basemax), &
               ew(basemax),         &
               conv(neigv),         &
               pivot(neigv) )
    
    call psi_link(phi(1:neigv),psi(1:neigv))
    call clonate(psi(1), phi(neigv+1:basemax))
    call clonate(psi(1), hphi(1:basemax))

    swap => hphi(basemax)   ! temporary swap wave-function needed when 
                            ! hphi(basemax) is not used

    ! orthonormalize the trial wave-functions
    
    call gram ( phi(1:neigv) )
    
    ! prepare the reduced hamiltonian for the first iteration         
    ! hp contains the projection of the hamiltonian onto the reduced space
    ! vc contains the eigenvectors of hp

    vc = cmplx(ZERO,ZERO)
    do n = 1, neigv
       call h_psi(phi(n),hphi(n))
       do m = n, neigv
          hp(m,n) = phi(m).dot.hphi(n)
          hp(n,m) = conjg(hp(m,n))
       enddo
       e(n)    = hp(n,n)
       vc(n,n) = cmplx(ONE,ZERO)
    enddo

    conv(:) = .false.
    nbase   = neigv
    
    ! ITERATE!

    iterate: do iter = 1, itermax
       np   = nbase
       sweep: do n = 1, neigv
          if ( .not.conv(n) ) then  
             ! this root not yet converged ... calculate the residual and 
             ! construct a new basis vector out of it
             np = np + 1
             phi(np) = ZERO
             do m=1, nbase                          !
                call axpy &                         ! phi(np)=(H-e(n)) * psi(n)
                     (      vc(m,n),hphi(m),phi(np))!
                call axpy &                         ! where
                     (-e(n)*vc(m,n), phi(m),phi(np))! psi(n) = 
             end do                                 !   sum_m phi(m)*vc(m,n)
             call g_psi(phi(np),e(n),phi(np))       ! phi(np) = (H-e(n)) phi(n)
                                                    !  / (Hd-e(n)) 
          end if                                    !          
       end do sweep

       pbase = nbase
       nbase = np
       call gram ( phi(1:nbase), exclude=pbase )    !  orthonormalize

       ! calculate h times the newly obtained basis-set elements
       ! and update the reduced hamiltonian

       do n = pbase+1, nbase
          call h_psi(phi(n),hphi(n))
          do m=1, n
             hp(m,n) = phi(m).dot.hphi(n)
             hp(n,m) = conjg(hp(m,n))
          end do
       enddo

       ! diagonalize the reduced hamiltonian

       call cdiagh( hp(:,1:nbase), ew(1:nbase), vc(:,1:nbase) )
       e(1:neigv) = ew(1:neigv)

       ! test for convergence
       ! the converge criterion is that the magnitude if the projection
       ! of each eigenvector onto the most recently generated basis
       ! vectors is smaller than thres

       conv(:) = (/ ( abs(dot_product(vc(pbase+1:nbase,n), &
            vc(pbase+1:nbase,n))) <= thresh, n=1, neigv ) /)
       notcnv     = neigv - count(conv(1:neigv))

       ! if overall convergence has been achieved or if the dimension of
       ! the reduced basis set is becoming too large, refresh the basis set. 
       ! i.e. replace its first neigv elements with the current estimates of 
       ! the eigenvectors and set the basis dimension to neigv
       
       if( notcnv==0 .or. nbase+notcnv > basemax )  then
          do n=1,neigv
             hphi(n) = ZERO
             do m=1,nbase
                call axpy(vc(m,n),phi(m),hphi(n))
             end do
          end do
          
          do n=1,neigv
             phi(n) = hphi(n)
             call h_psi(phi(n),hphi(n))
          end do

          nbase = neigv
          hp = (ZERO,ZERO)
          vc = (ZERO,ZERO)
          do n = 1, nbase
             hp(n,n) = cmplx(e(n),ZERO)
             vc(n,n) = cmplx(ONE,ZERO)
          end do
          
       end if

       ! all roots converged: return
       if( notcnv==0 ) exit iterate
       
    end do iterate

    do n = 1, neigv
       if(.not.conv(n)) write(*,"(a,i3,a,f10.5,a,1pe8.1,a,i3,a)") &
            "WARNING: e(",n,") =",e(n)," not converged to within ", &
            thresh, " in ",itermax," iterations"
    end do

    ! deallocate work arrays and exit

    swap => null()
    do n=1, basemax
       call destroy( phi(n))
       call destroy(hphi(n))
    end do
    deallocate ( phi, hphi, hp, vc, ew, conv, pivot )
    
    return
  end subroutine cegter
  
  subroutine gram(psi,exclude)

    use constants, only : dp, one, zero
    use pw_types,  only : axpy, scal, operator(.dot.), wave_function

    type(wave_function), intent(inout), dimension(:) :: psi
    integer, intent(in), optional                    :: exclude
    integer                                          :: start, m,n
    complex(DP)                                      :: ss
    
    if(present(exclude)) then
       start = exclude + 1
    else
       start = 1
    end if
    
    do n = start, size(psi)
       do m = 1, n
          ss = psi(m).dot.psi(n)
          if( m /= n ) then
             call axpy(-ss,psi(m),psi(n))
          else
             ! write(*,*) 'gram ',n,abs(ss)
             call scal( ONE / sqrt(abs(ss)) , psi(n) )
          end if
       end do
    end do

  end subroutine gram

end module iterative
