!=================================================================================
!
! Routines:
!
! (1) epscopy()                                 Last Modified 2014 (FHJ)
!
!     Copy dielectric matrices from eps0mat(.h5)/epsmat(.h5) files to memory.
!
!     This routine reads the eps0mat and epsmat files (units 10 and 11)
!     and copies the relevant information about those dielectric matrices
!     to memory (epsmpi). Processor 0 will actually read data and redistribute
!     them to other processors. The eps0mat can have more than one q-point,
!     which is useful when performing voronoi averages for the q->0 point.
!
!==================================================================================

#include "f_defs.h"

module epscopy_m

  use global_m
  use misc_m
  use io_utils_m
  use epsread_hdf5_m

  implicit none

  private

  public :: epscopy_init, epscopy

contains

!> Determines the number of q-points, the maximum rank of the dielectric matrix,
!! number of frequencies
subroutine epscopy_init(sig, neps)
  type (siginfo), intent(inout) :: sig
  integer, intent(out) :: neps !< Max. rank of dielectric matrix

  integer :: ngarbage1, nmtx, itrash, ig, igp, iq
  integer :: nFreq, freq_dep_flag
#ifdef HDF5
  integer :: ngarbage3(3), nmtxmax
  real(DP) :: dgarbage1
  logical :: file_exists
#endif

  PUSH_SUB(epscopy_init)

  if (sig%freq_dep>=0.and.sig%freq_dep<=3) then
#ifdef HDF5
    if (sig%use_hdf5) then
      if (peinf%inode==0) then
        call read_eps_grid_sizes_hdf5(ngarbage1, sig%nq0, dgarbage1, sig%nFreq, &
          nmtxmax, ngarbage3, 'eps0mat.h5')
        neps = nmtxmax
        ! FB: dirty trick, but it is really unlikely to have a full freq
        ! calculations with only 2 freqs
        select case(sig%nFreq)
        case(1)
          freq_dep_flag = 0
        case(2)
          freq_dep_flag = 3
        case default
          freq_dep_flag = 2
        end select

        ! Consistency check, before continuing
        if (freq_dep_flag==2.and.sig%freq_dep/=2) &
          call die('eps0mat is frequency-dependent, but this Sigma calculation is not.')
        if (freq_dep_flag==0.and.sig%freq_dep==2) then
          call die('This Sigma calculation is frequency-dependent, but eps0mat is not.')
        endif

        INQUIRE(FILE="epsmat.h5", EXIST=file_exists) 
        if (file_exists) then
          sig%igamma = 0
        else
          sig%igamma = 1
        endif

        if(sig%igamma/=0) then ! Gamma-only calculation
          sig%nq1 = 0
        else
          call read_eps_grid_sizes_hdf5(ngarbage1, sig%nq1, dgarbage1, nFreq, &
            nmtxmax, ngarbage3, 'epsmat.h5')
          if (nFreq/=sig%nFreq) then
            call die('nFreq mismatch between eps0mat.h5 and epsmat.h5')
          endif
          if (nmtxmax>neps) neps = nmtxmax
        endif
        write(6,*) "Found neps:", neps, sig%nq1
      endif !peinf%inode==0
    else
#endif
      if (peinf%inode==0) then
        call open_file(unit=10,file='eps0mat',form='unformatted',status='old')
        read(10)
        read(10) freq_dep_flag, sig%nFreq

        ! Consistency check, before continuing
        if (freq_dep_flag==2.and.sig%freq_dep/=2) &
          call die('eps0mat is frequency-dependent, but this Sigma calculation is not.')
        if (freq_dep_flag==0.and.sig%freq_dep==2) then
          call die('This Sigma calculation is frequency-dependent, but eps0mat is not.')
        endif

        read(10)
        read(10)
        read(10)
        read(10)
        read(10)
        read(10) sig%nq0
        read(10)
        read(10) itrash, neps
        call close_file(10)      
        call open_file(unit=10,file='eps0mat',form='unformatted',status='old')

        call open_file(unit=11,file='epsmat',form='unformatted',status='old',iostat=sig%igamma)
        if(sig%igamma/=0) then ! Gamma-only calculation
          sig%nq1 = 0
        else
          read(11)
          read(11) ngarbage1, nFreq
          if (nFreq/=sig%nFreq) then
            call die('nFreq mismatch between eps0mat and epsmat')
          endif
          read(11)
          read(11)
          read(11)
          read(11)
          read(11)
          read(11) sig%nq1
          read(11)
        
          do iq=1,sig%nq1
            read(11) itrash, nmtx
            read(11)
            if (sig%freq_dep==0.or.sig%freq_dep==1) then
              do ig=1,nmtx
                read(11)
              enddo
            endif
            if (sig%freq_dep==2.or.sig%freq_dep==3) then
              do igp=1,nmtx
                do ig=1,nmtx
                  read(11) ! For epsRDyn
                enddo
#ifdef CPLX
                do ig=1,nmtx
                  read(11) ! For epsADyn
                enddo
#endif
              enddo
            endif
            read(11)
            if(neps<nmtx) neps = nmtx
          enddo
          call close_file(11)
          call open_file(unit=11,file='epsmat',form='unformatted',status='old',iostat=sig%igamma)
        endif
      endif ! peinf%inode==0
#ifdef HDF5
    endif
#endif
#ifdef MPI
    call MPI_Bcast(sig%igamma, 1, MPI_INTEGER, 0 ,MPI_COMM_WORLD, mpierr)
    call MPI_Bcast(sig%nq0, 1, MPI_INTEGER, 0 ,MPI_COMM_WORLD, mpierr)
    call MPI_Bcast(sig%nq1, 1, MPI_INTEGER, 0 ,MPI_COMM_WORLD, mpierr)
    call MPI_Bcast(neps, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, mpierr)
    call MPI_Bcast(sig%nFreq, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, mpierr)
#endif
    if (sig%freq_dep==2.or.sig%freq_dep==3) then
      SAFE_ALLOCATE(sig%dFreqGrid,(sig%nFreq))
      SAFE_ALLOCATE(sig%dFreqBrd,(sig%nFreq))
    endif
  else ! sig%freq_dep>=0.and.sig_freq_dep<=2
    neps = 0
    sig%nFreq = 0
  endif
  if (sig%nq0/=1.and..not.sig%subsample) call die('epscopy_init: nq0/=1')

  POP_SUB(epscopy_init)

end subroutine epscopy_init


subroutine epscopy(crys,gvec,sig,neps,epsmpi,epshead,iunit_eps,fne)
  type (crystal), intent(in) :: crys
  type (gspace), intent(in) :: gvec
  type (siginfo), intent(inout) :: sig
  integer, intent(in) :: neps !< number of G-vectors up to epsilon cutoff defined in sigma.inp
  type (epsmpiinfo), intent(inout) :: epsmpi
  SCALAR, intent(out) :: epshead !< for full frequency, this is the retarded static head.
  integer, intent(in) :: iunit_eps
  character*20, intent(in) :: fne

  SCALAR, allocatable :: eps(:)
  complex(DPC), allocatable :: epsRDyn(:,:)
#ifdef CPLX
  complex(DPC), allocatable :: epsADyn(:,:)
#endif
  logical :: is_static

  PUSH_SUB(epscopy)

  is_static = sig%freq_dep/=2 .AND. sig%freq_dep/=3

  ! FHJ: Initialize files and buffers used by COMM_DISK and COMM_MPI
  if (sig%iwriteint==0) then ! COMM_DISK
    if (peinf%inode==0) then
      call open_file(iunit_eps,file=fne,form='unformatted',status='replace')
    endif
  else ! COMM_MPI
    if (is_static) then
      SAFE_ALLOCATE(epsmpi%eps, (neps,epsmpi%ngpown,sig%nq))
    else
      SAFE_ALLOCATE(epsmpi%epsR, (sig%nFreq,neps,epsmpi%ngpown,sig%nq))
#ifdef CPLX
      SAFE_ALLOCATE(epsmpi%epsA, (sig%nFreq,neps,epsmpi%ngpown,sig%nq))
#endif
    endif
  endif ! COMM_DISK

  ! FHJ: Temp buffers. Only used if it`s COMM_MPI or COMM_DISK and node 0
  if (sig%iwriteint==1.or.peinf%inode==0) then
    if (is_static) then
      SAFE_ALLOCATE(eps, (neps))
    else
      SAFE_ALLOCATE(epsRDyn, (sig%nFreq,neps))
#ifdef CPLX
      SAFE_ALLOCATE(epsADyn, (sig%nFreq,neps))
#endif
    endif
  endif

  !------------------------------------------------
  ! FHJ: Read dielectric matrices from eps0mat(.h5)
  !------------------------------------------------
  call read_epsmat(.true.)
#ifdef MPI
  call MPI_Bcast(epshead, 1, MPI_SCALAR, 0, MPI_COMM_WORLD, mpierr)
#endif
  if (dble(epshead)<1d-3.and.sig%iscreen==0.and.peinf%inode==0) then
    write(0,'(a)') 'WARNING: You are using semiconductor screening, but the'
    write(0,'(a)') 'head of epsilon inverse is very small and seems metallic.'
  endif
  !------------------------------------------------
  ! FHJ: Read dielectric matrices from epsmat(.h5)
  !------------------------------------------------
  if (sig%igamma==0) call read_epsmat(.false.)

#ifdef MPI
  call MPI_Bcast(sig%qgrid,3,MPI_INTEGER,0,MPI_COMM_WORLD,mpierr)
#endif
  if (sig%iwriteint==1) epsmpi%qk(:,:) = sig%qpt(:,:)

  ! FHJ: Free temp. buffers
  if (sig%iwriteint==1.or.peinf%inode==0) then
    if (is_static) then
      SAFE_DEALLOCATE(eps)
    else
      SAFE_DEALLOCATE(epsRDyn)
#ifdef CPLX
      SAFE_DEALLOCATE(epsADyn)
#endif
    endif
  endif
  if (sig%iwriteint==0.and.peinf%inode==0) call close_file(iunit_eps)

  POP_SUB(epscopy)

  return

contains


  subroutine read_epsmat(is_q0)
    logical, intent(in) :: is_q0

    character(len=16) :: fname
    integer :: qoffset, nq_file, iunit
    integer :: freq_dep_flag, nFreq, ng_old, ngq, nmtx
    integer :: ii, ig, igp, nq_tmp, iq, qgrid(3), gg(3), iout, iw, dest
    real(DP) :: ecuts, qk(3), ekin
    real(DP), allocatable :: ekold(:)
#ifdef HDF5
    integer :: nmtxmax
    integer, allocatable :: nmtx_of_q(:), isrtinvdummy(:)
#endif
    integer, allocatable :: isrtq(:), isrtqi(:), isrtold(:), gvecs_old(:,:)
    type(progress_info) :: prog_info !< a user-friendly progress report
    character(len=6) :: sname
    character(len=11) :: sdate

    PUSH_SUB(epscopy.read_epsmat)

    if (is_q0) then
      fname = 'eps0mat'
      qoffset = 0
      nq_file = sig%nq0
      iunit = 10
    else
      fname = 'epsmat'
      qoffset = sig%nq0
      nq_file = sig%nq1
      iunit = 11
    endif
    if (sig%use_hdf5) fname = TRUNC(fname) // '.h5'

    if (peinf%inode==0) then
      SAFE_ALLOCATE(ekold, (gvec%ng))
      SAFE_ALLOCATE(isrtold, (gvec%ng))
      SAFE_ALLOCATE(isrtq, (gvec%ng))
      SAFE_ALLOCATE(isrtqi, (gvec%ng))
    endif

!------------------------------------------------------------------------------
! Read q-grid, q-points, G-vectors and freq. grid
!------------------------------------------------------------------------------
    if (peinf%inode==0) then
#ifdef HDF5
      if (sig%use_hdf5) then
        sname = 'chiGG0'
        sdate = 'nodate'
        call read_eps_grid_sizes_hdf5(ng_old, nq_tmp, ecuts, nFreq, nmtxmax, qgrid, TRUNC(fname))
        if (.not.is_static.and.is_q0) then
          call read_eps_freqgrid_hdf5(nFreq, sig%dFreqGrid, sig%dFreqBrd, TRUNC(fname))
        endif

        SAFE_ALLOCATE(nmtx_of_q, (nq_file))
        call read_eps_qgrid_hdf5(nq_tmp, sig%qpt(:,qoffset+1:), nmtx_of_q, TRUNC(fname))
        ! Note: it seems like the old format stores isort up to ngq, and the
        ! HDF5 stores up to ng_old
        SAFE_ALLOCATE(gvecs_old, (3,ng_old))
        call read_eps_old_gvecs_hdf5(ng_old, gvecs_old, TRUNC(fname))
      else
#endif
        read(iunit) sname, sdate
        read(iunit) freq_dep_flag, nFreq
        read(iunit) qgrid(1:3)
        if (.not.is_static.and.is_q0) then
          read(iunit) sig%dFreqGrid, sig%dFreqBrd
        else
          read(iunit)
        endif
        read(iunit)
        read(iunit)
        read(iunit) ecuts
        read(iunit) nq_tmp, ((sig%qpt(ii,qoffset+iq), ii=1,3), iq=1,nq_tmp)
        read(iunit) ng_old
        backspace(iunit)
        SAFE_ALLOCATE(gvecs_old, (3, ng_old))
        read(iunit) ng_old, ((gvecs_old(ii,iq), ii=1,3), iq=1,ng_old)
#ifdef HDF5
      endif
#endif
      if (nq_tmp/=nq_file) call die('nq mismatch for '//TRUNC(fname))
      ! FHJ: only substitute sig%qgrid (used in voronoi averages) if sig%qgrid
      ! is empty and this is epsmat(.h5), if this file is available
      if (all(sig%qgrid(1:3)==0).and.((.not.is_q0).or.(sig%igamma==1))) then
        sig%qgrid(:) = qgrid(:)
      endif
    endif ! peinf%inode==0


!------------------------------------------------------------------------------
! Bcast q-points and allocate/bcast full-frequency and epsilon buffers
!------------------------------------------------------------------------------
#ifdef MPI
    call MPI_Bcast(sig%qpt(1,qoffset+1), 3*nq_file, MPI_DOUBLE_PRECISION, 0, MPI_COMM_WORLD, mpierr)
#endif
    if (is_q0) then
      sig%q0vec(:) = sig%qpt(:,1)
      ! We need to manually set the q0 point in sig%qpt to zero
      if (.not.sig%subsample) sig%qpt(:,:sig%nq0) = 0d0
      if (.not.is_static) then
#ifdef MPI
        call MPI_Bcast(sig%dFreqGrid,sig%nFreq,MPI_DOUBLE_PRECISION,0,MPI_COMM_WORLD,mpierr)
        call MPI_Bcast(sig%dFreqBrd,sig%nFreq,MPI_DOUBLE_COMPLEX,0,MPI_COMM_WORLD,mpierr)
#endif
      endif
    endif !is_q0


!------------------------------------------------------------------------------
! Loop over all q-points, map G-vectors and read/store dielectric matrices
!------------------------------------------------------------------------------
    call progress_init(prog_info, 'reading '//TRUNC(fname), 'q-point', nq_file)
    do iq=1,nq_file
      call progress_step(prog_info, iq)
      call logit('Storing eps to memory')

      ! FHJ: Map old isrt (from epsilon gspace) to new isrt (to gvec gspace)
      !--------------------------------------------------------=------------
      if (peinf%inode==0) then
#ifdef HDF5
        if (sig%use_hdf5) then
          SAFE_ALLOCATE(isrtinvdummy, (ng_old))
          call read_eps_gvecsofq_hdf5(ng_old,isrtold,isrtinvdummy,ekold,iq,TRUNC(fname))
          SAFE_DEALLOCATE(isrtinvdummy)
          nmtx = nmtx_of_q(iq)
          ngq = ng_old
        else
#endif
          read(iunit) ngq, nmtx, (isrtold(ig),ii,ig=1,ngq)
          read(iunit) (ekold(ig),ig=1,ngq)
          read(iunit) (qk(ii),ii=1,3)
#ifdef HDF5
        endif
#endif
        isrtq = 0
        isrtqi = 0
        qk = sig%qpt(:,iq+qoffset)
        if (is_q0) qk(:) = 0d0
        do ig=1,ngq
          if (ekold(isrtold(ig))<=sig%ecutb) then
            gg(:) = gvecs_old(:, isrtold(ig))
            call findvector(iout, gg, gvec)
            isrtq(ig) = iout
            isrtqi(iout) = ig
            if (ig==1) then
              ! just check the first so we do not waste time
              ekin = DOT_PRODUCT(gvec%components(:,iout)+qk(:), MATMUL(crys%bdot, gvec%components(:,iout)+qk(:)))
              if (abs(ekin-ekold(isrtold(ig))) > TOL_Zero) then
                write(0,*) 'iq = ', iq, ' ig = ',ig, ' qk = ', qk
                write(0,*) 'epsmat: ekold(isrtold(i)) = ', ekold(isrtold(ig)), ' ekin = ', ekin
                call die("Incorrect kinetic energies in epsmat.")
              endif
            endif
          endif  ! if (ekold(isrtold(i))<=sig%ecutb) 
        enddo  ! do ig=1,ngq
      endif  ! if (peinf%inode==0)
#ifdef MPI
      call MPI_Bcast(nmtx,1,MPI_INTEGER,0,MPI_COMM_WORLD,mpierr)
#endif
      if (sig%iwriteint==0) then
        if (peinf%inode==0) write(iunit_eps) gvec%ng, nmtx, &
          isrtq(1:gvec%ng), isrtqi(1:gvec%ng), (sig%qpt(ii,iq+qoffset),ii=1,3)
      else
        if (peinf%inode==0) then
          epsmpi%isrtq(:,iq+qoffset) = isrtq(:)
          epsmpi%isrtqi(:,iq+qoffset) = isrtqi(:)
        endif
        epsmpi%nmtx(iq+qoffset) = nmtx
      endif

#ifdef HDF5
      if (sig%use_hdf5) then
        ! FHJ: Read & store epsinv === HDF5 ===
        !--------------------------------------
        if (sig%iwriteint==1) then ! COMM_MPI below
          if (is_static) then
            call timacc(61,1)
            call read_eps_matrix_par_hdf5(epsmpi%eps(:,:,iq+qoffset), &
              epsmpi%ngpown, epsmpi%ngpown_max, epsmpi%inv_igp_index, nmtx, &
              iq, 1, fname)
            call timacc(61,2)
          else
            call timacc(61,1)
#ifdef CPLX
            call read_eps_matrix_par_f_hdf5(epsmpi%epsR(:,:,:,iq+qoffset), &
              epsmpi%epsA(:,:,:,iq+qoffset), epsmpi%ngpown, epsmpi%ngpown_max, &
              epsmpi%inv_igp_index, nmtx, sig%nFreq, iq, 1, fname)
#else
            call read_eps_matrix_par_f_hdf5(epsmpi%epsR(:,:,:,iq+qoffset), &
              epsmpi%ngpown, epsmpi%ngpown_max, epsmpi%inv_igp_index, nmtx, &
              sig%nFreq, iq, 1, fname)
#endif
            call timacc(61,2)
          endif
          if (is_q0.and.iq==1.and.peinf%inode==0) then
            if (is_static) then
              epshead = epsmpi%eps(1,1,1)
            else
              epshead = epsmpi%epsR(1,1,1,1)
            endif
          endif
        else ! COMM_DISK below
          if (peinf%inode==0) then
            do igp=1,nmtx
              if (is_static) then
                call timacc(61,1)
                call read_eps_matrix_col_hdf5(eps, igp, nmtx, iq, 1, fname)
                call timacc(61,2)
                if (is_q0.and.iq==1.and.igp==1) epshead = eps(1)
                call timacc(62,1)
                write(iunit) (eps(ig),ig=1,nmtx)
                call timacc(62,2)
              else
                call timacc(61,1)
#ifdef CPLX
                call read_eps_matrix_col_f_hdf5(epsRDyn, epsADyn, sig%nFreq, igp, nmtx, iq, 1, fname)
#else
                call read_eps_matrix_col_f_hdf5(epsRDyn, sig%nFreq, igp, nmtx, iq, 1, fname)
#endif
                call timacc(61,2)
                if (is_q0.and.iq==1.and.igp==1) epshead = epsRDyn(1,1)
                call timacc(62,1)
                if (sig%iwriteint==0) then
                  do ig=1,nmtx
                    write(iunit_eps) (epsRDyn(iw,ig),iw=1,sig%nFreq)
                  enddo
#ifdef CPLX
                  do ig=1,nmtx
                    write(iunit_eps) (epsADyn(iw,ig),iw=1,sig%nFreq)
                  enddo
#endif
                endif
                call timacc(62,2)
              endif
            enddo ! ig
          endif ! peinf%inode==0
        endif ! iwriteint==1
      else !HDF5
#endif
        ! FHJ: Read & store epsinv === Fortran binary  ===
        !-------------------------------------------------
        do igp=1,nmtx
          if (is_static) then
            if (peinf%inode==0) then
              call timacc(61,1)
              read(iunit) (eps(ig),ig=1,nmtx)
              call timacc(61,2)
              if (is_q0.and.iq==1.and.igp==1) epshead = eps(1)
              call timacc(62,1)
              if (sig%iwriteint==0) write(iunit_eps) (eps(ig),ig=1,nmtx)
              call timacc(62,2)
            endif

            if (sig%iwriteint==1) then ! COMM_MPI
              call timacc(62,1)
#ifdef MPI
              call MPI_Bcast(eps, neps, MPI_SCALAR, 0, MPI_COMM_WORLD, mpierr)
#endif
              call timacc(62,2)
            endif
          else ! is_static
            if (peinf%inode==0) then
              do ig=1,nmtx
                call timacc(61,1)
                read(iunit) (epsRDyn(iw,ig),iw=1,sig%nFreq) ! Retarded part
                call timacc(61,2)
                if (is_q0.and.iq==1.and.igp==1) epshead = epsRDyn(1,1)
                call timacc(62,1)
                if (sig%iwriteint==0) write(iunit_eps) (epsRDyn(iw,ig),iw=1,sig%nFreq)
                call timacc(62,2)
              enddo
#ifdef CPLX
              do ig=1,nmtx
                call timacc(61,1)
                read(iunit) (epsADyn(iw,ig),iw=1,sig%nFreq) ! Advanced part
                call timacc(61,2)
                call timacc(62,1)
                if (sig%iwriteint==0) write(iunit_eps) (epsADyn(iw,ig),iw=1,sig%nFreq)
                call timacc(62,2)
              enddo
#endif
            endif

            if (sig%iwriteint==1) then ! COMM_MPI
              call timacc(62,1)
#ifdef MPI
              call MPI_Bcast(epsRDyn, sig%nFreq*neps, MPI_COMPLEX_DPC, 0, MPI_COMM_WORLD, mpierr)
#ifdef CPLX
              call MPI_Bcast(epsADyn, sig%nFreq*neps, MPI_COMPLEX_DPC, 0, MPI_COMM_WORLD, mpierr)
#endif
#endif
              call timacc(62,2)
            endif
          endif ! is_static

          if (sig%iwriteint==1) then ! COMM_MPI
            dest = epsmpi%igp_owner(igp)
            call timacc(62,1)        
            if (peinf%pool_rank==dest) then
              if (is_static) then
                epsmpi%eps(:,epsmpi%igp_index(igp),iq+qoffset) = eps(:)
              else
                epsmpi%epsR(:,:,epsmpi%igp_index(igp),iq+qoffset) = epsRDyn(:,:)
#ifdef CPLX
                epsmpi%epsA(:,:,epsmpi%igp_index(igp),iq+qoffset) = epsADyn(:,:)
#endif
              endif ! is_static
            endif ! if (peinf%pool_rank==dest)
            call timacc(62,2)
          endif        
        enddo  ! do ig=1,nmtx
#ifdef HDF5
      endif !HDF5
#endif
    enddo ! iq
    call progress_free(prog_info)
    
!------------------------------------------------------------------------------
! Done!
!------------------------------------------------------------------------------

    if (peinf%inode==0) then
#ifdef HDF5
      if (sig%use_hdf5) then
        SAFE_DEALLOCATE(nmtx_of_q)
      else
#endif
        call close_file(iunit)
#ifdef HDF5
      endif
#endif
      SAFE_DEALLOCATE(ekold)
      SAFE_DEALLOCATE(isrtold)
      SAFE_DEALLOCATE(isrtq)
      SAFE_DEALLOCATE(isrtqi)
    endif

    POP_SUB(epscopy.read_epsmat)

  end subroutine read_epsmat

end subroutine epscopy

end module epscopy_m
