!> Preconditioned Conjugate Gradient (PCG) linear solver and operators.
!!
!! This module houses the pure numerical and linear algebra routines 
!! used by the LowMachReact-Hex framework. It manages the static Poisson 
!! matrix cache, temporary workspace allocations, Preconditioned Conjugate 
!! Gradient (PCG) solver loop, sparse matrix-vector products, and zero-mean 
!! gauge projections for pure Neumann systems.
module mod_linear_solver
   use mpi_f08
   use mod_precision, only : rk, zero, one, half, tiny_safe, fatal_error
   use mod_input, only : case_params_t
   use mod_mesh, only : mesh_t
   use mod_mpi_flow, only : flow_mpi_t, flow_exchange_cell_scalar, &
                            flow_global_dot_owned, flow_global_two_dots_owned
   use mod_bc, only : bc_set_t, bc_periodic, bc_neumann, bc_dirichlet, &
                      patch_type_for_face, boundary_pressure_type, &
                      face_effective_neighbor
   use mod_profiling, only : profiler_start, profiler_stop
   use mod_cantera, only : transport_properties_t
   implicit none

   private

   !> Cached sparse matrix coefficients for the Pressure Poisson operator.
   type, public :: pressure_operator_cache_t
      logical :: initialized = .false.      !< Initialization toggle.
      logical :: has_dirichlet_pressure = .false. !< True if at least one patch has a fixed pressure.
      logical :: has_neumann_outlet = .false. !< True if at least one boundary patch can absorb flux imbalance.
      integer :: ncells = 0                 !< Number of cells.
      integer :: max_faces = 0              !< Max faces per cell.
      integer, allocatable :: nb(:,:)       !< Neighbor indices for each cell face.
      real(rk), allocatable :: coeff(:,:)   !< Off-diagonal Laplacian coefficients
      real(rk), allocatable :: diag(:)      !< Diagonal Laplacian coefficients
      real(rk), allocatable :: var_coeff(:,:) !< Per-step variable-density coefficients
      real(rk), allocatable :: var_diag(:)  !< Per-step variable-density diagonal.
   end type pressure_operator_cache_t

   !> Temporary workspace for projection step calculations.
   type, public :: projection_workspace_t
      logical :: initialized = .false.
      integer :: ncells = 0
      integer :: nfaces = 0

      real(rk), allocatable :: local_vec(:,:)      !< Local momentum RHS.
      real(rk), allocatable :: local_vec_star(:,:) !< Local intermediate velocity.
      real(rk), allocatable :: local_scalar(:)     !< Local divergence scalar.
      real(rk), allocatable :: rhs_poisson(:)      !< Source term for Poisson solve.
      real(rk), allocatable :: local_face_flux(:)  !< Face fluxes on owned faces.
      real(rk), allocatable :: predicted_face_flux(:) !< Global intermediate face fluxes.

      real(rk), allocatable :: r(:)                !< PCG residual vector.
      real(rk), allocatable :: z(:)                !< PCG preconditioned residual.
      real(rk), allocatable :: pvec(:)             !< PCG search direction.
      real(rk), allocatable :: ap(:)               !< PCG operator-applied vector.
      real(rk), allocatable :: local_ap(:)         !< PCG local operator-applied vector.
   end type projection_workspace_t

   !> Dummy/diagnostic solver statistics structure mirroring the one in flow.
   !! Keeping structural declarations isolated or shared if needed.
   type, public :: solver_stats_t
      integer :: pressure_iterations = 0
      integer :: pressure_iterations_total = 0
      integer :: pressure_iterations_max = 0
      integer :: pressure_solve_count = 0
      real(rk) :: pressure_iterations_avg = zero
      real(rk) :: pressure_residual = zero      !< Relative PCG residual sqrt(r.r/r0.r0).
      real(rk) :: pressure_residual_abs = zero  !< Absolute RMS PCG residual sqrt(r.r/N).
      real(rk) :: max_divergence = zero
      real(rk) :: rms_divergence = zero
      real(rk) :: net_boundary_flux = zero
      real(rk) :: kinetic_energy = zero
      real(rk) :: cfl = zero
      real(rk) :: wall_time = zero
      real(rk) :: max_velocity = zero
      real(rk) :: total_mass = zero
      real(rk) :: min_species_y = zero
   end type solver_stats_t

   type(pressure_operator_cache_t), save, public :: pressure_cache
   type(projection_workspace_t), save, public :: projection_work

   !> Variable-density pure-Neumann pressure systems use an explicit
   !! zero-mean projection to control the constant nullspace.
   integer, parameter :: pcg_nullspace_projection_interval = 50

   public :: solve_pressure_correction, subtract_owned_mean_scalar
   public :: ensure_projection_workspace, ensure_pressure_operator_cache
   public :: finalize_flow_projection_workspace, update_variable_pressure_operator_cache
   public :: face_normal_distance, outward_normal
   public :: pressure_cache_has_dirichlet_pressure, pressure_cache_has_neumann_outlet
   public :: pressure_face_inverse_density

contains

   !> Expose Dirichlet flag from cache.
   logical function pressure_cache_has_dirichlet_pressure() result(flag)
      flag = pressure_cache%has_dirichlet_pressure
   end function pressure_cache_has_dirichlet_pressure

   !> Expose Neumann outlet flag from cache.
   logical function pressure_cache_has_neumann_outlet() result(flag)
      flag = pressure_cache%has_neumann_outlet
   end function pressure_cache_has_neumann_outlet

   !> Determines the outward unit normal relative to a specific cell.
   pure function outward_normal(mesh, face_id, cell_id) result(nvec)
      type(mesh_t), intent(in) :: mesh
      integer, intent(in) :: face_id
      integer, intent(in) :: cell_id
      real(rk) :: nvec(3)

      if (mesh%faces(face_id)%owner == cell_id) then
         nvec = mesh%faces(face_id)%normal
      else
         nvec = -mesh%faces(face_id)%normal
      end if
   end function outward_normal

   !> Calculates the distance between cell and neighbor center along face normal.
   function face_normal_distance(mesh, bc, face_id, cell_id, nb) result(dist)
      type(mesh_t), intent(in) :: mesh
      type(bc_set_t), intent(in) :: bc
      integer, intent(in) :: face_id
      integer, intent(in) :: cell_id
      integer, intent(in) :: nb
      real(rk) :: dist

      integer :: pair_face
      integer :: btype
      real(rk) :: nvec(3)

      nvec = outward_normal(mesh, face_id, cell_id)

      if (nb > 0) then
         if (mesh%faces(face_id)%neighbor == 0) then
            btype = patch_type_for_face(mesh, bc, face_id)

            if (btype == bc_periodic) then
               pair_face = mesh%faces(face_id)%periodic_face

               if (pair_face <= 0) then
                  call fatal_error('numerics', 'periodic face has no paired face')
               end if

               dist = abs(dot_product(mesh%faces(face_id)%center - &
                                      mesh%cells(cell_id)%center, nvec)) + &
                      abs(dot_product(mesh%cells(nb)%center - &
                                      mesh%faces(pair_face)%center, nvec))

               dist = max(dist, tiny_safe)
               return
            end if
         end if

         dist = abs(dot_product(mesh%cells(nb)%center - &
                                mesh%cells(cell_id)%center, nvec))
      else
         dist = abs(dot_product(mesh%faces(face_id)%center - &
                                mesh%cells(cell_id)%center, nvec))
      end if

      dist = max(dist, tiny_safe)
   end function face_normal_distance

   !> Allocate persistent solver workspace.
   subroutine ensure_projection_workspace(mesh)
      type(mesh_t), intent(in) :: mesh

      if (projection_work%initialized) then
         if (projection_work%ncells == mesh%ncells .and. &
             projection_work%nfaces == mesh%nfaces) return

         call finalize_flow_projection_workspace()
      end if

      projection_work%ncells = mesh%ncells
      projection_work%nfaces = mesh%nfaces

      allocate(projection_work%local_vec(3, mesh%ncells))
      allocate(projection_work%local_vec_star(3, mesh%ncells))
      allocate(projection_work%local_scalar(mesh%ncells))
      allocate(projection_work%rhs_poisson(mesh%ncells))
      allocate(projection_work%local_face_flux(mesh%nfaces))
      allocate(projection_work%predicted_face_flux(mesh%nfaces))

      allocate(projection_work%r(mesh%ncells))
      allocate(projection_work%z(mesh%ncells))
      allocate(projection_work%pvec(mesh%ncells))
      allocate(projection_work%ap(mesh%ncells))
      allocate(projection_work%local_ap(mesh%ncells))

      projection_work%local_vec = zero
      projection_work%local_vec_star = zero
      projection_work%local_scalar = zero
      projection_work%rhs_poisson = zero
      projection_work%local_face_flux = zero
      projection_work%predicted_face_flux = zero

      projection_work%r = zero
      projection_work%z = zero
      projection_work%pvec = zero
      projection_work%ap = zero
      projection_work%local_ap = zero

      projection_work%initialized = .true.
   end subroutine ensure_projection_workspace

   !> Deallocate temporary work arrays.
   subroutine finalize_flow_projection_workspace()
      if (allocated(projection_work%local_vec)) deallocate(projection_work%local_vec)
      if (allocated(projection_work%local_vec_star)) deallocate(projection_work%local_vec_star)
      if (allocated(projection_work%local_scalar)) deallocate(projection_work%local_scalar)
      if (allocated(projection_work%rhs_poisson)) deallocate(projection_work%rhs_poisson)
      if (allocated(projection_work%local_face_flux)) deallocate(projection_work%local_face_flux)
      if (allocated(projection_work%predicted_face_flux)) deallocate(projection_work%predicted_face_flux)

      if (allocated(projection_work%r)) deallocate(projection_work%r)
      if (allocated(projection_work%z)) deallocate(projection_work%z)
      if (allocated(projection_work%pvec)) deallocate(projection_work%pvec)
      if (allocated(projection_work%ap)) deallocate(projection_work%ap)
      if (allocated(projection_work%local_ap)) deallocate(projection_work%local_ap)

      projection_work%initialized = .false.
      projection_work%ncells = 0
      projection_work%nfaces = 0

      if (allocated(pressure_cache%nb)) deallocate(pressure_cache%nb)
      if (allocated(pressure_cache%coeff)) deallocate(pressure_cache%coeff)
      if (allocated(pressure_cache%diag)) deallocate(pressure_cache%diag)
      if (allocated(pressure_cache%var_coeff)) deallocate(pressure_cache%var_coeff)
      if (allocated(pressure_cache%var_diag)) deallocate(pressure_cache%var_diag)

      pressure_cache%initialized = .false.
      pressure_cache%has_dirichlet_pressure = .false.
      pressure_cache%has_neumann_outlet = .false.
      pressure_cache%ncells = 0
      pressure_cache%max_faces = 0
   end subroutine finalize_flow_projection_workspace

   !> Pre-computes Laplacian coefficients for the Poisson operator.
   subroutine ensure_pressure_operator_cache(mesh, flow, bc)
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(in) :: flow
      type(bc_set_t), intent(in) :: bc
      integer :: ierr, c, lf, f, nb
      real(rk) :: dist, coeff
      integer :: local_dirichlet_flag, global_dirichlet_flag

      if (pressure_cache%initialized) then
         if (pressure_cache%ncells == mesh%ncells) return
         call finalize_flow_projection_workspace()
      end if

      pressure_cache%ncells = mesh%ncells
      pressure_cache%max_faces = size(mesh%cell_faces, 1)

      allocate(pressure_cache%nb(pressure_cache%max_faces, mesh%ncells))
      allocate(pressure_cache%coeff(pressure_cache%max_faces, mesh%ncells))
      allocate(pressure_cache%diag(mesh%ncells))
      allocate(pressure_cache%var_coeff(pressure_cache%max_faces, mesh%ncells))
      allocate(pressure_cache%var_diag(mesh%ncells))

      pressure_cache%nb = 0
      pressure_cache%coeff = zero
      pressure_cache%diag = zero
      pressure_cache%var_coeff = zero
      pressure_cache%var_diag = zero
      pressure_cache%has_neumann_outlet = .false.

      local_dirichlet_flag = 0

      do c = 1, mesh%ncells
         do lf = 1, mesh%ncell_faces(c)
            f = mesh%cell_faces(lf, c)
            nb = face_effective_neighbor(mesh, bc, f, c)

            if (nb <= 0) then
               if (patch_type_for_face(mesh, bc, f) == bc_neumann) then
                  pressure_cache%has_neumann_outlet = .true.
               end if

               if (boundary_pressure_type(mesh, bc, f) == bc_dirichlet) then
                  dist = face_normal_distance(mesh, bc, f, c, 0)
                  coeff = mesh%faces(f)%area / dist
                  pressure_cache%diag(c) = pressure_cache%diag(c) + coeff
                  local_dirichlet_flag = 1
               end if
               cycle
            end if

            dist = face_normal_distance(mesh, bc, f, c, nb)
            coeff = mesh%faces(f)%area / dist

            pressure_cache%nb(lf, c) = nb
            pressure_cache%coeff(lf, c) = coeff
            pressure_cache%diag(c) = pressure_cache%diag(c) + coeff
         end do

         if (pressure_cache%diag(c) <= tiny_safe) then
            call fatal_error('numerics', 'non-positive cached pressure diagonal')
         end if
      end do

      call MPI_Allreduce(local_dirichlet_flag, global_dirichlet_flag, 1, MPI_INTEGER, MPI_MAX, flow%comm, ierr)
      pressure_cache%has_dirichlet_pressure = (global_dirichlet_flag > 0)

      pressure_cache%initialized = .true.
   end subroutine ensure_pressure_operator_cache

   !> Rebuild per-step variable-density pressure coefficients.
   subroutine update_variable_pressure_operator_cache(mesh, flow, bc, transport)
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(in) :: flow
      type(bc_set_t), intent(in) :: bc
      type(transport_properties_t), intent(in) :: transport

      integer :: c, lf, f, nb
      real(rk) :: dist, coeff, beta_face

      call profiler_start('Pressure_VarCoeff_Update')
      call ensure_pressure_operator_cache(mesh, flow, bc)

      if (.not. allocated(pressure_cache%var_coeff) .or. .not. allocated(pressure_cache%var_diag)) then
         call fatal_error('numerics', 'variable pressure coefficient cache is not allocated')
      end if

      do c = flow%first_cell, flow%last_cell
         pressure_cache%var_diag(c) = zero
         do lf = 1, mesh%ncell_faces(c)
            pressure_cache%var_coeff(lf, c) = zero
            f = mesh%cell_faces(lf, c)
            nb = face_effective_neighbor(mesh, bc, f, c)

            if (nb > 0) then
               dist = face_normal_distance(mesh, bc, f, c, nb)
               beta_face = pressure_face_inverse_density(transport, c, nb)
               coeff = beta_face * mesh%faces(f)%area / dist
               pressure_cache%nb(lf, c) = nb
               pressure_cache%var_coeff(lf, c) = coeff
               pressure_cache%var_diag(c) = pressure_cache%var_diag(c) + coeff
            else if (boundary_pressure_type(mesh, bc, f) == bc_dirichlet) then
               dist = face_normal_distance(mesh, bc, f, c, 0)
               beta_face = pressure_face_inverse_density(transport, c, 0)
               coeff = beta_face * mesh%faces(f)%area / dist
               pressure_cache%var_diag(c) = pressure_cache%var_diag(c) + coeff
            end if
         end do

         if (pressure_cache%var_diag(c) <= tiny_safe) then
            call fatal_error('numerics', 'non-positive variable-density pressure diagonal')
         end if
      end do

      call profiler_stop('Pressure_VarCoeff_Update')
   end subroutine update_variable_pressure_operator_cache

   !> Select the active diagonal used by the pressure PCG preconditioner.
   function pressure_preconditioner_diag(params, c) result(diag_value)
      type(case_params_t), intent(in) :: params
      integer, intent(in) :: c
      real(rk) :: diag_value

      if (params%enable_variable_density) then
         diag_value = pressure_cache%var_diag(c)
      else
         diag_value = pressure_cache%diag(c)
      end if
   end function pressure_preconditioner_diag

   !> Face-centered inverse density coefficient for variable-density projection.
   function pressure_face_inverse_density(transport, owner, nb) result(beta)
      type(transport_properties_t), intent(in) :: transport
      integer, intent(in) :: owner
      integer, intent(in) :: nb
      real(rk) :: beta

      if (.not. allocated(transport%rho)) then
         call fatal_error('numerics', 'transport rho must be allocated for variable-density projection')
      end if

      if (nb > 0) then
         beta = half * (one / max(transport%rho(owner), tiny_safe) + &
                        one / max(transport%rho(nb), tiny_safe))
      else
         beta = one / max(transport%rho(owner), tiny_safe)
      end if
   end function pressure_face_inverse_density

   !> Sparse Matrix-Vector multiplication for the Laplacian operator.
   subroutine pressure_matvec(mesh, flow, bc, params, x, local_ax)
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(in) :: flow
      type(bc_set_t), intent(in) :: bc
      type(case_params_t), intent(in) :: params
      real(rk), intent(in) :: x(:)
      real(rk), intent(out) :: local_ax(:)

      integer :: c, lf, nb

      call profiler_start('Pressure_Matvec')
      call ensure_pressure_operator_cache(mesh, flow, bc)

      local_ax = zero

      if (params%enable_variable_density) then
         do c = flow%first_cell, flow%last_cell
            local_ax(c) = pressure_cache%var_diag(c) * x(c)

            do lf = 1, mesh%ncell_faces(c)
               nb = pressure_cache%nb(lf, c)
               if (nb <= 0) cycle

               local_ax(c) = local_ax(c) - pressure_cache%var_coeff(lf, c) * x(nb)
            end do
         end do

         call profiler_stop('Pressure_Matvec')
         return
      end if

      do c = flow%first_cell, flow%last_cell
         if (c == 1 .and. .not. pressure_cache%has_dirichlet_pressure .and. &
                .not. params%enable_variable_density) then
            local_ax(c) = x(c)
            cycle
         end if

         local_ax(c) = pressure_cache%diag(c) * x(c)

         do lf = 1, mesh%ncell_faces(c)
            nb = pressure_cache%nb(lf, c)
            if (nb <= 0) cycle

            local_ax(c) = local_ax(c) - pressure_cache%coeff(lf, c) * x(nb)
         end do
      end do
      call profiler_stop('Pressure_Matvec')
   end subroutine pressure_matvec

   !> Remove the constant nullspace component from an owned-cell scalar.
   subroutine subtract_owned_mean_scalar(flow, scalar)
      type(flow_mpi_t), intent(in) :: flow
      real(rk), intent(inout) :: scalar(:)

      integer :: c, ierr
      integer :: global_count
      real(rk) :: local_sum, global_sum, mean_value

      local_sum = zero

      do c = flow%first_cell, flow%last_cell
         local_sum = local_sum + scalar(c)
      end do

      call profiler_start('Pressure_Gauge_Allreduce')
      call MPI_Allreduce(local_sum, global_sum, 1, MPI_DOUBLE_PRECISION, MPI_SUM, flow%comm, ierr)
      call profiler_stop('Pressure_Gauge_Allreduce')
      if (ierr /= MPI_SUCCESS) call fatal_error('numerics', 'MPI failure reducing zero-mean scalar sum')

      if (allocated(flow%gather_counts)) then
         global_count = sum(flow%gather_counts)
      else
         global_count = flow%nlocal
      end if

      if (global_count <= 0) return
      mean_value = global_sum / real(global_count, rk)

      do c = flow%first_cell, flow%last_cell
         scalar(c) = scalar(c) - mean_value
      end do
   end subroutine subtract_owned_mean_scalar

   !> Iteratively solves the Pressure Poisson system using PCG.
   subroutine solve_pressure_correction(mesh, flow, bc, params, transport, rhs, phi, stats)
      type(mesh_t), intent(in) :: mesh
      type(flow_mpi_t), intent(inout) :: flow
      type(bc_set_t), intent(in) :: bc
      type(case_params_t), intent(in) :: params
      type(transport_properties_t), intent(in) :: transport
      real(rk), intent(in) :: rhs(:)
      real(rk), intent(inout) :: phi(:)
      type(solver_stats_t), intent(inout) :: stats

      real(rk) :: rr, rr_new, rr_initial
      real(rk) :: rz, rz_new
      real(rk) :: pap
      real(rk) :: alpha, beta
      integer :: iter
      integer :: c
      logical :: use_zero_mean_gauge
      real(rk) :: abs_residual, rel_residual
      logical :: do_periodic_gauge

      call ensure_projection_workspace(mesh)
      call ensure_pressure_operator_cache(mesh, flow, bc)
      if (params%enable_variable_density) then
         call update_variable_pressure_operator_cache(mesh, flow, bc, transport)
      end if
      use_zero_mean_gauge = (.not. pressure_cache%has_dirichlet_pressure) .and. params%enable_variable_density

      associate( &
         r => projection_work%r, &
         z => projection_work%z, &
         pvec => projection_work%pvec, &
         ap => projection_work%ap, &
         local_ap => projection_work%local_ap)

         r = zero
         z = zero
         pvec = zero
         local_ap = zero

         ! The projection driver resets phi to zero before every pressure solve.
         ! Therefore A*phi is exactly zero for the initial PCG residual.
         local_ap = zero

         do c = flow%first_cell, flow%last_cell
            r(c) = rhs(c)
         end do

         if (.not. pressure_cache%has_dirichlet_pressure) then
            if (params%enable_variable_density) then
               call profiler_start('Pressure_Nullspace_Project')
               call subtract_owned_mean_scalar(flow, r)
               call subtract_owned_mean_scalar(flow, phi)
               call profiler_stop('Pressure_Nullspace_Project')
             else
               r(1) = zero
               phi(1) = zero
             end if
         end if

         call profiler_start('Pressure_Preconditioner')
         do c = flow%first_cell, flow%last_cell
            if (c == 1 .and. .not. pressure_cache%has_dirichlet_pressure .and. &
                .not. params%enable_variable_density) then
               z(c) = zero
            else
               z(c) = r(c) / max(pressure_preconditioner_diag(params, c), tiny_safe)
            end if
         end do
         call profiler_stop('Pressure_Preconditioner')

         if (use_zero_mean_gauge) then
            call profiler_start('Pressure_Nullspace_Project')
            call subtract_owned_mean_scalar(flow, z)
            call profiler_stop('Pressure_Nullspace_Project')
         end if
         block
            real(rk) :: res(2)
            call flow_global_two_dots_owned(flow, r, r, r, z, res)
            rr = res(1); rz = res(2)
         end block

         pvec = z
         if (.not. pressure_cache%has_dirichlet_pressure) then
            if (.not. params%enable_variable_density) then
               pvec(1) = zero
            end if
         end if

         stats%pressure_iterations = 0
         rr_initial = max(rr, tiny_safe)
         abs_residual = sqrt(max(rr, zero) / max(real(mesh%ncells, rk), one))
         rel_residual = sqrt(max(rr, zero) / rr_initial)
         stats%pressure_residual_abs = abs_residual
         stats%pressure_residual = rel_residual

         do iter = 1, params%pressure_max_iter
            if (stats%pressure_residual <= params%pressure_tol .or. &
                stats%pressure_residual_abs <= params%pressure_abs_tol) exit

            call profiler_start('Pressure_Exchange_Pvec')
            call flow_exchange_cell_scalar(flow, pvec)
            call profiler_stop('Pressure_Exchange_Pvec')
            call pressure_matvec(mesh, flow, bc, params, pvec, local_ap)

            call profiler_start('Pressure_Reduce_PAp')
            pap = flow_global_dot_owned(flow, pvec, local_ap)
            call profiler_stop('Pressure_Reduce_PAp')

            if (abs(pap) <= tiny_safe) exit

            alpha = rz / pap

            call profiler_start('Pressure_Vector_Update')
            do c = flow%first_cell, flow%last_cell
               phi(c) = phi(c) + alpha * pvec(c)
               r(c) = r(c) - alpha * local_ap(c)
            end do
            call profiler_stop('Pressure_Vector_Update')

            if (.not. pressure_cache%has_dirichlet_pressure) then
               if (params%enable_variable_density) then
                  do_periodic_gauge = (mod(iter, pcg_nullspace_projection_interval) == 0)
                  if (do_periodic_gauge) then
                     call profiler_start('Pressure_Nullspace_Project')
                     call subtract_owned_mean_scalar(flow, phi)
                     call subtract_owned_mean_scalar(flow, r)
                     call profiler_stop('Pressure_Nullspace_Project')
                  end if
               else
                  phi(1) = zero
                  r(1) = zero
               end if
            end if

            call profiler_start('Pressure_Preconditioner')
            do c = flow%first_cell, flow%last_cell
               if (c == 1 .and. .not. pressure_cache%has_dirichlet_pressure .and. &
                .not. params%enable_variable_density) then
                  z(c) = zero
               else
                  z(c) = r(c) / max(pressure_preconditioner_diag(params, c), tiny_safe)
               end if
            end do
            call profiler_stop('Pressure_Preconditioner')

            if (use_zero_mean_gauge) then
               call profiler_start('Pressure_Nullspace_Project')
               call subtract_owned_mean_scalar(flow, z)
               call profiler_stop('Pressure_Nullspace_Project')
            end if
            block
               real(rk) :: res(2)
               call profiler_start('Pressure_Reduce_RZ')
               call flow_global_two_dots_owned(flow, r, r, r, z, res)
               call profiler_stop('Pressure_Reduce_RZ')
               rr_new = res(1); rz_new = res(2)
            end block

            abs_residual = sqrt(max(rr_new, zero) / max(real(mesh%ncells, rk), one))
            rel_residual = sqrt(max(rr_new, zero) / rr_initial)
            stats%pressure_residual_abs = abs_residual
            stats%pressure_residual = rel_residual
            stats%pressure_iterations = iter

            if (stats%pressure_residual <= params%pressure_tol .or. &
                stats%pressure_residual_abs <= params%pressure_abs_tol) exit

            beta = rz_new / max(rz, tiny_safe)

            call profiler_start('Pressure_Vector_Update')
            do c = flow%first_cell, flow%last_cell
               pvec(c) = z(c) + beta * pvec(c)
            end do
            call profiler_stop('Pressure_Vector_Update')
            if (.not. pressure_cache%has_dirichlet_pressure) then
               if (.not. params%enable_variable_density) then
                  pvec(1) = zero
               end if
            end if

            rr = rr_new
            rz = rz_new
         end do

         if (.not. pressure_cache%has_dirichlet_pressure) then
            if (params%enable_variable_density) then
               call profiler_start('Pressure_Nullspace_Project')
               call subtract_owned_mean_scalar(flow, phi)
               call profiler_stop('Pressure_Nullspace_Project')
            else
               phi(1) = zero
            end if
         end if

         stats%pressure_solve_count = stats%pressure_solve_count + 1
         stats%pressure_iterations_total = stats%pressure_iterations_total + stats%pressure_iterations
         stats%pressure_iterations_max = max(stats%pressure_iterations_max, stats%pressure_iterations)
         stats%pressure_iterations_avg = real(stats%pressure_iterations_total, rk) / &
                                         max(real(stats%pressure_solve_count, rk), one)

         call profiler_start('Pressure_Exchange_FinalPhi')
         call flow_exchange_cell_scalar(flow, phi)
         call profiler_stop('Pressure_Exchange_FinalPhi')

      end associate
   end subroutine solve_pressure_correction

end module mod_linear_solver
