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

    use iso_c_binding
    use mpi
    use ChemShellGAModule

    implicit none

!    include 'mpif.h'

    private
    public :: taskfarmType

    integer, parameter :: intType = sizeof(MPI_INTEGER)

    type taskfarmType
        private
        integer(kind=intType), public  :: debug
        integer(kind=intType), public  :: ncolours
        integer(kind=intType), public  :: nkeys
        integer(kind=intType), public  :: nsubsetkeys
        integer(kind=intType), public  :: colour
        integer(kind=intType), public  :: key
        integer(kind=intType), public  :: key_subset
        integer(kind=intType), public  :: comm_colour
        integer(kind=intType), public  :: comm_key
        integer(kind=intType), public  :: comm_subset
        integer(kind=intType), public  :: ga_memory
        type(GAType), pointer, private :: ga => null()
        contains
        procedure, public  :: initialise  => private_initialise
        procedure, public  :: finalise    => private_finalise
        procedure, public  :: summarise   => private_summarise
        procedure, public  :: synchronise => private_synchronise
        procedure, public  :: print       => private_print
        procedure, public  :: is_root      => private_is_root
        procedure, public  :: is_subset    => private_is_subset
        procedure, private :: private_initialise
        procedure, private :: private_finalise
        procedure, private :: private_print
        procedure, private :: private_is_root
        procedure, private :: private_is_subset
        procedure, private :: private_summarise
        procedure, private :: private_synchronise

    endtype taskfarmType
    
    contains

    ! taskfarming initialisation: see tcg2mpi.f and parallel_ga.f in Tcl-ChemShell
    integer(kind=intType) function private_initialise(this, ncolours, nsubsetkeys, debug) result(ierror)

        class(taskfarmType)  , intent(inout) :: this
        integer(kind=intType), intent(in)    :: ncolours
        integer(kind=intType), intent(in)    :: nsubsetkeys
        integer(kind=intType), intent(in)    :: debug

        integer(kind=intType)                :: rank, nprocs, i
! YL FIXME: these have to be 32-bit (because of GA?)
        integer(kind=intType)                :: group_all, group_subset, rank_subset
        integer(kind=intType), pointer       :: subset(:)

        call MPI_COMM_SIZE(MPI_COMM_WORLD, nprocs, ierror)
        call MPI_COMM_RANK(MPI_COMM_WORLD, rank  , ierror)

        this%debug = debug

        ! initialise GA
        ! GA does not need sub-workgroup communicator as we only reduce MM nprocs
        allocate(this%ga)

        ierror = this%ga%initialise(ncolours, this%ga_memory, debug)

        this%ncolours = this%ga%ncolours
        this%nkeys    = this%ga%nkeys
        this%colour   = this%ga%colour
        this%key      = this%ga%key
!        this%comm_colour = this%ga%comm_colour

        ! nowhere GA communicator is needed, so we use the communicator from MPI parallelism
        call MPI_COMM_SPLIT(MPI_COMM_WORLD, this%colour, this%key   , this%comm_colour, ierror)
        ! called "counterpart" in Tcl-ChemShell
        call MPI_COMM_SPLIT(MPI_COMM_WORLD, this%key   , this%colour, this%comm_key   , ierror)

        if(ierror.ne.0) then
            print *, ">>> FATAL ERROR: MPI_COMM_SPLIT() failed."
            stop
        endif

        ! size of sub-group processes
        ! legal range: 1 <= nsubsetkeys <= nkeys
        if(nsubsetkeys.lt.1.or.nsubsetkeys.gt.this%nkeys) then
            this%nsubsetkeys = this%nkeys
        else
            this%nsubsetkeys = nsubsetkeys
        endif

        allocate(subset(this%nsubsetkeys))
        subset = (/(i, i=0,this%nsubsetkeys-1)/)

        ! create a "group" of the colour's (workgroup's) all keys (processes in this workgroup)
        call MPI_COMM_GROUP(this%comm_colour, group_all, ierror)

        if(ierror.ne.0) then
            print *, ">>> FATAL ERROR: MPI_COMM_GROUP() failed."
            stop
        endif

        ! create a new group containing only the subset (a subset of processes in the workgroup)
        call MPI_GROUP_INCL(group_all, this%nsubsetkeys, subset, group_subset, ierror)
!        call MPI_GROUP_EXCL(group_all, this%nkeys-this%nsubsetkeys, subset, group_rest  , ierror)

        if(ierror.ne.0) then
            print *, ">>> FATAL ERROR: MPI_GROUP_INCL() failed."
            stop
        endif

        ! create a communicator of the subset
        call MPI_COMM_CREATE(this%comm_colour, group_subset, this%comm_subset, ierror)

        if(ierror.ne.0) then
            print *, ">>> FATAL ERROR: MPI_COMM_CREATE() failed."
            stop
        endif

        call MPI_GROUP_RANK(group_subset, rank_subset, ierror)

        if(ierror.ne.0) then
            print *, ">>> FATAL ERROR: MPI_GROUP_RANK() failed."
            stop
        endif

        ! YL 10/11/2017: replaced all MPI_COMM_NULL with MPI_UNDEFINED as in Open MPI the former causes hanging
        if(rank_subset.eq.MPI_UNDEFINED) then
            this%comm_subset = MPI_UNDEFINED
            this%key_subset  = MPI_UNDEFINED
        else
            this%key_subset  = rank_subset
        endif

        deallocate(subset)

    endfunction private_initialise
    integer(kind=intType) function private_summarise(this, level) result(ierror)

        class(taskfarmType), intent(in) :: this
        integer            , intent(in) :: level

        ierror = this%ga%summarise(level)

    endfunction private_summarise
    integer(kind=intType) function private_synchronise(this) result(ierror)

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

        call this%ga%synchronise()
        call MPI_BARRIER(this%comm_colour, ierror)

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

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

!        call MPI_GROUP_FREE(group_all)
!        call MPI_GROUP_FREE(group_subset)
        if(associated(this%ga)) then
            deallocate(this%ga)
        endif

        ierror = 0

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

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

        ierror = this%ga%print()
        call flush(6)

    endfunction private_print
    logical function private_is_root(this) result(result)

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

        if(associated(this%ga)) then
            result = this%ga%is_root().or.this%key_subset.eq.0
        else
            print *, ">>> ERROR: taskfarm not yet initialised!"
            stop
        endif

    endfunction private_is_root
    logical function private_is_subset(this) result(result)

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

        ! YL 10/11/2017: replaced all MPI_COMM_NULL with MPI_UNDEFINED as in Open MPI the former causes hanging
        if(this%key_subset.ne.MPI_UNDEFINED.and.this%comm_subset.ne.MPI_UNDEFINED) then
            result = this%key_subset.lt.this%nsubsetkeys
        else
            result = .false.
        endif

    endfunction private_is_subset
    ! convert any kind=4 or kind=8 integer to intType
    integer(kind=intType) function getCastInt(input) result(result)

        class(*), intent(in) :: input

        selecttype(tmp=>input)
            type is (integer(kind=4))
                result = tmp
            type is (integer(kind=8))
                result = tmp
!            type is (integer(c_int))
!                result = tmp
        endselect

    endfunction getCastInt


endmodule ChemShellTaskfarmModule

