!-------------------------------------------------------------------------------
!
!   kgrid.f90
!   kgrid.x kgrid.in kgrid.out kgrid.log
!   Generates a list of k-points with a small q-shift for espresso input file
!   inspired by espresso-3.2.3/pwtools/kpoints.f
!   written by G. Samsonidze (March 2008)
!
!-------------------------------------------------------------------------------
!
!   See kgrid.inp for input file description.
!
!-------------------------------------------------------------------------------
!
!   kgrid generates k-points in the irreducible wedge of the Brillouin
!   zone with the symmetries of the space group of the crystal (for real
!   and complex versions of GW) and also with time-reversal symmetry
!   (for DFT) where real and complex stand for systems with and without
!   inversion symmetry.
!
!   Set trs to .true. for DFT and to .false. for complex version of GW
!   (it does not matter for real version of GW since inversion symmetry
!   takes k to -k same way as time-reversal symmetry).
!
!   Time-reversal symmetry is not currently implemented in BerkeleyGW.
!
!   If a non-zero q-shift is given, the irreducible wedge is unfolded with all
!   the symmetries of the space group of the crystal and folded back to the
!   irreducible wedge with the symmetries of a subgroup of the q-vector.
!
!   If some of the fractional translations are not commensurate with
!   the FFT grid, the corresponding symmetries of the space group of
!   the crystal will be dropped; you can skip FFT grid check by setting
!   (nr1 nr2 nr3) to (0 0 0).
!
!   Set number of atoms to 0 and skip atomic species and positions
!   to use the symmetries of the space group of the Bravais lattice.
!
!   Note that the units of dqj are nkj times the units of dkj (j=1,2,3)
!   according to convention inherited from PARATEC.
!
!   In the case of supercell calculation espresso generates additional
!   k-points; to use the list of k-points from espresso input file as is
!   set calculation to 'bands' rather than 'nscf' in namelist control in
!   espresso input file; note that in this case occupations in espresso
!   will be set to zero, so you will have to use flags wfng_occupation,
!   wfng_nvmin, and wfng_nvmax in pw2bgw input file,
!
!   If you set flag nosym in namelist system in espresso input file
!   to .true. there will be no symmetries written to WFN file which
!   means you will have to use the full Brillouin zone instead of
!   the irreducible wedge.
!
!   To use the k-point grid automatically generated by espresso
!   for complex version of GW set flag noinv in namelist system
!   in espresso input file to .true. (this disables the use of
!   time-reversal symmetry in k-point generation).
!
!-------------------------------------------------------------------------------

#include "f_defs.h"

program kgrid

  use global_m
  use kgrid_routines_m
  use misc_m
  use symmetries_m
  implicit none

  logical :: trs,trs_dq,dqflag,symr(48),syms(48),output_cartesian,output_octopus
  integer :: i,j,i1,i2,i3,ierr,ierr2,ia,na,ikf,nkf,nkfq,ikr,nkr,ikq, &
    nkq,nrb,irc,nrc,nrn,irr,nrr,nrq,cell_symmetry,nk(3),nfft(3),nfft_zero(3), &
    rb(3,3,48),rc(3,3,48),rn(3,3,48),rq(3,3,48),c2b(48),gpt(3),no_as(1),spacegroup_b,spacegroup_c,spacegroup_n
  real(DP) :: celvol,al,bl,dk(3),dq(3),a(3,3),b(3,3),no_ap(3,1),tnp_b(3,48),tnp_c(3,48),tnp_n(3,48)
  character*256 :: fi,fo,fl
  character*21 :: symbol_b, symbol_c, symbol_n

  integer, allocatable :: as(:)
  real(DP), allocatable :: ap(:,:),ap_lat(:,:)
  real(DP), allocatable :: kpoint(:,:)
  integer, allocatable :: kweight(:)
  integer, allocatable :: kfold(:)
  integer, allocatable :: ksymmetry(:)
  integer, allocatable :: kindex(:)
  real(DP), allocatable :: kr(:,:)
  real(DP), allocatable :: kw(:)
  real(DP), allocatable :: kf(:,:)
  real(DP), allocatable :: kfw(:)
  real(DP), allocatable :: kq(:,:)
  real(DP), allocatable :: kqw(:)

  output_cartesian = .false. ! default
  output_octopus = .false. ! default

  if (iargc().ne.3) then
    call die("USAGE: kgrid.x kgrid.in kgrid.out kgrid.log")
  else
    call getarg(1,fi)
    call getarg(2,fo)
    call getarg(3,fl)
  endif
  
  call open_file(21,file=TRUNC(fi),status='old')
  read(21,*,iostat=ierr)nk
  if (ierr.eq.0) read(21,*,iostat=ierr)dk
  if (ierr.eq.0) read(21,*,iostat=ierr)dq
  dqflag = (any(abs(dq(1:3)) > TOL_Zero))
  do i=1,3
    if (ierr.eq.0) read(21,*,iostat=ierr)(a(j,i),j=1,3)
  enddo
  na=0
  if (ierr.eq.0) read(21,*,iostat=ierr)na
  if (na.gt.0) then
    SAFE_ALLOCATE(as, (na))
    SAFE_ALLOCATE(ap, (3,na))
    SAFE_ALLOCATE(ap_lat, (3,na))
  endif
  do ia=1,na
    if (ierr.eq.0) read(21,*,iostat=ierr)as(ia),(ap(i,ia),i=1,3)
  enddo
  if (ierr.eq.0) read(21,*,iostat=ierr)(nfft(i),i=1,3)
  if (ierr.eq.0) read(21,*,iostat=ierr)trs
  if (ierr.eq.0) read(21,*,iostat=ierr2)output_cartesian ! this one is optional
  if (ierr2.eq.0) read(21,*,iostat=ierr2)output_octopus ! this one is optional
  if (ierr.eq.0) call close_file(21)
  if (ierr.ne.0) call die("Reading " // TRUNC(fi) // " failed.")
  
  if (any(nk(1:3) <= 0)) call die("kgrid must be > 0.")
  if (na < 0) call die("Number of atoms must be >= 0.")
  if (any(nfft(1:3) < 0)) call die("FFT grid must be >= 0.")
  
  nkf=nk(1)*nk(2)*nk(3)
  SAFE_ALLOCATE(kpoint, (3,nkf))
  SAFE_ALLOCATE(kweight, (nkf))
  SAFE_ALLOCATE(kfold, (nkf))
  SAFE_ALLOCATE(ksymmetry, (nkf))
  SAFE_ALLOCATE(kindex, (nkf))
  
  al = sqrt(a(1,1)**2 + a(2,1)**2 + a(3,1)**2)
  a(:,:) = a(:,:) / al
  celvol = a(1,1) * (a(2,2) * a(3,3) - a(2,3) * a(3,2)) - &
    a(2,1) * (a(1,2) * a(3,3) - a(1,3) * a(3,2)) + &
    a(3,1) * (a(1,2) * a(2,3) - a(1,3) * a(2,2))
  celvol = abs(celvol)
  if (celvol.le.TOL_Zero.or.celvol.ge.INF) call die("cell volume is <= 0 or infinite.")
  bl = 2.0d0 * PI_D / al
  b(1,1) = (a(2,2) * a(3,3) - a(3,2) * a(2,3)) / celvol
  b(2,1) = (a(3,2) * a(1,3) - a(1,2) * a(3,3)) / celvol
  b(3,1) = (a(1,2) * a(2,3) - a(2,2) * a(1,3)) / celvol
  b(1,2) = (a(2,3) * a(3,1) - a(3,3) * a(2,1)) / celvol
  b(2,2) = (a(3,3) * a(1,1) - a(1,3) * a(3,1)) / celvol
  b(3,2) = (a(1,3) * a(2,1) - a(2,3) * a(1,1)) / celvol
  b(1,3) = (a(2,1) * a(3,2) - a(3,1) * a(2,2)) / celvol
  b(2,3) = (a(3,1) * a(1,2) - a(1,1) * a(3,2)) / celvol
  b(3,3) = (a(1,1) * a(2,2) - a(2,1) * a(1,2)) / celvol
  celvol = abs(celvol) * al**3

  call open_file(23,file=TRUNC(fl),status='replace')

  ! Find symmetries of the empty Bravais lattice
  no_as(1) = 0        ! dummy species
  no_ap(1:3, 1) = 0d0 ! dummy atom at origin
  call get_symmetries(1, no_as, no_ap, a, nfft, cell_symmetry, nrb, rb, tnp_b, spacegroup_b, symbol_b)

  ! Find symmetries of the actual crystal with atoms
  if (na.gt.0) then
    ap_lat = matmul(transpose(b),ap) / al ! convert to crystal coords

    ! first, with FFT grid = 0
    nfft_zero(1:3) = 0
    call get_symmetries(na, as, ap_lat, a, nfft_zero, cell_symmetry, nrn, rn, tnp_n, spacegroup_n, symbol_n)

    ! then, with FFT grid as given in input
    call get_symmetries(na, as, ap_lat, a, nfft, cell_symmetry, nrc, rc, tnp_c, spacegroup_c, symbol_c)
  else
    nrc = nrb
    rc = rb
  endif

  ! Originally, c2b was used to map crystal symmetries to Bravais symmetries
  ! for writing into the log file. Here the mapping is disabled, hence the
  ! log file contains the indices of crystal symmetries instead of Bravais.
  do irc = 1, nrc
    c2b(irc) = irc
  enddo

  if(output_cartesian) then
    write(23,'(a)') 'k-points are in Cartesian coordinates'
  else
    write(23,'(a)') 'k-points are in crystal coordinates'
  endif
  if(trs) then
    write(23,'(a)') 'time-reversal symmetry is being used'
    write(0,'(a)') 'WARNING: you will need to run the wfn_time_reversal.x utility to unfold'
    write(0,'(a)') 'a WFN file generated with this k-grid before you can run BerkeleyGW.'
  endif
  write(23,'("nk(3) =",1x,3i5)')nk
  write(23,'("dk(3) =",1x,3f13.9)')dk
  if(dqflag) then
    write(23,'("dq(3) =",1x,3f13.9)')dq
  endif
  write(23,'(a)')
  write(23,'("direct lattice vectors")')
  do i=1,3
    write(23,'("a",i1," =",1x,3f13.9)')i,(al*a(j,i),j=1,3)
  enddo
  write(23,'(a)')
  write(23,'("unit cell volume =",1x,e16.9)')celvol
  write(23,'(a)')
  write(23,'("reciprocal lattice vectors")')
  do i=1,3
    write(23,'("b",i1," =",1x,3f13.9)')i,(bl*b(j,i),j=1,3)
  enddo
  call write_sym("of the Bravais lattice", symbol_b, nrb, rb, tnp_b, spacegroup = spacegroup_b)
  if (na.gt.0) then
    call write_sym("of the crystal without FFT grid", symbol_n, nrn, rn, tnp_n, spacegroup = spacegroup_n)
    if(any(nfft(1:3) /= 0)) then
      write(23,'(/,a,3i6)') 'FFT grid: ', nfft(1:3)
      call write_sym("of the crystal with FFT grid", symbol_c, nrc, rc, tnp_c, spacegroup = spacegroup_c)
    endif
  endif

  ikf=0
  do i1=0,nk(1)-1
    do i2=0,nk(2)-1
      do i3=0,nk(3)-1
        ikf=ikf+1
        kpoint(1,ikf)=(dble(i1)+dk(1))/dble(nk(1))
        kpoint(2,ikf)=(dble(i2)+dk(2))/dble(nk(2))
        kpoint(3,ikf)=(dble(i3)+dk(3))/dble(nk(3))
        call k_range(kpoint(:,ikf),gpt,TOL_Zero)
        kweight(ikf)=1
        kfold(ikf)=0
        ksymmetry(ikf)=0
      enddo
    enddo
  enddo
  nkf=ikf
  call fold(nkf,kpoint,kweight,kfold,ksymmetry,nrc,rc,trs,TOL_Zero)
  ikr=0
  do ikf=1,nkf
    if (kweight(ikf).gt.0) ikr=ikr+1
  enddo
  nkr=ikr
  call sort(nkf,nkr,kpoint,kweight,kindex,TOL_Zero)
  call reduce(nkf,kweight,kfold)
  SAFE_ALLOCATE(kr, (3,nkr))
  SAFE_ALLOCATE(kw, (nkr))
  do ikr=1,nkr
    kr(:, ikr) = kpoint(:, kindex(ikr))
    kw(ikr) = kweight(kindex(ikr))
  enddo

  irr=0
  do irc=1,nrc
    symr(irc)=.false.
    do ikf=1,nkf
      if (ksymmetry(ikf).eq.irc) symr(irc)=.true.
    enddo
    if (symr(irc)) irr=irr+1
  enddo
  nrr=irr

  if (nkr.lt.nkf) then
    write(23,'(a)')
    write(23,'("symmetries that reduce the k-points")')
    do irc=1,nrc
      if (symr(irc)) then
        write(23,'("r",i2.2," =",2x,3i4,2x,3i4,2x,3i4)') &
          c2b(irc),((rc(j,i,irc),j=1,3),i=1,3)
      endif
    enddo
    if (nrr.eq.0) write(23,'("none")')
  endif
  
  if (dqflag) then
    ! allocate to a simple upper bound of the number of k-points to unfold to
    nkfq=nkr*nrc
    ! even with time-reversal symmetry, there can be a maximum of 48 symmetries
    if(trs) nkfq = min(nkfq * 2, nkr * 48)
    SAFE_ALLOCATE(kf, (3,nkfq))
    SAFE_ALLOCATE(kfw, (nkfq))

    call dqunfold(nkr,kr,nkfq,kf,kfw,nrc,rc,trs,TOL_Zero)
    call dqsubgrp(dq,nrc,rc,nrq,rq,syms,trs_dq,TOL_Zero)
    if(.not. trs) trs_dq = .false.

    write(23,'(a)')
    write(23,'("symmetries of a subgroup of the q-vector")')
    do irc=1,nrc
      if (syms(irc)) then
        write(23,'("r",i2.2," =",2x,3i4,2x,3i4,2x,3i4)') &
          c2b(irc),((rc(j,i,irc),j=1,3),i=1,3)
      endif
    enddo
    if (nrq.eq.0) write(23,'("none")')

    ! another upper bound: the folded number can be no more than the unfolded number
    nkq=nkfq
    SAFE_ALLOCATE(kq, (3,nkq))
    SAFE_ALLOCATE(kqw, (nkq))

    call dqfold(nkfq,kf,kfw,nkq,kq,kqw,nrq,rq,trs_dq,TOL_Zero)
    call dqsort(nkfq,kf,kfw,TOL_Zero)
    call dqsort(nkq,kq,kqw,TOL_Zero)
    do ikf=1,nkfq
      kf(1:3,ikf)=kf(1:3,ikf)+dq(1:3)
    enddo
    do ikq=1,nkq
      kq(1:3,ikq)=kq(1:3,ikq)+dq(1:3)
    enddo

    call open_file(22,file=TRUNC(fo),status='replace')
    if(output_octopus) then
      call write_octopus(22, nkq, kq, kqw, nk, dk + dq * nk)
    else
      call write_espresso(22, nkq, kq, kqw)
    endif
  else
    call open_file(22,file=TRUNC(fo),status='replace')
    if(output_octopus) then
      call write_octopus(22, nkr, kr, kw, nk, dk)
    else
      call write_espresso(22, nkr, kr, kw)
    endif
  endif
  call close_file(22)
  
  write(23,'(a)')
  write(23,'("k-points in the original uniform grid")')
  write(23,'(i7)')nkf
  do ikf=1,nkf
    if(output_cartesian) kpoint(1:3, ikf) = bl * matmul(b(1:3, 1:3), kpoint(1:3, ikf))

    if (kfold(ikf).eq.0) then
      write(23,'(i7,3f13.9,f6.1,1x,i7,2x,"---")')ikf, &
        (kpoint(j,ikf),j=1,3),dble(kweight(ikf)),0
    else
      if (ksymmetry(ikf).eq.0) then
        write(23,'(i7,3f13.9,f6.1,1x,i7,2x,"r",i2.2)')ikf, &
          (kpoint(j,ikf),j=1,3),dble(kweight(ikf)), &
          kfold(ikf),c2b(ksymmetry(kfold(ikf)))
      else
        write(23,'(i7,3f13.9,f6.1,1x,i7,2x,"r",i2.2)')ikf, &
          (kpoint(j,ikf),j=1,3),dble(kweight(ikf)), &
          kfold(ikf),c2b(ksymmetry(ikf))
      endif
    endif
  enddo
  write(23,'(a)')
  write(23,'("k-points in the irreducible wedge")')
  write(23,'(i7)')nkr
  do ikr=1,nkr
    if(output_cartesian) kr(1:3, ikr) = bl * matmul(b(1:3, 1:3), kr(1:3, ikr))
    write(23,'(i7,3f13.9,f6.1,1x,i7)')ikr,kr(1:3,ikr),dble(kw(ikr)),kindex(ikr)
  enddo
  if (dqflag) then
    write(23,'(a)')
    write(23,'("k-points in the unfolded Brillouin Zone")')
    write(23,'(i7)')nkfq
    do ikf=1,nkfq
      if(output_cartesian) kf(1:3, ikf) = bl * matmul(b(1:3, 1:3), kf(1:3, ikf))
      write(23,'(i7,3f13.9,f6.1)')ikf,(kf(j,ikf),j=1,3),1.0d0
    enddo
    write(23,'(a)')
    write(23,'("k-points folded with the symmetries of a subgroup of the q-vector")')
    write(23,'(i7)')nkq
    do ikq=1,nkq
      if(output_cartesian) kq(1:3, ikq) = bl * matmul(b(1:3, 1:3), kq(1:3, ikq))
      write(23,'(i7,3f13.9,f6.1)')ikq,(kq(j,ikq),j=1,3),kqw(ikq)
    enddo
  endif
  write(23,'(a)')
  call close_file(23)
  
  if (na.gt.0) then
    SAFE_DEALLOCATE(as)
    SAFE_DEALLOCATE(ap)
  endif
  SAFE_DEALLOCATE(kpoint)
  SAFE_DEALLOCATE(kweight)
  SAFE_DEALLOCATE(kfold)
  SAFE_DEALLOCATE(ksymmetry)
  SAFE_DEALLOCATE(kindex)
  if (dqflag) then
    SAFE_DEALLOCATE(kr)
    SAFE_DEALLOCATE(kf)
    SAFE_DEALLOCATE(kfw)
    SAFE_DEALLOCATE(kq)
    SAFE_DEALLOCATE(kqw)
  endif
  
end program kgrid
