! This MODULE contains the requests for transfert_mpi
MODULE transfert_request_mod
   USE domain_mod, ONLY: ndomain, ndomain_glo, domain, domain_glo, domloc_glo_ind, domglo_rank, domglo_loc_ind
   USE transfert_types_mod
   IMPLICIT NONE
   PRIVATE

   PUBLIC ::          init_all_requests

CONTAINS

   ! Initialize requests
   SUBROUTINE init_all_requests
      USE field_mod, ONLY: t_field, field_T, field_U, field_Z
      USE dimensions, ONLY: swap_dimensions, ii_begin, ii_end, jj_begin, jj_end
      USE metric, ONLY: right, rdown, ldown, left, lup, rup, vdown, vlup, vrdown, vrup, vup
      INTEGER :: ind, i, j

      ! Create requests req_* see corresponding MODULE variables for description
      CALL request_init(req_i0, field_T)
      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         DO i = ii_begin, ii_end
            CALL request_add_point(req_i0, ind, i, jj_begin)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_i0, ind, ii_begin, j)
            CALL request_add_point(req_i0, ind, ii_end, j)
         END DO
         DO i = ii_begin, ii_end
            CALL request_add_point(req_i0, ind, i, jj_end)
         END DO
      END DO
      CALL request_exchange(req_i0)

      CALL request_init(req_i1, field_T)
      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         DO i = ii_begin, ii_end + 1
            CALL request_add_point(req_i1, ind, i, jj_begin - 1)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_i1, ind, ii_begin - 1, j)
            CALL request_add_point(req_i1, ind, ii_end + 1, j)
         END DO
         CALL request_add_point(req_i1, ind, ii_begin - 1, jj_end + 1)
         DO i = ii_begin, ii_end
            CALL request_add_point(req_i1, ind, i, jj_end + 1)
         END DO
      END DO
      CALL request_exchange(req_i1)

      CALL request_init(req_e0_scal, field_U)
      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         DO i = ii_begin + 1, ii_end - 1
            CALL request_add_point(req_e0_scal, ind, i, jj_begin, right)
            CALL request_add_point(req_e0_scal, ind, i, jj_end, right)
         END DO
         DO j = jj_begin + 1, jj_end - 1
            CALL request_add_point(req_e0_scal, ind, ii_begin, j, rup)
            CALL request_add_point(req_e0_scal, ind, ii_end, j, rup)
         END DO
         CALL request_add_point(req_e0_scal, ind, ii_begin + 1, jj_begin, left)
         CALL request_add_point(req_e0_scal, ind, ii_begin, jj_begin + 1, ldown)
         CALL request_add_point(req_e0_scal, ind, ii_begin + 1, jj_end, left)
         CALL request_add_point(req_e0_scal, ind, ii_end, jj_begin + 1, ldown)
      END DO
      CALL request_exchange(req_e0_scal)

      ! req_e0_vect is the same request than req_e0_scal but with the vector flag to true
      ALLOCATE (req_e0_vect(ndomain))
      req_e0_vect = req_e0_scal
      req_e0_vect%vector = .true.

      CALL request_init(req_e1_scal, field_U)
      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         DO i = ii_begin, ii_end
            CALL request_add_point(req_e1_scal, ind, i, jj_begin - 1, rup)
            CALL request_add_point(req_e1_scal, ind, i + 1, jj_begin - 1, lup)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_e1_scal, ind, ii_end + 1, j, left)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_e1_scal, ind, ii_end + 1, j - 1, lup)
         END DO
         DO i = ii_begin, ii_end
            CALL request_add_point(req_e1_scal, ind, i, jj_end + 1, ldown)
            CALL request_add_point(req_e1_scal, ind, i - 1, jj_end + 1, rdown)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_e1_scal, ind, ii_begin - 1, j, right)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_e1_scal, ind, ii_begin - 1, j + 1, rdown)
         END DO
      END DO
      CALL request_exchange(req_e1_scal)

      ALLOCATE (req_e1_vect(ndomain))
      req_e1_vect = req_e1_scal
      req_e1_vect%vector = .true.

      CALL request_init(req_z1_scal, field_z)
      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         DO i = ii_begin, ii_end
            CALL request_add_point(req_z1_scal, ind, i, jj_begin - 1, vrup)
            CALL request_add_point(req_z1_scal, ind, i + 1, jj_begin - 1, vup)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_z1_scal, ind, ii_end + 1, j, vlup)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_z1_scal, ind, ii_end + 1, j - 1, vup)
         END DO
         DO i = ii_begin, ii_end
            CALL request_add_point(req_z1_scal, ind, i, jj_end + 1, vdown)
            CALL request_add_point(req_z1_scal, ind, i - 1, jj_end + 1, vrdown)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_z1_scal, ind, ii_begin - 1, j, vrup)
         END DO
         DO j = jj_begin, jj_end
            CALL request_add_point(req_z1_scal, ind, ii_begin - 1, j, vrdown)
         END DO
      END DO
      CALL request_exchange(req_z1_scal)

   END SUBROUTINE init_all_requests

   ! ---- Points ----
   ! Initialize (ALLOCATE) points
   SUBROUTINE points_init(points, has_elt)
      TYPE(t_points):: points
      LOGICAL :: has_elt
      INTEGER, PARAMETER :: INITIAL_ALLOC_SIZE = 10

      ALLOCATE (points%i(INITIAL_ALLOC_SIZE))
      ALLOCATE (points%j(INITIAL_ALLOC_SIZE))
      IF (has_elt) ALLOCATE (points%elt(INITIAL_ALLOC_SIZE))
      points%npoints = 0
   END SUBROUTINE points_init

   ! Append point (i, j [,elt]) to point list
   SUBROUTINE points_append(points, i, j, elt)
      TYPE(t_points), INTENT(INOUT):: points
      INTEGER, INTENT(IN) :: i, j
      INTEGER, INTENT(IN), OPTIONAL :: elt
      INTEGER :: begin_size

      begin_size = points%npoints
      CALL array_append(points%i, begin_size, i)
      begin_size = points%npoints
      CALL array_append(points%j, begin_size, j)
      IF (present(elt)) THEN
         begin_size = points%npoints
         CALL array_append(points%elt, begin_size, elt)
      END IF

      points%npoints = points%npoints + 1
   END SUBROUTINE points_append

   ! Add element to array, and reALLOCATE IF necessary
   SUBROUTINE array_append(a, a_size, elt)
      INTEGER, allocatable, intent(inout) :: a(:)
      INTEGER, intent(inout) :: a_size
      INTEGER, intent(in) :: elt
      INTEGER, 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

   ! Sort and remove duplicates in points
   ! dimensions%iim must be set to the right value with swap_dimensions before
   ! After this CALL, size(points%i) == points%npoints
   SUBROUTINE points_sort(points)
      USE abort_mod, ONLY: dynamico_abort
      USE dimensions, ONLY: iim, jjm
      TYPE(t_points):: points
      INTEGER :: displs(points%npoints)
      INTEGER :: k, i, last, i_min, tmp

      displs = (points%i(1:points%npoints) - 1) + (points%j(1:points%npoints) - 1)*iim
      IF (ALLOCATED(points%elt)) displs = displs + (points%elt(1:points%npoints) - 1)*iim*jjm

      last = 0
      DO i = 1, points%npoints
         i_min = minloc(displs(i:), 1) + i - 1
         tmp = displs(i_min)
         displs(i_min) = displs(i)
         IF (last == 0) THEN
            last = last + 1
         ELSE IF (displs(i_min) /= displs(last)) THEN
            last = last + 1
         END IF
         displs(last) = tmp
      END DO

      points%npoints = last
      DEALLOCATE (points%i); ALLOCATE (points%i(last))
      DEALLOCATE (points%j); ALLOCATE (points%j(last))
      DO k = 1, last
         points%i(k) = mod(displs(k), iim) + 1
         points%j(k) = mod((displs(k) - (points%i(k) - 1))/iim, jjm) + 1
      END DO
      IF (ALLOCATEd(points%elt)) THEN
         DEALLOCATE (points%elt); ALLOCATE (points%elt(last))
         points%elt(:) = displs(1:last)/(iim*jjm) + 1
      END IF

      IF (size(points%i) /= points%npoints) CALL dynamico_abort("Storage too big for points")
   END SUBROUTINE points_sort

   ! ---- Requests ----
   ! Initialize request
   SUBROUTINE request_init(request, field_type)
      USE field_mod, ONLY: t_field, field_T, field_U, field_Z
      TYPE(t_request), pointer, intent(out) :: request(:)
      INTEGER, intent(in) :: field_type
      INTEGER :: ind, remote_ind_glo

      ALLOCATE (request(ndomain))
      DO ind = 1, ndomain
         request(ind)%field_type = field_type
         ALLOCATE (request(ind)%points_BtoH(ndomain_glo))
         ALLOCATE (request(ind)%points_HtoB(ndomain_glo))
         DO remote_ind_glo = 1, ndomain_glo
            CALL points_init(request(ind)%points_BtoH(remote_ind_glo), field_TYPE /= field_T)
            !CALL points_init(request(ind)%points_HtoB(remote_ind_glo)) ! Initialized before MPI communication in request_exchange
         END DO
      END DO
   END SUBROUTINE request_init

   ! Append local point (ind_glo,i,j) to the point list of domain ind, where ind_glo is the domain owning the point
   SUBROUTINE request_add_point(req, ind, i, j, elt)
      USE field_mod, ONLY: t_field, field_T, field_U, field_Z
      TYPE(t_request), intent(inout) :: req(:)
      INTEGER :: ind, i, j
      INTEGER, optional :: elt

      IF (req(ind)%field_TYPE == field_U) THEN
         IF (domain(ind)%edge_assign_domain(elt - 1, i, j) /= domloc_glo_ind(ind)) & ! Add ONLY IF remote domain != local domain
            CALL points_append(req(ind)%points_BtoH(domain(ind)%edge_assign_domain(elt - 1, i, j)), i, j, elt)
      ELSE IF (req(ind)%field_TYPE == field_Z) THEN
         IF (domain(ind)%vertex_assign_domain(elt - 1, i, j) /= domloc_glo_ind(ind)) & ! Add ONLY IF remote domain != local domain
            CALL points_append(req(ind)%points_BtoH(domain(ind)%vertex_assign_domain(elt - 1, i, j)), i, j, elt)
      ELSE ! T field
         IF (domain(ind)%assign_domain(i, j) /= domloc_glo_ind(ind)) &
            CALL points_append(req(ind)%points_BtoH(domain(ind)%assign_domain(i, j)), i, j)
      END IF
   END SUBROUTINE request_add_point

   ! Each domain send the cell index it needs from other domains
   SUBROUTINE request_exchange(req)
      USE field_mod, ONLY: t_field, field_T, field_U, field_Z
      USE mpi_mod
      USE dimensions, ONLY: swap_dimensions
      USE mpipara, ONLY: comm_icosa
      TYPE(t_request), target, intent(in) :: req(:)
      TYPE(t_points) :: points_send(ndomain, ndomain_glo)
      TYPE(t_points), pointer :: points

      INTEGER :: ind_loc, remote_ind_glo, k
      INTEGER :: i, j, elt
      INTEGER :: reqs(ndomain, ndomain_glo, 2), reqs2(ndomain, ndomain_glo, 2), reqs3(ndomain, ndomain_glo, 3), ierr

      ! Transform position in local domain to position in remote domain pos : (i,j) -> pos_send : (assign_i(i,j), assign_j(i,j))
      DO ind_loc = 1, ndomain
         CALL swap_dimensions(ind_loc)
         DO remote_ind_glo = 1, ndomain_glo
            points => req(ind_loc)%points_BtoH(remote_ind_glo)
            CALL points_sort(points)
            CALL points_init(points_send(ind_loc, remote_ind_glo), req(ind_loc)%field_TYPE /= field_T)
            DO k = 1, points%npoints
               i = points%i(k)
               j = points%j(k)
               IF (req(ind_loc)%field_TYPE == field_U) THEN
                  elt = points%elt(k) - 1 ! WARNING : domain%edge_assign_* USE index in 0:5 instead of 1:6 everywhere ELSE
                  CALL points_append(points_send(ind_loc, remote_ind_glo), domain(ind_loc)%edge_assign_i(elt, i, j), &
                                     DOmain(ind_loc)%edge_assign_j(elt, i, j), &
                                     domain(ind_loc)%edge_assign_pos(elt, i, j) + 1)
               ELSE IF (req(ind_loc)%field_TYPE == field_Z) THEN
                  elt = points%elt(k) - 1 ! WARNING : domain%vertex_assign_* USE index in 0:5 instead of 1:6 everywhere ELSE
                  CALL points_append(points_send(ind_loc, remote_ind_glo), domain(ind_loc)%vertex_assign_i(elt, i, j), &
                                     domain(ind_loc)%vertex_assign_j(elt, i, j), &
                                     domain(ind_loc)%vertex_assign_pos(elt, i, j) + 1)
               ELSE ! T field
                  CALL points_append(points_send(ind_loc, remote_ind_glo), domain(ind_loc)%assign_i(i, j), &
                                     domain(ind_loc)%assign_j(i, j))
               END IF
            END DO
         END DO
      END DO

      ! Exchange sizes
      DO ind_loc = 1, ndomain
         DO remote_ind_glo = 1, ndomain_glo
            CALL MPI_Isend(points_send(ind_loc, remote_ind_glo)%npoints, 1, MPI_INT, domglo_rank(remote_ind_glo), &
                           remote_ind_glo + ndomain_glo*domloc_glo_ind(ind_loc), comm_icosa, reqs(ind_loc, remote_ind_glo, 1), ierr)
            CALL MPI_Irecv(req(ind_loc)%points_HtoB(remote_ind_glo)%npoints, 1, MPI_INT, domglo_rank(remote_ind_glo), &
                           domloc_glo_ind(ind_loc) + ndomain_glo*remote_ind_glo, comm_icosa, reqs(ind_loc, remote_ind_glo, 2), ierr)
         END DO
      END DO
      CALL waitall(reqs, ndomain*ndomain_glo*2)

      ! Exchange points_send -> points_recv
      DO ind_loc = 1, ndomain
         DO remote_ind_glo = 1, ndomain_glo
            CALL MPI_Isend(points_send(ind_loc, remote_ind_glo)%i, points_send(ind_loc, remote_ind_glo)%npoints, MPI_INT, &
                           domglo_rank(remote_ind_glo), remote_ind_glo + ndomain_glo*domloc_glo_ind(ind_loc), &
                           comm_icosa, reqs(ind_loc, remote_ind_glo, 1), ierr)
            CALL MPI_Isend(points_send(ind_loc, remote_ind_glo)%j, points_send(ind_loc, remote_ind_glo)%npoints, MPI_INT, &
                           domglo_rank(remote_ind_glo), remote_ind_glo + ndomain_glo*domloc_glo_ind(ind_loc), &
                           comm_icosa, reqs2(ind_loc, remote_ind_glo, 1), ierr)

            points => req(ind_loc)%points_HtoB(remote_ind_glo) ! Points to recieve (HtoB)
            ALLOCATE (points%i(points%npoints))
            ALLOCATE (points%j(points%npoints))
            CALL MPI_Irecv(points%i, points%npoints, MPI_INT, domglo_rank(remote_ind_glo), &
                           domloc_glo_ind(ind_loc) + ndomain_glo*remote_ind_glo, &
                           comm_icosa, reqs(ind_loc, remote_ind_glo, 2), ierr)
            CALL MPI_Irecv(points%j, points%npoints, MPI_INT, domglo_rank(remote_ind_glo), &
                           domloc_glo_ind(ind_loc) + ndomain_glo*remote_ind_glo, &
                           comm_icosa, reqs2(ind_loc, remote_ind_glo, 2), ierr)

            IF (req(1)%field_type /= field_T) THEN
               CALL MPI_Isend(points_send(ind_loc, remote_ind_glo)%elt, points_send(ind_loc, remote_ind_glo)%npoints, MPI_INT, &
                              domglo_rank(remote_ind_glo), remote_ind_glo + ndomain_glo*domloc_glo_ind(ind_loc), &
                              comm_icosa, reqs3(ind_loc, remote_ind_glo, 1), ierr)
               ALLOCATE (points%elt(points%npoints))
               CALL MPI_Irecv(points%elt, points%npoints, MPI_INT, &
                              domglo_rank(remote_ind_glo), domloc_glo_ind(ind_loc) + ndomain_glo*remote_ind_glo, &
                              comm_icosa, reqs3(ind_loc, remote_ind_glo, 2), ierr)
            END IF
         END DO
      END DO
      CALL waitall(reqs, ndomain*ndomain_glo*2)
      CALL waitall(reqs2, ndomain*ndomain_glo*2)
      IF (req(1)%field_TYPE /= field_T) CALL waitall(reqs3, ndomain*ndomain_glo*2)
   END SUBROUTINE request_exchange

   ! MPI_Waitall, but with a N-Dim array of requests
   SUBROUTINE waitall(reqs, size)
      USE mpi_mod
      INTEGER :: size, reqs(size), ierr
      CALL MPI_Waitall(size, reqs, MPI_STATUSES_IGNORE, ierr)
   END SUBROUTINE waitall
END MODULE transfert_request_mod
