! This module defines the derived types involved in halo exchanges
! Instances of this types are manipulated by modules transfert_requests_mod, ...
! This module uses Fortran 2003 features : allocatable type fields

module transfert_types_mod
  USE field_mod, ONLY : t_field
  USE abort_mod, ONLY : dynamico_abort
  USE profiling_mod, ONLY : enter_profile, exit_profile

  implicit none
  PUBLIC
  SAVE

  ! Describes how to pack/unpack a message from a local domain to another
  TYPE t_local_submessage
    INTEGER   :: src_ind_loc, dest_ind_loc ! index of local and remote domain
    INTEGER   :: npoints                   ! Number of cells to transfer (dim12)
    INTEGER, ALLOCATABLE :: displ_src(:)  ! List of indexes to copy from domain src_ind_loc
    INTEGER, ALLOCATABLE :: displ_dest(:) ! List of indexes to copy to domain dest_ind_loc
    INTEGER, ALLOCATABLE :: sign(:)       ! Sign change to be applied for vector requests
  END TYPE t_local_submessage

  ! Describes how to pack/unpack a message from a domain to another
  TYPE t_submessage
    INTEGER              :: ind_loc, remote_ind_glo, remote_rank ! index of local and remote domain
    INTEGER              :: npoints                              ! Number of cells to transfer (dim12)
    INTEGER, ALLOCATABLE :: displs(:)                            ! List of indexes to copy from field to buffer for each level
    INTEGER, ALLOCATABLE :: sign(:)                              ! Sign change to be applied for vector requests
    INTEGER              :: mpi_buffer_displ = -1
  END TYPE t_submessage

  ! Points for t_request : list of points coordinates in a domain
  TYPE t_points
    INTEGER, ALLOCATABLE :: i(:), j(:)
    INTEGER, ALLOCATABLE :: elt(:) ! Element : cell edge or cell point
    INTEGER :: npoints
  END TYPE t_points

  ! Describes how to exchange data from/to one domain to/from other domains
  TYPE t_request

     INTEGER :: field_type
     TYPE(t_points), ALLOCATABLE :: points_BtoH(:) ! req(local_ind)%points_BtoH(remote_ind) are the points of
     ! domain local_ind to unpack from mpi buffer recieved from remote_ind
     TYPE(t_points), ALLOCATABLE :: points_HtoB(:) ! req(local_ind)%points_HtoB(remote_ind) are the points of
     ! domain local_ind to pack to mpi buffer to send to remote_ind
     LOGICAL :: vector = .false. ! Is a vector request (true if sign change needed)
     
     ! the above information describes what data must be exchanged
     ! but not the details of how to copy data to/from buffers
     ! detailed information is contained in the arrays below
     ! ICO : init_all_requests leaves has_message at .FALSE. Later, init_message calls init_message_in_out_local
     !       which fills the detailed information and sets has_message to .TRUE.
     ! UNST : init_all_requests_unst fills the detailed information and sets has_message to .TRUE.
     
     LOGICAL :: has_message = .false.
     TYPE (t_local_submessage), ALLOCATABLE :: message_local(:)     ! Local halo copies
     TYPE (t_submessage),       ALLOCATABLE :: message_in(:)        ! Messages to receive from remote ranks and to copy back to the field
     TYPE (t_submessage),       ALLOCATABLE :: message_out(:)       ! Halos to copy to MPI buffer and to send to remote ranks
  END TYPE t_request

  TYPE t_compact_submessages                     ! Describes MPI halo exchange, as in copy_BtoH and copy_HtoB
     INTEGER :: npoints                          ! number of halo points to send/receive
     INTEGER, ALLOCATABLE :: field_ind(:),    &  ! domain index, i.e. use field(ind)%rval4d
          &                  field_displ(:),  &  ! horizontal index ij in field
          &                  sign(:),         &  ! sign by which incoming data must be multiplied
          &                  remote_rank(:),  &  ! MPI rank for send/receive
          &                  mpi_displ(:),    &  ! horizontal index in MPI buffer
          &                  level_offset(:)     ! horizontal size of submessage, used as stride when field is 3D or 4! horizontal size of submessage, used as stride when field is 3D or 4D
  END TYPE t_compact_submessages

  TYPE t_compact_local_submessages                 ! Describes in-memory halo exchange, as in copy_HtoH
     INTEGER :: npoints                            ! number of halo points to copy
     INTEGER, ALLOCATABLE :: field_ind_src(:),   & ! source : domain index, horizontal index, sign
          &                  field_displ_src(:), &
          &                  sign(:),            &
          &                  field_ind_dest(:),  & ! destination : domain index, horizontal index
          &                  field_displ_dest(:)
  END TYPE t_compact_local_submessages

  TYPE mpi_buffer_t
    INTEGER :: n
    REAL, ALLOCATABLE :: buff(:)
  END TYPE mpi_buffer_t

  ! Describes how to exchange data for a field, and contains MPI buffers.
  TYPE t_message
    TYPE (t_field), POINTER                     :: field(:) => null()   ! Field to exchange
    TYPE (t_request), POINTER                   :: request(:) => null() ! Type of message to send
    TYPE (t_local_submessage), ALLOCATABLE          :: message_local(:)     ! Local halo copies
    TYPE (t_submessage), ALLOCATABLE                :: message_in(:)  ! Messages to recieve from remote ranks and to copy back to the field
    TYPE (t_submessage), ALLOCATABLE                :: message_out(:) ! Halos to copy to MPI buffer and to send to remote ranks
    TYPE (mpi_buffer_t), POINTER                :: mpi_buffer_in(:)
    TYPE (mpi_buffer_t), POINTER                :: mpi_buffer_out(:)
    TYPE (t_compact_submessages), POINTER       :: message_in_compact
    TYPE (t_compact_submessages), POINTER       :: message_out_compact
    TYPE (t_compact_local_submessages), POINTER :: message_local_compact

    ! NOTE : requests are persistant requests initialized in init_message.
    !           MPI_Start and MPI_Wait are then used to initiate and complete communications.
    !           ex : Give mpi_requests_in(i) to MPI_Start to send the buffer contained in message_in(i)
    INTEGER, POINTER  :: mpi_requests_in(:)  ! MPI requests used for message_in.
    INTEGER, POINTER  :: mpi_requests_out(:) ! MPI requests used for message_out.

    INTEGER  :: send_seq ! Sequence number : send_seq is incremented each time send_message is called
    INTEGER  :: wait_seq ! Sequence number : wait_seq is incremented each time wait_message is called
    LOGICAL  :: ondevice ! Ready to transfer ondevice field
  END TYPE t_message

  ! NOTE : create_request stores halo exchange information in fields points_BtoH and point_HtoB
  ! init_message rearranges the same information into fields  message_in_compact, message_out_compact, message_local_compact

  TYPE(t_request), POINTER :: req_i1(:)      ! Halos for T fields
  TYPE(t_request), POINTER :: req_e1_scal(:) ! Halos for U fields
  TYPE(t_request), POINTER :: req_e1_vect(:) ! Halos for U fields (with sign change)
  TYPE(t_request), POINTER :: req_z1_scal(:) ! Halos for Z fields

  TYPE(t_request), POINTER :: req_i0(:)      ! Duplicated cells for T fields
  TYPE(t_request), POINTER :: req_e0_scal(:)
  TYPE(t_request), POINTER :: req_e0_vect(:)

  ! Profiling id for mpi
  LOGICAL, PARAMETER :: profile_mpi_detail = .true.
  INTEGER :: profile_mpi, profile_mpi_copies, profile_mpi_waitall, profile_mpi_barrier

CONTAINS

  ! We put in this module only routines that do not depend on icosahedral data structures (domain_mod)
  ! TODO : store communicator and/or mpi_size in t_message rather than using module variables (=> dependency)
  
  !================================= Auxiliary routines for init_message =================================

  SUBROUTINE message_alloc_mpi_buf(message)
    ! Allocates MPI buffers and computes offsets within buffers
    ! There is one buffer_in and buffer_out for each remote_rank
    ! IN :  message_in(:)%remote_rank
    !       message_in(:)%npoints
    !       message_out(:)%remote_rank
    !       message_out(:)%npoints
    !       These arrays have the same size, equal to the number of submessages (halo from one domain to another domain)
    !       In copy_BtoH and copy_HtoB, submessages exchanged with the same remote_rank are copied from/to a single buffer
    !       => for each submessage an offset inside the buffer is computed here
    !       if each MPI rank manages a single domain all offsets are zero
    ! OUT : allocate(mpi_buffer_in(:)), allocate(mpi_buffer_out(:)) 
    !       mpi_buffer_out%n       ! total size of buffer, used to compute indices when copy-in / copy-out
    !       message_out(:)%displ   ! offset of outgoing submessage within MPI buffer
    !       message_in(:)%displ    ! offset of ingoing  submessage within MPI buffer

    USE mpipara, ONLY : mpi_size
    TYPE(t_message), TARGET, INTENT(INOUT) :: message
    INTEGER :: buffer_in_size, buffer_out_size, remote_rank
    INTEGER :: i, dim34

    dim34 = message%field(1)%dim3 * message%field(1)%dim4

    ! Allocate MPI buffers
    ALLOCATE( message%mpi_buffer_in(0:mpi_size-1) )
    ALLOCATE( message%mpi_buffer_out(0:mpi_size-1) )
    DO i = 0, mpi_size-1
       buffer_in_size = dim34*sum( message%message_in(:)%npoints, message%message_in(:)%remote_rank == i )
       buffer_out_size = dim34*sum( message%message_out(:)%npoints, message%message_out(:)%remote_rank == i )
       !TODO : what if size == 0 ?
       ALLOCATE( message%mpi_buffer_in(i)%buff( buffer_in_size ) )
       ALLOCATE( message%mpi_buffer_out(i)%buff( buffer_out_size ) )
       message%mpi_buffer_in(i)%n=0
       message%mpi_buffer_out(i)%n=0
    END DO

    ! Set offsets in submessages
    DO i=1, size(message%message_out)
      remote_rank = message%message_out(i)%remote_rank
      message%message_out(i)%mpi_buffer_displ = message%mpi_buffer_out(remote_rank)%n
      message%mpi_buffer_out(remote_rank)%n = message%mpi_buffer_out(remote_rank)%n + message%message_out(i)%npoints*dim34
    END DO
    DO i=1, size(message%message_in)
      remote_rank = message%message_in(i)%remote_rank
      message%message_in(i)%mpi_buffer_displ = message%mpi_buffer_in(remote_rank)%n
      message%mpi_buffer_in(remote_rank)%n = message%mpi_buffer_in(remote_rank)%n + message%message_in(i)%npoints*dim34
    END DO

  END SUBROUTINE message_alloc_mpi_buf

  SUBROUTINE message_create_MPI_request(message)
    ! Create persistent MPI requests
    USE mpipara, ONLY : mpi_size, comm_icosa
    USE mpi_mod, ONLY : MPI_REQUEST_NULL, MPI_REAL8
    TYPE(t_message), TARGET, INTENT(INOUT) :: message
    INTEGER :: i, ierr

    ALLOCATE( message%mpi_requests_in(0:mpi_size-1) )
    ALLOCATE( message%mpi_requests_out(0:mpi_size-1) )
    message%mpi_requests_in(0:mpi_size-1) = MPI_REQUEST_NULL
    message%mpi_requests_out(0:mpi_size-1) = MPI_REQUEST_NULL
    DO i = 0, mpi_size-1
       IF(  size(message%mpi_buffer_in(i)%buff) /= message%mpi_buffer_in(i)%n &
            .or. size(message%mpi_buffer_out(i)%buff) /= message%mpi_buffer_out(i)%n)&
            CALL dynamico_abort("Internal error in transfert_mpi : mpi buffer size different than expected")
       IF( message%mpi_buffer_out(i)%n > 0) THEN
          CALL MPI_Send_Init( message%mpi_buffer_out(i)%buff, message%mpi_buffer_out(i)%n, MPI_REAL8, i,&
               100, comm_icosa, message%mpi_requests_out(i), ierr )
       ENDIF
       IF( message%mpi_buffer_in(i)%n > 0) THEN
          CALL MPI_Recv_Init( message%mpi_buffer_in(i)%buff, message%mpi_buffer_in(i)%n, MPI_REAL8, i,&
               100, comm_icosa, message%mpi_requests_in(i), ierr )
       ENDIF
    END DO
  END SUBROUTINE message_create_MPI_request

  SUBROUTINE message_compactify(message)
    ! Rearranges halo exchange data
    ! IN : halo exchange is described as a number of submessages (message_in(:), message_out(:))
    !      each submessage relates to a remote_rank and a domain index and has size npoints
    !      within a submessage for each point to exchange we have the index in the field (field_dipsl=ij) and the index in the MPI buffer (displ)
    ! OUT : halo exchange is described as a flat list of points
    !       for each point we store (remote_rank, ind, level_offset=npoints) in addition to (field_displ, displ)
    !       level_offset is used during copy_HtoB and copy_BtoH to compute index in MPI buffer and keep each submessage contiguous

    USE mpipara, ONLY : mpi_size
    USE layout_mod, ONLY : layout, order_ij_l_p, order_l_ij_p
    TYPE(t_message), TARGET, INTENT(INOUT) :: message
    INTEGER :: field_type, index_order, dim34, npoints, last_point, i, k

    field_type = message%field(1)%field_type
    index_order = layout(1, field_type)%index_order
    dim34 = message%field(1)%dim3 * message%field(1)%dim4

    ALLOCATE(message%message_in_compact)
    message%message_in_compact%npoints = sum(message%message_in(:)%npoints)
    npoints = message%message_in_compact%npoints
    ALLOCATE(message%message_in_compact%field_ind(npoints))
    ALLOCATE(message%message_in_compact%field_displ(npoints))
    ALLOCATE(message%message_in_compact%sign(npoints))
    ALLOCATE(message%message_in_compact%remote_rank(npoints))
    ALLOCATE(message%message_in_compact%mpi_displ(npoints))
    ALLOCATE(message%message_in_compact%level_offset(npoints))

    last_point=0
    DO i = 1, size( message%message_in )
       DO k = 1, message%message_in(i)%npoints
          last_point = last_point+1
          message%message_in_compact%field_ind(last_point)   = message%message_in(i)%ind_loc
          message%message_in_compact%field_displ(last_point) = message%message_in(i)%displs(k)
          message%message_in_compact%sign(last_point)        = message%message_in(i)%sign(k)
          message%message_in_compact%remote_rank(last_point) = message%message_in(i)%remote_rank
          SELECT CASE(index_order)
          CASE(order_ij_l_p)
             message%message_in_compact%mpi_displ(last_point)   = message%message_in(i)%mpi_buffer_displ + k
             message%message_in_compact%level_offset(last_point)= message%message_in(i)%npoints
          CASE(order_l_ij_p)
             message%message_in_compact%mpi_displ(last_point)   = 1+dim34*(k-1)
             message%message_in_compact%level_offset(last_point)= 1
          END SELECT
       END DO
    END DO

    ALLOCATE(message%message_out_compact)
    message%message_out_compact%npoints = sum(message%message_out(:)%npoints)
    npoints = message%message_out_compact%npoints
    ALLOCATE(message%message_out_compact%field_ind(npoints))
    ALLOCATE(message%message_out_compact%field_displ(npoints))
    ALLOCATE(message%message_out_compact%sign(npoints))
    ALLOCATE(message%message_out_compact%remote_rank(npoints))
    ALLOCATE(message%message_out_compact%mpi_displ(npoints))
    ALLOCATE(message%message_out_compact%level_offset(npoints))

    last_point=0
    DO i = 1, size( message%message_out )
       DO k = 1, message%message_out(i)%npoints
          last_point = last_point+1
          message%message_out_compact%field_ind(last_point)   = message%message_out(i)%ind_loc
          message%message_out_compact%field_displ(last_point) = message%message_out(i)%displs(k)
          message%message_out_compact%sign(last_point)        = message%message_out(i)%sign(k)
          message%message_out_compact%remote_rank(last_point) = message%message_out(i)%remote_rank
          SELECT CASE(index_order)
          CASE(order_ij_l_p)
             message%message_out_compact%mpi_displ(last_point)   = message%message_out(i)%mpi_buffer_displ + k
             message%message_out_compact%level_offset(last_point)= message%message_out(i)%npoints
          CASE(order_l_ij_p)
             message%message_out_compact%mpi_displ(last_point)   = 1+dim34*(k-1)
             message%message_out_compact%level_offset(last_point)= 1
          END SELECT
       END DO
    END DO
    
    ALLOCATE(message%message_local_compact)
    message%message_local_compact%npoints = sum(message%message_local(:)%npoints)
    npoints = message%message_local_compact%npoints
    ALLOCATE(message%message_local_compact%field_ind_src(npoints))
    ALLOCATE(message%message_local_compact%field_displ_src(npoints))
    ALLOCATE(message%message_local_compact%sign(npoints))
    ALLOCATE(message%message_local_compact%field_ind_dest(npoints))
    ALLOCATE(message%message_local_compact%field_displ_dest(npoints))
    
    last_point=0
    DO i = 1, size( message%message_local )
       DO k = 1, message%message_local(i)%npoints
          last_point = last_point+1
          message%message_local_compact%field_ind_src(last_point)   = message%message_local(i)%src_ind_loc
          message%message_local_compact%field_displ_src(last_point) = message%message_local(i)%displ_src(k)
          message%message_local_compact%sign(last_point)            = message%message_local(i)%sign(k)
          message%message_local_compact%field_ind_dest(last_point)  = message%message_local(i)%dest_ind_loc
          message%message_local_compact%field_displ_dest(last_point)= message%message_local(i)%displ_dest(k)
       END DO
    END DO

  END SUBROUTINE message_compactify

  
  !====================== Halo copy : field to buffer, field to field, buffer to field =====================

  ! Halo to Buffer : copy outbound message to MPI buffers
  SUBROUTINE copy_HtoB(msg)
    USE omp_para, ONLY : distrib_level
    USE layout_mod, ONLY : layout, order_ij_l_p, order_l_ij_p

    TYPE(t_message), INTENT(INOUT) :: msg
    INTEGER :: field_type, index_order, dim3, dim4, d3_begin, d3_end
    INTEGER :: d3, d4, d34, i, ij, ind, rank, displ, stride

    field_type = msg%field(1)%field_type
    index_order = layout(1, field_type)%index_order
    dim3 = msg%field(1)%dim3
    dim4 = msg%field(1)%dim4

    SELECT CASE( index_order )
       
    CASE( order_ij_l_p )
       CALL distrib_level( 1, dim3, d3_begin, d3_end )
       !$acc parallel loop collapse(3) present(msg) default(present) async if(msg%ondevice)
       DO d4 = 1, dim4
          DO d3 = d3_begin, d3_end
             DO i=1, msg%message_out_compact%npoints
                d34 = d3-1 + dim3*(d4-1) ! linearize index for non-horizontal dimensions
                ! source
                ind    = msg%message_out_compact%field_ind(i)    ! ind (domain) to read from
                ij     = msg%message_out_compact%field_displ(i)  ! horizontal index in field
                ! dest
                rank   = msg%message_out_compact%remote_rank(i)  ! MPI rank to which data will be sent
                displ  = msg%message_out_compact%mpi_displ(i)    ! position in MPI buffer
                stride = msg%message_out_compact%level_offset(i)
                !
                msg%mpi_buffer_out(rank)%buff(displ + stride*d34) = msg%field(ind)%rval4d(ij, d3, d4)
             END DO
          END DO
       END DO
       
    CASE( order_l_ij_p )
       !$OMP DO SCHEDULE(STATIC)
       DO i=1, msg%message_out_compact%npoints
          ! source
          ind    = msg%message_out_compact%field_ind(i)    ! ind (domain) to read from
          ij     = msg%message_out_compact%field_displ(i)  ! horizontal index in field
          ! dest
          rank   = msg%message_out_compact%remote_rank(i)  ! MPI rank to which data will be sent
          displ  = msg%message_out_compact%mpi_displ(i)    ! position in MPI buffer
          stride = msg%message_out_compact%level_offset(i)
          
          DO d4 = 1, dim4
             DO d3 = 1, dim3
                d34 = d3-1 + dim3*(d4-1) ! linearize index for non-horizontal dimensions
                msg%mpi_buffer_out(rank)%buff(displ + stride*d34) = msg%field(ind)%rval4d(d3, ij, d4)
             END DO
          END DO
          
       END DO
       
    END SELECT
    
  END SUBROUTINE copy_HtoB

  ! Halo to Halo : copy local messages from source field to destination field
  SUBROUTINE copy_HtoH(msg)
    USE omp_para, ONLY : distrib_level
    USE layout_mod, ONLY : layout, order_ij_l_p, order_l_ij_p
    TYPE(t_message), INTENT(INOUT) :: msg
    INTEGER :: dim3, dim4, d3_begin, d3_end, d3, d4
    INTEGER ::  i, ind_src, ind_dest, ij_src, ij_dest, sign, index_order
    
    dim4 = msg%field(1)%dim4
    dim3 = msg%field(1)%dim3
    index_order = layout(1, 1)%index_order
    
    ! TODO : too many copies when tiles are distributed among threads
    SELECT CASE( index_order )
       
    CASE( order_ij_l_p )
       CALL distrib_level( 1, dim3, d3_begin, d3_end )
       !$acc parallel loop collapse(3) present(msg) default(present) async if(msg%ondevice)
       DO d4 = 1, dim4
          DO d3 = d3_begin, d3_end
             DO i=1, msg%message_local_compact%npoints
                ! src
                ind_src = msg%message_local_compact%field_ind_src(i)
                ij_src  = msg%message_local_compact%field_displ_src(i)
                sign    = msg%message_local_compact%sign(i)
                ! dest
                ind_dest = msg%message_local_compact%field_ind_dest(i)
                ij_dest  = msg%message_local_compact%field_displ_dest(i)
                !
                msg%field(ind_dest)%rval4d(ij_dest,d3,d4) = sign * msg%field(ind_src)%rval4d(ij_src,d3,d4)
             END DO
          END DO
       END DO
       
    CASE( order_l_ij_p )
       IF(msg%message_local_compact%npoints /= 0) THEN
          PRINT *, 'npoints = ', msg%message_local_compact%npoints
          STOP 'There should never be local halo copies when unstructured'
       END IF
       DO i=1, msg%message_local_compact%npoints
          ! src
          ind_src = msg%message_local_compact%field_ind_src(i)
          ij_src  = msg%message_local_compact%field_displ_src(i)
          sign    = msg%message_local_compact%sign(i)
          ! dest
          ind_dest = msg%message_local_compact%field_ind_dest(i)
          ij_dest  = msg%message_local_compact%field_displ_dest(i)
          !
          DO d4 = 1, dim4
             DO d3 = d3_begin, d3_end
                msg%field(ind_dest)%rval4d(d3, ij_dest, d4) = sign * msg%field(ind_src)%rval4d(d3, ij_src, d4)
             END DO
          END DO
       END DO
       
    END SELECT
    
  END SUBROUTINE copy_HtoH

  ! Buffer to Halo : copy inbound message to field
  SUBROUTINE copy_BtoH(msg)
    USE omp_para, ONLY : distrib_level
    USE layout_mod, ONLY : layout, order_ij_l_p, order_l_ij_p

    TYPE(t_message), INTENT(INOUT) :: msg
    INTEGER :: field_type, index_order, dim3, dim4, d3_begin, d3_end
    INTEGER :: d4, d3, d34, i, ij, ind, rank, displ, stride, sign 

    field_type = msg%field(1)%field_type
    index_order = layout(1, field_type)%index_order
    dim3 = msg%field(1)%dim3
    dim4 = msg%field(1)%dim4

    SELECT CASE( index_order )

       CASE( order_ij_l_p )
         CALL distrib_level( 1, dim3, d3_begin, d3_end )
         !$acc parallel loop collapse(3) present(msg) default(present) async if(msg%ondevice)
         DO d4 = 1, dim4
            DO d3 = d3_begin, d3_end
               DO i=1, msg%message_in_compact%npoints
                  d34 = d3-1 + dim3*(d4-1) ! linearized index for non-horizontal dimensions
                  ! src
                  rank   = msg%message_in_compact%remote_rank(i)  ! MPI rank from which data has been received
                  displ  = msg%message_in_compact%mpi_displ(i)    ! position in MPI buffer
                  stride = msg%message_in_compact%level_offset(i)
                  sign   = msg%message_in_compact%sign(i)
                  ! dest
                  ind    = msg%message_in_compact%field_ind(i)    ! ind (domain) to write to
                  ij     = msg%message_in_compact%field_displ(i)  ! horizontal index in field
                  !
                   msg%field(ind)%rval4d(ij,d3,d4) = sign * msg%mpi_buffer_in(rank)%buff(displ+stride*d34)
               END DO
            END DO
         END DO

       CASE( order_l_ij_p )
         !$OMP DO SCHEDULE(static)
         DO i=1, msg%message_in_compact%npoints
            ! src
            rank   = msg%message_in_compact%remote_rank(i)  ! MPI rank from which data has been received
            displ  = msg%message_in_compact%mpi_displ(i)    ! position in MPI buffer
            stride = msg%message_in_compact%level_offset(i)
            sign   = msg%message_in_compact%sign(i)
            ! dest
            ind    = msg%message_in_compact%field_ind(i)    ! ind (domain) to write to
            ij     = msg%message_in_compact%field_displ(i)  ! horizontal index in field
            !
            DO d4 = 1, dim4
               DO d3 = 1, dim3
                  d34 = d3-1 + dim3*(d4-1) ! linearized index for non-horizontal dimensions
                 msg%field(ind)%rval4d(d3, ij, d4) = sign * msg%mpi_buffer_in(rank)%buff(displ+stride*d34)
               END DO
            END DO
         END DO
    END SELECT

  END SUBROUTINE copy_BtoH


  !================================== MPI buffer send/receive ================================


  SUBROUTINE send_message(field, message)
    USE mpi_mod
    USE mpipara , ONLY : mpi_size
    TYPE(t_field),POINTER :: field(:)
    TYPE(t_message), TARGET :: message
    INTEGER :: ierr, i

    CALL enter_profile(profile_mpi)

    ! Needed because rval4d is set in init_message
    IF( .not. associated( message%field, field ) ) then
         CALL dynamico_abort("send_message must be called with the same field used in init_message")
    endif

    !Prepare 'message' for on-device copies if field is on device
    !$omp master
    IF( field(1)%ondevice .and. .not. message%ondevice ) CALL message_create_ondevice(message)
    IF( field(1)%ondevice .neqv. message%ondevice ) CALL dynamico_abort("send_message : internal device/host memory synchronization error")
    ! Check if previous message has been waited
    IF(message%send_seq /= message%wait_seq) &
         CALL dynamico_abort("No matching wait_message before new send_message")
    message%send_seq = message%send_seq + 1
    !$omp end master

    CALL enter_profile(profile_mpi_barrier)
    !$omp barrier
    CALL exit_profile(profile_mpi_barrier)

    CALL enter_profile(profile_mpi_copies)
    CALL copy_HtoB(message)
    CALL exit_profile(profile_mpi_copies)

    CALL enter_profile(profile_mpi_barrier)
    !$acc wait
    CALL exit_profile(profile_mpi_barrier)

    !$omp master
    DO i=0, mpi_size-1
       IF(message%mpi_requests_in(i) /= MPI_REQUEST_NULL) CALL MPI_Start( message%mpi_requests_in(i), ierr )
    END DO
    !$omp end master

    CALL enter_profile(profile_mpi_barrier)
    !$omp barrier
    CALL exit_profile(profile_mpi_barrier)

    !$omp master
    DO i=0, mpi_size-1
       IF(message%mpi_requests_out(i) /= MPI_REQUEST_NULL) CALL MPI_Start( message%mpi_requests_out(i), ierr )
    END DO
    !$omp end master

    CALL enter_profile(profile_mpi_copies)
    CALL copy_HtoH(message)
    CALL exit_profile(profile_mpi_copies)

    CALL exit_profile(profile_mpi)
  END SUBROUTINE send_message

  SUBROUTINE test_message(message)
    USE mpi_mod
    TYPE(t_message) :: message
    INTEGER :: ierr
    LOGICAL :: completed

    ! !$omp master
    ! CALL MPI_Testall( size(message%mpi_requests_out), message%mpi_requests_out, completed, MPI_STATUSES_IGNORE, ierr )
    ! CALL MPI_Testall( size(message%mpi_requests_in), message%mpi_requests_in, completed, MPI_STATUSES_IGNORE, ierr )
    ! !$omp end master
  END SUBROUTINE test_message

  SUBROUTINE wait_message(message)
    USE mpi_mod
    TYPE(t_message), TARGET :: message
    INTEGER :: ierr

    ! Check if message has been sent and not recieved yet
    ! note : barrier needed between this and send_seq increment, and this and wait_seq increment
    ! note : watch out for integer overflow a = b+1 doesn't imply a>b
    IF(message%send_seq /= message%wait_seq+1) THEN
       PRINT*, "WARNING : wait_message called multiple times for one send_message, skipping"
       RETURN ! Don't recieve message if already recieved
    END IF

    CALL enter_profile(profile_mpi)

    CALL enter_profile(profile_mpi_waitall)
    !$omp master
    CALL MPI_Waitall( size(message%mpi_requests_out), message%mpi_requests_out, MPI_STATUSES_IGNORE, ierr )
    CALL MPI_Waitall( size(message%mpi_requests_in), message%mpi_requests_in, MPI_STATUSES_IGNORE, ierr )
    !$omp end master
    CALL exit_profile(profile_mpi_waitall)

    CALL enter_profile(profile_mpi_barrier)
    !$omp barrier
    CALL exit_profile(profile_mpi_barrier)

    CALL enter_profile(profile_mpi_copies)
    CALL copy_BtoH(message)
    CALL exit_profile(profile_mpi_copies)

    !$omp master
    message%wait_seq = message%wait_seq + 1
    !$omp end master

    CALL enter_profile(profile_mpi_barrier)
    !$omp barrier
    CALL exit_profile(profile_mpi_barrier)

    CALL exit_profile(profile_mpi)
  END SUBROUTINE wait_message

  !====== GPU allocation/deallocation of halo exchange information and buffers =====

    SUBROUTINE message_create_ondevice(message)
    USE mpi_mod
    USE mpipara, ONLY : mpi_size, comm_icosa
    TYPE(t_message), INTENT(INOUT) :: message
    INTEGER :: i, ierr

    IF( message%ondevice ) CALL dynamico_abort("Message already on device")

    !$acc enter data copyin(message) async
    !$acc enter data copyin(message%mpi_buffer_in(:)) async
    !$acc enter data copyin(message%mpi_buffer_out(:)) async
    DO i = 0, mpi_size-1
       !$acc enter data copyin(message%mpi_buffer_in(i)%buff(:)) async
       !$acc enter data copyin(message%mpi_buffer_out(i)%buff(:)) async
    END DO

    ! !$acc enter data copyin(message%message_in(:)) async
    ! DO i = 1, size( message%message_in )
    !   !$acc enter data copyin(message%message_in(i)%displs(:)) async
    !   !$acc enter data copyin(message%message_in(i)%sign(:)) async
    ! END DO
    ! !$acc enter data copyin(message%message_out(:)) async
    ! DO i = 1, size( message%message_out )
    !   !$acc enter data copyin(message%message_out(i)%displs(:)) async
    !  !$acc enter data copyin(message%message_out(i)%sign(:)) async
    ! END DO
    ! !$acc enter data copyin(message%message_local(:)) async
    ! DO i = 1, size( message%message_local )
    !   !$acc enter data copyin(message%message_local(i)%displ_src(:)) async
    !   !$acc enter data copyin(message%message_local(i)%displ_dest(:)) async
    !   !$acc enter data copyin(message%message_local(i)%sign(:)) async
    ! END DO

    !$acc enter data copyin(message%field(:)) async
    DO i = 1, SIZE(message%field)
       !$acc enter data copyin(message%field(i)%rval4d(:,:,:)) async
    END DO

    !$acc enter data copyin(message%message_in_compact) async
    !$acc enter data copyin(message%message_in_compact%field_ind(:)) async
    !$acc enter data copyin(message%message_in_compact%field_displ(:)) async
    !$acc enter data copyin(message%message_in_compact%sign(:)) async
    !$acc enter data copyin(message%message_in_compact%remote_rank(:)) async
    !$acc enter data copyin(message%message_in_compact%mpi_displ(:)) async
    !$acc enter data copyin(message%message_in_compact%level_offset(:)) async

    !$acc enter data copyin(message%message_out_compact) async
    !$acc enter data copyin(message%message_out_compact%field_ind(:)) async
    !$acc enter data copyin(message%message_out_compact%field_displ(:)) async
    !$acc enter data copyin(message%message_out_compact%sign(:)) async
    !$acc enter data copyin(message%message_out_compact%remote_rank(:)) async
    !$acc enter data copyin(message%message_out_compact%mpi_displ(:)) async
    !$acc enter data copyin(message%message_out_compact%level_offset(:)) async

    !$acc enter data copyin(message%message_local_compact) async
    !$acc enter data copyin(message%message_local_compact%field_ind_src(:)) async
    !$acc enter data copyin(message%message_local_compact%field_displ_src(:)) async
    !$acc enter data copyin(message%message_local_compact%sign(:)) async
    !$acc enter data copyin(message%message_local_compact%field_ind_dest(:)) async
    !$acc enter data copyin(message%message_local_compact%field_displ_dest(:)) async

    !$acc wait
    DO i = 0, mpi_size-1
       IF( message%mpi_requests_out(i) /= MPI_REQUEST_NULL ) THEN
          CALL MPI_Request_free(message%mpi_requests_out(i), ierr)
          !$acc host_data use_device(message%mpi_buffer_out(i)%buff)
          ! /!\ buff(1) is important for PGI to avoid temporary array copy
          CALL MPI_Send_Init( message%mpi_buffer_out(i)%buff(1), message%mpi_buffer_out(i)%n, MPI_REAL8, i,&
               0, comm_icosa, message%mpi_requests_out(i), ierr )
          !$acc end host_data
       END IF
       IF( message%mpi_requests_in(i) /= MPI_REQUEST_NULL ) THEN
          CALL MPI_Request_free(message%mpi_requests_in(i), ierr)
          !$acc host_data use_device(message%mpi_buffer_in(i)%buff)
          CALL MPI_Recv_Init( message%mpi_buffer_in(i)%buff(1), message%mpi_buffer_in(i)%n, MPI_REAL8, i,&
               0, comm_icosa, message%mpi_requests_in(i), ierr )
          !$acc end host_data
       ENDIF
    END DO
    message%ondevice=.true.
    ! !$acc update device(message%ondevice)
  END SUBROUTINE message_create_ondevice

  SUBROUTINE message_delete_ondevice(message)
    USE mpipara, ONLY : mpi_size
    TYPE(t_message), INTENT(INOUT) :: message
    INTEGER :: i

    IF( .not. message%ondevice ) CALL dynamico_abort("Message not on device")

    ! DO i = 1, size( message%message_in )
    !  !$acc exit data delete(message%message_in(i)%displs(:)) async
    !  !$acc exit data delete(message%message_in(i)%sign(:)) async
    ! END DO
    ! !$acc exit data delete(message%message_in(:)) async
    ! DO i = 1, size( message%message_out )
    !   !$acc exit data delete(message%message_out(i)%displs(:)) async
    !   !!$acc exit data delete(message%message_out(i)%sign(:)) async
    ! END DO
    ! !$acc exit data delete(message%message_out(:)) async
    ! DO i = 1, size( message%message_local )
    !   !$acc exit data delete(message%message_local(i)%displ_src(:)) async
    !   !$acc exit data delete(message%message_local(i)%displ_dest(:)) async
    !   !$acc exit data delete(message%message_local(i)%sign(:)) async
    ! END DO
    ! !$acc exit data delete(message%message_local(:)) async

    DO i = 0, mpi_size-1
       !$acc exit data delete(message%mpi_buffer_in(i)%buff(:)) async
       !$acc exit data delete(message%mpi_buffer_out(i)%buff(:)) async
    END DO
    !$acc exit data delete(message%mpi_buffer_in(:)) async
    !$acc exit data delete(message%mpi_buffer_out(:)) async
    DO i = 1, SIZE(message%field)
       !$acc exit data delete(message%field(i)%rval4d(:,:,:)) async
    END DO
    !$acc exit data delete(message%field(:)) async
    !$acc exit data delete(message) async

    !$acc exit data delete(message%message_in_compact%field_ind(:)) async
    !$acc exit data delete(message%message_in_compact%field_displ(:)) async
    !$acc exit data delete(message%message_in_compact%sign(:)) async
    !$acc exit data delete(message%message_in_compact%remote_rank(:)) async
    !$acc exit data delete(message%message_in_compact%mpi_displ(:)) async
    !$acc exit data delete(message%message_in_compact%level_offset(:)) async
    !$acc exit data delete(message%message_in_compact) async

    !$acc exit data delete(message%message_out_compact%field_ind(:)) async
    !$acc exit data delete(message%message_out_compact%field_displ(:)) async
    !$acc exit data delete(message%message_out_compact%sign(:)) async
    !$acc exit data delete(message%message_out_compact%remote_rank(:)) async
    !$acc exit data delete(message%message_out_compact%mpi_displ(:)) async
    !$acc exit data delete(message%message_out_compact%level_offset(:)) async
    !$acc exit data delete(message%message_out_compact) async

    !$acc exit data delete(message%message_local_compact%field_ind_src(:)) async
    !$acc exit data delete(message%message_local_compact%field_displ_src(:)) async
    !$acc exit data delete(message%message_local_compact%sign(:)) async
    !$acc exit data delete(message%message_local_compact%field_ind_dest(:)) async
    !$acc exit data delete(message%message_local_compact%field_displ_dest(:)) async
    !$acc exit data delete(message%message_local_compact) async

    message%ondevice=.false.
  END SUBROUTINE message_delete_ondevice

END MODULE transfert_types_mod
