MODULE mpipara
   IMPLICIT NONE
   SAVE

   INTEGER :: mpi_rank
   INTEGER :: mpi_size
   INTEGER :: mpi_threading_mode

   INTEGER :: comm_icosa
   INTEGER :: ierr
   LOGICAL :: using_mpi
   LOGICAL :: is_mpi_root
   LOGICAL, BIND(C) :: is_mpi_master = .TRUE.

   CHARACTER(LEN=256) :: using_oasis

   INTEGER :: mpi_master

   INTEGER :: id_mpi ! id for profiling

   INTERFACE allocate_mpi_buffer
      MODULE PROCEDURE allocate_mpi_buffer_r2, allocate_mpi_buffer_r3, allocate_mpi_buffer_r4
   END INTERFACE allocate_mpi_buffer

   INTERFACE free_mpi_buffer
      MODULE PROCEDURE free_mpi_buffer_r2, free_mpi_buffer_r3, free_mpi_buffer_r4
   END INTERFACE free_mpi_buffer

   PRIVATE :: getin

CONTAINS

   SUBROUTINE getin(name, value) ! Copied from getin.f90 to avoid circular dependency
      USE ioipsl, ONLY: getin_ => getin
      USE transfert_omp_mod
      USE omp_para
      IMPLICIT NONE
      CHARACTER(LEN=*) :: name
      CHARACTER(LEN=*) :: value

!$OMP MASTER
      CALL getin_(name, value)
      IF (is_mpi_root) PRINT *, 'GETIN ', TRIM(name), ' = ', TRIM(value)
!$OMP END MASTER
      IF (omp_in_parallel()) CALL bcast_omp(value)
   END SUBROUTINE getin

   SUBROUTINE init_mpipara
      USE mpi_mod
#ifdef CPP_USING_XIOS
      USE xios
#endif
      USE abort_mod
      USE openacc_mod
      CHARACTER(LEN=256) :: required_mode_str
      INTEGER :: required_mode

      using_mpi = .FALSE.
#ifdef CPP_USING_MPI
      using_mpi = .TRUE.
#endif

      CALL set_openacc_device

      IF (using_mpi) THEN

         using_oasis = 'n'
         CALL getin('using_oasis', using_oasis)

         IF (using_oasis == 'n') THEN

            required_mode_str = 'funneled'
            CALL getin('mpi_threading_mode', required_mode_str)

            SELECT CASE (TRIM(required_mode_str))
            CASE ('single')
               required_mode = MPI_THREAD_SINGLE
            CASE ('funneled')
               required_mode = MPI_THREAD_FUNNELED
            CASE ('serialized')
               required_mode = MPI_THREAD_SERIALIZED
            CASE ('multiple')
               required_mode = MPI_THREAD_MULTIPLE
            CASE DEFAULT
               PRINT *, 'Bad selector for variable mpi_threading_mode  : <', TRIM(required_mode_str), &
                  '>  => options are <single>, <funneled>, <serialized>, <multiple>'
               STOP
            END SELECT

            IF (required_mode == MPI_THREAD_SERIALIZED .OR. required_mode == MPI_THREAD_MULTIPLE) THEN
               CALL abort_acc("mpi_threading_mode /= 'single' .AND. mpi_threading_mode /= 'funneled'")
            END IF

            IF (required_mode == MPI_THREAD_SINGLE) PRINT *, 'MPI_INIT_THREAD : MPI_SINGLE_THREAD required'
            IF (required_mode == MPI_THREAD_FUNNELED) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_FUNNELED required'
            IF (required_mode == MPI_THREAD_SERIALIZED) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_SERIALIZED required'
            IF (required_mode == MPI_THREAD_MULTIPLE) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_MULTIPLE required'

            CALL MPI_INIT_THREAD(required_mode, mpi_threading_mode, ierr)

            IF (mpi_threading_mode == MPI_THREAD_SINGLE) PRINT *, 'MPI_INIT_THREAD : MPI_SINGLE_THREAD provided'
            IF (mpi_threading_mode == MPI_THREAD_FUNNELED) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_FUNNELED provided'
            IF (mpi_threading_mode == MPI_THREAD_SERIALIZED) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_SERIALIZED provided'
            IF (mpi_threading_mode == MPI_THREAD_MULTIPLE) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_MULTIPLE provided'

            IF (mpi_threading_mode > required_mode) mpi_threading_mode = required_mode

            IF (mpi_threading_mode == MPI_THREAD_SINGLE) THEN
               PRINT *, 'MPI_INIT_THREAD : MPI_SINGLE_THREAD used : Warning : openMP is not garanted to work'
            END IF
            IF (mpi_threading_mode == MPI_THREAD_FUNNELED) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_FUNNELED used'
            IF (mpi_threading_mode == MPI_THREAD_SERIALIZED) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_SERIALIZED used'
            IF (mpi_threading_mode == MPI_THREAD_MULTIPLE) PRINT *, 'MPI_INIT_THREAD : MPI_THREAD_MULTIPLE used'

         END IF

#ifdef CPP_USING_XIOS
         CALL xios_initialize("icosagcm", return_comm=comm_icosa)
#else
         comm_icosa = MPI_COMM_WORLD
#endif
         CALL MPI_COMM_SIZE(comm_icosa, mpi_size, ierr)
         CALL MPI_COMM_RANK(comm_icosa, mpi_rank, ierr)
         PRINT *, 'MPI Process ', mpi_rank, '/', mpi_size
      ELSE
         comm_icosa = -1
         mpi_size = 1
         mpi_rank = 0
      END IF

      mpi_master = 0
      IF (mpi_rank == 0) THEN
         is_mpi_root = .TRUE.
         is_mpi_master = .TRUE.
      ELSE
         is_mpi_root = .FALSE.
         is_mpi_master = .FALSE.
      END IF

      CALL set_openacc_device(mpi_size, mpi_rank)
      PRINT *, 'GPU device ', openacc_device_id

   END SUBROUTINE init_mpipara

   SUBROUTINE finalize_mpipara
      USE mpi_mod
#ifdef CPP_USING_XIOS
      USE xios
#endif
      IMPLICIT NONE

#ifdef CPP_USING_XIOS
      CALL xios_finalize
#endif
      IF (using_mpi) THEN
         IF (using_oasis == 'n') CALL MPI_FINALIZE(ierr)
      END IF

   END SUBROUTINE finalize_mpipara

   SUBROUTINE allocate_mpi_buffer_r2(buffer, length)
      USE ISO_C_BINDING
      USE mpi_mod
      USE prec
      IMPLICIT NONE
      REAL(rstd), POINTER :: buffer(:)
      INTEGER, INTENT(IN)  :: length

      TYPE(C_PTR)         :: base_ptr
      INTEGER(KIND=MPI_ADDRESS_KIND) :: real_size, lb, size
      INTEGER :: ierr

      CALL MPI_Type_get_extent(MPI_REAL8, lb, real_size, ierr)
      size = length*real_size

      CALL MPI_ALLOC_MEM(size, MPI_INFO_NULL, base_ptr, ierr)
      CALL C_F_POINTER(base_ptr, buffer, (/length/))

   END SUBROUTINE allocate_mpi_buffer_r2

   SUBROUTINE free_mpi_buffer_r2(buffer)
      USE ISO_C_BINDING
      USE mpi_mod
      USE prec
      IMPLICIT NONE
      REAL(rstd), POINTER :: buffer(:)

      CALL MPI_FREE_MEM(buffer, ierr)

   END SUBROUTINE free_mpi_buffer_r2

   SUBROUTINE allocate_mpi_buffer_r3(buffer, length, dim3)
      USE ISO_C_BINDING
      USE mpi_mod
      USE prec
      IMPLICIT NONE
      REAL(rstd), POINTER :: buffer(:, :)
      INTEGER, INTENT(IN)  :: length
      INTEGER, INTENT(IN)  :: dim3

      TYPE(C_PTR)         :: base_ptr
      INTEGER(KIND=MPI_ADDRESS_KIND) :: real_size, lb, size
      INTEGER :: ierr

      CALL MPI_Type_get_extent(MPI_REAL8, lb, real_size, ierr)
      size = length*real_size*dim3

      CALL MPI_ALLOC_MEM(size, MPI_INFO_NULL, base_ptr, ierr)
      CALL C_F_POINTER(base_ptr, buffer, (/length, dim3/))

   END SUBROUTINE allocate_mpi_buffer_r3

   SUBROUTINE free_mpi_buffer_r3(buffer)
      USE ISO_C_BINDING
      USE mpi_mod
      USE prec
      IMPLICIT NONE
      REAL(rstd), POINTER :: buffer(:, :)

      CALL MPI_FREE_MEM(buffer, ierr)

   END SUBROUTINE free_mpi_buffer_r3

   SUBROUTINE allocate_mpi_buffer_r4(buffer, length, dim3, dim4)
      USE ISO_C_BINDING
      USE mpi_mod
      USE prec
      IMPLICIT NONE
      REAL(rstd), POINTER :: buffer(:, :, :)
      INTEGER, INTENT(IN)  :: length
      INTEGER, INTENT(IN)  :: dim3
      INTEGER, INTENT(IN)  :: dim4

      TYPE(C_PTR)         :: base_ptr
      INTEGER(KIND=MPI_ADDRESS_KIND) :: real_size, lb, size
      INTEGER :: ierr

      CALL MPI_Type_get_extent(MPI_REAL8, lb, real_size, ierr)
      size = length*real_size*dim3*dim4

      CALL MPI_ALLOC_MEM(size, MPI_INFO_NULL, base_ptr, ierr)
      CALL C_F_POINTER(base_ptr, buffer, (/length, dim3, dim4/))

   END SUBROUTINE allocate_mpi_buffer_r4

   SUBROUTINE free_mpi_buffer_r4(buffer)
      USE ISO_C_BINDING
      USE mpi_mod
      USE prec
      IMPLICIT NONE
      REAL(rstd), POINTER :: buffer(:, :, :)

      CALL MPI_FREE_MEM(buffer, ierr)

   END SUBROUTINE free_mpi_buffer_r4

END MODULE mpipara
