!  Copyright (C) 2023 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 ChemShellCastepModule
    
    use, intrinsic :: iso_c_binding
    
    !Rajany: ptrtodict-function convert PyType ptr to dictionary
    !Rajany: PyType is a bind-c structure in fortran
    
    use DL_PY2F, only : PyType, dictType, ptr2dict
    use ChemShellParallelModule, only : chemsh_parallel

    private
   
    public register_castep, interface_castep, castep_wrap
    public natoms
    
    !! These pointers are saved across all function calls in module
    type(dictType)  , pointer   :: PyOptions
    type(dictType)  , pointer   :: PyFrag
    type(dictType)  , pointer   :: PyResult
    type(dictType)  , pointer   :: PyCell
    integer(c_int)              :: c_debug
    integer(c_int)              :: castep_handle

    logical           , save    :: initialise = .true.

    contains

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !Rajany: executed by the startup main function 

    subroutine register_castep(debug) bind(c, name="f_register_castep")

        integer(c_int), intent(in) :: debug
        type(c_funptr)             :: castep_pointer
        c_debug        = debug
        castep_pointer = c_funloc(castep_wrap) 
        castep_handle  = chemsh_parallel%get_function_handle(castep_pointer)
    
    endsubroutine register_castep
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 
    ! called by master process only
    integer(c_int) function interface_castep(fragPtr, optionsPtr, resultPtr,&
                                            cellPtr) result(ierror) bind(c)
 
        implicit none
        type(PyType)  , intent(in) :: fragPtr
        type(PyType)  , intent(in) :: optionsPtr
        type(PyType)  , intent(in) :: ResultPtr
        type(PyType)  , intent(in) :: cellPtr

        !!!!!!!!!!!!!!!! declarations end !!!!!!!!!!!!!!!!!!
        ierror = 1
        ! set up dict of fragObj
        allocate(PyFrag, source=ptr2dict(fragPtr))
        allocate(PyOptions, source=ptr2dict(optionsPtr))
        allocate(PyResult, source=ptr2dict(resultPtr))
        allocate(PyCell, source=ptr2dict(cellPtr))

        !get's arguments:metaobj,key,value(out)
        !call PyOptions%get("input", ptr_input) !!error,keyword "input"
        !.. not found in dictionary
        !calls the wrapper
        if(chemsh_parallel%is_parallel()) then
            ierror = chemsh_parallel%exec_function(castep_handle, c_debug)
        else
            ierror = castep_wrap()
        endif

        if(associated(PyFrag)) then
            call PyFrag%finalise()
            deallocate(PyFrag)
        endif

       if(associated(PyResult)) then
            call PyResult%finalise()
            deallocate(PyResult)
       endif
       if(associated(PyOptions)) then
            call PyOptions%finalise()
            deallocate(PyOptions)
        endif
 
       if(associated(PyCell)) then
            call PyCell%finalise()
            deallocate(PyCell)
        endif
    endfunction interface_castep
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    ! also called by slave processes
    integer function castep_wrap() result(ierror)

        use ChemShellParallelModule, only : chemsh_parallel
        
        implicit none
        
        character(len=255) :: castep_input
        integer         :: linput
        integer(kind=4) :: communicator
        !for mpi error message
        integer         :: error
        logical         :: isroot
        integer, parameter :: dp=selected_real_kind(15,300)
        !Rajany: Not initialising natoms here caused overflow
        integer         :: natoms=0
        real(kind=dp)   :: energy
        real(kind=dp)   :: new_lattice(3,3)
        real(kind=dp), allocatable, dimension(:,:)  :: new_positions
        real(kind=dp), allocatable, dimension(:,:)  :: forces

        !!!! Rajany: new_positions are passed as dummy array, so interface included.

        interface
             subroutine castepfunc(input_seed, initial, is_root, communicator, &
                        energy, forces, new_lattice, new_positions, natoms)

              integer, parameter :: dp=selected_real_kind(15,300)
              logical, save :: initialised 

              character(len=*), intent(in)         :: input_seed
              logical         , intent(in)         :: initial
              logical         , intent(in)         :: is_root
              integer(kind=4) , intent(in)         :: communicator 
              real(kind=dp)   , intent(inout)      :: energy
              real(kind=dp)   , dimension(:,:),   allocatable,  intent(inout):: forces
              real(kind=dp)   , dimension(3,3),   intent(inout)              :: new_lattice
              real(kind=dp)   , dimension(:,:),   intent(in)                 :: new_positions
              integer(kind=4) , intent(in)         :: natoms
              integer                              :: linput
              character(len=:), allocatable        :: input_seedname
             end subroutine
         end interface

        !!!!!!!!!!!!!!!!!!!!!!!!   declarations end !!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        ierror = 1
        castep_input = repeat(' ',255)
        castep_input = "_castep"
        linput  = len(trim(castep_input))
        !ierror = chemsh_parallel%broadcast(linput , 1, 0)
        !ierror = chemsh_parallel%broadcast(castep_input(1:linput), linput , 0)
        !call flush(6)
        
        communicator = chemsh_parallel%comm_workgroup
        isroot=chemsh_parallel%is_root() 

        if(isroot) then
        call PyFrag%get('npts'   , natoms)
        call PyCell%get('vectors'   , new_lattice)
        endif

        ierror = chemsh_parallel%broadcast(natoms,1,0)
        ierror = chemsh_parallel%broadcast(new_lattice,sizeof(new_lattice),0)
        !print *,'natoms',natoms
        !print *,'new_lattice',new_lattice
        !call flush(6)
        !Rajany: fortran column major
        allocate(new_positions(3,natoms), source=0.0_dp)

        if(isroot) then
        call PyFrag%get('xyz',new_positions)
        endif
        ierror = chemsh_parallel%broadcast(new_positions,sizeof(new_positions),0)
        !print *,'new_positions', new_positions
        !call flush(6)

        !Rajany:lattice vector components is across row in castep and column in chemsh
        !tanspose is not necessary here because of the conversion to fortran matrix
        !in column order, but should transpose in input file

        call castepfunc(castep_input(1:linput), &
                        initialise, isroot, communicator, energy, forces, &
                        new_lattice, new_positions, natoms )
        initialise = .false.
        ierror = 0

        !Rajany:update the results with the returned values      
        ! update, self._results with the energy and forces
        if(isroot) then
        call PyResult%set('energy', [energy])
        endif

        if(isroot) then
        call PyResult%set('gradients', forces(1:3, 1:natoms))
        endif

        deallocate(new_positions)

    endfunction castep_wrap
endmodule ChemShellCastepModule
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
