module reciprocal_lattice
  use constants, only: dp
  use utilities, only: hpsort, error
  implicit none

  REAL(DP)                      :: gcutm    ! squared-wavevector cutoff (in
                                            ! units of tpiba2 = (2*pi/alat)^2)
                                            ! only wavevectors whose squared 
                                            ! norm does not exceed gcutm are 
                                            ! kept

  real(DP), allocatable, target :: g(:,:),& ! array of reciprocal-space 
                                            ! vectors arranged in order of
                                            ! increasing magnitude (unit is
                                            ! tpiba = (2*pi/alat) )
                                   gg(:)    ! squared norms of the g's

  integer                       :: ngm      ! current # of g's

  ! first two days up to here

  integer                       :: nr1=0, & !
                                   nr2=0, & ! current fft dimensions
                                   nr3=0    !
  
  integer, allocatable, target  :: ig1(:),& ! ig1(i), ig2(i), ig3(i)
                                   ig2(:),& ! are the fft indeces of
                                   ig3(:)   ! g(:,i)

  ! third day  up to here

  complex(dp), allocatable, target :: structure_factor(:,:)

  ! fourth day up to here

contains

  subroutine ggen(gcut)
    use structure, only: at, bg

    real(DP), intent(in)       :: gcut

    real(DP)                   :: t(3), tt
    integer, allocatable       :: nl(:)
    integer                    :: i,j,k, ng, nli, nx1,nx2,nx3

    logical, save              :: first=.true.

    gcutm = gcut

    ! count the total number of g-vectors such that |g|^2 <= gcutm

    if(.not. first) call error('ggen', &
         'this routine should be called only once',1)

    nx1=int(2*sqrt(gcutm)*sqrt(at(1,1)**2+at(2,1)**2+at(3,1)**2))+1
    nx2=int(2*sqrt(gcutm)*sqrt(at(1,2)**2+at(2,2)**2+at(3,2)**2))+1
    nx3=int(2*sqrt(gcutm)*sqrt(at(1,3)**2+at(2,3)**2+at(3,3)**2))+1

    ngm = 0
    do i = -(nx1+1), nx1+1
       do j = -(nx2+1), nx2+1
          do k = -(nx3+1), nx3+1
             t(:) = i*bg(:,1) + j*bg(:,2) + k*bg(:,3)
             tt   = t(1)*t(1) + t(2)*t(2) + t(3)*t(3)
             if ( tt <= gcutm ) ngm = ngm + 1
          end do
       end do
    end do

    ! allocate memory for g-vectors and their squared norms

    allocate(g(3,ngm))
    allocate(gg(ngm))
    allocate(nl(ngm))

    ! re-calculate the g-vectors and their squared norms and store them

    ngm = 0
    do i = -(nx1+1), nx1+1
       do j = -(nx2+1), nx2+1
          do k = -(nx3+1), nx3+1
             t(:) = i*bg(:,1) + j*bg(:,2) + k*bg(:,3)
             tt   = t(1)*t(1) + t(2)*t(2) + t(3)*t(3)
             if ( tt <= gcutm ) then
                ngm      = ngm + 1
                g(:,ngm) = t(:)
                gg(ngm)  = tt
             end if
          end do
       end do
    end do

    ! sort the g's in order of increasing norm

    nl(1) = 0 ! this is an input value to the hpsort routine to 
              ! indicate that the nl array is not initialized
    call hpsort( gg(1:ngm), nl )

    do ng = 1, ngm-1
       i = ng
       t(:) = g(:,i)
       do while ( nl(i) /= ng )
          nli    = nl(i)
          g(:,i) = g(:,nli)
          nl(i)  = i
          i      = nli
       end do
       g(:,i) = t(:)
       nl(i) = i
    end do

    deallocate (nl)
    first = .false.

    return
  end subroutine ggen

  ! first two days up to here

  subroutine set_fft(mr1,mr2,mr3)
    
    ! sets fft dimensions (nr1,nr2,nr3), pivot arrays (ig1,ig2,ig3,nl), 
    ! and current number of g vectors (|g|<=gcutm)

    use structure, only: at

    integer, intent(in), optional :: mr1, mr2, mr3
    integer :: ng, n1, n2, n3, minr1, minr2, minr3, maxr1, maxr2, maxr3, ngx

    nr1 = 0
    if(present(mr1)) nr1 = mr1
    nr2 = 0
    if(present(mr2)) nr2 = mr2
    nr1 = 0
    if(present(mr3)) nr3 = mr3

    ngx = ngm
    
    if(.not.allocated(ig1)) allocate ( ig1(ngx) )
    if(.not.allocated(ig2)) allocate ( ig2(ngx) )
    if(.not.allocated(ig3)) allocate ( ig3(ngx) )

    maxr1 = -100000
    maxr2 = -100000
    maxr3 = -100000
    minr1 =  100000
    minr2 =  100000
    minr3 =  100000

    do ng =1, ngx

       if( dot_product(g(:,ng),g(:,ng)) > gcutm ) exit
       
       n1      = nint( dot_product(g(:,ng),at(:,1)) )
       maxr1   = max(n1,maxr1)
       minr1   = min(n1,minr1)
       ig1(ng) = n1
       
       n2      = nint( dot_product(g(:,ng),at(:,2)) )
       maxr2   = max(n2,maxr2)
       minr2   = min(n2,minr2)
       ig2(ng) = n2

       n3      = nint( dot_product(g(:,ng),at(:,3)) )
       maxr3   = max(n3,maxr3)
       minr3   = min(n3,minr3)
       ig3(ng) = n3

    end do

    if( minr1 /= -maxr1 .or. minr2 /= -maxr2 .or. minr3 /= -maxr3 ) &
         call error("set_fft","wrong input data: nrmin /= -nrmax",1)
    
    if( 2*maxr1+1 > nr1 .or. 2*maxr2+1 > nr2 .or. 2*maxr3+1 > nr3 ) then
       nr1 = 2*maxr1 + 1
       nr2 = 2*maxr2 + 1
       nr3 = 2*maxr3 + 1
    end if

    do
       if( allowed(nr1) ) exit
       nr1 = nr1 + 1
    end do

    do 
       if( allowed(nr2) ) exit
       nr2 = nr2 + 1
    end do

    do 
       if( allowed(nr3) ) exit
       nr3 = nr3 + 1
    end do

    do ng = 1, ngm
       
       if( ig1(ng) < 0 ) ig1(ng) = ig1(ng) + nr1
       ig1(ng) = ig1(ng) + 1
       
       if( ig2(ng) < 0 ) ig2(ng) = ig2(ng) + nr2
       ig2(ng) = ig2(ng) + 1
       
       if( ig3(ng) < 0 ) ig3(ng) = ig3(ng) + nr3
       ig3(ng) = ig3(ng) + 1
       
    end do
    
    return
    
  end subroutine set_fft
  
  function allowed ( nr ) result(is_allowed)
    
    ! find if the fft dimension is a good one
    ! a "bad one" is either not implemented (as on IBM with ESSL)
    ! or implemented but with awful performances (most other cases)

    integer, intent(in)         :: nr
    logical                     :: is_allowed
    
    integer, dimension(5)            :: pwr
    integer, dimension(5), parameter :: factors = (/ 2, 3, 5, 7, 11 /)
    
    integer                     :: mr, i, fac, p, maxpwr
    
    ! find the factors of the input fft dimension
    
    mr = nr
    
    pwr(:) = 0
    
    for_each_factor: do i=1,5
       fac=factors(i)
       maxpwr = nint(log(real(mr))/log(real(fac)))+1
       do p=1, maxpwr
          if ( mr==1 ) then
             exit for_each_factor
          else if ( modulo (mr,fac) == 0 ) then
             mr=mr/fac
             pwr(i)=pwr(i)+1
          end if
       end do
    end do for_each_factor
    
    is_allowed = ( mr == 1 )
    
    return
    
  end function allowed
  
!> Subroutine {\tt set\_structure\_factor} sets up the atomic structure 
!> factor for each atomic {\tt species}: \begin{equation} {1\over \Omega} 
!> S({\bf G}) = \sum_{\tau} {\rm e}^{-i{\bf G}\cdot\tau}, \end{equation} 
!> where the sum is extended to all the atoms in the unit cell of the given 
!> species, and $\Omega$ is the volume of the unit cell.

  subroutine set_structure_factor
    use constants, only: zero, tpi
    use structure, only: element, species, atom, atoms, ntyp, nats, omega
    implicit none
    integer        :: na, nt, k
    real(DP)       :: phase
    
    if (.not.allocated(structure_factor)) then
       allocate(structure_factor(ngm,ntyp))
    else
       if (ngm/=size(structure_factor,1)) call error ('set_structure_factor', &
            '1st structure_factor dimension has been changed ',ngm)
       if (ntyp/=size(structure_factor,2)) call error ('set_structure_factor',&
            '2nd structure_factor dimension has been changed ',ngm)
    end if

    do k=1, ngm
       structure_factor(k,:) = cmplx(ZERO,ZERO)
       do na=1, nats
          nt = atoms(na)%itype 
          phase = dot_product(g(:,k),atoms(na)%tau) * TPI
          structure_factor(k,nt) = structure_factor(k,nt) &
               + cmplx(cos(phase),-sin(phase))
       end do
    end do
    structure_factor = structure_factor / omega

    do nt=1,ntyp
       species(nt)%s => structure_factor(:,nt)
    end do

    return
    
  end subroutine set_structure_factor
  
end module reciprocal_lattice
