!
! Copyright (C) 2017- Orbital-free DFT group at University of Florida
!
! This program is free software; you can redistribute it and/or modify 
! it under the terms of the GNU General Public License as published by 
! the Free Software Foundation; either version 2 of the License, or 
! any later version. See the file LICENSE in the root directory of the 
! present distribution, or http://www.gnu.org/copyleft/gpl.txt ,
! or contact developers via e-mail: calderin@lcalderin.net, or paper mail:
!
! Quantum Theory Project
! University of Florida
! P.O. Box 118435
! Gainesville, FL 32611
!
! This program is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
! 
! Authors: Lazaro Calderin and Valentin Karasiev
!-----------------------------------------------------------------------
! 

program kgec

! Kubo Greenwood Electron Conductivity (KGEC) code
! Post-processing tool for Quantum Espresso

use qevars
use qe_p_psi, ONLY: p_psi

use apawgm,   ONLY: use_real_spherical_harm, awfc_grad_awfc, &
                    check_p_pswfc, diff_awfc_grad_awfc,       &
                    psi_norm_correction_paw_at,               &
                    agm_free_memory
use kgecvars
use kgecchecks
use kgecpaw
use kgecgradient
use kgsigma1
use kgsigma2
use kgsigma2_w
use parallel
use kgecout
use kgepsilon1
use kgepsilon2
use kgrrae

implicit none

   NAMELIST / kgecpp / outdir, prefix, calculation,          &
                       sigma1_exact, decomposition,           &
                       calculate_sigma2,                      &
                       deltarep, deltawidth, wmin, wmax, nw,  &
                       use_real_spherical_harm,               &
                       check_wfc, check_delta,                &
                       non_local, yessym, yesinv, readgm, de, &
                       npwrecover, writegm, dc, ac,           &
                       sigma1_notintra, sigma2_notintra,      &
                       calculate_epsilon1, calculate_epsilon2,& 
                       calculate_rrae

   ! Initialization
   call mp_startup ( )
   ! 
   ! Get the starting header according to QE
   ! 
   
   call environment_start ( 'KGEC FOR QE' )
   
   
   write(stdout,'(A)') ''
   write(stdout,'(A)') &
'=------------------------------------------------------------------------------='
   write(stdout,*) &
     '                                  KGEC v.2.3.5'
   write(stdout,'(A)') &
'=------------------------------------------------------------------------------='


   ! Main reference
   write(stdout,'(A)') ''
   write(stdout,'(A)') &
   "     Please cite:"
   write(stdout,'(A)') &
   "                L. Calderin, V. V. Karasiev and S. B. Trickey," 
   write(stdout,'(A)') &
   "                 Comput Phys Commun 221  (2017), p. 118-142."  
   
   ! 
   ! Read the KGEC input file
   ! 
   ios=0
   if ( ionode )  then
      !call input_from_file TODO: this has not been needed for a while
      read(5, kgecpp, IOSTAT=ios)
      tmp_dir = trim(outdir)
   endif

   call mp_bcast( ios, ionode_id, world_comm )
   if( ios/=0 ) call errore('kgec', 'reading namelist KGECPP', ABS(ios))

   ! Broadcast and check KGEC specific input variables
   call kgecvars_check

   ! 
   ! Read QE files
   !
   if ( .not. readgm ) call read_file
   if (       readgm ) call read_file_nowfc

   ! wf_collect is now set to false always even if it was true
   ! twfcollect is set to true if that is the case during read_file but is
   ! defined back to false in openfil_pp
   ! so we have to check for collection of wavefunctions right here
#if defined(__QE64)
   if( ( .not. wf_collect ) .and. ( .not. readgm ) ) &
#else
   if( ( .not. twfcollect ) .and. ( .not. readgm ) ) &
#endif
      call errore( 'kgec:','Wave functions not collected',           3 )

   if ( .not. readgm ) call openfil_pp

   ! Check if KGEC requirements for the QE calculation were met and print out
   call kgec_check_requirements

   ! Print KGEC variables
   call kgecvars_print

   ! Get the frequency grid in Ha
   call frequency_grid()

   ! Initialization for bands parallelization
   call parallel_bands_init

   write(stdout,*) ' '
   write(stdout,*) &
         '    --------------------------------------------------------------'
   write(stdout,*) &
         "                           Calculations "
   write(stdout,*) &
         '    --------------------------------------------------------------'

   READGMATRIX: SELECT CASE ( readgm )
   CASE ( .false. ) 

      !
      ! Calculate wave function re-normalization factors in case of PAW sets
      !
      select case ( only_paw  )
      case( .true. )

         call kgec_paw_renorm

      end select

      !
      ! Checking ortho-normalization of the (pseudo) wave functions
      !
      CHECKORTHO: SELECT CASE ( check_wfc )
      CASE (.true.)

         select case( only_paw ) ! According to pseudopotentials
         case( .false. )         ! Norm conserving pseudos

            call kgec_check_orthonormalization_nc

         case(.true.) ! PAW sets

            call kgec_check_orthonormalization_paw

         end select 

      END SELECT CHECKORTHO

      !============================================================================
      ! Gradient matrix elements over the 'solid' pseudo wave function \Psi_{n,k}
      !                                    psi_grad_psi
      !============================================================================

      call kgec_gradient

      !=========================================================================
      !           Calculate PAW correction to the gradient matrix elements
      !                          ( Grad_paw_correction )
      !=========================================================================

      select case ( only_paw )
      case( .true. )

         call kgec_gradient_paw_correction

      end select ! PAW correction

   CASE ( .true. )

      call kgec_readgm

   END SELECT READGMATRIX


   ! Write gradient matrix elements
   select case ( writegm )
   case( .true. )
      
     call kgec_writegm

   end select


   call mp_barrier( inter_pool_comm )

   ! There is no more parallelization over G from this point on, so
   ! the -np cores would be wasted. Instead we recover them for use 
   ! in the bands parallelization with two caveat: the name of the comm
   ! must be inter_bgrp_comm but it must be an intra comm. The first 
   ! condition is to keep the same name for the comm of bands 
   ! parallelization and the second to be
   ! able to communicate with all process of one band group.

   select case( npwrecover )
   case(.true.)

     call recover_processes

   end select

   ! 
   ! Calculate the exact sum rule ( no dirac Delta function involved ) 
   ! 
   if ( .not. readgm ) call sumrule_exact

   ! check delta function representation
   
   select case ( check_delta )
   case( .true. )

      write(stdout,*) ''
      write(stdout,'(A)',advance='no') &
                          '     Calculating checks for delta function rep. ... '
      call check_delta_sub

      write(stdout,*) 'done.'

   end select

   write(stdout,*) ' '
   write(stdout,'(A)',advance='no') '     Calculating sigma1 ... '

   SIGMA1AC: SELECT CASE ( ac )
   CASE(.true.)  

   SIGMA1TENSORORTRACE: SELECT CASE ( calculation )
   CASE ( 'tensor' )

      SIGMAONEEXACT: SELECT CASE ( sigma1_exact )

      CASE ( .false. ) ! Approximated formulas

         select case ( deltarep )
         case ( '1l' )
            call sigma1_approx_1lorentzian
         case ( '2l' )
            call sigma1_approx_2lorentzian
         case ( '1g'  )
            call sigma1_approx_1gaussian
         case ( '2g'  )
            call sigma1_approx_2gaussian
         end select

      CASE ( .true. ) ! Exact formulas

         select case ( decomposition )
         case( .false. )

            select case ( deltarep )  ! Exact formulas
            case ('1l')
               call sigma1_exact_1lorentzian
            case ('2l')
               call sigma1_exact_2lorentzian
            case ('1g')
               call sigma1_exact_1gaussian
            case ('2g')
               call sigma1_exact_2gaussian
            end select

         case(.true.) ! Separate degenerate, intra and inter band contributions

            select case ( deltarep )  ! Exact formula
            case ('1l')
               call sigma1_exact_decomp_1lorentzian
            case ('2l')
               call sigma1_exact_decomp_2lorentzian
            case ('1g')
               call sigma1_exact_decomp_1gaussian
            case ('2g')
               call sigma1_exact_decomp_2gaussian
            end select

         end select

      END SELECT SIGMAONEEXACT

    CASE( 'atrace' )

      SIGMAONETRACEEXACT: SELECT CASE ( sigma1_exact )
      CASE ( .false. ) ! Approximated formula

         select case ( deltarep )
         case ('1l')
            call sigma1_atrace_approx_1lorentzian
         case ('2l')
            call sigma1_atrace_approx_2lorentzian
         case ('1g')
            call sigma1_atrace_approx_1gaussian
         case ('2g')
            call sigma1_atrace_approx_2gaussian
         end select

      CASE ( .true. ) ! Exact formula

         select case ( decomposition )
         case( .false. )

            select case ( deltarep )
            case ('1l')
               call sigma1_atrace_exact_1lorentzian
            case ('2l')
               call sigma1_atrace_exact_2lorentzian
            case ('1g')
               call sigma1_atrace_exact_1gaussian
            case ('2g')
               call sigma1_atrace_exact_2gaussian
            end select ! Delta function representation 

         case(.true.) ! Separate degenerate, intra and inter band contributions

            select case ( deltarep )  ! Exact formula
            case ('1l')
               call sigma1_atrace_exact_decomp_1lorentzian
            case ('2l')
               call sigma1_atrace_exact_decomp_2lorentzian
            case ('1g')
               call sigma1_atrace_exact_decomp_1gaussian   
            case ('2g')
               call sigma1_atrace_exact_decomp_2gaussian   
            end select ! Delta function representation 

         end select ! Decomposition

      END SELECT SIGMAONETRACEEXACT

   END SELECT SIGMA1TENSORORTRACE

   END SELECT SIGMA1AC
   
   write(stdout,'(A)') 'done.'


   SIGMA1DC: SELECT CASE ( dc )
   CASE( .true. )

   SIGMA1TENSORORTRACEDC: SELECT CASE ( calculation )
   CASE ( 'tensor' )

      SIGMAONEEXACTDC: SELECT CASE ( sigma1_exact )

      CASE ( .false. ) ! Approximated formulas

         select case ( deltarep )
         case ( '1l' )
            write(stdout, '(A)') &
            '    WARNING: Not possible to calculate the DC component with & 
            &the approximated formula and one Lorentzian, using two instead.'
            call sigma1_approx_2lorentzian_DC
         case ( '2l' )
            call sigma1_approx_2lorentzian_DC
         case ( '1g'  )
            write(stdout, '(A)') &
            '    WARNING: Not possible to calculate the DC component with & 
            &the approximated formula and one Gaussian, using two instead.'
            call sigma1_approx_2gaussian_DC
         case ( '2g'  )
            call sigma1_approx_2gaussian_DC
         end select

      CASE ( .true. ) ! Exact formulas

         select case ( decomposition )
         case( .false. )

            select case ( deltarep )  ! Exact formulas
            case ('1l','2l')
               call sigma1_exact_lorentzian_DC
            case ('1g','2g')
               call sigma1_exact_gaussian_DC
            end select

         case(.true.) ! Separate degenerate, intra and inter band contributions

            select case ( deltarep )  ! Exact formula
            case ('1l', '2l')
               call sigma1_exact_decomp_lorentzian_DC
            case ('1g', '2g')
               call sigma1_exact_decomp_gaussian_DC
            end select

         end select

      END SELECT SIGMAONEEXACTDC

    CASE( 'atrace' )

      SIGMAONETRACEEXACTDC: SELECT CASE ( sigma1_exact )
      CASE ( .false. ) ! Approximated formula

         select case ( deltarep )
         case ('1l')
            write(stdout, '(A)') &
            '    WARNING: Not possible to calculate the DC component with & 
            &the approximated formula and one Lorentzian, using two instead.'
            call sigma1_atrace_approx_2lorentzian_DC
         case ('2l')
            call sigma1_atrace_approx_2lorentzian_DC
         case ('1g')
            write(stdout, '(A)') &
            '    WARNING: Not possible to calculate the DC component with & 
            &the approximated formula and one Gaussian, using two instead.'
            call sigma1_atrace_approx_2gaussian_DC
         case ('2g')
            call sigma1_atrace_approx_2gaussian_DC
         end select

      CASE ( .true. ) ! Exact formula

         select case ( decomposition )
         case( .false. )

            select case ( deltarep )
            case ('1l','2l')
               call sigma1_atrace_exact_lorentzian_DC
            case ('1g','2g')
               call sigma1_atrace_exact_gaussian_DC
            end select 

         case(.true.) ! Separate degenerate, intra and inter band contributions

            select case ( deltarep )  ! Exact formula
            case ('1l','2l')
               call sigma1_atrace_exact_decomp_lorentzian_DC
            case ('1g','2g')
               call sigma1_atrace_exact_decomp_gaussian_DC
            end select 
            
         end select ! Decomposition

      END SELECT SIGMAONETRACEEXACTDC

   END SELECT SIGMA1TENSORORTRACEDC
   
   
   END SELECT SIGMA1DC


! To implement
!   SELECT CASE ( dc_all )
!   CASE( 'w', 'f', 'a' )
!      write(stdout,*) ' '
!      write(stdout,'(A)',advance='no') &
!   '     Calculating the DC using both the exact and approx sigma1 formula ...'
!!      call sigma1_dc_all_deltarep
!      write(stdout,'(A)') 'done.'
!   END SELECT

   !===========================================================================
   !                          Conductivity tensor (sigma2)
   !===========================================================================

   SIGMA2CALCULATION: SELECT CASE ( calculate_sigma2 )
   CASE(.true.)
   write(stdout,'(A)',advance='no') '     Calculating sigma2 ... '

      SIGMA2TENSORORTRACE: SELECT CASE ( calculation )
      CASE ( 'tensor' )
 
         select case ( decomposition )
         case( .false. )
            call sigma2_sub
         case(.true.) ! Separate degenerate, intra and inter-band contributions
            call sigma2_decomp
         end select 

      CASE( 'atrace' )

         select case ( decomposition )
         case( .false. )
            call sigma2_atrace_sub
         case(.true.) ! Separate degenerate, intra and inter-band contributions
            call sigma2_atrace_decomp
         end select ! Decomposition

      END SELECT SIGMA2TENSORORTRACE

      write(stdout,'(A)') 'done.'

   END SELECT SIGMA2CALCULATION
   
   IF(IONODE) THEN

   !
   ! Sum rules
   !

   ! Sum rule (no delta function)
   if ( .not. readgm ) then
      write(stdout,*) ' ' 
      write(stdout,'(A,f10.3)') "     Sum rule (exact)          :", sumrulef
   endif

   ! Sum rule (with delta function)
   if ( AC ) then
     call sumrule_integral
     write(stdout,'(A,f10.3)') "     Sum rule (int. over freq.):", sumruleintegral
     write(stdout,*)
   endif

   !
   ! SIGMA OUTPUT
   !
   if ( .not. sigma2_via_eps1mod ) call kgec_output_sigma

   ENDIF ! IONODE

   !===========================================================================
   !                Real Part of Dielectric tensor (epsilon_1)
   !===========================================================================
   EPSILON1CALCULATION: SELECT CASE ( calculate_epsilon1 )
   CASE(.true.)
      write(stdout,'(A)') ' '
      if( .not. sigma2_via_eps1mod ) then
           write(stdout,'(A)',advance='no') '     Calculating epsilon1 ... '
      else
           write(stdout,'(A)',advance='no') '     Calculating sigma2 and epsilon1 ... '
      endif
 
      EPSILON1TENSORORTRACE: SELECT CASE ( calculation )
      CASE ( 'tensor' )

         select case ( ac )
         case(.true.) 
            select case ( decomposition )
            case( .false. )
               call epsilon1_tensor
               if ( sigma2_via_eps1mod ) call epsilon1_tensor_sub_sigma2
            case(.true.) ! Separate degenerate, intra and inter-band contributions
               call epsilon1_tensor_decomp
               if ( sigma2_via_eps1mod ) call epsilon1_tensor_decomp_sigma2
            end select 

         end select ! ac

         select case ( dc )
         case(.true.) 

            select case ( decomposition )
            case( .false. )
               call epsilon1_dc_tensor
            case(.true.) ! Separate degenerate, intra and inter-band contributions
               call epsilon1_dc_tensor_decomp
            end select 

         end select ! dc

      CASE( 'atrace' )

         select case ( ac )
         case( .true. )

            select case ( decomposition )
            case( .false. )
               call epsilon1_atrace_sub
               if ( sigma2_via_eps1mod ) call epsilon1_atrace_sub_sigma2
            case(.true.) ! Separate degenerate, intra and inter-band contributions
               call epsilon1_atrace_decomp
               if ( sigma2_via_eps1mod ) call epsilon1_atrace_decomp_sigma2
            end select ! Decomposition

         end select ! ac

         select case ( dc )
         case( .true. )

            select case ( decomposition )
            case( .false. )
print*,"KGEC"
               call epsilon1_dc_atrace_sub
            case(.true.) ! Separate degenerate, intra and inter-band contributions
               call epsilon1_dc_atrace_decomp
            end select ! Decomposition

         end select ! dc

      END SELECT EPSILON1TENSORORTRACE

      write(stdout,'(A)') 'done.'

      if ( IONODE .and. sigma2_via_eps1mod ) then
         calculate_sigma2=.true. ! This is to print out sigma2
         call kgec_output_sigma
      endif

       
   END SELECT EPSILON1CALCULATION


   !===========================================================================
   !               Imaginary Part of Dielectric tensor (epsilon_2)
   !===========================================================================

   EPSILON2CALCULATION: SELECT CASE ( calculate_epsilon2 )
   CASE(.true.)
   write(stdout,*) ' '
   write(stdout,'(A)',advance='no') '     Calculating epsilon2 ... '

      EPSILON2TENSORORTRACE: SELECT CASE ( calculation )
      CASE ( 'tensor' )
 
         select case ( ac )
         case( .true. )

            select case ( decomposition )
            case( .false. )
               call epsilon2_tensor
            case( .true. ) ! Separate degenerate, intra and inter-band contributions
               call epsilon2_tensor_decomp
            end select 
         end select ! ac

         ! Epsilon2 DC is infinity for all components
 
      CASE( 'atrace' )

         select case ( ac )
         case( .true. )
 
            select case ( decomposition )
            case( .false. )
               call epsilon2_atrace_sub
            case( .true. ) ! Separate degenerate, intra and inter-band contributions
               call epsilon2_atrace_decomp
            end select     ! Decomposition

         end select ! ac

         ! Epsilon2 DC is infinity for all components

      END SELECT EPSILON2TENSORORTRACE

      write(stdout,'(A)') 'done.'

   END SELECT EPSILON2CALCULATION

   !
   ! EPSILON OUTPUT
   !

   if ( IONODE .and. calculate_epsilon1 .or. calculate_epsilon2 ) &
       call kgec_output_epsilon

   !
   ! Refraction, Reflection and Absorption Coefficients 
   !

   select case( calculate_rrae )
   case( .true. )
      write(stdout,*) ''
      write(stdout,'(A)',advance='no') &
      '     Calculating refraction, reflection and absorption coeffs. ... '

      if ( AC ) call rrae

      write(stdout,'(A)') 'done.'

      call kgec_output_rrae

   end select




   write(stdout,*) ""
   write(stdout,*) &
          '    ----------------------------------------------------------------'
   write(stdout,*) &
          '           Kubo Greenwood conductivity calculation finished'
   write(stdout,*) &
          '    ----------------------------------------------------------------'

   call environment_end("KGEC FOR QE")

   call mp_global_end()

end program kgec
