module scf
  use constants, only : dp, one
  use utilities, only : error
  use pw_types,  only : k_point, field, direct_space, available,  &
                        direct_to_reciprocal_space, reciprocal_to_direct_space,&
                        destroy, construct, axpy, field_ft, scal, &
                        is_idle, operator(.dot.), assignment(=)

  implicit none

  type(field)            :: rho, Vscf, delta_v
  real(DP)               :: eband, exc, evxc, ehar, eewald, etot
  real(DP)               :: V0
  logical                :: add_nscf_energy_correction=.true.
  real(DP), allocatable  :: force(:,:),forcelc(:,:),forceion(:,:)
  logical                :: new_atomic_configuration = .true.
  logical                :: verbose = .true.

contains

!> Subroutine {\tt v\_of\_rho} calculates the electrostatic and
!> exchange-correlation potential, $V({\bf r})$, generated by the
!> electron charge-density distribution, $\rho({\bf r})$: 
!> \[ V({\bf r}) =e^2 \int \frac{\rho({\bf r}')} {|{\bf r} - {\bf r'} |} + 
!> \mu_{xc}[\rho({\bf r})], \], 
!> where $\mu_{xc}[\rho({\bf r})]$ is the functional derivative
!> of the exchange-correlation energy in the LDA approximation.
  subroutine v_of_rho(rho,v) 
    use constants, only : e2, zero, fpi, tpiba2
    use structure,  only: omega
    use functionals, only: xc

    implicit none
 
    integer                     :: i,j,k
    type(field), intent(in)     :: rho
    type(field), intent(inout)  :: v
    complex(DP), allocatable    :: aux(:,:,:)
    complex(DP)                 :: car
    real(DP)                    :: ex,ec, vx,vc, fac
    integer                     :: nr1, nr2, nr3
    integer,     pointer        :: ig1(:), ig2(:), ig3(:) 
    real(DP),    pointer        :: gg(:)


    ! check that on input rho is in real space

    if (.not.available(rho)) &
         call error('v_of_rho','rho should already exists on input',1)

    nr1 = rho%nr1
    nr2 = rho%nr2
    nr3 = rho%nr3

    ig1 => rho%ig1
    ig2 => rho%ig2
    ig3 => rho%ig3

    gg  => rho%gg

    if(rho%state /= DIRECT_SPACE) &
         call error('v_of_rho','on input, rho should be in real space',1)

    ! check that v is available and that its dimensions match those of rho

    if( available(v) ) then
       if( v%nr1 /= nr1 .or. v%nr2 /= nr2 .or. v%nr3 /= nr3 ) &
            call error('v_of_rho', &
             'input and output field dimensions do not match',2)
    else
       call error('v_of_rho','output field is not available',3)
    end if

    v = rho
    call field_ft(v,DIRECT_TO_RECIPROCAL_SPACE)

    allocate(aux(nr1,nr2,nr3))

    aux(:,:,:) = (ZERO,ZERO)

    ehar=ZERO
    do i=2, v%ngm
       fac = (FPI * E2 / gg(i) / tpiba2 )  
       aux(ig1(i),ig2(i),ig3(i)) =  fac * v%val(ig1(i),ig2(i),ig3(i))
       ehar = ehar + fac * abs(v%val(ig1(i),ig2(i),ig3(i)))**2
    end do

    ehar=ehar*omega/2.d0
    v%val(:,:,:) = aux(:,:,:)
    
    call field_ft(v,RECIPROCAL_TO_DIRECT_SPACE)

    exc=ZERO
    evxc=ZERO
    do i=1, nr1
       do j=1, nr2
          do k=1, nr3
             call xc(real(rho%val(i,j,k)),ex,ec,vx,vc) 
             v%val(i,j,k) = v%val(i,j,k) + e2*(vx + vc)
             exc = exc + E2 * (ex+ec) * rho%val(i,j,k)
             evxc=evxc + E2 * (vx+vc) * rho%val(i,j,k)
          end do
       end do
    end do
    exc  = exc  * omega/(nr1*nr2*nr3)
    evxc = evxc * omega/(nr1*nr2*nr3)
    
    deallocate(aux)
    
    return
  end subroutine v_of_rho

  logical function linear_mixing(vin,vout,alpha,tol,tr2)

    ! simple potential linear mixing of potential
    ! On output: the mixed potential is in vout, vin=(vin-vout)

    implicit none
    real(DP), intent(in)    :: alpha, tol
    real(DP), intent(out)   :: tr2
    type(field), intent(inout) :: vin, vout
    
    ! vin <-  vin - vout
    call axpy(-ONE,vout,vin)
    tr2 = (vin .dot. vin)
    
    linear_mixing = (tr2 < tol)
    ! vout <- vout + ( ONE - alpha ) * vin
    call axpy((ONE-alpha),vin,vout)

    return

  end function linear_mixing

  logical function broyden (vin,vout,alpha,iter,n_iter,tol,tr2)

    ! Modified Broyden's method for potential mixing
    !             d.d. johnson prb 38, 12807 (1988)
    ! On output: the mixed potential is in vout, vin=(vin-vout)

    implicit none
    integer, intent(in)     :: iter, n_iter
    real(DP), intent(in)    :: alpha, tol
    real(DP), intent(out)   :: tr2
    type(field), intent(inout) :: vin, vout

    integer                 :: i, j, iwork(n_iter), info, iter_used, &
                               ipos, inext
    real(DP)                :: beta(n_iter,n_iter), gamma, work(n_iter),&
                               norm, DDOT, DNRM2, w0, w(n_iter)
    type(field)             :: vsave
    type(field), allocatable, save :: df(:), dv(:)
    external DSYTRF, DSYTRI

    w0=0.01_DP
    w =ONE
    
    if (iter < 1)   call error('broyden','iteration number wrong',1)
    if (n_iter < 1) call error('broyden','n_iter too small',2)
    if (vin%state /= vout%state .or. vin%status == IS_IDLE .or. &
         vout%status == IS_IDLE) call error('broyden','ma de che?',3)

    ! vin <-  vin - vout
    call axpy(-ONE,vout,vin)
    tr2 = (vin .dot. vin)

    broyden = (tr2 < tol)

    if (iter == 1) then
       allocate(df(n_iter))
       allocate(dv(n_iter))
       do i=1,n_iter
          call construct(df(i),vin%state)
          call construct(dv(i),vin%state)
       end do
    end if
    if (broyden) then
       do i=1,n_iter
          call destroy(df(i))
          call destroy(dv(i))
       end do
       deallocate(dv)
       deallocate(df)
       return
    end if

    ! iter_used = iter-1  if iter <= n_iter
    ! iter_used = n_iter  if iter >  n_iter
    
    iter_used=min(iter-1,n_iter)
    
    ! ipos is the position in which results from the present iteraction 
    ! are stored. ipos=iter-1 until ipos=n_iter, then back to 1,2,...
    
    ipos =iter-1-((iter-2)/n_iter)*n_iter
    
    if (iter > 1) then
       ! df(ipos) <-  vin - df(ipos)
       call scal(-ONE,df(ipos)) 
       call axpy(ONE,vin,df(ipos))
       ! dv(ipos) <-  vout- dv(ipos)
       call scal(-ONE,dv(ipos))
       call axpy(ONE,vout,dv(ipos))

       norm = sqrt((df(ipos) .dot. df(ipos)))
       call scal(ONE/norm,df(ipos))
       call scal(ONE/norm,dv(ipos))
    end if
    
    do i=1,iter_used
       do j=i+1,iter_used
          beta(i,j) = w(i)*w(j)*(df(j).dot.df(i))
       end do
       beta(i,i)=w0**2+w(i)**2
    end do
    
    call DSYTRF ('u',iter_used,beta,n_iter,iwork,work,n_iter,info)
    call error('broyden','factorization',info)
    call DSYTRI ('u',iter_used,beta,n_iter,iwork,work,info)
    call error('broyden','DSYTRI',info)
    
    do i=1,iter_used
       do j=i+1,iter_used
          beta(j,i)=beta(i,j)
       end do
    end do
    
    do i=1,iter_used
       work(i) = (df(i).dot.vin)
    end do
    
    call construct(vsave,vin%state)
    vsave = vout
    ! vout <- vout + (ONE-alpha) * vin
    call axpy(ONE-alpha,vin,vout)
    
    do i=1,iter_used
       gamma=0.d0
       do j=1,iter_used
          gamma = gamma + beta(j,i)*w(j)*work(j)
       end do
       ! vout <- vout - w(i) * gamma * ( alpha * df(i) + dv(i) )
       call axpy(-w(i)*gamma*(ONE-alpha), df(i), vout)
       call axpy(-w(i)*gamma            , dv(i), vout)
    end do
    
    inext=iter-((iter-1)/n_iter)*n_iter
    df(inext) = vin
    dv(inext) = vsave
    call destroy(vsave)
    
    return
  end function broyden

  subroutine tot_ene(conv)
    use structure, only : omega
    use constants, only : zero
    use ewald, only : ewald_energy
    use pw_types, only : operator(.dot.)
    implicit none
    logical, intent(in)    :: conv
    real(DP)               :: delta_e = ZERO
    logical                :: first = .true.
    integer, allocatable   :: ityp(:)
    integer                :: na

    if (first .or. new_atomic_configuration) then
       eewald= ewald_energy()
       first = .false.
       new_atomic_configuration = .false.
    end if

    if (add_nscf_energy_correction) delta_e = - (delta_v.dot.rho) * omega

    etot = eband-ehar+(exc-evxc)+eewald + delta_e

    write(*,*)
    if (conv) then
       write(*,'("!",4x,''Total   energy: '',f15.6)') etot
    else
       write(*,'(5x,''Total   energy: '',f15.6)') etot
    end if
    if (verbose.or.conv) then
       write(*,'(5x,''Band    energy: '',f15.6)') eband
       write(*,'(5x,''Hartree energy: '',f15.6)') ehar
       write(*,'(5x,''XC      energy: '',f15.6)') (exc-evxc)
       write(*,'(5x,''Ewald   energy: '',f15.6)') eewald
       if (add_nscf_energy_correction) &
            write(*,'(5x,''dscf    energy: '',f15.6)') delta_e
    end if

    return
  end subroutine tot_ene

  !> \large \newpage 
  !> The subroutine {\tt forces} computes the atomic forces using the
  !> Hellmann-Feynman  theorem 
  !> \[ {\bf F}_i = - \frac{\partial E}{\partial {\bf R}_i} =  
  !> - \left\langle \psi_{{\bf k}\nu} \mid \frac{\partial {\cal H}}
  !> {\partial {\bf R}_i} \mid \psi_{{\bf k}\nu} \right\rangle. \] 

  subroutine forces
    use constants, only : tpi, tpiba, tpiba2, zero
    use structure, only : nats, species, atoms
    use ewald, only: ewald_force
    use pw_types, only : available, operator(.dot.), destroy, &
                         direct_to_reciprocal_space, field_ft, &
                         assignment(=)
    implicit none

    real(DP)              :: valr, vali, arg
    integer               :: na, ig, ipol
    type(field)           :: aux
    complex(DP), pointer  :: vg(:)
    integer               :: nr1,nr2,nr3
    integer,     pointer  :: ig1(:), ig2(:), ig3(:)
    real(DP),    pointer  :: g(:,:)

    integer,  allocatable :: ityp(:)
    
    if (.not.allocated(forceion)) allocate(forceion(3,nats))
    if (.not.allocated(forcelc))  allocate(forcelc(3,nats))
    if (.not.allocated(force))    allocate(force(3,nats))

    !> The ionic contribution ${\bf F_i^{\rm ion}}$ arising from the 
    !> electrostatic potential on site $i$ due to all other ions at 
    !> sites $j \ne i$ is conditionally convergent with the distance 
    !> $R_{ij}=\mid {\bf R_{ij}} \mid$. 
    !> (notation: ${\bf R_{ij}}={\bf R_i - R_j}$). 
    !> To improve convergence the \hyperlink{ewald}{Ewald technique} is 
    !> employed (for a reference see M.P. Allen and D.J.
    !> Tildesley, "Computer Simulation of Liquids", Clarendon Press, Oxford 
    !> UK.). The ionic contribution to the total force 
    !> \[ -\frac{1}{2} \sum_{i,j} Z_i \, Z_j
    !> \frac{\partial}{\partial{\bf R_i}} \frac{1} {R_{ij}} = 
    !> -\frac{1}{2} \sum_{i,j} \frac{Z_i \, Z_j}{R_{ij}^{2}}  \, \,
    !> \frac{\bf R_{ij}}{R_{ij}} \]
    !> is calculated in \hyperlink{ewaldforce}{\tt ewald\_force}.
    !> \hypertarget{htforceion}{}

    call ewald_force  ( forceion )

    !> The electronic contribution ${\bf F}_j^{\rm el}$:
    !> \[ {\bf F}_j^{\rm el} = - \int \rho_{\bf R}({\bf r}) 
    !> \frac{\partial V_{\bf R}({\bf r})}{\partial{\bf R_i}} {\rm d \bf r}. \]
    !> is calculated in  reciprocal space:
    !> \[ {\bf F}_j^{\rm el} = \sum_{\bf G} i \, {\bf G} \, 
    !> V_{\bf G}^{\rm ion} \, \rho_{\bf G} \, e^{i \, {\bf G \cdot \tau}} \]
    !> To begin, get a reciprocal-space representation of the charge density
    !> $\rho_{\bf G}$ via a fast fourier transform:
    aux=rho
    call field_ft(aux,DIRECT_TO_RECIPROCAL_SPACE)

    nr1 = rho%nr1
    nr2 = rho%nr2
    nr3 = rho%nr3

    ig1 => rho%ig1
    ig2 => rho%ig2
    ig3 => rho%ig3
 
    g   => rho%g

    forcelc(:,:)=ZERO
    do na=1, nats    ! loop over atoms
       vg => species(atoms(na)%itype)%ps%vg
       do ig=1, aux%ngm   
          arg = ( g(1,ig)*atoms(na)%tau(1) + &
                  g(2,ig)*atoms(na)%tau(2) + &
                  g(3,ig)*atoms(na)%tau(3)  ) * TPI
          valr = real( aux%val(ig1(ig),ig2(ig),ig3(ig)) )
          vali = aimag(aux%val(ig1(ig),ig2(ig),ig3(ig)) )
          !> \[ \sum_{\bf G} {\bf G} \, V_{\bf G}^{\rm ion} \, \rho_{\bf G} \,
          !> \left[ \sin({\bf G \cdot \tau}){\rm Re}(\rho_{\bf G}) +
          !> \cos({\bf G \cdot \tau}){\rm Im}(\rho_{\bf G}) \right] \]
          do ipol=1,3
             forcelc(ipol,na) = forcelc(ipol,na) + g(ipol,ig) * vg(ig) * &
                  ( sin(arg) * valr + cos(arg) * vali )
          enddo
       enddo
       do ipol=1,3
          forcelc(ipol,na) = forcelc(ipol,na)*tpiba
       enddo
    enddo
  
    !> Sum the ionic and electronic contributions:
    !> \[ {\bf F_i} = {\bf F_i^{\rm ion}} + {\bf F_i^{\rm el}} \]
    do na=1,nats
       do ipol=1,3
          force(ipol,na)=forcelc(ipol,na)+forceion(ipol,na) 
       enddo
    enddo

    !> Print out

    write(*,*)  'Total forces'
    write(*,"(5x,'atom ',i3,' type ',i2,'   force = ',3f14.8)") &
         (na, atoms(na)%itype, force(:,na), na=1,nats)
    if (verbose) then
       write(*,*)  'Local forces'
       write(*,"(5x,'atom ',i3,' type ',i2,'   force = ',3f14.8)") &
            (na, atoms(na)%itype, forcelc(:,na), na=1, nats)
       write(*,*)  'Ionic forces'
       write(*,"(5x,'atom ',i3,' type ',i2,'   force = ',3f14.8)") &
            (na, atoms(na)%itype, forceion(:,na), na=1, nats)
    end if
    
    call destroy(aux)

  end subroutine forces

  !> \hyperlink{htendforces}{\red BACK} to the main program.
  
end module scf

