! This module contains the requests for transfert_mpi
module transfert_request_mod
  use abort_mod
  use field_mod, only : t_field, field_T, field_U, field_Z
  use domain_mod, only : ndomain, ndomain_glo, domain, domain_glo, domloc_glo_ind, domglo_rank, domglo_loc_ind
  implicit none
  private

  ! 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

  ! 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)
  end type t_request

  type(t_request),save,pointer :: req_i1(:) ! Halos for T fields
  type(t_request),save,pointer :: req_e1_scal(:) ! Halos for U fields
  type(t_request),save,pointer :: req_e1_vect(:) ! Halos for U fields (with sign change)
  type(t_request),save,pointer :: req_z1_scal(:) ! Halos for Z fields

  type(t_request),save,pointer :: req_i0(:) ! Duplicated cells for T fields
  type(t_request),save,pointer :: req_e0_scal(:)
  type(t_request),save,pointer :: req_e0_vect(:)

  public :: t_points, t_request, &
            req_i1, req_e1_scal, req_e1_vect, &
            req_i0, req_e0_scal, req_e0_vect, &
            req_z1_scal, &
            init_all_requests
contains

  ! Initialize requests
  subroutine init_all_requests
    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)
      ENDDO
      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)
      ENDDO
      DO i=ii_begin,ii_end
        CALL request_add_point(req_i0,ind,i,jj_end)
      ENDDO
    ENDDO
    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)
      enddo
      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)
      enddo
      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)
      enddo
    enddo
    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)
      ENDDO
      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)
      ENDDO
      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)
    ENDDO
    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)
      ENDDO
      DO j=jj_begin,jj_end
        CALL request_add_point(req_e1_scal,ind,ii_end+1,j,left)
      ENDDO
      DO j=jj_begin,jj_end
        CALL request_add_point(req_e1_scal,ind,ii_end+1,j-1,lup)
      ENDDO
      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)
      ENDDO
      DO j=jj_begin,jj_end
        CALL request_add_point(req_e1_scal,ind,ii_begin-1,j,right)
      ENDDO
      DO j=jj_begin,jj_end
        CALL request_add_point(req_e1_scal,ind,ii_begin-1,j+1,rdown)
      ENDDO
    ENDDO
    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)
      enddo
      do j=jj_begin,jj_end
        call request_add_point(req_z1_scal,ind,ii_end+1,j,vlup)
      enddo
      do j=jj_begin,jj_end
        call request_add_point(req_z1_scal,ind,ii_end+1,j-1,vup)
      enddo
      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)
      enddo
      do j=jj_begin,jj_end
        call request_add_point(req_z1_scal,ind,ii_begin-1,j,vrup)
      enddo
      do j=jj_begin,jj_end
        call request_add_point(req_z1_scal,ind,ii_begin-1,j,vrdown)
      enddo
    enddo
    call request_exchange(req_z1_scal)
  contains

    ! ---- 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

    ! 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

    ! 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

    ! 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 dimensions, only : iim, jjm
      type(t_points):: points
      integer :: displs(points%npoints)
      integer :: k, i, last, i_min, tmp

      displs = (points%i-1) + (points%j-1)*iim
      if(allocated(points%elt)) displs = displs + (points%elt-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 .or. displs(i_min) /= displs(last)) last=last+1
        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

    ! ---- Requests ----
    ! Initialize request
    subroutine request_init( request, field_type )
      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

    ! 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)
      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 )
      endif
    end subroutine

    ! Each domain send the cell index it needs from other domains
    subroutine request_exchange(req)
      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))
            endif
          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 )
          endif
        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

    ! 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
  end subroutine
end module