! Adapted from parallel_ga.F of Tcl-ChemShell
! you.lu@stfc.ac.uk
! Feb 2017

!  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 ChemShellGAModule

    implicit none

!DEC$ NOFREEFORM
#include "mafdecls.fh"
#include "global.fh"
!DEC$ FREEFORM

    private
    public :: GAType

    integer, parameter :: intType = sizeof(0)

    type GAType
        private
        integer             , public :: debug
        integer             , public :: ncolours
        integer             , public :: nkeys
        integer             , public :: colour
        integer             , public :: key
        integer             , public :: comm_colour
        integer             , public :: comm_key
        integer, allocatable, public :: procgroup(:)
        integer             , private :: ga_memory
        contains
        procedure, public :: initialise  => private_initialise
        procedure, public :: finalise    => private_finalise
        procedure, public :: synchronise => private_synchronise
        procedure, public :: summarise   => private_summarise
        procedure, public :: print       => private_print
        procedure, public :: is_root      => private_is_root
        procedure, public :: setProcGroup => private_set_processors_group
        procedure, private :: private_initialise
        procedure, private :: private_finalise
        procedure, private :: private_synchronise
        procedure, private :: private_summarise
        procedure, private :: private_print
        procedure, private :: private_set_processors_group
        procedure, private :: private_is_root
    endtype GAType

    contains

    integer(kind=intType) function private_initialise(this, ncolours, ga_memory, debug) result(ierror)

        use iso_c_binding

        implicit none

!#include "mafdecls.fh"

        class(GAType)              , intent(inout) :: this
        class(*)                   , intent(in)    :: ncolours
        class(*)                   , intent(in)    :: debug
        class(*)                   , intent(in)    :: ga_memory

        integer                                    :: ga_nprocs
        integer                                    :: i, icolour
        integer                                    :: stack, heap, global
        integer                                    :: stack_auto, heap_auto, global_auto
        integer         , parameter                :: MB2Byte  = 1024*1024
        integer         , parameter                :: MB2ndbls = 1024*1024/8
        character(len=8)                           :: sbuffh = "", sbuffs = "", sbuffg = ""


        ! always compile libga.so 64-bit while libparallel.so could be 32-bit
        selecttype(ncolours)
            type is(integer(kind=4))
                this%ncolours = ncolours
            type is(integer(kind=8))
                this%ncolours = ncolours
        endselect
        selecttype(ga_memory)
            type is(integer(kind=4))
                this%ga_memory = ga_memory
            type is(integer(kind=8))
                this%ga_memory = ga_memory
        endselect
        selecttype(debug)
            type is(integer(kind=4))
                this%debug = debug
            type is(integer(kind=8))
                this%debug = debug
        endselect

        ! TD 22/04/2024: I didn't find documentation on this point, but by looking at GA
        ! examples I think calling both MPI_INIT and PBEGINF is an error
        ! This sometimes causes hangs on archer2
        !call PBEGINF

!        ! get the environment variables of HEAP/STACK/GLOBAL in MByte
!        ! TODO: review if this reduces the modularity of parallel lib
!        call get_environment_variable("HEAP  ", sbuffh)
!        read(sbuffh, '(i8)') heap
!        call get_environment_variable("STACK ", sbuffs)
!        read(sbuffs, '(i8)') stack
!        call get_environment_variable("GLOBAL", sbuffg)
!        read(sbuffg, '(i8)') global

        ! default size (MB)
!        if(any((/heap, stack, global/).le.0)) then
        global = this%ga_memory
        heap   = global/2
        stack  = global/2
!        endif

        ! convert MByte to number of doubles
!            heap   = MA_SIZEOF(MT_BYTE, MB2Byte*heap  , MT_DBL)
!            stack  = MA_SIZEOF(MT_BYTE, MB2Byte*stack , MT_DBL)
!            global = MA_SIZEOF(MT_BYTE, MB2Byte*global, MT_DBL)
        
        ! TODO this the TCS method
!        mt_dbl_i8  = int(MT_DBL , kind=i8)
!        mt_byte_i8 = int(MT_BYTE, kind=i8)

!        size_i8 = int(MA_SIZEOF(mt_dbl_i8, global, mt_byte_i8), kind=i8)

        ! amount of memory in bytes per process (this is a collective operation)
        call GA_INITIALIZE_LTD(global*MB2Byte)
!        call GA_INITIALIZE_LTD(size_i8)

        ! GA uses local memory
        if(GA_USES_MA()) then
            if(.not.MA_INIT(MT_DBL, -1, -1)) then
            !if(.not.MA_INIT(mt_dbl_i8, -1, -1)) then
                print *, ">>> GA module: MA_INIT() failed (local memory mode)"
                stop
            endif
        ! GA uses shared memory
        else
            if(.not.MA_INIT(MT_DBL, stack*MB2ndbls, heap*MB2ndbls)) then
            !if(.not.MA_INIT(mt_dbl_i8, stack, heap+global)) then
                print *, ">>> GA module: MA_INIT() failed (shared memory mode)"
                stop
            endif
        endif

        ! hardware memory limit
        global_auto = GA_MEMORY_AVAIL()
        heap_auto   = MA_INQUIRE_HEAP(MT_DBL)
        stack_auto  = MA_INQUIRE_STACK(MT_DBL)

        ! shouldn't exceed the hardware limit
        global = min(global, global_auto)
        heap   = min(heap  , heap_auto)
        stack  = min(stack , stack_auto)

        ! TCS: WG_size
        ga_nprocs = GA_NNODES()
        if(mod(ga_nprocs, this%ncolours).eq.0) then
            this%nkeys = ga_nprocs / this%ncolours
        else
            print *, ">>> FATAL ERROR: number of processors (-np) not divisible by number of workgroups (-nwg)"
            stop
        endif

        ! GA_NODEID() returns MPI_WORLD_RANK
        this%colour   = GA_NODEID() / this%nkeys
        this%key      = mod(GA_NODEID(), this%nkeys)

        allocate(this%procgroup(this%ncolours))

        do icolour = 1, this%ncolours
            ! for example we have 6 processes in 3 workgroups
            ! proclist used for GA_PGROUP_CREATE is (/0,1/), (/2,3/), and (/4,5/), respectively
            this%procgroup(icolour) = GA_PGROUP_CREATE((/(i, i=this%nkeys*(icolour-1), this%nkeys*icolour-1)/), this%nkeys)
        enddo

        call GA_PGROUP_SET_DEFAULT(this%procgroup(this%colour+1))

        call GA_SYNC()

        call GA_MPI_COMM_PGROUP_DEFAULT(this%comm_colour)

        if(this%colour.eq.0.and.this%key.eq.0) then
            write(6,'(//1x,a)') repeat("=", 48)
            write(6,'(/a/)') " Initialising Parallelism driven by Global Arrays (GA)"
            write(6,'(1x,a,1x,i5,1x,a,i7,a)') "Heap   size:", heap  , "MB (", heap_auto  /1024, " GB available)"
            write(6,'(1x,a,1x,i5,1x,a,i7,a)') "Stack  size:", stack , "MB (", stack_auto /1024, " GB available)"
            write(6,'(1x,a,1x,i5,1x,a,i7,a)') "Global size:", global, "MB (", global_auto/1024, " GB available)"
            write(6,'(/1x,a/)') repeat("=", 48)
            if(this%debug.eq.5) then
                call MA_SUMMARIZE_ALLOCATED_BLOCKS()
                call MA_PRINT_STATS(.true.)
            endif
        endif

        ierror = 0

    endfunction private_initialise
    integer(kind=intType) function private_finalise(this) result(ierror)

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

        deallocate(this%procgroup)
        call GA_TERMINATE()
        call PEND

        ierror = 0

    endfunction private_finalise
    integer function private_summarise(this, level) result(ierror)

        class(GAType), intent(in) :: this
        class(*)     , intent(in) :: level

        selecttype(level)
            type is (integer(kind=4))
                call GA_SUMMARIZE(int8(level))
            type is (integer(kind=8))
                call GA_SUMMARIZE(level)
        endselect

        ierror = 0

    endfunction private_summarise
    subroutine private_synchronise(this)

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

        call GA_SYNC()

    endsubroutine private_synchronise
    integer function private_print(this) result(ierror)

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

        if(this%ncolours.gt.1.and.this%key.eq.0.and.this%debug.eq.5) then
            call GA_PRINT_STATS()
        endif

        ierror = 0

    endfunction private_print
    subroutine private_set_processors_group(this, proc_group)

        class(GAType), intent(inout) :: this
        integer      , intent(in)    :: proc_group(this%nkeys)

        call GA_PGROUP_SET_DEFAULT(proc_group(this%colour))

    endsubroutine private_set_processors_group
    logical function private_is_root(this) result(result)

        class(GAType), intent(in) :: this

!        result = this%key.eq.0
        result = GA_NODEID().eq.0

    endfunction private_is_root

endmodule ChemShellGAModule

