!> @file halo_exchange_2d.F90
!> @brief Two-axis non-blocking halo exchange for the 2D block decomposition.
!!
!! Exchanges +/-x then +/-y halos sequentially; no corner exchange. Each axis
!! uses four non-blocking operations (2 Irecv + 2 Isend) followed by a single
!! MPI_Waitall, matching the 1D halo_exchange pattern.
!!
!! Neighbours stored as MPI_PROC_NULL (non-periodic physical boundaries) make
!! the matching Irecv/Isend calls silent no-ops per MPI standard §6.6. The
!! conditional halo-copy guards that follow are still needed so stale (-1)
!! interior values are not written into those halo positions.
!!
!! The Cartesian communicator is obtained via cart_comm_2d(); call
!! setup_decomp_2d(state) before exchange_halos_2d.
!!
!! Legacy-MPI note (CFD_SOLVER_LEGACY_MPI): MPI_Request is a bare integer under
!! the legacy `mpi` module; the mpi_f08 path uses derived-type handles.
!! MPI_STATUSES_IGNORE is provided by both but requires the same guard to stay
!! type-consistent. MPI_DOUBLE_PRECISION is an integer constant in both paths.
module halo_exchange_2d
#ifdef CFD_SOLVER_LEGACY_MPI
  use mpi, only: MPI_Isend, MPI_Irecv, MPI_Waitall, MPI_DOUBLE_PRECISION, &
                 MPI_STATUSES_IGNORE
#else
  use mpi_f08, only: MPI_Isend, MPI_Irecv, MPI_Waitall, MPI_DOUBLE_PRECISION, &
                     MPI_STATUSES_IGNORE, MPI_Request, MPI_Comm
#endif
  use precision, only: wp
  use solver_state_2d, only: solver_state_2d_t, neq2d
  use domain_decomposition_2d, only: decomp_2d_t
  use mpi_cart_2d, only: cart_comm_2d
  use mpi_runtime, only: mpi_proc_null_value, parallel_fatal
  use option_registry, only: method_fdm
  implicit none
  private
  public :: exchange_halos_2d

  ! Use distinct lo/hi tags on each axis so that when a periodic boundary
  ! causes x_lo_neighbour == x_hi_neighbour (e.g. np=2 with dims=[2,1]) the
  ! two concurrent Irecv/Isend pairs targeting the same rank are
  ! unambiguously matched. Mirrors the TAG_HALO_L / TAG_HALO_R convention
  ! from the 1D halo_exchange module.
  integer, parameter :: TAG_X_LO = 2001, TAG_X_HI = 2002
  integer, parameter :: TAG_Y_LO = 2003, TAG_Y_HI = 2004

contains

  !> Fill state%ub halo cells from MPI neighbours in both x and y directions.
  !!
  !! @param state  Solver state (ub must already be allocated via allocate_work_arrays_2d).
  !! @param d      Per-rank 2D decomposition (neighbour ranks, local extents).
  subroutine exchange_halos_2d(state, d)
    type(solver_state_2d_t), intent(inout) :: state
    type(decomp_2d_t), intent(in) :: d

#ifdef CFD_SOLVER_LEGACY_MPI
    integer :: req(4)
    integer :: comm
#else
    type(MPI_Request) :: req(4)
    type(MPI_Comm) :: comm
#endif

    integer :: h, nxl, nyl, cnt, ierr, k, j, i
    integer :: xlo_s, xhi_s, ylo_s, yhi_s
    logical :: nodal_skip
    real(wp), allocatable :: sx_lo(:,:,:), sx_hi(:,:,:), rx_lo(:,:,:), rx_hi(:,:,:)
    real(wp), allocatable :: sy_lo(:,:,:), sy_hi(:,:,:), ry_lo(:,:,:), ry_hi(:,:,:)

    h   = state % halo_width
    nxl = state % nx_local
    nyl = state % ny_local
    comm = cart_comm_2d()

    ! Periodic-wrap node-skip (NODAL/FDM grid only).  The FDM nodal grid has a
    ! DUPLICATED periodic endpoint: global node 1 coincides with global node
    ! n_global (x_R == x_L under periodicity).  The serial wrap collapses this
    ! duplicate (period (n-1)*dx); the cross-rank wrap must match.  The two
    ! periodic-edge ranks therefore SKIP their duplicated boundary node when
    ! sending across the seam, shifting the sent slice by one node: the lo-edge
    ! rank (coord==0, owns global node 1) skips its first interior node, and the
    ! hi-edge rank (coord==dim-1, owns global node n_global) skips its last.
    ! Interior ranks are unshifted.  When dim==1 (axis not split) the single rank
    ! is BOTH edges and wraps to itself, so both shifts apply and reproduce the
    ! serial nodal wrap exactly.  This mirrors the verified 1D halo_exchange
    ! correction.  The FVM (cell-centred) grid has NO duplicated endpoint, so the
    ! skip MUST NOT apply there; gate strictly on method so the FVM path stays
    ! byte-identical.  Each guarded coord test is in its own if (Fortran .and.
    ! does not short-circuit; nodal_skip is the only compound condition and has
    ! no index operands).
    nodal_skip = (trim(state % blocks(1) % method) == method_fdm)

    ! ---- x-axis exchange ----
    ! Pack: send leftmost h interior cols to x_lo neighbour (TAG_X recv);
    !       send rightmost h interior cols to x_hi neighbour (TAG_X recv).
    allocate(sx_lo(neq2d, h, nyl), sx_hi(neq2d, h, nyl), &
             rx_lo(neq2d, h, nyl), rx_hi(neq2d, h, nyl), stat=ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: x-axis halo buffer allocation failed')
    xlo_s = 1; xhi_s = nxl - h + 1
    if (nodal_skip) then
      if (d % periodic_x) then
        if (d % coord_x == 0) xlo_s = 2                 ! skip global node 1
        if (d % coord_x == d % dim_x - 1) xhi_s = nxl - h   ! skip global node n_global
      end if
    end if
    sx_lo = state % ub(:, xlo_s:xlo_s+h-1, 1:nyl)
    sx_hi = state % ub(:, xhi_s:xhi_s+h-1, 1:nyl)
    cnt = neq2d * h * nyl

    ! Post receives before sends to avoid potential deadlock on small MPIs.
    ! Recv from lo-neighbour with TAG_X_HI: lo-neighbour sends its hi-side
    ! data to us using TAG_X_HI. Recv from hi-neighbour with TAG_X_LO.
    call MPI_Irecv(rx_lo, cnt, MPI_DOUBLE_PRECISION, d % x_lo_neighbour, TAG_X_HI, comm, req(1), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Irecv x_lo failed')
    call MPI_Irecv(rx_hi, cnt, MPI_DOUBLE_PRECISION, d % x_hi_neighbour, TAG_X_LO, comm, req(2), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Irecv x_hi failed')
    ! Send our lo-side data to lo-neighbour with TAG_X_LO (they recv as rx_hi).
    ! Send our hi-side data to hi-neighbour with TAG_X_HI (they recv as rx_lo).
    call MPI_Isend(sx_lo, cnt, MPI_DOUBLE_PRECISION, d % x_lo_neighbour, TAG_X_LO, comm, req(3), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Isend x_lo failed')
    call MPI_Isend(sx_hi, cnt, MPI_DOUBLE_PRECISION, d % x_hi_neighbour, TAG_X_HI, comm, req(4), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Isend x_hi failed')
    call MPI_Waitall(4, req, MPI_STATUSES_IGNORE, ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Waitall x failed')

    ! Unpack x halos. x_lo_neighbour's rightmost cols -> our lo halo (col 0, -1, ..., 1-h).
    ! x_hi_neighbour's leftmost cols -> our hi halo (col nxl+1, ..., nxl+h).
    if (d % x_lo_neighbour /= mpi_proc_null_value()) then
      do j = 1, nyl
        do k = 1, h
          state % ub(:, 1 - k, j) = rx_lo(:, h - k + 1, j)
        end do
      end do
    end if
    if (d % x_hi_neighbour /= mpi_proc_null_value()) then
      do j = 1, nyl
        do k = 1, h
          state % ub(:, nxl + k, j) = rx_hi(:, k, j)
        end do
      end do
    end if

    deallocate(sx_lo, sx_hi, rx_lo, rx_hi, stat=ierr)

    ! ---- y-axis exchange ----
    ! Pack: send bottom h interior rows to y_lo neighbour;
    !       send top h interior rows to y_hi neighbour.
    allocate(sy_lo(neq2d, nxl, h), sy_hi(neq2d, nxl, h), &
             ry_lo(neq2d, nxl, h), ry_hi(neq2d, nxl, h), stat=ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: y-axis halo buffer allocation failed')
    ylo_s = 1; yhi_s = nyl - h + 1
    if (nodal_skip) then
      if (d % periodic_y) then
        if (d % coord_y == 0) ylo_s = 2                 ! skip global node 1
        if (d % coord_y == d % dim_y - 1) yhi_s = nyl - h   ! skip global node n_global
      end if
    end if
    sy_lo = state % ub(:, 1:nxl, ylo_s:ylo_s+h-1)
    sy_hi = state % ub(:, 1:nxl, yhi_s:yhi_s+h-1)
    cnt = neq2d * nxl * h

    ! Same tag disambiguation for y-axis.
    call MPI_Irecv(ry_lo, cnt, MPI_DOUBLE_PRECISION, d % y_lo_neighbour, TAG_Y_HI, comm, req(1), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Irecv y_lo failed')
    call MPI_Irecv(ry_hi, cnt, MPI_DOUBLE_PRECISION, d % y_hi_neighbour, TAG_Y_LO, comm, req(2), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Irecv y_hi failed')
    call MPI_Isend(sy_lo, cnt, MPI_DOUBLE_PRECISION, d % y_lo_neighbour, TAG_Y_LO, comm, req(3), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Isend y_lo failed')
    call MPI_Isend(sy_hi, cnt, MPI_DOUBLE_PRECISION, d % y_hi_neighbour, TAG_Y_HI, comm, req(4), ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Isend y_hi failed')
    call MPI_Waitall(4, req, MPI_STATUSES_IGNORE, ierr)
    if (ierr /= 0) call parallel_fatal('halo_exchange_2d: Waitall y failed')

    ! Unpack y halos. y_lo_neighbour's top rows -> our bottom halo (row 0, -1, ..., 1-h).
    ! y_hi_neighbour's bottom rows -> our top halo (row nyl+1, ..., nyl+h).
    if (d % y_lo_neighbour /= mpi_proc_null_value()) then
      do k = 1, h
        do i = 1, nxl
          state % ub(:, i, 1 - k) = ry_lo(:, i, h - k + 1)
        end do
      end do
    end if
    if (d % y_hi_neighbour /= mpi_proc_null_value()) then
      do k = 1, h
        do i = 1, nxl
          state % ub(:, i, nyl + k) = ry_hi(:, i, k)
        end do
      end do
    end if

    deallocate(sy_lo, sy_hi, ry_lo, ry_hi, stat=ierr)
  end subroutine exchange_halos_2d

end module halo_exchange_2d