module psb_comm_rma_mod
  use psb_const_mod
  use psb_desc_const_mod, only: psb_proc_id_, psb_n_elem_recv_, psb_elem_recv_, &
       & psb_n_elem_send_, psb_elem_send_
  use psb_error_mod
#ifdef PSB_MPI_MOD
  use mpi
#endif
    use psb_comm_schemes_mod, only: psb_comm_handle_type, psb_comm_rma_pull_, psb_comm_rma_push_, &
      & psb_comm_unknown_
  implicit none
  integer(psb_mpk_), parameter :: psb_rma_meta_tag = 913
#ifdef PSB_MPI_H
  include 'mpif.h'
#endif

  type, extends(psb_comm_handle_type) :: psb_comm_rma_handle
    integer(psb_mpk_) :: win = mpi_win_null
    logical :: window_ready = .false.
    logical :: window_open = .false.
    logical :: layout_ready = .false.
    integer(psb_ipk_) :: layout_nnbr = -1
    integer(psb_ipk_) :: layout_send = -1
    integer(psb_ipk_) :: layout_recv = -1
    integer(psb_ipk_), allocatable :: peer_proc(:)
    integer(psb_ipk_), allocatable :: peer_send_counts(:)
    integer(psb_ipk_), allocatable :: peer_recv_counts(:)
    integer(psb_ipk_), allocatable :: peer_send_displs(:)
    integer(psb_ipk_), allocatable :: peer_recv_displs(:)
    integer(psb_ipk_), allocatable :: peer_mpi_rank(:)
    integer(psb_ipk_), allocatable :: peer_remote_send_displs(:)
    integer(psb_ipk_), allocatable :: peer_remote_recv_displs(:)
    integer(psb_ipk_), allocatable :: peer_send_indexes(:)
    integer(psb_ipk_), allocatable :: peer_recv_indexes(:)
  contains
    procedure, pass :: init => psb_comm_rma_init
    procedure, pass :: free => psb_comm_rma_free
    procedure, pass :: clear_memory_buffer_layout => psb_comm_rma_clear_memory_buffer_layout
    procedure, pass :: init_memory_buffer_layout => psb_comm_rma_ini_memory_buffer_layout
    procedure, pass :: set_swap_status => psb_comm_rma_set_swap_status
    procedure, pass :: get_swap_status => psb_comm_rma_get_swap_status
  end type psb_comm_rma_handle

contains

  subroutine psb_comm_rma_init(this, info)
    class(psb_comm_rma_handle), intent(inout) :: this
    integer(psb_ipk_), intent(out) :: info
    info = 0
    this%comm_type = psb_comm_unknown_
    this%id = 0
    this%swap_status = 0
    this%win = mpi_win_null
    this%window_ready = .false.
    this%window_open = .false.
    this%layout_ready = .false.
    this%layout_nnbr = -1
    this%layout_send = -1
    this%layout_recv = -1
    call this%clear_memory_buffer_layout(info)
  end subroutine psb_comm_rma_init

  subroutine psb_comm_rma_clear_memory_buffer_layout(this, info)
    class(psb_comm_rma_handle), intent(inout) :: this
    integer(psb_ipk_), intent(out) :: info

    info = psb_success_
    if (allocated(this%peer_proc)) deallocate(this%peer_proc)
    if (allocated(this%peer_send_counts)) deallocate(this%peer_send_counts)
    if (allocated(this%peer_recv_counts)) deallocate(this%peer_recv_counts)
    if (allocated(this%peer_send_displs)) deallocate(this%peer_send_displs)
    if (allocated(this%peer_recv_displs)) deallocate(this%peer_recv_displs)
    if (allocated(this%peer_mpi_rank)) deallocate(this%peer_mpi_rank)
    if (allocated(this%peer_remote_send_displs)) deallocate(this%peer_remote_send_displs)
    if (allocated(this%peer_remote_recv_displs)) deallocate(this%peer_remote_recv_displs)
    if (allocated(this%peer_send_indexes)) deallocate(this%peer_send_indexes)
    if (allocated(this%peer_recv_indexes)) deallocate(this%peer_recv_indexes)

    this%layout_ready = .false.
    this%layout_nnbr = -1
    this%layout_send = -1
    this%layout_recv = -1
  end subroutine psb_comm_rma_clear_memory_buffer_layout

    subroutine psb_comm_rma_ini_memory_buffer_layout(this, info, comm_list, peer_mpi_rank, &
      & num_neighbors, total_send, total_recv, my_rank, icomm)
    class(psb_comm_rma_handle), intent(inout) :: this
    integer(psb_ipk_), intent(out) :: info
    integer(psb_ipk_), intent(in) :: comm_list(:)
    integer(psb_ipk_), intent(in) :: peer_mpi_rank(:)
    integer(psb_ipk_), intent(in) :: num_neighbors, total_send, total_recv
    integer(psb_ipk_), intent(in) :: my_rank
    integer(psb_mpk_), intent(in) :: icomm

    integer(psb_mpk_) :: iret, p2pstat(mpi_status_size), prc_rank
    integer(psb_ipk_) :: n_neighbors, send_total, recv_total
    integer(psb_ipk_) :: neighbor_idx, item_idx
    integer(psb_ipk_) :: list_pos, send_offset, recv_offset
    integer(psb_ipk_) :: proc_to_comm, recv_count, send_count
    integer(psb_ipk_) :: local_meta(4), remote_meta(4)

    call this%clear_memory_buffer_layout(info)
    if (info /= psb_success_) return

    n_neighbors = num_neighbors
    send_total = total_send
    recv_total = total_recv

    if (n_neighbors > 0) then
       allocate(this%peer_proc(n_neighbors), this%peer_send_counts(n_neighbors), &
         & this%peer_recv_counts(n_neighbors), this%peer_send_displs(n_neighbors), &
         & this%peer_recv_displs(n_neighbors), this%peer_mpi_rank(n_neighbors), &
         & this%peer_remote_send_displs(n_neighbors), this%peer_remote_recv_displs(n_neighbors), stat=iret)
      if (iret /= 0) then
        info = psb_err_alloc_dealloc_
        return
      end if
    end if

    if (send_total > 0) then
      allocate(this%peer_send_indexes(send_total), stat=iret)
      if (iret /= 0) then
        info = psb_err_alloc_dealloc_
        return
      end if
    end if

    if (recv_total > 0) then
      allocate(this%peer_recv_indexes(recv_total), stat=iret)
      if (iret /= 0) then
        info = psb_err_alloc_dealloc_
        return
      end if
    end if

    list_pos = 1
    send_offset = 0
    recv_offset = 0
    do neighbor_idx = 1, n_neighbors
      proc_to_comm = comm_list(list_pos + psb_proc_id_)
      recv_count = comm_list(list_pos + psb_n_elem_recv_)
      send_count = comm_list(list_pos + recv_count + psb_n_elem_send_)

      this%peer_proc(neighbor_idx) = proc_to_comm
      this%peer_recv_counts(neighbor_idx) = recv_count
      this%peer_send_counts(neighbor_idx) = send_count
      this%peer_recv_displs(neighbor_idx) = recv_offset
      this%peer_send_displs(neighbor_idx) = send_offset
      this%peer_mpi_rank(neighbor_idx) = peer_mpi_rank(neighbor_idx)

      if (recv_count > 0) then
        do item_idx = 1, recv_count
          this%peer_recv_indexes(recv_offset + item_idx) = comm_list(list_pos + psb_elem_recv_ + item_idx - 1)
        end do
      end if

      if (send_count > 0) then
        do item_idx = 1, send_count
          this%peer_send_indexes(send_offset + item_idx) = comm_list(list_pos + recv_count + psb_elem_send_ + item_idx - 1)
        end do
      end if

      recv_offset = recv_offset + recv_count
      send_offset = send_offset + send_count
      list_pos = list_pos + recv_count + send_count + 3
    end do

    do neighbor_idx = 1, n_neighbors
      proc_to_comm = this%peer_proc(neighbor_idx)
      if (proc_to_comm /= my_rank) then
        prc_rank = this%peer_mpi_rank(neighbor_idx)
        local_meta = (/ this%peer_send_displs(neighbor_idx)+1, this%peer_send_counts(neighbor_idx), &
             & this%peer_recv_displs(neighbor_idx)+1+send_total, this%peer_recv_counts(neighbor_idx) /)
           call mpi_sendrecv(local_meta, 4, psb_mpi_mpk_, prc_rank, psb_rma_meta_tag, &
             & remote_meta, 4, psb_mpi_mpk_, prc_rank, psb_rma_meta_tag, icomm, p2pstat, iret)
        if (iret /= mpi_success) then
          info = psb_err_mpi_error_
          return
        end if
        this%peer_remote_send_displs(neighbor_idx) = remote_meta(1)
        this%peer_remote_recv_displs(neighbor_idx) = remote_meta(3)
      else
        this%peer_remote_send_displs(neighbor_idx) = this%peer_send_displs(neighbor_idx)+1
        this%peer_remote_recv_displs(neighbor_idx) = this%peer_recv_displs(neighbor_idx)+1+send_total
      end if
    end do

    if ((send_offset /= send_total) .or. (recv_offset /= recv_total)) then
      info = psb_err_internal_error_
      return
    end if

    this%layout_nnbr = n_neighbors
    this%layout_send = send_total
    this%layout_recv = recv_total
    this%layout_ready = .true.
  end subroutine psb_comm_rma_ini_memory_buffer_layout

  subroutine psb_comm_rma_free(this, info)
#ifdef PSB_MPI_MOD
    use mpi
#endif
    class(psb_comm_rma_handle), intent(inout) :: this
    integer(psb_ipk_), intent(out) :: info
    integer(psb_mpk_) :: iret

    info = 0
    if (this%window_open) then
      call mpi_win_unlock_all(this%win, iret)
      this%window_open = .false.
    end if
    if (this%win /= mpi_win_null) then
      call mpi_win_free(this%win, iret)
      this%win = mpi_win_null
    end if
    this%window_ready = .false.
    this%layout_ready = .false.
    this%layout_nnbr = -1
    this%layout_send = -1
    this%layout_recv = -1
    call this%clear_memory_buffer_layout(info)
  end subroutine psb_comm_rma_free

  subroutine psb_comm_rma_set_swap_status(this, flag, info)
    class(psb_comm_rma_handle), intent(inout) :: this
    integer(psb_ipk_), intent(in) :: flag
    integer(psb_ipk_), intent(out) :: info
    info = 0
    this%swap_status = flag
  end subroutine psb_comm_rma_set_swap_status

  subroutine psb_comm_rma_get_swap_status(this, flag, info)
    class(psb_comm_rma_handle), intent(in) :: this
    integer(psb_ipk_), intent(out) :: flag
    integer(psb_ipk_), intent(out) :: info
    info = 0
    flag = this%swap_status
  end subroutine psb_comm_rma_get_swap_status

end module psb_comm_rma_mod