! Interface to the MPI module
module MW_parallel
#ifndef MW_SERIAL
      use mpi
      use MW_errors, only: MW_errors_mpilib_error => mpilib_error
#endif
#ifdef _OPENACC
   use openacc
#endif
   implicit none
   private

   ! public types
   public :: MW_parallel_t

   ! public subroutines
   public :: initialize
   public :: initialize_python
   public :: finalize
   public :: define_type
   public :: void_type
   public :: print_type

   ! A data structure to hold comm, rank and group for a process
   type MW_parallel_t
      integer :: comm   ! Parallel communicator
      integer :: comm_rank   ! Parallel rank of the process on this communicator
      integer :: size   ! Size of the parallel communicator
   end type MW_parallel_t

   integer, public, save :: MW_COMM_WORLD = 0

contains
   !================================================================================
   ! Initialize MPI and set MW_COMM_WORLD communicator
   subroutine initialize_python(comm)
      implicit none
      integer, intent(in) :: comm
#ifndef MW_SERIAL
      MW_COMM_WORLD = comm
#else
      integer, intent(in) :: comm
      MW_COMM_WORLD = 0
#endif
   end subroutine initialize_python

   !================================================================================
   ! Initialize MPI and set MW_COMM_WORLD communicator
   subroutine initialize()
      implicit none
#ifndef MW_SERIAL
      integer :: ierr
      call MPI_Init(ierr)
      if (ierr /= MPI_SUCCESS) then
         call MW_errors_mpilib_error("initialize", "parallel.f90", "MPI_Init", ierr)
      end if
      MW_COMM_WORLD = MPI_COMM_WORLD
#else
      MW_COMM_WORLD = 0
#endif
   end subroutine initialize


   !================================================================================
   ! Finalize MPI and free MW_COMM_WORLD communicator
   subroutine finalize()
      implicit none
#ifndef MW_SERIAL
      integer :: ierr
      call MPI_Finalize(ierr)
      if (ierr /= MPI_SUCCESS) then
         call MW_errors_mpilib_error("finalize", "parallel.f90", "MPI_Finalize", ierr)
      end if
#else
      continue
#endif
   end subroutine finalize

   !================================================================================
   ! Define a parallel data structure
   subroutine define_type(this, comm)
      implicit none
      type(MW_parallel_t), intent(inout) :: this
      integer, intent(in) :: comm
#ifndef MW_SERIAL
      integer :: ierr
      integer :: rank
      integer :: size
#endif
#ifdef _OPENACC
      integer(acc_device_kind) :: devicetype
      integer :: num_devices, devicenum
#endif
#ifndef MW_SERIAL
      call MPI_Comm_size(comm, size, ierr)
      if (ierr /= MPI_SUCCESS) then
         call MW_errors_mpilib_error("define_type", "parallel.f90", "MPI_Comm_size", ierr)
      end if

      call MPI_Comm_rank(comm, rank, ierr)
      if (ierr /= MPI_SUCCESS) then
         call MW_errors_mpilib_error("define_type", "parallel.f90", "MPI_Comm_rank", ierr)
      end if

      this%comm = comm
      this%comm_rank = rank
      this%size = size
#else
      this%comm = comm
      this%comm_rank = 0
      this%size = 1
#endif
#ifdef _OPENACC
      devicetype = acc_get_device_type() ! hopefully this is the GPUs
      num_devices = acc_get_num_devices(devicetype)
      devicenum = mod(this%comm_rank, num_devices)
      call acc_set_device_num(devicenum, devicetype)
      call acc_init(devicetype)
#endif

   end subroutine define_type

   !================================================================================
   ! Void a parallel data structure
   subroutine void_type(this)
      implicit none
      type(MW_parallel_t), intent(inout) :: this

      this%comm = 0
      this%comm_rank = 0
      this%size = 0

   end subroutine void_type

   ! ================================================================================
   ! Print parallel comm characteristics
   subroutine print_type(this, ounit)
      implicit none
      type(MW_parallel_t), intent(in) :: this
      integer, intent(in) :: ounit

      write(ounit, *)
      write(ounit, '("Parallel setup")')
      write(ounit, '("==============")')
      write(ounit, *)
      write(ounit, '("MPI communicator id:   ",i8)') this%comm
      write(ounit, '("MPI communicator size: ",i8)') this%size
      write(ounit, '("MPI rank:              ",i8)') this%comm_rank

   end subroutine print_type
end module MW_parallel
