!  Copyright (C) 2017 The authors of Py-ChemShell
!
!  This file is part of Py-ChemShell.
!
!  Py-ChemShell is free software: you can redistribute it and/or modify
!  it under the terms of the GNU Lesser General Public License as
!  published by the Free Software Foundation, either version 3 of the
!  License, or (at your option) any later version.
!
!  Py-ChemShell 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 Lesser General Public License for more details.
!
!  You should have received a copy of the GNU Lesser General Public
!  License along with Py-ChemShell.  If not, see
!  <http://www.gnu.org/licenses/>.


! you.lu@stfc.ac.uk
! YL: rewritten 2021 according to NWChem 7.0.2

module ChemShellNWChemModule

    use iso_c_binding
    use DL_PY2F

    implicit none

    private

    public :: interface_nwchem, register_nwchem
    public :: chemsh_nwchem_input, chemsh_nwchem_output
    public :: fragmovecs, NWChemObj

    integer(c_int) :: c_debug

    type NWChemType
        private
        integer                   :: debug
        integer                   :: dft
        integer                   :: rtdb
        integer                   :: natoms
        integer                   :: nbqs
        integer                   :: necps
        integer                   :: nstates
        character(len=:), pointer :: name
        real(kind=8), allocatable, public :: excitation_energies(:)
        real(kind=8), allocatable, public :: excitation_oscillator_s(:)
        real(kind=8), allocatable, public :: excitation_wavelengths(:)
        character(len=1), allocatable, public :: excitation_symmetries(:)
        character(len=1), allocatable, public :: excitation_spins(:)
        contains
            procedure, public  :: run => private_run
            procedure, private :: private_run
    endtype NWChemType

    type(NWChemType)          :: NWChemObj
    type(dictType)  , pointer :: PyTheory, PyFrag, PyBQs, PyResult, PyECPs
    integer(c_int)            :: nwchem_handle
    integer                   :: nwchem_task

    logical                   :: fragmovecs
    logical                   :: calc_grads

    character(len=:), pointer :: chemsh_nwchem_input, chemsh_nwchem_output

    contains

    subroutine register_nwchem(debug) bind(c, name='f_register_nwchem')

        use ChemShellParallelModule, only : chemsh_parallel

        implicit none

        integer(c_int), intent(in) :: debug

        type(c_funptr) :: nwchem_pointer

        c_debug = debug

        nwchem_pointer = c_funloc(wrap_nwchem)
        nwchem_handle  = chemsh_parallel%get_function_handle(nwchem_pointer)

    endsubroutine register_nwchem
    integer(c_int) function interface_nwchem(theoryPtr, resultPtr, calcGradients, debug) result(ierror) bind(c)

        implicit none

        type(PyType)              , intent(in) :: theoryPtr
        type(PyType)              , intent(in) :: resultPtr
        logical(c_bool)           , intent(in) :: calcGradients
        integer(c_int)            , intent(in) :: debug

        type(PyType)    , pointer              :: fragPtr, bqsPtr, ecpsPtr
        integer                                :: natoms, nbqs, necps, nstates
        integer                                :: code
        integer         , external             :: nodeid
        character(len=:), pointer              :: method
        character(len=8), pointer              :: names(:)

        character(len=:), pointer              :: cbuff
        logical(c_bool)                        :: clbuff

        ierror = 1

        ! YL 06/08/2019: sometimes we need to switch it off e.g. if numerical
        !                gradients are required in Py-ChemShell
        calc_grads = calcGradients

        ! initialise dict of theory object
        allocate(PyTheory, source=ptr2dict(theoryPtr))
        
        ! initialise dict of resultObj
        allocate(PyResult, source=ptr2dict(resultPtr))

        ! set up dict of fragObj
        call PyTheory%get('frag' , fragPtr)
        call PyTheory%get('necps', necps)
        call PyTheory%get('nstates', nstates)
        if(necps.gt.0) then
            call PyTheory%get('ecps', ecpsPtr)
            allocate(PyECPs, source=ptr2dict(ecpsPtr))
        endif
        allocate(PyFrag, source=ptr2dict(fragPtr))

        ! number of atoms
        ! YL: changed meaning of natoms and nbqs to stay consistent with PCS definitions:
        ! natoms does NOT contain nbqs; and in PCS QM nbqs already contains nshells
        call PyFrag%get('npts'   , natoms)
        call PyFrag%get('nbqs'   , nbqs)

        ! N.B.: natoms include separate ECPs in NWChem
        NWChemObj%natoms  = natoms + necps
        NWChemObj%nbqs    = nbqs
        NWChemObj%necps   = necps
        NWChemObj%nstates = nstates
        NWChemObj%debug   = debug
        allocate(NWChemObj%excitation_energies(nstates))
        allocate(NWChemObj%excitation_oscillator_s(nstates))
        allocate(NWChemObj%excitation_wavelengths(nstates))
        allocate(NWChemObj%excitation_symmetries(nstates))
        allocate(NWChemObj%excitation_spins(nstates))

        call PyFrag%get('bqs', bqsPtr)
        allocate(PyBQs, source=ptr2dict(bqsPtr))

        ! input/output filenames
        ! YL NOTE: the filenames are unavailable for the slave nodes
        call PyTheory%get('input' , chemsh_nwchem_input)
        call PyTheory%get('output', chemsh_nwchem_output)

        ! fragment molecular orbitals
        call PyTheory%get('fragmovecs', cbuff)
        ! not string value given to key "fragmovecs"
        if(len_trim(cbuff).eq.255) then
            call PyTheory%get('fragmovecs', clbuff)
            fragmovecs = clbuff
        ! a string value, e.g. "atoms"
        else
            fragmovecs = .true.
        endif

! YL: bytes array are still unavailable in py2f
!        allocate(names(natoms))
!        call PyFrag%get('names', names)

        call PyTheory%get('method', method)

        ierror = NWChemObj%run()

        deallocate(PyTheory, PyResult, PyFrag, PyBQs)
        ! YL 26/07/2021: on some systems, e.g., ARCHER2, it crashes randomly if PyECPs isn't deallocated
        if(necps.gt.0) then
            deallocate(PyECPs)
        endif

    endfunction interface_nwchem
    integer function private_run(this) result(ierror)

        use ChemShellParallelModule, only : chemsh_parallel

        implicit none

        class(NWChemType), intent(inout) :: this

        if (chemsh_parallel%is_parallel()) then
            ierror = chemsh_parallel%exec_function(nwchem_handle, c_debug)
        else
            ierror = wrap_nwchem()
        endif

    endfunction private_run
    ! also called by replica processes
    integer function wrap_nwchem() result(ierror)! bind(c)

        use ChemShellParallelModule, only : chemsh_parallel

        implicit none

        real(kind=8)            :: energy
        real(kind=8), pointer   :: coords(:,:)     => null()
        real(kind=8), pointer   :: charges(:)      => null()
        real(kind=8), pointer   :: bq_coords(:,:)  => null()
        real(kind=8), pointer   :: bq_charges(:)   => null()
        real(kind=8), pointer   :: bq_grads(:,:)  => null()
!        real(kind=8), pointer   :: ecp_coords(:,:) => null()
!        real(kind=8), pointer   :: ecp_charges(:)  => null()
        real(kind=8), pointer   :: gradients(:,:)  => null()
        integer                 :: rtdb, i, j
        integer                 :: natoms, nbqs, necps, nstates, debug
        integer                 :: ncoords, ncharges
        integer     , external  :: chemshell_nwchem

        ! size of integer
        integer     , parameter :: nbytes = 8
        real(kind=8), allocatable :: dummies(:)

        ierror = 1

        if(chemsh_parallel%is_root()) then
            rtdb    = NWChemObj%rtdb
            ! N.B.: natoms includes separate ECPs in NWChem
            natoms  = NWChemObj%natoms
            nbqs    = NWChemObj%nbqs
            necps   = NWChemObj%necps
            nstates = NWChemObj%nstates
            debug   = NWChemObj%debug
        endif

        ! broadcast the code and rtdb label
        ! NB: in ChemShell we only pass in number of *bytes* rather than number of data
        ierror = chemsh_parallel%broadcast(natoms, nbytes, 0)
        ierror = chemsh_parallel%broadcast(nbqs  , nbytes, 0)
        ierror = chemsh_parallel%broadcast(debug , nbytes, 0)
        ! YL 09/03/2021: not every compiler initialise integers as 0 so we must broadcast necps, too
        ierror = chemsh_parallel%broadcast(necps , nbytes, 0)
        call flush(6)

        ! initialise
        ! set geometry
        ! N.B.: natoms includes separate ECPs in NWChem
        allocate(coords(3, natoms))
        allocate(charges(natoms))
        allocate(bq_coords(3, nbqs))
        allocate(bq_charges(nbqs))
!        allocate(ecp_coords(3, necps))
!        allocate(ecp_charges(necps))

        ncoords = nbytes*natoms*3

        ! broadcast coords, bq_coords, charges, and bq_charges
        if(chemsh_parallel%is_root()) then
            call PyFrag%get('xyz'    , coords(:,1:natoms-necps))
            call PyFrag%get('charges', charges(1:natoms-necps))
            call PyBQs%get('npts', nbqs)
            if(nbqs.gt.0) then
                call PyBQs%get('xyz'    , bq_coords)
                call PyBQs%get('charges', bq_charges)
            endif
            if(necps.gt.0) then
                call PyECPs%get('xyz'    , coords(:,natoms-necps+1:natoms))
                call PyECPs%get('charges', charges(natoms-necps+1:natoms))
            endif
        endif
        ierror = chemsh_parallel%broadcast(coords , nbytes*natoms*3, 0)
        ierror = chemsh_parallel%broadcast(charges, nbytes*natoms  , 0)
        call flush(6)
        if(nbqs.gt.0) then
            ierror = chemsh_parallel%broadcast(bq_coords , nbytes*nbqs*3  , 0)
            ierror = chemsh_parallel%broadcast(bq_charges, nbytes*nbqs    , 0)
        endif

        ! gradients and BQ gradients
        allocate(gradients(3, natoms))
        allocate(bq_grads(3, nbqs))

        ierror = chemshell_nwchem(natoms,     &
                                  nbqs,       &
                                  coords,     &
                                  charges,    &
                                  bq_coords,  &
                                  bq_charges, &
                                  bq_grads,   &
                                  energy,     &
                                  gradients,  &
                                  rtdb,       &
                                  debug)

        ! NWChem returned OK
        if(ierror.ne.0) then
            print *, ">>> NWChem has failed: error code =", ierror
            call flush(6)
        endif

        ierror = chemsh_parallel%broadcast(rtdb, nbytes, 0)
        NWChemObj%rtdb = rtdb
        ! unset geometry
        deallocate(coords, charges, bq_coords, bq_charges)
        if(chemsh_parallel%is_root()) then
            ! energy
            ! YL 01/08/2021: cast to a precision of 1D-12 to minimise numerical errors which lead
            !                to unstable geometry optimisations (for gfortran, e.g., on ARCHER2)
            energy = anint(energy*(10.d0**12), 8)/(10.d0**12)
            call PyResult%set('energy', (/ energy /))
            ! gradients
            if(calc_grads) then
                ! YL 01/08/2021: cast to a precision of 1D-12 to minimise numerical errors which lead
                !                to unstable geometry optimisations (for gfortran, e.g., on ARCHER2)
                gradients = anint(gradients*(10.d0**12), 8)/(10.d0**12)
                ! YL 26/07/2021: on some systems it is safer than call PyResult%set('gradients', gradients)
                call PyResult%set('gradients', gradients(1:3,1:natoms-necps))
            endif
            ! BQ gradients
            if(nbqs.gt.0) then
                ! YL 01/08/2021: cast to a precision of 1D-12 to minimise numerical errors which lead
                !                to unstable geometry optimisations (for gfortran, e.g., on ARCHER2)
                bq_grads = anint(bq_grads*(10.d0**12), 8)/(10.d0**12)
                call PyBQs%set('gradients', bq_grads)
            endif
            if(necps.gt.0) then
!                do i = 1, 3
!                    do j = natoms-necps+1, natoms
!                        gradients(i,j) = anint(gradients(i,j)*(10.d0**12), 8)/(10.d0**12)
!                    enddo
!                enddo
                call PyECPs%set('gradients', gradients(1:3,natoms-necps+1:natoms))
            endif
            if(nstates.gt.0) then
                allocate(dummies(nstates))
                dummies = 0.0
                ! YL 16/12/2023: we use some dummy data to align with the Python result datatype
                call PyResult%set('excitations', transpose(reshape((/ NWChemObj%excitation_energies    , &
                                                                      NWChemObj%excitation_oscillator_s, &
                                                                      NWChemObj%excitation_wavelengths , &
                                                                      dummies, dummies /),               &
                                                                   (/nstates,5/))))
                deallocate(dummies)
            endif
        endif
        deallocate(gradients, bq_grads)
        if(allocated(NWChemObj%excitation_energies)) then
            deallocate(NWChemObj%excitation_energies)
        endif
        if(allocated(NWChemObj%excitation_oscillator_s)) then
            deallocate(NWChemObj%excitation_oscillator_s)
        endif
        if(allocated(NWChemObj%excitation_wavelengths)) then
            deallocate(NWChemObj%excitation_wavelengths)
        endif
        if(allocated(NWChemObj%excitation_symmetries)) then
            deallocate(NWChemObj%excitation_symmetries)
        endif
        if(allocated(NWChemObj%excitation_spins)) then
            deallocate(NWChemObj%excitation_spins)
        endif
!        deallocate(ecp_coords, ecp_charges)

    endfunction wrap_nwchem

endmodule ChemShellNWChemModule

