!> @file domain_decomposition.f90
!> @brief Slab partition of the 1D grid across MPI ranks.
!!
!! Owns the per-rank `decomp_t` derived type and the pure `decompose()`
!! function that, given (my_rank, n_ranks, n_global, halo_width, is_periodic),
!! returns a fully-populated decomp_t with neighbour rank ids and the
!! global-index range owned by this rank. No MPI calls live here — the
!! function is callable from unit tests with synthetic rank ids.
!!
!! Validation of the (n_local >= halo_width) feasibility constraint is done
!! by `validate_decomp()`, which aborts via parallel_fatal when violated.

module domain_decomposition
  use mpi_runtime, only: mpi_proc_null_value, parallel_fatal
  implicit none
  private

  public :: decomp_t
  public :: decompose
  public :: is_decomp_feasible
  public :: validate_decomp
  public :: rank_local_count
  public :: rank_first_global

  !> Per-rank domain decomposition descriptor. All integer fields are
  !! local to the rank that holds it; only `n_ranks`, `n_global`,
  !! `halo_width`, and `is_periodic` are constant across ranks.
  type :: decomp_t
    integer :: my_rank = -1
    integer :: n_ranks = -1
    integer :: n_global = -1
    integer :: n_local = -1
    integer :: i_first_global = -1   !< global index of local cell 1
    integer :: i_last_global = -1   !< global index of local cell n_local
    integer :: halo_width = -1
    integer :: left_neighbour = -1   !< MPI_PROC_NULL on rank 0 (non-periodic)
    integer :: right_neighbour = -1   !< MPI_PROC_NULL on rank N-1 (non-periodic)
    logical :: is_periodic = .false.
  end type decomp_t

contains

  !> Single source of truth for the slab/block partition: the number of global
  !! cells owned by `rank_id` (0-based) when `n_global` cells are split across
  !! `n_ranks`. floor(n_global/n_ranks) per rank, with the remainder distributed
  !! one extra cell each to the first `mod(n_global,n_ranks)` ranks. Pure — used
  !! by decompose(), the gather/scatter count builders, and the 2D axis split so
  !! the partition arithmetic exists in exactly one place.
  pure function rank_local_count(n_global, n_ranks, rank_id) result(n_local)
    integer, intent(in) :: n_global, n_ranks, rank_id
    integer :: n_local
    integer :: base

    base = n_global / n_ranks
    if (rank_id < mod(n_global, n_ranks)) then
      n_local = base + 1
    else
      n_local = base
    end if
  end function rank_local_count

  !> 1-based global index of `rank_id`'s first owned cell under the same
  !! partition as `rank_local_count`. Equals the sum of all lower ranks'
  !! counts, plus 1. Pure — single source of truth (see rank_local_count).
  pure function rank_first_global(n_global, n_ranks, rank_id) result(i_first)
    integer, intent(in) :: n_global, n_ranks, rank_id
    integer :: i_first
    integer :: base, rem

    base = n_global / n_ranks
    rem = mod(n_global, n_ranks)
    if (rank_id < rem) then
      ! All lower ranks are "fat" (base+1 cells each).
      i_first = rank_id * (base + 1) + 1
    else
      ! `rem` fat ranks then (rank_id-rem) "thin" ranks below this one.
      ! Algebraically equal to rem*(base+1) + (rank_id-rem)*base + 1.
      i_first = rank_id * base + rem + 1
    end if
  end function rank_first_global

  !> Compute this rank's decomposition. Pure: no MPI calls, no global state.
  !!
  !! Slab partition: floor(n_global / n_ranks) cells per rank, with the
  !! remainder distributed to the first (n_global mod n_ranks) ranks. This
  !! is the standard "block" partition used by virtually every 1D MPI
  !! decomposition; it minimises the cell-count spread.
  pure function decompose(my_rank, n_ranks, n_global, halo_width, is_periodic) result(d)
    integer, intent(in) :: my_rank, n_ranks, n_global, halo_width
    logical, intent(in) :: is_periodic
    type(decomp_t) :: d

    d % my_rank = my_rank
    d % n_ranks = n_ranks
    d % n_global = n_global
    d % halo_width = halo_width
    d % is_periodic = is_periodic

    d % n_local = rank_local_count(n_global, n_ranks, my_rank)
    d % i_first_global = rank_first_global(n_global, n_ranks, my_rank)
    d % i_last_global = d % i_first_global + d % n_local - 1

    if (my_rank == 0) then
      if (is_periodic) then
        d % left_neighbour = n_ranks - 1
      else
        d % left_neighbour = mpi_proc_null_value()
      end if
    else
      d % left_neighbour = my_rank - 1
    end if

    if (my_rank == n_ranks - 1) then
      if (is_periodic) then
        d % right_neighbour = 0
      else
        d % right_neighbour = mpi_proc_null_value()
      end if
    else
      d % right_neighbour = my_rank + 1
    end if
  end function decompose

  !> Minimum interior cell count this rank must hold. A non-periodic rank needs
  !! at least `halo_width` cells. A periodic decomposition needs one more: the
  !! wrap send skips the duplicated endpoint node (global node n_global == node
  !! 1), so the periodic edge ranks send cells [2 .. h+1] / [n_local-h ..
  !! n_local-1], which require n_local >= halo_width + 1. The bound is applied to
  !! every rank (not just the edges) so the feasibility guarantee is uniform.
  pure function required_local(d) result(needed)
    type(decomp_t), intent(in) :: d
    integer :: needed

    needed = d % halo_width
    if (d % is_periodic) needed = needed + 1
  end function required_local

  !> Return .true. iff this rank's n_local is large enough to accommodate the
  !! halo (and, for periodic grids, the wrap send's one-cell offset).
  pure function is_decomp_feasible(d) result(ok)
    type(decomp_t), intent(in) :: d
    logical :: ok

    ok = d % n_local >= required_local(d)
  end function is_decomp_feasible

  !> Abort with a clear message if this rank's n_local is too small to hold
  !! the halo cells the chosen reconstruction scheme requires.
  subroutine validate_decomp(d)
    type(decomp_t), intent(in) :: d
    character(len=256) :: msg

    if (.not. is_decomp_feasible(d)) then
      write (msg, '(A,I0,A,I0,A,I0,A,I0,A)') &
        'decomp: n_local (', d % n_local, &
        ') < required (', required_local(d), &
        ') on rank ', d % my_rank, &
        ' of ', d % n_ranks, &
        '; use fewer ranks or a larger n_cell'// &
        ' (periodic grids need halo_width+1 per rank)'
      call parallel_fatal(trim(msg))
    end if
  end subroutine validate_decomp

end module domain_decomposition