!  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/>.

module ChemShellDLPOLYModule

    use, intrinsic :: iso_c_binding
    use DL_PY2F 
    use ChemShellParallelModule
    use DLPOLYModule

    implicit none
    private

    public register_dl_poly, interface_dl_poly, dl_poly_wrap
    public PyFrag, PyShells, PyOptions, nshells

    type(dictType) , pointer :: PyFrag, PyShells, PyResult, PyOptions
    integer                  :: natoms, nshells

    integer        :: dl_poly_task
    integer(c_int) :: c_debug
    integer(c_int) :: dl_poly_handle

    contains

      
    subroutine register_dl_poly(debug) bind(c, name="f_register_dl_poly")

        use ChemShellParallelModule, only : chemsh_parallel

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

        type(c_funptr) :: dl_poly_pointer

        c_debug = debug

        dl_poly_pointer = c_funloc(dl_poly_wrap)
        dl_poly_handle = chemsh_parallel%get_function_handle(dl_poly_pointer)

    endsubroutine register_dl_poly
    ! called by master process only
    subroutine interface_dl_poly(fragPtr, resultPtr, optionsPtr) bind(c)
 
        use ChemShellParallelModule, only : chemsh_parallel

        implicit none
 
        type(PyType)  , intent(in) :: fragPtr
        type(PyType)  , intent(in) :: resultPtr
        type(PyType)  , intent(in) :: optionsPtr

        integer                    :: ierror
        type(PyType)  , pointer    :: shellsPtr

        ! set up dict of fragObj
        allocate(PyFrag, source=ptr2dict(fragPtr))
        call PyFrag%get('npts', natoms)
        call PyFrag%get('shells', shellsPtr)
        allocate(PyResult, source=ptr2dict(resultPtr))
        allocate(PyOptions, source=ptr2dict(optionsPtr))

        ! set up dict of shellsObj and save nshells
        allocate(PyShells, source=ptr2dict(shellsPtr))
        call PyShells%get('npts', nshells)
 
        if (chemsh_parallel%is_parallel()) then
           ierror = chemsh_parallel%exec_function(dl_poly_handle, c_debug)
        else
           ierror = dl_poly_wrap()
        endif
 
        if(associated(PyFrag)) then
            deallocate(PyFrag)
        endif
        if(associated(PyShells)) then
            deallocate(PyShells)
        endif
        if(associated(PyResult)) then
            deallocate(PyResult)
        endif
        if(associated(PyOptions)) then
            deallocate(PyOptions)
        endif
 
    endsubroutine interface_dl_poly
    ! also called by slave processes
    integer function dl_poly_wrap() result(ierror)

        use ChemShellParallelModule, only : chemsh_parallel
        use DLPOLYModule           , only : dl_poly_gradients, dl_poly_global_indices, control_filename, output_filename

        integer                   :: itype, lenbuf, ifrom, itag
        logical                   :: opened
        integer                   :: i, shapebuff(2), nnodes, node
        integer(kind=4) , pointer :: natoms_per_domain(:), global_indices_per_node(:)
        real(kind=8)    , pointer :: gradients(:,:), gbuff(:,:)
        character(len=:), pointer :: strbuff

        ! size of int for GULP
        lenbuf = 4
        ifrom  = 0
        itag   = 66

        if(chemsh_parallel%is_root()) then
            call PyOptions%get('output', strbuff)
            output_filename  = trim(strbuff)
! YL 17/08/2021: we can't yet get values out from a tuple so it's hardcoded for the time being
!            call PyOptions%get('input' , strbuff)
            control_filename = '_dl_poly.inp'
        endif

        ! number of nodes per subset (rather than workgroup!)
        nnodes = chemsh_parallel%nnodes_subset
        node   = chemsh_parallel%node_subset

        ierror = chemsh_parallel%broadcast(natoms, lenbuf, ifrom)

!        include "dl_poly_routines.f90"
        allocate(gradients(3,natoms))
        allocate(gbuff(3,natoms))
        allocate(global_indices_per_node(natoms))
        allocate(natoms_per_domain(nnodes))
        gradients = 0.0
        natoms_per_domain = 0
        
        ! share task control flag
        itype = 402

!        ierror = parallel%broadcast(itype, dl_poly_task, lenbuf, ifrom, 0)

        ! YL 30/11/2018: dl_poly_gradients is allocated by the DL_POLY main routine
!        allocate(dl_poly_gradients(3,natoms), stat=ierror)

        if(chemsh_parallel%is_subset()) then
            ! cast to system integer kind
            call dl_poly(int(chemsh_parallel%comm_subset))
        endif

        ! YL 24/11/2018: I don't know how it worked before, but in DL_POLY 4 the forces are held per
        !                domain (the fxx, fyy, & fzz arrays) and one domain has no idea about others'
        !                sizes, so we must creat a full list of natoms/domain to gather the gradients
        shapebuff = shape(dl_poly_gradients)
        if(chemsh_parallel%is_root()) then
            natoms_per_domain(1) = shapebuff(2)
            do i = 1, nnodes-1
                ierror = chemsh_parallel%receive(natoms_per_domain(i+1), lenbuf, i, itag)
            enddo
        else
            ! YL 21/11/2021: restrained to the MM subset otherwise Cray's MPICH (GNU>=10.2) segfaults
            !                on the remaining processes
            if(chemsh_parallel%is_subset()) then
                ierror = chemsh_parallel%send(shapebuff(2), lenbuf, 0, itag)
            endif
        endif

        ! get the whole array of indices per node
        if(chemsh_parallel%is_root()) then
            global_indices_per_node(1:natoms_per_domain(1)) = dl_poly_global_indices
            do i = 1, nnodes-1
                ierror = chemsh_parallel%receive(global_indices_per_node(sum(natoms_per_domain(1:i))+1:sum(natoms_per_domain(1:i+1))), 4*natoms_per_domain(i+1), i, itag)
            enddo
        else
            ! YL 21/11/2021: restrained to the MM subset otherwise Cray's MPICH (GNU>=10.2) segfaults
            !                on the remaining processes
            if(chemsh_parallel%is_subset()) then
                ierror = chemsh_parallel%send(dl_poly_global_indices, 4*shapebuff(2), 0, itag)
            endif
        endif

        ierror = chemsh_parallel%broadcast(natoms_per_domain, lenbuf*nnodes, 0)

        ! copy the domain's gradients to the buffer in the nodes order
        ! YL 21/11/2021: we restrain this operation to the MPI subset because otherwise GNU-compiled
        !                code throws out "Fortran runtime error: Array bound mismatch for dimension 1 of array 'gbuff' (1/3)"
        !                for the non-subset MPI processes; but do not use -fcheck=no-bounds in compile flags which
        !                causes more problems
        if(chemsh_parallel%is_subset()) then
            gbuff(1:3,sum(natoms_per_domain(1:node))+1:sum(natoms_per_domain(1:node+1))) = dl_poly_gradients
        endif

        ! the root gathers the gradients
        if(chemsh_parallel%is_root()) then
            do i = 1, nnodes-1
                ierror = chemsh_parallel%receive(gbuff(1:3,sum(natoms_per_domain(1:i))+1:sum(natoms_per_domain(1:i+1))), 8*3*natoms_per_domain(i+1), i, itag)
            enddo
        else
            ! YL 21/11/2021: restrained to the MM subset otherwise Cray's MPICH (GNU>=10.2) segfaults
            !                on the remaining processes
            if(chemsh_parallel%is_subset()) then
                ierror = chemsh_parallel%send(dl_poly_gradients, 8*3*shapebuff(2), 0, itag)
            endif
        endif

        ! sort the gradients according to the global indices
        if(chemsh_parallel%is_root()) then
            gradients(1:3,dl_poly_global_indices) = dl_poly_gradients
            do i = 1, nnodes-1
                gradients(1:3,global_indices_per_node(sum(natoms_per_domain(1:i))+1:sum(natoms_per_domain(1:i+1)))) = gbuff(1:3,sum(natoms_per_domain(1:i))+1:sum(natoms_per_domain(1:i+1)))
            enddo
        endif

        if(chemsh_parallel%is_root()) then
            call PyResult%set('energy'   , (/dl_poly_energy/))
            call PyResult%set('gradients', gradients)
        endif
        
        deallocate(gradients, gbuff, natoms_per_domain, global_indices_per_node)
        ! YL 14/02/2019: these public arrays should only be deallocated by
        !                ChemShell/DL_POLY4 interface
        if(associated(dl_poly_global_indices)) deallocate(dl_poly_global_indices)
        if(associated(dl_poly_gradients))      deallocate(dl_poly_gradients)

    endfunction dl_poly_wrap

endmodule ChemShellDLPOLYModule

