! Module for MPI communication of field halos
! This module uses Fortran 2003 features : move_alloc intrinsic
MODULE transfert_mpi_mod
   USE abort_mod, ONLY: dynamico_abort
   USE profiling_mod, ONLY: register_id
   USE field_mod, ONLY: t_field
   USE transfert_types_mod
   USE transfert_request_mod, ONLY: init_all_requests_ico => init_all_requests
   USE transfert_request_unst_mod, ONLY: init_all_requests_unst

   IMPLICIT NONE
   PRIVATE

   public :: t_message, t_request, &
             req_i1, req_e1_scal, req_e1_vect, &
             req_i0, req_e0_scal, req_e0_vect, &
             req_z1_scal, &
             send_message, &
             wait_message, &
             test_message, &
             init_transfert, &
             init_message, &
             finalize_message

CONTAINS

   ! Initialize transfert : must be called before any other transfert_mpi routines
   SUBROUTINE init_transfert
      USE mpi_mod, ONLY: MPI_THREAD_SINGLE, MPI_THREAD_FUNNELED
      USE mpipara, ONLY: mpi_threading_mode
      USE profiling_mod, ONLY: register_id
      USE grid_param

      !$omp master
      ! Check requested threads support
      IF (mpi_threading_mode /= MPI_THREAD_SINGLE .and. mpi_threading_mode /= MPI_THREAD_FUNNELED) &
         CALL dynamico_abort("Only single and funneled threading mode are supported.")

      ! Register profiling ids
      CALL register_id("MPI", profile_mpi)
      IF (profile_mpi_detail) THEN
         CALL register_id("MPI_copies", profile_mpi_copies)
         CALL register_id("MPI_waitall", profile_mpi_waitall)
         CALL register_id("MPI_omp_barrier", profile_mpi_barrier)
      ELSE
         profile_mpi_copies = profile_mpi
         profile_mpi_waitall = profile_mpi
         profile_mpi_barrier = profile_mpi
      END IF

      ! Initialize requests
      SELECT CASE (grid_type)
      CASE (grid_ico)
         CALL init_all_requests_ico()
      CASE (grid_unst)
         CALL init_all_requests_unst(pmesh)
      END SELECT
      !$omp end master
      !$omp barrier
   END SUBROUTINE init_transfert

   SUBROUTINE init_message(field, request, message, name)
      USE domain_mod, ONLY: ndomain_glo, domloc_glo_ind, domglo_loc_ind, domglo_rank
      USE mpi_mod
      USE mpipara
      USE grid_param

      TYPE(t_field), POINTER, INTENT(IN) :: field(:)
      TYPE(t_request), POINTER, INTENT(INOUT) :: request(:)
      TYPE(t_message), TARGET, INTENT(OUT) :: message ! Needs intent out for CALL to finalize_message
      CHARACTER(LEN=*), INTENT(IN), OPTIONAL :: name

      !$omp barrier
      !$omp master
      !init off-device
      message%ondevice = .false.
      message%send_seq = 0
      message%wait_seq = 0

      IF (request(1)%field_type /= field(1)%field_type) CALL dynamico_abort("init_message : field_type/request mismatch")

      message%field => field
      message%request => request
      ! Fortran 2003 : assignment to allocatable array allocates the array to fit r.h.s
      IF (request(1)%has_message) THEN
         message%message_in = request(1)%message_in
         message%message_out = request(1)%message_out
         message%message_local = request(1)%message_local
      ELSE
         CALL init_message_in_out_local(message)
         request(1)%message_in = message%message_in
         request(1)%message_out = message%message_out
         request(1)%message_local = message%message_local
         request(1)%has_message = .TRUE.
      END IF

      CALL message_alloc_mpi_buf(message)
      CALL message_create_MPI_request(message)
      CALL message_compactify(message)

      !$omp end master
      !$omp barrier
   END SUBROUTINE init_message

   SUBROUTINE init_message_in_out_local(message)
      ! Create list of inbound/outbound/local messages
      USE mpi_mod
      USE mpipara
      USE domain_mod, ONLY: ndomain_glo, domloc_glo_ind, domglo_loc_ind, domglo_rank
      TYPE(t_message), TARGET, INTENT(INOUT) :: message

      TYPE(t_local_submessage), ALLOCATABLE :: message_local_tmp(:)
      TYPE(t_submessage), ALLOCATABLE :: message_in_tmp(:), message_out_tmp(:)
      TYPE(t_request), POINTER :: request(:)
      TYPE(t_submessage) :: submessage_in, submessage_out
      TYPE(t_local_submessage) :: submessage_local
      INTEGER, PARAMETER :: INITIAL_ALLOC_SIZE = 10, GROW_FACTOR = 2
      INTEGER                  :: field_type, ind, ind_loc, remote_ind_glo, loc_ind_glo, i, k
      INTEGER                  :: message_in_size, message_out_size, message_local_size

      request => message%request
      field_type = request(1)%field_type

      ALLOCATE (message_in_tmp(INITIAL_ALLOC_SIZE))
      message_in_size = 0
      ALLOCATE (message_out_tmp(INITIAL_ALLOC_SIZE))
      message_out_size = 0
      ALLOCATE (message_local_tmp(INITIAL_ALLOC_SIZE))
      message_local_size = 0
      DO loc_ind_glo = 1, ndomain_glo
         DO remote_ind_glo = 1, ndomain_glo
            IF (domglo_rank(loc_ind_glo) == mpi_rank) THEN
               ind_loc = domglo_loc_ind(loc_ind_glo)
               IF (domglo_rank(remote_ind_glo) == mpi_rank) THEN ! If sending to local domain
                  IF (request(ind_loc)%points_HtoB(remote_ind_glo)%npoints > 0) THEN ! Add only non-empty messages
                     ! Add local message ind_loc -> remote_ind_glo, aggregarting submessage_in and submessage_out into submessage_local
                     submessage_out = make_submessage(field_type, &
                                                      request(ind_loc)%points_HtoB(remote_ind_glo), &
                                                      ind_loc, remote_ind_glo, request(1)%vector)
                     submessage_in = make_submessage(field_type, &
                                                     request(domglo_loc_ind(remote_ind_glo))%points_BtoH(domloc_glo_ind(ind_loc)), &
                                                     domglo_loc_ind(remote_ind_glo), domloc_glo_ind(ind_loc), request(1)%vector)
                     submessage_local%src_ind_loc = ind_loc
                     submessage_local%dest_ind_loc = domglo_loc_ind(remote_ind_glo)
                     submessage_local%npoints = submessage_out%npoints
                     submessage_local%displ_src = submessage_out%displs
                     submessage_local%displ_dest = submessage_in%displs
                     submessage_local%sign = submessage_in%sign
                     ! Add to local message list
                     CALL array_append_local_submessage(message_local_tmp, message_local_size, submessage_local)
                  END IF
               ELSE ! If remote domain
                  ! When data to send to remote_domain, add submessage in message%message_out
                  IF (request(ind_loc)%points_HtoB(remote_ind_glo)%npoints > 0) THEN
                     submessage_out = make_submessage(field_type, request(ind_loc)%points_HtoB(remote_ind_glo), &
                                                      ind_loc, remote_ind_glo, request(1)%vector)
                     CALL array_append_submessage(message_out_tmp, message_out_size, submessage_out)
                  END IF
               END IF
            END IF
         END DO
      END DO
      ! Recv and Send submessages are transposed to recieve and send in same order
      ! We iterate over global domain index to match sends with recieves (local domains are not ordered like global domains)
      DO remote_ind_glo = 1, ndomain_glo
         DO loc_ind_glo = 1, ndomain_glo
            IF ((domglo_rank(loc_ind_glo) == mpi_rank) .and. (domglo_rank(remote_ind_glo) /= mpi_rank)) THEN
               ind_loc = domglo_loc_ind(loc_ind_glo)
               IF (request(ind_loc)%points_BtoH(remote_ind_glo)%npoints > 0) THEN
                  submessage_in = make_submessage(field_type, request(ind_loc)%points_BtoH(remote_ind_glo), &
                                                  ind_loc, remote_ind_glo, request(1)%vector)
                  CALL array_append_submessage(message_in_tmp, message_in_size, submessage_in)
               END IF
            END IF
         END DO
      END DO

      ! Trim message_xx_tmp and put it in message%message_xx
      ! Fortran 2003 : assignment to allocatable array allocates the array to fit r.h.s
      message%message_in = message_in_tmp(:message_in_size)
      message%message_out = message_out_tmp(:message_out_size)
      message%message_local = message_local_tmp(:message_local_size)
   END SUBROUTINE init_message_in_out_local

   ! Generate submessage from points
   FUNCTION make_submessage(field_type, points, ind_loc, remote_ind_glo, vector) result(submessage)
      USE dimensions, ONLY: swap_dimensions, iim, u_pos, z_pos
      USE domain_mod, ONLY: domain, domglo_rank
      USE field_mod, ONLY: field_T, field_U, field_Z
      INTEGER, INTENT(IN) :: field_type
      TYPE(t_points), INTENT(IN) :: points
      INTEGER, INTENT(IN) :: ind_loc, remote_ind_glo
      LOGICAL, INTENT(IN) :: vector
      INTEGER :: k
      TYPE(t_submessage) :: submessage

      CALL swap_dimensions(ind_loc)
      submessage%ind_loc = ind_loc
      submessage%remote_ind_glo = remote_ind_glo
      submessage%remote_rank = domglo_rank(remote_ind_glo)
      submessage%npoints = points%npoints
      submessage%mpi_buffer_displ = -1 ! Buffers not allocated yet
      ALLOCATE (submessage%displs(points%npoints))
      submessage%displs(:) = points%i + (points%j - 1)*iim
      IF (field_type == field_U) submessage%displs = submessage%displs + u_pos(points%elt)
      IF (field_type == field_Z) submessage%displs = submessage%displs + z_pos(points%elt)
      ALLOCATE (submessage%sign(points%npoints))
      IF (vector) THEN ! For U fields only
         submessage%sign(:) = (/(domain(ind_loc)%edge_assign_sign(points%elt(k) - 1, &
                                                                  points%i(k), points%j(k)), k=1, points%npoints)/)
      ELSE
         submessage%sign(:) = 1
      END IF
   END FUNCTION make_submessage

   ! Add element to array, and reallocate if necessary
   SUBROUTINE array_append_submessage(a, a_size, elt)
      TYPE(t_submessage), ALLOCATABLE, INTENT(INOUT) :: a(:)
      INTEGER, INTENT(INOUT) :: a_size
      TYPE(t_submessage), INTENT(IN) :: elt
      TYPE(t_submessage), ALLOCATABLE :: a_tmp(:)
      INTEGER, PARAMETER :: GROW_FACTOR = 2

      IF (size(a) <= a_size) THEN
         ALLOCATE (a_tmp(a_size*GROW_FACTOR))
         a_tmp(1:a_size) = a(1:a_size)
         CALL move_alloc(a_tmp, a)
      end if
      a_size = a_size + 1
      a(a_size) = elt
   END SUBROUTINE array_append_submessage

   ! Add element to array, and reallocate if necessary
   SUBROUTINE array_append_local_submessage(a, a_size, elt)
      TYPE(t_local_submessage), ALLOCATABLE, INTENT(INOUT) :: a(:)
      INTEGER, INTENT(INOUT) :: a_size
      TYPE(t_local_submessage), INTENT(IN) :: elt
      TYPE(t_local_submessage), ALLOCATABLE :: a_tmp(:)
      INTEGER, PARAMETER :: GROW_FACTOR = 2

      IF (size(a) <= a_size) THEN
         ALLOCATE (a_tmp(a_size*GROW_FACTOR))
         a_tmp(1:a_size) = a(1:a_size)
         CALL move_alloc(a_tmp, a)
      end if
      a_size = a_size + 1
      a(a_size) = elt
   END SUBROUTINE array_append_local_submessage

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

      !$omp barrier
      !$omp master
      IF (message%send_seq /= message%wait_seq) CALL dynamico_abort("No matching wait_message before finalization")

      IF (message%ondevice) CALL message_delete_ondevice(message)
      DEALLOCATE (message%message_in)
      DEALLOCATE (message%message_out)
      DEALLOCATE (message%message_local)
      DO i = 0, mpi_size - 1
         IF (message%mpi_requests_in(i) /= MPI_REQUEST_NULL) CALL MPI_Request_free(message%mpi_requests_in(i), ierr)
         IF (message%mpi_requests_out(i) /= MPI_REQUEST_NULL) CALL MPI_Request_free(message%mpi_requests_out(i), ierr)
         DEALLOCATE (message%mpi_buffer_in(i)%buff)
         DEALLOCATE (message%mpi_buffer_out(i)%buff)
      END DO
      DEALLOCATE (message%mpi_buffer_in)
      DEALLOCATE (message%mpi_buffer_out)
      DEALLOCATE (message%mpi_requests_in)
      DEALLOCATE (message%mpi_requests_out)
      DEALLOCATE (message%message_in_compact)
      DEALLOCATE (message%message_out_compact)
      DEALLOCATE (message%message_local_compact)
      !$omp END master
      !$omp barrier
   END SUBROUTINE finalize_message

END MODULE transfert_mpi_mod
