!> @file solution_gather_2d.F90
!> @brief Rank-0 global gather + ASCII write for the 2D Euler solver.
!!
!! At np>1 every rank owns a tile of the global domain. `gather_and_write_2d`
!! collects all tiles onto rank 0 via point-to-point MPI (simpler than
!! MPI_Gatherv for irregular 2D blocks), assembles the full (neq2d,nx,ny) field
!! on rank 0, and writes the same "x y rho u v p" ASCII format that
!! write_result_2d produces at np=1.
!!
!! At np=1 this is identical to write_result_2d (root receives nothing, just
!! writes its own block which covers the whole domain).
!!
!! The Cartesian communicator (from mpi_cart_2d) is used consistently so that
!! rank renumbering by MPI_Cart_create(reorder=.true.) cannot cause a mismatch
!! between gather source/destination and decomp placement indices.
!!
!! Legacy-MPI note: MS-MPI ships only the legacy `use mpi` module; OpenMPI
!! ships both. We guard with CFD_SOLVER_LEGACY_MPI in the same way as the
!! other Phase-2D MPI source files.

module solution_gather_2d
  use precision, only: wp, safe_vel
  use solver_state_2d, only: solver_state_2d_t, neq2d
  use domain_decomposition_2d, only: decomp_2d_t
  use domain_decomposition, only: rank_local_count, rank_first_global
  use mpi_cart_2d, only: cart_comm_2d
  use mpi_runtime, only: n_ranks, parallel_fatal
  use parallel_reductions, only: par_lor
  use logger, only: log_warn
  use option_registry, only: method_fdm
  use tecplot_writer, only: write_tecplot_ascii, write_tecplot_binary
  use output_format_list, only: parse_format_list, derived_filename, max_formats
#ifdef CFD_SOLVER_LEGACY_MPI
  use mpi, only: MPI_Send, MPI_Recv, MPI_Comm_rank, MPI_Cart_coords, &
                 MPI_DOUBLE_PRECISION, MPI_INTEGER, MPI_STATUS_IGNORE
#else
  use mpi_f08, only: MPI_Send, MPI_Recv, MPI_Comm_rank, MPI_Cart_coords, &
                     MPI_DOUBLE_PRECISION, MPI_INTEGER, MPI_STATUS_IGNORE, MPI_Comm
#endif
  implicit none
  private
  public :: gather_and_write_2d
  public :: gather_solution_to_root_2d
  public :: gather_boundary_ghosts_to_root, scatter_boundary_ghosts_from_root

  integer, parameter :: TAG_META  = 3001
  integer, parameter :: TAG_BLOCK = 3002
  integer, parameter :: TAG_COORD = 3003  !< coordinate block (curvilinear only): FVM xc/yc, FDM x_node/y_node
                                          !! (TAG_COORD+1 = 3004 carries the y block)
  !> Boundary-ghost frame slabs (checkpoint v6; see gather_boundary_ghosts_to_root).
  integer, parameter :: TAG_GHOST_L = 3005
  integer, parameter :: TAG_GHOST_R = 3006
  integer, parameter :: TAG_GHOST_B = 3007
  integer, parameter :: TAG_GHOST_T = 3008

contains

  !> Gather all per-rank tiles to rank 0 and write the global field as ASCII.
  !!
  !! Format per line: x y rho u v p (6 × ES24.12E3), i (x) fastest.
  !! A blank line is written after each constant-j row so gnuplot/ParaView
  !! treat the output as a structured grid — identical to write_result_2d.
  !!
  !! @param state     Solver state (ub(:,1:nx_local,1:ny_local) = interior cells).
  !! @param filename  Output path (written by rank 0 only).
  !! @param is_ok     .true. on success, .false. on I/O or MPI error.
  !! @param message   Human-readable error description on failure.
  subroutine gather_and_write_2d(state, filename, is_ok, message)
    type(solver_state_2d_t), intent(in) :: state
    character(len=*),        intent(in) :: filename
    logical,                 intent(out) :: is_ok
    character(len=*),        intent(out) :: message

#ifdef CFD_SOLVER_LEGACY_MPI
    integer :: comm
#else
    type(MPI_Comm) :: comm
#endif
    integer :: my_cart_rank, ierr
    real(wp), allocatable :: global_field(:,:,:)
    real(wp), allocatable :: global_xc(:,:)   !< gathered x coords (root, curvilinear: FVM centroids / FDM nodes)
    real(wp), allocatable :: global_yc(:,:)   !< gathered y coords (root, curvilinear: FVM centroids / FDM nodes)
    logical :: local_ok, failed
    logical :: curvi                          !< .true. when mesh is curvilinear

    is_ok = .false.
    message = ''
    ! Per-rank success; reduced across ranks via par_lor before returning so
    ! every rank agrees on the outcome (mirrors write_solution_file). EVERY
    ! rank must reach the par_lor below — no early `return` from an error path.
    local_ok = .true.

    curvi = .not. state % mesh % uniform

    if (curvi) then
      call gather_solution_to_root_2d(state, global_field, local_ok, message, global_xc, global_yc)
    else
      call gather_solution_to_root_2d(state, global_field, local_ok, message)
    end if

    ! --- Write global field on root only (one file per requested format) ---
    comm         = cart_comm_2d()
    my_cart_rank = -1
    call MPI_Comm_rank(comm, my_cart_rank, ierr)
    if (ierr /= 0) then
      call log_warn('gather_and_write_2d: MPI_Comm_rank failed')
      local_ok = .false.
    end if

    if (local_ok .and. my_cart_rank == 0) then
      block
        character(len=8) :: tokens(max_formats)
        integer :: n_fmt, kf, nx_g2, ny_g2, ii, jj, pp
        logical :: fmt_ok
        character(len=256) :: fmt_msg, w_msg
        character(len=:), allocatable :: fname
        real(wp), allocatable :: tp_data(:, :)
        character(len=8) :: tp_names(6)
        real(wp) :: rr, uu, vvel, pp_pres, gam2
        logical :: w_ok

        call parse_format_list(state % cfg % output_format, &
               [character(len=8) :: 'dat', 'vtk', 'tec', 'plt'], &
               tokens, n_fmt, fmt_ok, fmt_msg)
        if (.not. fmt_ok) then
          call log_warn('gather_and_write_2d: '//trim(fmt_msg))
          local_ok = .false.
        end if

        if (local_ok) then
        nx_g2 = state % decomp_2d % nx_global
        ny_g2 = state % decomp_2d % ny_global
        gam2  = state % cfg % gam

        if (any(tokens(1:n_fmt) == 'tec') .or. any(tokens(1:n_fmt) == 'plt')) then
          tp_names = [character(len=8) :: 'X', 'Y', 'rho', 'u', 'v', 'p']
          allocate(tp_data(6, nx_g2 * ny_g2), stat=ierr)
          if (ierr /= 0) then
            call log_warn('solution_gather_2d: tecplot buffer allocation failed')
            local_ok = .false.
          end if
          if (local_ok) then
            pp = 0
            do jj = 1, ny_g2
              do ii = 1, nx_g2
                pp = pp + 1
                rr = global_field(1, ii, jj)
                ! safe_vel emits 0 velocity (no division) for a vacuum cell (rho <= 0)
                ! so output cannot raise a SIGFPE; no-op for rho > 0.
                uu = safe_vel(global_field(2, ii, jj), rr)
                vvel = safe_vel(global_field(3, ii, jj), rr)
                pp_pres = (gam2 - 1.0_wp) * (global_field(4, ii, jj) &
                           - 0.5_wp * rr * (uu*uu + vvel*vvel))
                if (curvi) then
                  tp_data(1, pp) = global_xc(ii, jj)
                  tp_data(2, pp) = global_yc(ii, jj)
                else if (trim(state % blocks(1) % method) == method_fdm) then
                  ! FDM: nodes start at x_left; node index (1,1) is at (x_left,y_left).
                  tp_data(1, pp) = state % cfg % x_left + real(ii - 1, wp) * state % dx
                  tp_data(2, pp) = state % cfg % y_left + real(jj - 1, wp) * state % dy
                else
                  tp_data(1, pp) = state % cfg % x_left + (real(ii, wp) - 0.5_wp) * state % dx
                  tp_data(2, pp) = state % cfg % y_left + (real(jj, wp) - 0.5_wp) * state % dy
                end if
                tp_data(3, pp) = rr
                tp_data(4, pp) = uu
                tp_data(5, pp) = vvel
                tp_data(6, pp) = pp_pres
              end do
            end do
          end if
        end if

        if (local_ok) then
        do kf = 1, n_fmt
          fname = derived_filename(trim(filename), tokens(kf), n_fmt)
          w_ok = .false.
          w_msg = ''
          select case (trim(tokens(kf)))
          case ('vtk')
            call write_global_field_vtk(global_field, state, fname, w_ok, w_msg)
          case ('tec')
            call write_tecplot_ascii(fname, 'euler_2d', tp_names, nx_g2, ny_g2, &
                                      tp_data, w_ok, w_msg)
          case ('plt')
            call write_tecplot_binary(fname, 'euler_2d', tp_names, nx_g2, ny_g2, &
                                       tp_data, w_ok, w_msg)
          case default   ! 'dat'
            if (curvi) then
              call write_global_field(global_field, state, fname, w_ok, w_msg, &
                                      global_xc, global_yc)
            else
              call write_global_field(global_field, state, fname, w_ok, w_msg)
            end if
          end select
          if (.not. w_ok) then
            call log_warn('gather_and_write_2d: '//trim(w_msg))
            local_ok = .false.
            exit
          end if
        end do
        end if

        if (allocated(tp_data)) deallocate(tp_data, stat=ierr)
        end if
      end block
    end if
    if (allocated(global_field)) deallocate(global_field, stat=ierr)
    if (allocated(global_xc))    deallocate(global_xc, stat=ierr)
    if (allocated(global_yc))    deallocate(global_yc, stat=ierr)

    ! Collective failure agreement: every rank reaches this reduction so all
    ! ranks return the SAME outcome, even when only rank 0's I/O failed. Each
    ! failing rank has already emitted its specific diagnostic via log_warn;
    ! the returned `message` is a single uniform string set identically on every
    ! rank (mirrors write_solution_file) so is_ok and message agree everywhere.
    failed = par_lor(.not. local_ok)
    if (failed) then
      is_ok = .false.
      message = 'gather_and_write_2d: write/gather failed on some rank'
    else
      is_ok = .true.
      message = ''
    end if
  end subroutine gather_and_write_2d

  !> Gather all per-rank tiles to rank 0, assembling the full global solution field.
  !!
  !! This is a pure data-move routine: on rank 0 it assembles the global
  !! (neq2d, nx_g, ny_g) field from all ranks' tiles via point-to-point MPI.
  !! On non-root ranks it sends the local tile and allocates a 0-size placeholder.
  !! No cross-rank reduction of is_ok is performed — the caller is responsible
  !! for any par_lor aggregation, matching how 1D's gather_solution_to_root
  !! is a pure data move.
  !!
  !! When global_xc/global_yc are present (curvilinear grids), the mesh
  !! coordinates are also gathered to rank 0: FVM cell centroids (mesh%xc/yc)
  !! or, for curvilinear FDM (mesh%fdm_curvilinear), node coordinates
  !! (mesh%x_node/y_node) — xc/yc are never allocated on the FDM path.
  !!
  !! @param state        Solver state (ub(:,1:nx_local,1:ny_local) = interior cells).
  !! @param global_field On root: (neq2d,nx_g,ny_g) assembled field. Off-root: (neq2d,0,0).
  !! @param is_ok        .true. on local success; caller reduces if needed.
  !! @param message      Human-readable error on failure.
  !! @param global_xc    Optional: gathered x coords (root only, curvilinear).
  !! @param global_yc    Optional: gathered y coords (root only, curvilinear).
  subroutine gather_solution_to_root_2d(state, global_field, is_ok, message, global_xc, global_yc)
    type(solver_state_2d_t), intent(in)  :: state
    real(wp), allocatable,   intent(out) :: global_field(:,:,:)  !< (neq2d,nx_g,ny_g) on root; 0-size off-root
    logical,                 intent(out) :: is_ok                !< rank-local; caller reduces if needed
    character(len=*),        intent(out) :: message
    real(wp), allocatable, optional, intent(out) :: global_xc(:,:)  !< curvilinear, root only
    real(wp), allocatable, optional, intent(out) :: global_yc(:,:)  !< curvilinear, root only

#ifdef CFD_SOLVER_LEGACY_MPI
    integer :: comm
#else
    type(MPI_Comm) :: comm
#endif
    integer :: my_cart_rank, nr, ierr, astat
    integer :: nx_g, ny_g, nxl, nyl, ix1, iy1
    integer :: meta(4)
    real(wp), allocatable :: send_buf(:,:,:)
    real(wp), allocatable :: send_coord(:,:)  !< xc or yc send scratch (curvilinear)
    integer :: src
    logical :: curvi                          !< .true. when mesh is curvilinear
    logical :: want_coords                    !< .true. when BOTH global_xc and global_yc are present

    ! --- meta variables for received blocks (root only) ---
    integer :: r_ix1, r_iy1, r_nxl, r_nyl
    integer :: r_meta(4)
    real(wp), allocatable :: recv_buf(:,:,:)
    real(wp), allocatable :: recv_coord(:,:)  !< coordinate recv scratch (curvilinear, root)

    is_ok   = .true.
    message = ''

    ! Both coord optionals must be present together; requiring either alone is
    ! illegal Fortran (root would assign to an absent dummy).
    want_coords = present(global_xc) .and. present(global_yc)

    comm          = cart_comm_2d()
    ! Safe sentinel so my_cart_rank is always defined when read below. Fortran
    ! does NOT guarantee short-circuit `.and.`, so `if (is_ok .and.
    ! my_cart_rank /= 0)` may evaluate my_cart_rank even when is_ok is
    ! .false. (e.g. MPI_Comm_rank failed and left my_cart_rank undefined).
    my_cart_rank  = -1
    call MPI_Comm_rank(comm, my_cart_rank, ierr)
    if (ierr /= 0) then
      call log_warn('gather_solution_to_root_2d: MPI_Comm_rank failed')
      is_ok = .false.
    end if

    nr   = n_ranks()
    nxl  = state % nx_local
    nyl  = state % ny_local
    nx_g = state % decomp_2d % nx_global
    ny_g = state % decomp_2d % ny_global
    ix1  = state % decomp_2d % ix_first_global
    iy1  = state % decomp_2d % iy_first_global
    curvi = .not. state % mesh % uniform

    ! --- Non-root: send metadata then block buffer to rank 0 ---
    if (is_ok .and. my_cart_rank /= 0) then
      ! Allocate ALL send buffers up front, BEFORE announcing meta. A mid-gather
      ! allocation failure is then signalled to root via a sentinel meta
      ! (meta(3) = -1) instead of a skipped MPI_Send that would deadlock root's
      ! matching MPI_Recv.
      allocate(send_buf(neq2d, nxl, nyl), stat=astat)
      if (astat /= 0) then
        call log_warn('gather_solution_to_root_2d: allocate send_buf failed')
        is_ok = .false.
      end if
      if (is_ok .and. curvi .and. want_coords) then
        allocate(send_coord(nxl, nyl), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate send_coord failed')
          is_ok = .false.
        end if
      end if

      meta(1) = ix1
      meta(2) = iy1
      meta(3) = nxl
      meta(4) = nyl
      if (.not. is_ok) meta(3) = -1   ! sentinel: root skips this rank's block/coord recv
      call MPI_Send(meta, 4, MPI_INTEGER, 0, TAG_META, comm, ierr)
      if (ierr /= 0) then
        call log_warn('gather_solution_to_root_2d: MPI_Send meta failed')
        is_ok = .false.
      end if

      if (is_ok) then
        send_buf = state % ub(:, 1:nxl, 1:nyl)
        call MPI_Send(send_buf, neq2d*nxl*nyl, MPI_DOUBLE_PRECISION, 0, TAG_BLOCK, comm, ierr)
        if (ierr /= 0) then
          call log_warn('gather_solution_to_root_2d: MPI_Send block failed')
          is_ok = .false.
        end if
      end if
      ! Curvilinear: also send the local x then y coordinates to root (same block
      ! index meta). Curvilinear FDM stores NODE coordinates (x_node/y_node,
      ! sliced rank-local by build_mesh_2d_local); FVM stores cell centroids
      ! (xc/yc — never allocated on the FDM path). Both are (nx_local, ny_local).
      if (is_ok .and. curvi .and. want_coords) then
        if (state % mesh % fdm_curvilinear) then
          send_coord = state % mesh % x_node(1:nxl, 1:nyl)
        else
          send_coord = state % mesh % xc(1:nxl, 1:nyl)
        end if
        call MPI_Send(send_coord, nxl*nyl, MPI_DOUBLE_PRECISION, 0, TAG_COORD, comm, ierr)
        if (ierr /= 0) then
          call log_warn('gather_solution_to_root_2d: MPI_Send x-coord failed')
          is_ok = .false.
        end if
        if (is_ok) then
          if (state % mesh % fdm_curvilinear) then
            send_coord = state % mesh % y_node(1:nxl, 1:nyl)
          else
            send_coord = state % mesh % yc(1:nxl, 1:nyl)
          end if
          call MPI_Send(send_coord, nxl*nyl, MPI_DOUBLE_PRECISION, 0, TAG_COORD+1, comm, ierr)
          if (ierr /= 0) then
            call log_warn('gather_solution_to_root_2d: MPI_Send y-coord failed')
            is_ok = .false.
          end if
        end if
      end if
      if (allocated(send_buf)) then
        deallocate(send_buf, stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: deallocate send_buf failed')
          is_ok = .false.
        end if
      end if
      if (allocated(send_coord)) then
        deallocate(send_coord, stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: deallocate send_coord failed')
          is_ok = .false.
        end if
      end if
      ! Off-root: always allocate a 0-size placeholder so global_field is always allocated.
      allocate(global_field(neq2d, 0, 0), stat=astat)
      if (astat /= 0) then
        call log_warn('gather_solution_to_root_2d: allocate global_field placeholder failed')
        is_ok = .false.
      end if
      if (want_coords) then
        allocate(global_xc(0, 0), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate global_xc placeholder failed')
          is_ok = .false.
        end if
        allocate(global_yc(0, 0), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate global_yc placeholder failed')
          is_ok = .false.
        end if
      end if
    else if (is_ok) then
      ! --- Root: assemble global field ---
      allocate(global_field(neq2d, nx_g, ny_g), stat=astat)
      if (astat /= 0) then
        call log_warn('gather_solution_to_root_2d: allocate global_field failed')
        is_ok = .false.
      end if

      ! Place our own tile first.
      if (is_ok) global_field(:, ix1:ix1+nxl-1, iy1:iy1+nyl-1) = state % ub(:, 1:nxl, 1:nyl)

      ! Curvilinear: root builds global_xc/global_yc from its own tile + received tiles.
      if (curvi .and. want_coords) then
        allocate(global_xc(nx_g, ny_g), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate global_xc failed')
          is_ok = .false.
        end if
        allocate(global_yc(nx_g, ny_g), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate global_yc failed')
          is_ok = .false.
        end if
        ! Coordinate source mirrors the non-root send: FDM nodal x_node/y_node,
        ! FVM cell-centre xc/yc.
        if (is_ok) then
          if (state % mesh % fdm_curvilinear) then
            global_xc(ix1:ix1+nxl-1, iy1:iy1+nyl-1) = state % mesh % x_node(1:nxl, 1:nyl)
            global_yc(ix1:ix1+nxl-1, iy1:iy1+nyl-1) = state % mesh % y_node(1:nxl, 1:nyl)
          else
            global_xc(ix1:ix1+nxl-1, iy1:iy1+nyl-1) = state % mesh % xc(1:nxl, 1:nyl)
            global_yc(ix1:ix1+nxl-1, iy1:iy1+nyl-1) = state % mesh % yc(1:nxl, 1:nyl)
          end if
        end if
      end if

      ! Receive each other cart rank's tile.
      do src = 1, nr - 1
        call MPI_Recv(r_meta, 4, MPI_INTEGER, src, TAG_META, comm, MPI_STATUS_IGNORE, ierr)
        if (ierr /= 0) then
          call log_warn('gather_solution_to_root_2d: MPI_Recv meta failed')
          is_ok = .false.
          exit
        end if
        r_ix1 = r_meta(1)
        r_iy1 = r_meta(2)
        r_nxl = r_meta(3)
        r_nyl = r_meta(4)

        if (r_nxl < 0) then
          ! Sentinel: this source failed to allocate its send buffers and sent
          ! no block/coord. Record the failure and skip its recvs so root does
          ! not block on a matching MPI_Recv that will never be posted.
          call log_warn('gather_solution_to_root_2d: source reported allocation failure')
          is_ok = .false.
          cycle
        end if

        allocate(recv_buf(neq2d, r_nxl, r_nyl), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate recv_buf failed')
          is_ok = .false.
          exit
        end if
        call MPI_Recv(recv_buf, neq2d*r_nxl*r_nyl, MPI_DOUBLE_PRECISION, &
                      src, TAG_BLOCK, comm, MPI_STATUS_IGNORE, ierr)
        if (ierr /= 0) then
          call log_warn('gather_solution_to_root_2d: MPI_Recv block failed')
          is_ok = .false.
          deallocate(recv_buf, stat=astat)
          exit
        end if
        global_field(:, r_ix1:r_ix1+r_nxl-1, r_iy1:r_iy1+r_nyl-1) = recv_buf
        deallocate(recv_buf, stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: deallocate recv_buf failed')
          is_ok = .false.
          exit
        end if

        ! Curvilinear: receive xc then yc from this rank.
        if (curvi .and. want_coords) then
          allocate(recv_coord(r_nxl, r_nyl), stat=astat)
          if (astat /= 0) then
            call log_warn('gather_solution_to_root_2d: allocate recv_coord failed')
            is_ok = .false.
            exit
          end if
          call MPI_Recv(recv_coord, r_nxl*r_nyl, MPI_DOUBLE_PRECISION, &
                        src, TAG_COORD, comm, MPI_STATUS_IGNORE, ierr)
          if (ierr /= 0) then
            call log_warn('gather_solution_to_root_2d: MPI_Recv x-coord failed')
            is_ok = .false.
            deallocate(recv_coord, stat=astat)
            exit
          end if
          global_xc(r_ix1:r_ix1+r_nxl-1, r_iy1:r_iy1+r_nyl-1) = recv_coord
          call MPI_Recv(recv_coord, r_nxl*r_nyl, MPI_DOUBLE_PRECISION, &
                        src, TAG_COORD+1, comm, MPI_STATUS_IGNORE, ierr)
          if (ierr /= 0) then
            call log_warn('gather_solution_to_root_2d: MPI_Recv y-coord failed')
            is_ok = .false.
            deallocate(recv_coord, stat=astat)
            exit
          end if
          global_yc(r_ix1:r_ix1+r_nxl-1, r_iy1:r_iy1+r_nyl-1) = recv_coord
          deallocate(recv_coord, stat=astat)
          if (astat /= 0) then
            call log_warn('gather_solution_to_root_2d: deallocate recv_coord failed')
            is_ok = .false.
            exit
          end if
        end if
      end do
    else
      ! MPI_Comm_rank failed — still allocate 0-size placeholders to satisfy intent(out).
      allocate(global_field(neq2d, 0, 0), stat=astat)
      if (astat /= 0) then
        call log_warn('gather_solution_to_root_2d: allocate global_field placeholder failed')
        is_ok = .false.
      end if
      if (want_coords) then
        allocate(global_xc(0, 0), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate global_xc placeholder failed')
          is_ok = .false.
        end if
        allocate(global_yc(0, 0), stat=astat)
        if (astat /= 0) then
          call log_warn('gather_solution_to_root_2d: allocate global_yc placeholder failed')
          is_ok = .false.
        end if
      end if
    end if
  end subroutine gather_solution_to_root_2d

  !> Write the assembled global field to `filename` in the standard ASCII format.
  !!
  !! i (x) is fastest; a blank line separates each constant-j row so
  !! gnuplot/ParaView parse the output as a structured 2D grid.
  !! When `xc`/`yc` are present (curvilinear), those cell centroids are written
  !! instead of the uniform `x_left+(i-0.5)*dx` formula. The uniform branch
  !! (no optional args) is byte-identical to the prior implementation.
  subroutine write_global_field(global_field, state, filename, is_ok, message, xc, yc)
    real(wp),                intent(in)  :: global_field(:,:,:)  !< (neq2d, nx_g, ny_g)
    type(solver_state_2d_t), intent(in)  :: state
    character(len=*),        intent(in)  :: filename
    logical,                 intent(out) :: is_ok
    character(len=*),        intent(out) :: message
    real(wp), optional,      intent(in)  :: xc(:,:)  !< gathered cell-centre x (curvilinear)
    real(wp), optional,      intent(in)  :: yc(:,:)  !< gathered cell-centre y (curvilinear)

    integer :: u, ios, i, j
    real(wp) :: x, y, rho, vu, vv, p, gam
    integer :: nx_g, ny_g

    is_ok = .false.
    message = ''
    gam  = state % cfg % gam
    nx_g = state % decomp_2d % nx_global
    ny_g = state % decomp_2d % ny_global

    open(newunit=u, file=filename, status='replace', action='write', iostat=ios)
    if (ios /= 0) then
      message = 'gather_and_write_2d: cannot open "'//trim(filename)//'"'
      return
    end if

    do j = 1, ny_g
      do i = 1, nx_g
        ! Coordinate: curvilinear → gathered centroids; FDM uniform → node formula
        ! (node i=1 at x_left, spacing dx); FVM uniform → cell-centre formula.
        if (present(xc) .and. present(yc)) then
          x = xc(i, j)
          y = yc(i, j)
        else if (trim(state % blocks(1) % method) == method_fdm) then
          x = state % cfg % x_left + real(i - 1, wp) * state % dx
          y = state % cfg % y_left + real(j - 1, wp) * state % dy
        else
          x = state % cfg % x_left + (real(i, wp) - 0.5_wp) * state % dx
          y = state % cfg % y_left + (real(j, wp) - 0.5_wp) * state % dy
        end if
        rho = global_field(1, i, j)
        ! safe_vel emits 0 velocity (no division) for a vacuum cell (rho <= 0)
        ! so output cannot raise a SIGFPE; no-op for rho > 0.
        vu  = safe_vel(global_field(2, i, j), rho)
        vv  = safe_vel(global_field(3, i, j), rho)
        p   = (gam - 1.0_wp) * (global_field(4, i, j) - 0.5_wp * rho * (vu*vu + vv*vv))
        ! ES24.12E3: explicit 3-digit exponent field ensures gfortran always
        ! emits 'E±NNN' even for |exp| >= 100 (e.g. 1.23E-105), making the
        ! file parseable by np.loadtxt without clamping values.
        write(u, '(6ES24.12E3)', iostat=ios) x, y, rho, vu, vv, p
        if (ios /= 0) then
          message = 'gather_and_write_2d: write error'
          close(u, iostat=ios)
          return
        end if
      end do
      write(u, '(a)', iostat=ios) ''   ! blank line between rows (gnuplot structured grid)
      if (ios /= 0) then
        message = 'gather_and_write_2d: write error'
        close(u, iostat=ios)
        return
      end if
    end do

    close(u, iostat=ios)
    if (ios /= 0) then
      message = 'gather_and_write_2d: close failed for "'//trim(filename)//'"'
      return
    end if
    is_ok = .true.
  end subroutine write_global_field

  !> Write the assembled global field as legacy-ASCII VTK STRUCTURED_POINTS.
  !! Cell-centred: ORIGIN at the first cell centre, SPACING=(dx,dy). Emits
  !! SCALARS rho, SCALARS pressure, VECTORS velocity (i fastest, then j).
  subroutine write_global_field_vtk(global_field, state, filename, is_ok, message)
    real(wp),                intent(in)  :: global_field(:,:,:)
    type(solver_state_2d_t), intent(in)  :: state
    character(len=*),        intent(in)  :: filename
    logical,                 intent(out) :: is_ok
    character(len=*),        intent(out) :: message
    integer :: u, ios, i, j, nx_g, ny_g
    real(wp) :: rho, vu, vv, p, gam, x0, y0

    is_ok = .false.; message = ''
    if (.not. state % mesh % uniform) then
      message = 'gather_and_write_2d: VTK output (STRUCTURED_POINTS) is uniform-only; '// &
                'use output_format=dat or tec for curvilinear grids'
      return
    end if
    gam = state % cfg % gam
    nx_g = state % decomp_2d % nx_global
    ny_g = state % decomp_2d % ny_global
    x0 = state % cfg % x_left + 0.5_wp * state % dx
    y0 = state % cfg % y_left + 0.5_wp * state % dy

    open(newunit=u, file=filename, status='replace', action='write', iostat=ios)
    if (ios /= 0) then
      message = 'gather_and_write_2d: cannot open "'//trim(filename)//'"'
      return
    end if
    ! ios accumulates the first non-zero write status; checked once before close
    ! so a failed write aborts via the status-return API without changing the
    ! success path (on failure the partial file is discarded by the caller).
    write(u,'(a)', iostat=ios) '# vtk DataFile Version 3.0'
    if (ios == 0) write(u,'(a)', iostat=ios) 'euler_2d output'
    if (ios == 0) write(u,'(a)', iostat=ios) 'ASCII'
    if (ios == 0) write(u,'(a)', iostat=ios) 'DATASET STRUCTURED_POINTS'
    if (ios == 0) write(u,'(a,i0,1x,i0,1x,i0)', iostat=ios) 'DIMENSIONS ', nx_g, ny_g, 1
    if (ios == 0) write(u,'(a,3ES20.12)', iostat=ios) 'ORIGIN ', x0, y0, 0.0_wp
    if (ios == 0) write(u,'(a,3ES20.12)', iostat=ios) 'SPACING ', state % dx, state % dy, 1.0_wp
    if (ios == 0) write(u,'(a,i0)', iostat=ios) 'POINT_DATA ', nx_g * ny_g
    if (ios == 0) write(u,'(a)', iostat=ios) 'SCALARS rho double 1'
    if (ios == 0) write(u,'(a)', iostat=ios) 'LOOKUP_TABLE default'
    do j = 1, ny_g; do i = 1, nx_g
      if (ios == 0) write(u,'(ES24.12E3)', iostat=ios) global_field(1, i, j)
    end do; end do
    if (ios == 0) write(u,'(a)', iostat=ios) 'SCALARS pressure double 1'
    if (ios == 0) write(u,'(a)', iostat=ios) 'LOOKUP_TABLE default'
    do j = 1, ny_g; do i = 1, nx_g
      rho = global_field(1,i,j)
      ! safe_vel emits 0 velocity (no division) for a vacuum cell (rho <= 0); no-op for rho > 0.
      vu = safe_vel(global_field(2,i,j), rho); vv = safe_vel(global_field(3,i,j), rho)
      p = (gam - 1.0_wp) * (global_field(4,i,j) - 0.5_wp*rho*(vu*vu+vv*vv))
      if (ios == 0) write(u,'(ES24.12E3)', iostat=ios) p
    end do; end do
    if (ios == 0) write(u,'(a)', iostat=ios) 'VECTORS velocity double'
    do j = 1, ny_g; do i = 1, nx_g
      rho = global_field(1,i,j)
      ! safe_vel emits 0 velocity (no division) for a vacuum cell (rho <= 0); no-op for rho > 0.
      if (ios == 0) write(u,'(3ES24.12E3)', iostat=ios) safe_vel(global_field(2,i,j), rho), &
                              safe_vel(global_field(3,i,j), rho), 0.0_wp
    end do; end do
    if (ios /= 0) then
      message = 'gather_and_write_2d: write error'
      close(u, iostat=ios)
      return
    end if
    close(u, iostat=ios)
    if (ios /= 0) then
      message = 'gather_and_write_2d: close failed for "'//trim(filename)//'"'
      return
    end if
    is_ok = .true.
  end subroutine write_global_field_vtk

  !> Slab geometry of one rank's owned piece of the GLOBAL boundary-ghost
  !! frame (the four width-h bands outside the global interior, corners
  !! included). Ownership is disjoint and exhaustive: the left/right bands
  !! cover global rows 1..ny_g only; the bottom/top bands span the FULL
  !! padded width 1-h..nx_g+h and therefore own the four corners. A rank
  !! owns a band segment iff its tile touches that global edge; a
  !! bottom/top-edge rank's segment extends into the corner columns iff it
  !! is also a left/right-edge rank (`ilo`/`ihi` below, in LOCAL i). Ghost
  !! cells of a rank that are NOT global-boundary cells are interior MPI-halo
  !! words: they are refreshed from neighbour interiors by the next halo
  !! exchange and are deliberately not part of the frame.
  pure subroutine frame_slab_extents(ix1, iy1, nxl, nyl, nx_g, ny_g, h, &
                                     has_l, has_r, has_b, has_t, ilo, ihi)
    integer, intent(in) :: ix1, iy1, nxl, nyl, nx_g, ny_g, h
    logical, intent(out) :: has_l, has_r, has_b, has_t
    integer, intent(out) :: ilo, ihi

    has_l = (ix1 == 1)
    has_r = (ix1 + nxl - 1 == nx_g)
    has_b = (iy1 == 1)
    has_t = (iy1 + nyl - 1 == ny_g)
    ilo = 1
    if (has_l) ilo = 1 - h
    ihi = nxl
    if (has_r) ihi = nxl + h
  end subroutine frame_slab_extents

  !> Gather the GLOBAL boundary-ghost frame to cart rank 0: the four width-h
  !! ghost bands outside the global interior, corners included — every ghost
  !! cell of every rank that is NOT an interior MPI-halo cell. The 2D twin of
  !! 1D's `gather_edge_ghosts_to_root` (checkpoint v6; 1D's is v5/F2): the
  !! checkpoint must persist these words because the multi-stage whole-array
  !! RK kernels leave a stage-snapshot blend in them that is not
  !! reconstructable from the interior (see checkpoint_2d.f90's record-block
  !! comment for the full root cause).
  !!
  !! Band layout (all 1-based; h = halo width, nx_g/ny_g = global cells):
  !!   lft/rgt (neq2d, h,          ny_g): column k <-> global i = k - h (left)
  !!                                      or nx_g + k (right); j = 1..ny_g.
  !!   bot/top (neq2d, nx_g + 2*h, h):    column c <-> global i = c - h
  !!                                      (spanning 1-h..nx_g+h, corners
  !!                                      included); row k <-> global j = k - h
  !!                                      (bottom) or ny_g + k (top).
  !!
  !! Edge ownership needs no handshake (mirroring the 1D twin's argument):
  !! `axis_split` delegates to the pure `rank_local_count`/`rank_first_global`
  !! partition, so rank 0 recomputes every source rank's tile from its cart
  !! coordinates (`MPI_Cart_coords` is a local call) — deterministic, no
  !! rank-divergent I/O. At np=1 all four edges are rank 0's and this is a
  !! pure local copy — no MPI traffic. Collective discipline: EVERY rank must
  !! call, at a fixed position; MPI/allocation failures are parallel_fatal
  !! (matching the 1D twin), never silent.
  !!
  !! @param state  This rank's solver state (`ub` must be allocated/populated).
  !! @param lft    inout (neq2d, h, ny_g): valid on cart rank 0 after the call.
  !! @param rgt    inout (neq2d, h, ny_g): valid on cart rank 0 after the call.
  !! @param bot    inout (neq2d, nx_g+2h, h): valid on cart rank 0 after the call.
  !! @param top    inout (neq2d, nx_g+2h, h): valid on cart rank 0 after the call.
  subroutine gather_boundary_ghosts_to_root(state, lft, rgt, bot, top)
    type(solver_state_2d_t), intent(in)    :: state
    real(wp),                intent(inout) :: lft(:, :, :)
    real(wp),                intent(inout) :: rgt(:, :, :)
    real(wp),                intent(inout) :: bot(:, :, :)
    real(wp),                intent(inout) :: top(:, :, :)

#ifdef CFD_SOLVER_LEGACY_MPI
    integer :: comm
#else
    type(MPI_Comm) :: comm
#endif
    integer :: me, ierr, astat, nr, src
    integer :: h, nx_g, ny_g
    integer :: ix1, iy1, nxl, nyl, ilo, ihi
    logical :: has_l, has_r, has_b, has_t
    integer :: coords(2)
    real(wp), allocatable :: buf(:, :, :)

    comm = cart_comm_2d()
    me = -1
    call MPI_Comm_rank(comm, me, ierr)
    if (ierr /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: MPI_Comm_rank failed')

    h    = state % halo_width
    nx_g = state % decomp_2d % nx_global
    ny_g = state % decomp_2d % ny_global
    nr   = n_ranks()

    if (me == 0) then
      ! Own slabs first (at np=1 this alone covers the whole frame).
      ix1 = state % decomp_2d % ix_first_global
      iy1 = state % decomp_2d % iy_first_global
      nxl = state % nx_local
      nyl = state % ny_local
      call frame_slab_extents(ix1, iy1, nxl, nyl, nx_g, ny_g, h, &
                              has_l, has_r, has_b, has_t, ilo, ihi)
      if (has_l) lft(:, 1:h, iy1:iy1 + nyl - 1) = state % ub(:, 1 - h:0, 1:nyl)
      if (has_r) rgt(:, 1:h, iy1:iy1 + nyl - 1) = state % ub(:, nxl + 1:nxl + h, 1:nyl)
      if (has_b) bot(:, ix1 + ilo - 1 + h:ix1 + ihi - 1 + h, 1:h) = state % ub(:, ilo:ihi, 1 - h:0)
      if (has_t) top(:, ix1 + ilo - 1 + h:ix1 + ihi - 1 + h, 1:h) = state % ub(:, ilo:ihi, nyl + 1:nyl + h)

      ! Then every other cart rank's slabs, in cart-rank order, each in the
      ! fixed L, R, B, T slab order both sides derive from the same pure
      ! partition functions — send/recv counts always agree by construction.
      do src = 1, nr - 1
        call MPI_Cart_coords(comm, src, 2, coords, ierr)
        if (ierr /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: MPI_Cart_coords failed')
        nxl = rank_local_count(nx_g, state % decomp_2d % dim_x, coords(1))
        ix1 = rank_first_global(nx_g, state % decomp_2d % dim_x, coords(1))
        nyl = rank_local_count(ny_g, state % decomp_2d % dim_y, coords(2))
        iy1 = rank_first_global(ny_g, state % decomp_2d % dim_y, coords(2))
        call frame_slab_extents(ix1, iy1, nxl, nyl, nx_g, ny_g, h, &
                                has_l, has_r, has_b, has_t, ilo, ihi)
        if (has_l) then
          call recv_slab(buf, neq2d, h, nyl, src, TAG_GHOST_L, comm, 'left')
          lft(:, 1:h, iy1:iy1 + nyl - 1) = buf
          call drop_slab(buf, 'left')
        end if
        if (has_r) then
          call recv_slab(buf, neq2d, h, nyl, src, TAG_GHOST_R, comm, 'right')
          rgt(:, 1:h, iy1:iy1 + nyl - 1) = buf
          call drop_slab(buf, 'right')
        end if
        if (has_b) then
          call recv_slab(buf, neq2d, ihi - ilo + 1, h, src, TAG_GHOST_B, comm, 'bottom')
          bot(:, ix1 + ilo - 1 + h:ix1 + ihi - 1 + h, 1:h) = buf
          call drop_slab(buf, 'bottom')
        end if
        if (has_t) then
          call recv_slab(buf, neq2d, ihi - ilo + 1, h, src, TAG_GHOST_T, comm, 'top')
          top(:, ix1 + ilo - 1 + h:ix1 + ihi - 1 + h, 1:h) = buf
          call drop_slab(buf, 'top')
        end if
      end do
    else
      ! Non-root: send our owned slabs to root in the same fixed L, R, B, T
      ! order (contiguous copies — the ub sections are strided).
      ix1 = state % decomp_2d % ix_first_global
      iy1 = state % decomp_2d % iy_first_global
      nxl = state % nx_local
      nyl = state % ny_local
      call frame_slab_extents(ix1, iy1, nxl, nyl, nx_g, ny_g, h, &
                              has_l, has_r, has_b, has_t, ilo, ihi)
      if (has_l) then
        allocate (buf(neq2d, h, nyl), stat=astat)
        if (astat /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: allocate left slab failed')
        buf = state % ub(:, 1 - h:0, 1:nyl)
        call MPI_Send(buf, neq2d * h * nyl, MPI_DOUBLE_PRECISION, 0, TAG_GHOST_L, comm, ierr)
        if (ierr /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: MPI_Send (left) failed')
        call drop_slab(buf, 'left')
      end if
      if (has_r) then
        allocate (buf(neq2d, h, nyl), stat=astat)
        if (astat /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: allocate right slab failed')
        buf = state % ub(:, nxl + 1:nxl + h, 1:nyl)
        call MPI_Send(buf, neq2d * h * nyl, MPI_DOUBLE_PRECISION, 0, TAG_GHOST_R, comm, ierr)
        if (ierr /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: MPI_Send (right) failed')
        call drop_slab(buf, 'right')
      end if
      if (has_b) then
        allocate (buf(neq2d, ihi - ilo + 1, h), stat=astat)
        if (astat /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: allocate bottom slab failed')
        buf = state % ub(:, ilo:ihi, 1 - h:0)
        call MPI_Send(buf, neq2d * (ihi - ilo + 1) * h, MPI_DOUBLE_PRECISION, 0, TAG_GHOST_B, comm, ierr)
        if (ierr /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: MPI_Send (bottom) failed')
        call drop_slab(buf, 'bottom')
      end if
      if (has_t) then
        allocate (buf(neq2d, ihi - ilo + 1, h), stat=astat)
        if (astat /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: allocate top slab failed')
        buf = state % ub(:, ilo:ihi, nyl + 1:nyl + h)
        call MPI_Send(buf, neq2d * (ihi - ilo + 1) * h, MPI_DOUBLE_PRECISION, 0, TAG_GHOST_T, comm, ierr)
        if (ierr /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: MPI_Send (top) failed')
        call drop_slab(buf, 'top')
      end if
    end if
  end subroutine gather_boundary_ghosts_to_root

  !> Allocate `buf` to (n1, n2, n3) and blocking-receive one frame slab into
  !! it from cart rank `src` (gather_boundary_ghosts_to_root's root side).
  subroutine recv_slab(buf, n1, n2, n3, src, tag, comm, which)
    real(wp), allocatable, intent(inout) :: buf(:, :, :)
    integer,               intent(in)    :: n1, n2, n3, src, tag
#ifdef CFD_SOLVER_LEGACY_MPI
    integer,               intent(in)    :: comm
#else
    type(MPI_Comm),        intent(in)    :: comm
#endif
    character(len=*),      intent(in)    :: which
    integer :: astat, ierr

    allocate (buf(n1, n2, n3), stat=astat)
    if (astat /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: allocate '//which//' recv slab failed')
    call MPI_Recv(buf, n1 * n2 * n3, MPI_DOUBLE_PRECISION, src, tag, comm, MPI_STATUS_IGNORE, ierr)
    if (ierr /= 0) call parallel_fatal('gather_boundary_ghosts_to_root: MPI_Recv ('//which//') failed')
  end subroutine recv_slab

  !> Deallocate a frame slab buffer with the mandatory stat check.
  subroutine drop_slab(buf, which)
    real(wp), allocatable, intent(inout) :: buf(:, :, :)
    character(len=*),      intent(in)    :: which
    integer :: astat

    deallocate (buf, stat=astat)
    if (astat /= 0) call parallel_fatal('boundary ghost frame: deallocate '//which//' slab failed')
  end subroutine drop_slab

  !> Inverse of `gather_boundary_ghosts_to_root`: deliver the checkpoint's
  !! GLOBAL boundary-ghost frame bands into this rank's `ub` ghost cells.
  !! Same band layout and disjoint ownership as the gather (the slab geometry
  !! lives ONLY in `frame_slab_extents`, so gather and scatter cannot drift).
  !!
  !! Unlike the 1D twin (`scatter_edge_ghosts_from_root`, where only rank 0
  !! reads the checkpoint and MPI-delivers the edges), 2D's read path opens
  !! the global checkpoint file on EVERY rank (checkpoint_2d.f90's documented
  !! contract) — each rank already holds the root-authored band buffers it
  !! read itself, so delivery is a pure local slice: no MPI traffic, exactly
  !! like the per-rank `glob` interior slice. The `_from_root` name mirrors
  !! the 1D twin and the bands' root-gathered layout. Collective discipline
  !! still applies: every rank must call, at a fixed position, with the bands
  !! it read (zero-filled if its read failed — harmless, the caller's
  !! collective conclude aborts the restart in that case).
  !!
  !! @param state  This rank's solver state (`ub` receives the restored ghosts).
  !! @param lft    in (neq2d, h, ny_g): left band, as read from the checkpoint.
  !! @param rgt    in (neq2d, h, ny_g): right band.
  !! @param bot    in (neq2d, nx_g+2h, h): bottom band (owns the corners).
  !! @param top    in (neq2d, nx_g+2h, h): top band (owns the corners).
  subroutine scatter_boundary_ghosts_from_root(state, lft, rgt, bot, top)
    type(solver_state_2d_t), intent(inout) :: state
    real(wp),                intent(in)    :: lft(:, :, :)
    real(wp),                intent(in)    :: rgt(:, :, :)
    real(wp),                intent(in)    :: bot(:, :, :)
    real(wp),                intent(in)    :: top(:, :, :)

    integer :: h, nx_g, ny_g
    integer :: ix1, iy1, nxl, nyl, ilo, ihi
    logical :: has_l, has_r, has_b, has_t

    h    = state % halo_width
    nx_g = state % decomp_2d % nx_global
    ny_g = state % decomp_2d % ny_global
    ix1  = state % decomp_2d % ix_first_global
    iy1  = state % decomp_2d % iy_first_global
    nxl  = state % nx_local
    nyl  = state % ny_local
    call frame_slab_extents(ix1, iy1, nxl, nyl, nx_g, ny_g, h, &
                            has_l, has_r, has_b, has_t, ilo, ihi)
    if (has_l) state % ub(:, 1 - h:0, 1:nyl) = lft(:, 1:h, iy1:iy1 + nyl - 1)
    if (has_r) state % ub(:, nxl + 1:nxl + h, 1:nyl) = rgt(:, 1:h, iy1:iy1 + nyl - 1)
    if (has_b) state % ub(:, ilo:ihi, 1 - h:0) = bot(:, ix1 + ilo - 1 + h:ix1 + ihi - 1 + h, 1:h)
    if (has_t) state % ub(:, ilo:ihi, nyl + 1:nyl + h) = top(:, ix1 + ilo - 1 + h:ix1 + ihi - 1 + h, 1:h)
  end subroutine scatter_boundary_ghosts_from_root

end module solution_gather_2d
