!> @file mpi_cart_2d.F90
!> @brief Build a 2D Cartesian MPI communicator and the per-rank 2D decomposition.
!!
!! `cart_setup_2d` calls MPI_Dims_create to auto-factor n_ranks() into a
!! 2D grid, creates the Cartesian communicator (with optional periodicity),
!! queries this rank's coords, and derives the axis-neighbour ranks via
!! MPI_Cart_shift before delegating the actual tile sizing to decompose_2d.
!!
!! Boundary ranks at non-periodic edges receive MPI_PROC_NULL from
!! MPI_Cart_shift, which propagates into decomp_2d_t%x_lo_neighbour etc.
!! halo_exchange_2d (Task 3) relies on MPI standard §6.6: Isend/Irecv to
!! MPI_PROC_NULL are silent no-ops, so no special-casing is needed there.
!!
!! The Cartesian communicator handle is stored in the module-level save
!! variable g_cart and exposed via cart_comm_2d() for use by the halo
!! exchange and any other collective that needs it.
!!
!! Legacy-MPI note: MPI_Cart_create, MPI_Cart_coords, MPI_Cart_shift, and
!! MPI_Dims_create are not used in the 1D path, so they are not imported in
!! mpi_runtime.F90. We import them directly here. On Windows/MS-MPI
!! (CFD_SOLVER_LEGACY_MPI) the `mpi` module provides the same symbols with
!! bare-integer handles; the mpi_f08 path uses derived-type handles.
module mpi_cart_2d
#ifdef CFD_SOLVER_LEGACY_MPI
  use mpi, only: MPI_Dims_create, MPI_Cart_create, MPI_Comm_rank, &
                 MPI_Cart_coords, MPI_Cart_shift, MPI_COMM_NULL
#else
  use mpi_f08, only: MPI_Dims_create, MPI_Cart_create, MPI_Comm_rank, &
                     MPI_Cart_coords, MPI_Cart_shift, MPI_Comm, MPI_COMM_NULL
#endif
  use mpi_runtime, only: n_ranks, mpi_world
  use domain_decomposition_2d, only: decomp_2d_t, decompose_2d
  use solver_state_2d, only: solver_state_2d_t
  use option_registry, only: bc_periodic
  implicit none
  private
  public :: cart_setup_2d, cart_comm_2d, setup_decomp_2d, cart_is_setup, cart_rank_2d

#ifdef CFD_SOLVER_LEGACY_MPI
  integer, save :: g_cart = MPI_COMM_NULL
#else
  type(MPI_Comm), save :: g_cart = MPI_COMM_NULL
#endif
  logical, save :: g_cart_setup = .false.

contains

  !> Create the Cartesian communicator (dims auto-chosen by MPI_Dims_create)
  !! and return this rank's per-rank decomposition.
  !!
  !! @param nx_global  Total interior cells in x.
  !! @param ny_global  Total interior cells in y.
  !! @param halo_width Ghost-cell layer width (scheme-derived).
  !! @param periodic_x Wrap in x?
  !! @param periodic_y Wrap in y?
  !! @param d          Output: per-rank decomposition including local extents
  !!                   and all four axis-neighbour ranks.
  subroutine cart_setup_2d(nx_global, ny_global, halo_width, periodic_x, periodic_y, d)
    integer, intent(in) :: nx_global, ny_global, halo_width
    logical, intent(in) :: periodic_x, periodic_y
    type(decomp_2d_t), intent(out) :: d

    integer :: dims(2), coords(2), ierr, my_cart_rank
    logical :: periods(2)
    integer :: x_lo, x_hi, y_lo, y_hi

    ! Auto-factor n_ranks into a 2D grid. dims must be initialised to 0 so
    ! MPI_Dims_create treats both dimensions as unconstrained.
    dims = 0
    call MPI_Dims_create(n_ranks(), 2, dims, ierr)
    if (ierr /= 0) error stop 'mpi_cart_2d: MPI_Dims_create failed'

    ! Build the Cartesian communicator. reorder=.true. allows MPI to reorder
    ! ranks for topology locality (safe; we query our rank immediately after).
    periods = [periodic_x, periodic_y]
    call MPI_Cart_create(mpi_world(), 2, dims, periods, .true., g_cart, ierr)
    if (ierr /= 0) error stop 'mpi_cart_2d: MPI_Cart_create failed'
    g_cart_setup = .true.

    ! Discover which rank we are in the new communicator (may differ from
    ! COMM_WORLD rank if reorder was honoured).
    call MPI_Comm_rank(g_cart, my_cart_rank, ierr)
    if (ierr /= 0) error stop 'mpi_cart_2d: MPI_Comm_rank failed'

    ! Map Cartesian rank -> (coord_x, coord_y).
    call MPI_Cart_coords(g_cart, my_cart_rank, 2, coords, ierr)
    if (ierr /= 0) error stop 'mpi_cart_2d: MPI_Cart_coords failed'

    ! Axis-0 (x) neighbours: direction=0, displacement=1.
    ! rank_source = x_lo  (the rank that sends to us from the low-x side)
    ! rank_dest   = x_hi  (the rank we send to on the high-x side)
    ! At non-periodic boundaries these are MPI_PROC_NULL.
    call MPI_Cart_shift(g_cart, 0, 1, x_lo, x_hi, ierr)
    if (ierr /= 0) error stop 'mpi_cart_2d: MPI_Cart_shift (x) failed'

    ! Axis-1 (y) neighbours.
    call MPI_Cart_shift(g_cart, 1, 1, y_lo, y_hi, ierr)
    if (ierr /= 0) error stop 'mpi_cart_2d: MPI_Cart_shift (y) failed'

    ! Delegate tile sizing and field population to the pure decompose_2d.
    d = decompose_2d(coords, dims, nx_global, ny_global, halo_width, &
                     periodic_x, periodic_y, x_lo, x_hi, y_lo, y_hi)
  end subroutine cart_setup_2d

  !> Build the MPI Cartesian decomposition and store it on the state (MPI runs).
  !!
  !! Derives periodicity from state % cfg boundary condition strings, calls
  !! cart_setup_2d, and updates state % nx_local / state % ny_local to the
  !! per-rank local sizes so allocate_work_arrays_2d sizes correctly.
  subroutine setup_decomp_2d(state)
    type(solver_state_2d_t), intent(inout) :: state
    logical :: periodic_x, periodic_y
    periodic_x = (trim(state % cfg % bc_left) == bc_periodic .and. &
                  trim(state % cfg % bc_right) == bc_periodic)
    periodic_y = (trim(state % cfg % bc_bottom) == bc_periodic .and. &
                  trim(state % cfg % bc_top) == bc_periodic)
    call cart_setup_2d(state % nx, state % ny, state % halo_width, &
                       periodic_x, periodic_y, state % decomp_2d)
    state % nx_local = state % decomp_2d % nx_local
    state % ny_local = state % decomp_2d % ny_local
  end subroutine setup_decomp_2d

  !> True once cart_setup_2d / setup_decomp_2d has created the Cartesian comm.
  !!
  !! Uses a save flag rather than comparing g_cart /= MPI_COMM_NULL to avoid
  !! relying on derived-type comparison operators in the mpi_f08 path.
  function cart_is_setup() result(yes)
    logical :: yes
    yes = g_cart_setup
  end function cart_is_setup

  !> Return the Cartesian communicator handle.
  !!
  !! Returns MPI_COMM_NULL if cart_setup_2d has not been called yet.
  !! halo_exchange_2d and any collective that needs the 2D topology should
  !! call this rather than referencing MPI_COMM_WORLD directly.
  function cart_comm_2d() result(c)
#ifdef CFD_SOLVER_LEGACY_MPI
    integer :: c
#else
    type(MPI_Comm) :: c
#endif
    c = g_cart
  end function cart_comm_2d

  !> This rank's index within the Cartesian communicator.
  !!
  !! Returns -1 if cart_setup_2d has not been called yet or MPI_Comm_rank
  !! fails. Callers that gate file I/O on rank 0 must use THIS, not the
  !! COMM_WORLD rank: MPI_Cart_create(reorder=.true.) may renumber ranks, so
  !! the cart-rank-0 process (which holds gathered data) can differ from
  !! world-rank-0.
  function cart_rank_2d() result(r)
    integer :: r
    integer :: ierr
    r = -1
    if (.not. g_cart_setup) return
    call MPI_Comm_rank(g_cart, r, ierr)
    if (ierr /= 0) r = -1
  end function cart_rank_2d

end module mpi_cart_2d