!> @file solution_gather.F90
!> @brief Gather a per-rank primitives slice into a rank-0-owned global
!!        buffer via MPI_Gatherv.
!!
!! Phase D helper. Reused by write_solution_file and write_snapshot to
!! produce byte-identical multi-rank output.
!!
!! Hybrid cpp split mirrors src/parallel_reductions.F90: MS-MPI ships only
!! the legacy `use mpi` module; OpenMPI ships both and we prefer `mpi_f08`
!! for the strong typing.

module solution_gather
  use precision, only: wp
  use domain_decomposition, only: decomp_t, rank_local_count
  use solver_state, only: solver_state_t, neq
#ifdef CFD_SOLVER_LEGACY_MPI
  use mpi, only: MPI_Gatherv, MPI_Send, MPI_Recv, MPI_DOUBLE_PRECISION, MPI_STATUS_IGNORE
#else
  use mpi_f08, only: MPI_Gatherv, MPI_Send, MPI_Recv, MPI_DOUBLE_PRECISION, MPI_STATUS_IGNORE
#endif
  use mpi_runtime, only: mpi_world, my_rank, n_ranks, parallel_fatal
  implicit none
  private

  public :: gather_solution_to_root
  public :: gather_edge_ghosts_to_root, scatter_edge_ghosts_from_root

  !> Point-to-point tag for the right-edge ghost exchange (F2). The left edge
  !! never needs a tag: `decompose()`'s contiguous block partition always
  !! gives rank 0 global index 1, so the left edge is always already local to
  !! rank 0 (see the two subroutines below for the full ownership argument).
  integer, parameter :: TAG_EDGE_RIGHT = 4001

contains

  !> Gather (neq, n_local) primitives from every rank into (neq, n_global)
  !! on rank 0 via MPI_Gatherv.
  !!
  !! @param local_prim   per-rank send buffer, shape (neq, n_local)
  !! @param n_local      this rank's interior cell count
  !! @param decomp       per-rank decomposition descriptor (needed for n_global, n_ranks)
  !! @param gather_buf   inout: rank 0 receives (neq, decomp%n_global);
  !!                     other ranks must pass a (neq, 0) buffer (untouched).
  subroutine gather_solution_to_root(local_prim, n_local, decomp, gather_buf)
    real(wp),       intent(in)    :: local_prim(:, :)
    integer,        intent(in)    :: n_local
    type(decomp_t), intent(in)    :: decomp
    real(wp),       intent(inout) :: gather_buf(:, :)

    integer, allocatable :: recv_counts(:), displs(:)
    integer :: rk, send_count, ierr
    integer :: nr

    nr = n_ranks()
    send_count = neq * n_local

    if (my_rank() == 0) then
      allocate (recv_counts(0:nr - 1), displs(0:nr - 1), stat=ierr)
      if (ierr /= 0) error stop 'solution_gather: recv_counts/displs allocation failed'
      do rk = 0, nr - 1
        recv_counts(rk) = neq * rank_local_count(decomp % n_global, nr, rk)
      end do
      displs(0) = 0
      do rk = 1, nr - 1
        displs(rk) = displs(rk - 1) + recv_counts(rk - 1)
      end do
    else
      ! Non-root must still pass valid pointers to MPI_Gatherv; allocate
      ! zero-length so the addresses are valid even though they are unused.
      allocate (recv_counts(0:0), displs(0:0), stat=ierr)
      if (ierr /= 0) error stop 'solution_gather: non-root recv_counts/displs allocation failed'
      recv_counts = 0
      displs = 0
    end if

    call MPI_Gatherv(local_prim, send_count, MPI_DOUBLE_PRECISION, &
                     gather_buf, recv_counts, displs, MPI_DOUBLE_PRECISION, &
                     0, mpi_world(), ierr)
    if (ierr /= 0) call parallel_fatal('gather_solution_to_root: MPI_Gatherv failed')

    deallocate (recv_counts, displs, stat=ierr)
  end subroutine gather_solution_to_root

  !> Gather the boundary halo/ghost cells owned by the global left- and
  !! right-edge ranks to rank 0: `left(neq,h)` = the left-edge rank's
  !! `ub(:,1-h:0)`, `right(neq,h)` = the right-edge rank's
  !! `ub(:,n_local+1:n_local+h)`. Needed because checkpoint/restart must
  !! persist these cells for "prescribed" BCs (dirichlet/inflow/
  !! supersonic_inlet) whose ghost values are never recomputed by apply_bcs
  !! after the initial seed, yet are NOT actually constant across a run (see
  !! checkpoint.f90's module header for the full story) — F2.
  !!
  !! Edge ownership follows `decompose()`'s deterministic contiguous-block
  !! partition: rank 0 always owns global index 1 (the left edge). The right
  !! edge is owned by rank `decomp%n_ranks - 1`, a value every rank computes
  !! identically from its own `decomp` (n_ranks is constant across ranks) —
  !! no rank-divergent I/O or extra handshake is needed to discover it. At
  !! np=1 (or whenever n_ranks-1 == 0) both edges are local to rank 0 and this
  !! is a pure local copy — no MPI traffic.
  !!
  !! @param state  This rank's solver state (`ub` must be allocated/populated).
  !! @param left   inout (neq, h): valid on rank 0 after the call.
  !! @param right  inout (neq, h): valid on rank 0 after the call.
  subroutine gather_edge_ghosts_to_root(state, left, right)
    type(solver_state_t), intent(in)    :: state
    real(wp),              intent(inout) :: left(:, :)
    real(wp),              intent(inout) :: right(:, :)

    integer :: h, n_pt, right_rank, ierr

    ! W3: derive the halo width from `state % halo_width` -- the SAME field
    ! `ub` is dimensioned from, and the one checkpoint.f90 sizes its transfer
    ! buffers with. `state % decomp % halo_width` is a copy made at decompose()
    ! time and is equal only by init ordering; sizing the `ub` slices and the
    ! MPI counts from the copy while the caller sized its buffer from the
    ! original would silently become an OOB copy / MPI count mismatch if that
    ! ordering ever changed. One field, everywhere `ub` is touched (this also
    ! matches the 2D twin, which uses `state % halo_width` throughout).
    h = state % halo_width
    n_pt = state % n_pt
    right_rank = state % decomp % n_ranks - 1

    if (my_rank() == 0) then
      left = state % ub(:, 1 - h:0)
      if (right_rank == 0) then
        right = state % ub(:, n_pt + 1:n_pt + h)
      else
        call MPI_Recv(right, neq * h, MPI_DOUBLE_PRECISION, right_rank, TAG_EDGE_RIGHT, &
                     mpi_world(), MPI_STATUS_IGNORE, ierr)
        if (ierr /= 0) call parallel_fatal('gather_edge_ghosts_to_root: MPI_Recv (right) failed')
      end if
    else if (my_rank() == right_rank) then
      call MPI_Send(state % ub(:, n_pt + 1:n_pt + h), neq * h, MPI_DOUBLE_PRECISION, 0, &
                   TAG_EDGE_RIGHT, mpi_world(), ierr)
      if (ierr /= 0) call parallel_fatal('gather_edge_ghosts_to_root: MPI_Send (right) failed')
    end if
  end subroutine gather_edge_ghosts_to_root

  !> Inverse of `gather_edge_ghosts_to_root`: rank 0 delivers the `left`/
  !! `right` ghost blocks (as just read from a checkpoint) to the ranks that
  !! own those edges, copying directly into `state % ub`'s halo columns.
  !! Same ownership derivation and np=1 pure-local-copy behaviour as the
  !! gather — F2.
  !!
  !! @param state  This rank's solver state (`ub` receives the restored ghosts).
  !! @param left   in (neq, h): the gathered left-edge ghost block (rank 0's).
  !! @param right  in (neq, h): the gathered right-edge ghost block (rank 0's).
  subroutine scatter_edge_ghosts_from_root(state, left, right)
    type(solver_state_t), intent(inout) :: state
    real(wp),              intent(in)    :: left(:, :)
    real(wp),              intent(in)    :: right(:, :)

    integer :: h, n_pt, right_rank, ierr

    ! W3: same single-source halo width as the gather above.
    h = state % halo_width
    n_pt = state % n_pt
    right_rank = state % decomp % n_ranks - 1

    if (my_rank() == 0) then
      state % ub(:, 1 - h:0) = left
      if (right_rank == 0) then
        state % ub(:, n_pt + 1:n_pt + h) = right
      else
        call MPI_Send(right, neq * h, MPI_DOUBLE_PRECISION, right_rank, TAG_EDGE_RIGHT, &
                     mpi_world(), ierr)
        if (ierr /= 0) call parallel_fatal('scatter_edge_ghosts_from_root: MPI_Send (right) failed')
      end if
    else if (my_rank() == right_rank) then
      call MPI_Recv(state % ub(:, n_pt + 1:n_pt + h), neq * h, MPI_DOUBLE_PRECISION, 0, &
                   TAG_EDGE_RIGHT, mpi_world(), MPI_STATUS_IGNORE, ierr)
      if (ierr /= 0) call parallel_fatal('scatter_edge_ghosts_from_root: MPI_Recv (right) failed')
    end if
  end subroutine scatter_edge_ghosts_from_root

end module solution_gather
