//! SIMPLE (Semi-Implicit Method for Pressure Linked Equations) algorithm
//!
//! The SIMPLE algorithm is a widely-used iterative solution method for
//! the incompressible Navier-Stokes equations. It uses a pressure-velocity
//! coupling approach to handle the incompressibility constraint.
//!
//! Algorithm steps:
//! 1. Solve momentum equations with guessed pressure field → u*, v*
//! 2. Solve pressure correction equation → p'
//! 3. Correct velocities and pressure
//! 4. Check convergence and iterate

use super::{BoundaryConditions, FlowField, IncompressibleSolver, SolverResult};
use crate::turbulence::{KEpsilonModel, KEpsilonVariant, TurbulenceModel, TurbulenceState};
use crate::{CfdConfig, CfdError, CfdResult};
use async_trait::async_trait;
use nalgebra::{DMatrix, DVector, Vector3};
use std::time::Instant;

/// Discretisation of the convective term in the momentum equations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConvectionScheme {
    /// First-order upwind. Unconditionally bounded, but carries a numerical
    /// viscosity of about `|u| dx / 2`, which caps the observed order of the
    /// whole discretisation at 1 whenever convection matters.
    Upwind,
    /// Deferred-correction TVD with the van Albada limiter
    /// `psi(r) = (r^2 + r) / (r^2 + 1)` (0 for `r <= 0`).
    ///
    /// The upwind operator stays implicit, so `a_p = sum(a_nb)` and diagonal
    /// dominance survive unconditionally; the limited high-order-minus-upwind
    /// flux difference is added explicitly to the source, evaluated at the
    /// current iterate. At a converged state the two agree, so the fixed
    /// point is the TVD discretisation — relaxation changes the path, never
    /// the answer. Faces whose far-upwind node lies outside the domain fall
    /// back to pure upwind, the standard TVD boundary treatment.
    TvdVanAlbada,
    /// Deferred-correction TVD with the van Leer limiter
    /// `psi(r) = (r + |r|) / (1 + |r|)`. Same construction as
    /// [`ConvectionScheme::TvdVanAlbada`].
    TvdVanLeer,
}

impl ConvectionScheme {
    /// Flux limiter `psi(r)`. Zero recovers pure upwind, one recovers central
    /// differencing; both TVD limiters satisfy `psi(1) = 1`, which is what
    /// makes them second order in smooth regions.
    fn limiter(self, r: f64) -> f64 {
        match self {
            Self::Upwind => 0.0,
            Self::TvdVanAlbada => {
                if r > 0.0 {
                    (r * r + r) / (r * r + 1.0)
                } else {
                    0.0
                }
            }
            Self::TvdVanLeer => (r + r.abs()) / (1.0 + r.abs()),
        }
    }

    /// The limited correction `u_face_HO - u_face_upwind` for one face, given
    /// the far-upwind, upwind and downwind values along the flow direction.
    /// `None` for the far-upwind value means it lies outside the domain, and
    /// the face falls back to pure upwind.
    fn face_correction(self, far_upwind: Option<f64>, upwind: f64, downwind: f64) -> f64 {
        let Some(far) = far_upwind else {
            return 0.0;
        };
        let denominator = downwind - upwind;
        if denominator.abs() < 1e-300 {
            return 0.0;
        }
        let r = (upwind - far) / denominator;
        0.5 * self.limiter(r) * denominator
    }
}

/// Parameters for SIMPLE algorithm
#[derive(Debug, Clone)]
pub struct SimpleParameters {
    /// Under-relaxation factor for pressure (typically 0.2-0.8)
    pub pressure_relaxation: f64,
    /// Under-relaxation factor for velocity (typically 0.5-0.8)
    pub velocity_relaxation: f64,
    /// Convection discretisation. Defaults to first-order upwind.
    pub convection_scheme: ConvectionScheme,
    /// Maximum number of iterations
    pub max_iterations: usize,
    /// Convergence tolerance for residuals
    pub tolerance: f64,
    /// Time step for transient problems
    pub time_step: f64,
    /// Maximum Courant number for stability
    pub max_courant: f64,
    /// Enable turbulence modeling
    pub use_turbulence: bool,
    /// Drop the transient term and solve for the steady state directly.
    ///
    /// Standard SIMPLE is a steady-state algorithm: it has no pseudo-time
    /// term, and stability comes from under-relaxation folded implicitly into
    /// the momentum coefficients. Keeping a false-transient term instead makes
    /// the converged answer depend on `time_step`, which a steady state cannot
    /// legitimately do. Set false only for genuinely transient problems, where
    /// `time_step` is a physical time step rather than a relaxation knob.
    pub steady: bool,
}

impl SimpleParameters {
    /// Create new SIMPLE parameters with default values
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Set pressure under-relaxation factor
    #[must_use]
    pub fn with_pressure_relaxation(mut self, factor: f64) -> Self {
        self.pressure_relaxation = factor.max(0.0); // Ensure non-negative
        self
    }

    /// Set velocity under-relaxation factor
    #[must_use]
    pub fn with_velocity_relaxation(mut self, factor: f64) -> Self {
        self.velocity_relaxation = factor.max(0.0);
        self
    }

    /// Set the convection scheme
    #[must_use]
    pub fn with_convection_scheme(mut self, scheme: ConvectionScheme) -> Self {
        self.convection_scheme = scheme;
        self
    }

    /// Set maximum iterations
    #[must_use]
    pub fn with_max_iterations(mut self, max_iter: usize) -> Self {
        self.max_iterations = max_iter;
        self
    }

    /// Set convergence tolerance
    #[must_use]
    pub fn with_tolerance(mut self, tol: f64) -> Self {
        self.tolerance = tol.abs();
        self
    }

    /// Set time step
    #[must_use]
    pub fn with_time_step(mut self, dt: f64) -> Self {
        self.time_step = dt.abs();
        self
    }

    /// Validate parameters
    pub fn validate(&self) -> CfdResult<()> {
        if self.pressure_relaxation <= 0.0 || self.pressure_relaxation > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Pressure relaxation factor must be in (0, 1]",
            ));
        }

        if self.velocity_relaxation <= 0.0 || self.velocity_relaxation > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Velocity relaxation factor must be in (0, 1]",
            ));
        }

        if self.tolerance <= 0.0 {
            return Err(CfdError::invalid_parameter("Tolerance must be positive"));
        }

        if self.time_step <= 0.0 {
            return Err(CfdError::invalid_parameter("Time step must be positive"));
        }

        Ok(())
    }
}

impl Default for SimpleParameters {
    fn default() -> Self {
        Self {
            pressure_relaxation: 0.3,
            velocity_relaxation: 0.7,
            convection_scheme: ConvectionScheme::Upwind,
            max_iterations: 1000,
            tolerance: 1e-6,
            time_step: 0.001,
            max_courant: 1.0,
            use_turbulence: false,
            steady: true,
        }
    }
}

/// Result of SIMPLE algorithm execution
#[derive(Debug, Clone)]
pub struct SimpleResult {
    /// Base solver result information
    pub solver_result: SolverResult,
    /// Pressure correction iterations per SIMPLE iteration
    pub pressure_iterations: Vec<usize>,
    /// Final mass residual
    pub mass_residual: f64,
    /// Final momentum residual
    pub momentum_residual: f64,
}

/// SIMPLE algorithm implementation
pub struct SimpleSolver {
    /// CFD configuration
    config: CfdConfig,
    /// SIMPLE parameters
    parameters: SimpleParameters,
    /// Linear algebra workspace
    workspace: LinearAlgebraWorkspace,
    /// Turbulence model (optional)
    turbulence_model: Option<KEpsilonModel>,
    /// Optional volumetric momentum source `f(x, y) -> (f_x, f_y)`, per unit
    /// volume.
    ///
    /// Exists so a manufactured solution can be imposed: given any velocity
    /// and pressure field, the residual of the momentum equations *is* the
    /// body force that makes that field exact, and applying it turns the
    /// solver into something whose exact answer is known in closed form.
    #[allow(clippy::type_complexity)]
    momentum_source: Option<Box<dyn Fn(f64, f64) -> (f64, f64) + Send + Sync>>,
    /// Optional wall velocity `f(x, y) -> (u_wall, v_wall)`, sampled at the
    /// wall face position.
    ///
    /// The near-wall momentum control volumes need the *tangential* velocity of
    /// the wall that bounds them, and on this staggered layout there is nowhere
    /// to store it: no u node lies on the bottom or top wall, and no v node
    /// lies on the left or right wall. Every u row is at `y = (j + 0.5) dy`,
    /// strictly interior. Supplying it as a function of position is the only
    /// way a spatially varying wall — a manufactured solution's, for instance —
    /// can reach the discretisation at all.
    ///
    /// When unset the solver falls back to the value stored on the near-wall
    /// line itself, which is exactly what `BoundaryConditions` writes there, so
    /// an existing configuration keeps the wall it always had.
    #[allow(clippy::type_complexity)]
    wall_velocity: Option<Box<dyn Fn(f64, f64) -> (f64, f64) + Send + Sync>>,
}

/// Workspace for linear algebra operations
struct LinearAlgebraWorkspace {
    /// Matrix for pressure correction equation
    pressure_matrix: Option<DMatrix<f64>>,
    /// RHS vector for pressure correction
    pressure_rhs: Option<DVector<f64>>,
    /// Solution vector for pressure correction
    pressure_solution: Option<DVector<f64>>,
    /// Momentum equation coefficients
    momentum_coefficients: Option<MomentumCoefficients>,
}

/// Coefficients for momentum equations discretization
#[derive(Debug, Clone)]
struct MomentumCoefficients {
    /// Central coefficient (diagonal)
    pub ap: DMatrix<f64>,
    /// East neighbor coefficient
    pub ae: DMatrix<f64>,
    /// West neighbor coefficient
    pub aw: DMatrix<f64>,
    /// North neighbor coefficient
    pub an: DMatrix<f64>,
    /// South neighbor coefficient
    pub as_: DMatrix<f64>,
    /// Source term
    pub su: DMatrix<f64>,
}

impl SimpleSolver {
    /// How far the residual may rise above its best value before the solve is
    /// declared divergent.
    const DIVERGENCE_GROWTH: f64 = 1e6;

    /// Create new SIMPLE solver
    pub fn new(config: CfdConfig, parameters: SimpleParameters) -> CfdResult<Self> {
        config.validate()?;
        parameters.validate()?;

        // Initialize turbulence model if enabled (will be properly sized when flow field is available)
        let turbulence_model = if parameters.use_turbulence {
            None // Will be initialized when flow field dimensions are known
        } else {
            None
        };

        Ok(Self {
            config,
            parameters,
            workspace: LinearAlgebraWorkspace {
                pressure_matrix: None,
                pressure_rhs: None,
                pressure_solution: None,
                momentum_coefficients: None,
            },
            turbulence_model,
            momentum_source: None,
            wall_velocity: None,
        })
    }

    /// Set a volumetric momentum source. See [`Self::momentum_source`].
    pub fn set_momentum_source<F>(&mut self, source: F)
    where
        F: Fn(f64, f64) -> (f64, f64) + Send + Sync + 'static,
    {
        self.momentum_source = Some(Box::new(source));
    }

    /// Set the wall velocity as a function of position. See
    /// [`Self::wall_velocity`].
    pub fn set_wall_velocity<F>(&mut self, f: F)
    where
        F: Fn(f64, f64) -> (f64, f64) + Send + Sync + 'static,
    {
        self.wall_velocity = Some(Box::new(f));
    }

    /// Tangential `u` of the horizontal wall bounding the near-wall u control
    /// volume at face `i`, row `j`, with the wall itself at `y_wall`.
    ///
    /// Falls back to the value stored on the near-wall line — see
    /// [`Self::wall_velocity`].
    fn u_wall(&self, flow_field: &FlowField, i: usize, j: usize, y_wall: f64, dx: f64) -> f64 {
        self.wall_velocity
            .as_ref()
            .map_or(flow_field.u[(j, i)], |f| f(i as f64 * dx, y_wall).0)
    }

    /// Tangential `v` of the vertical wall bounding the near-wall v control
    /// volume at column `i`, face `j`, with the wall itself at `x_wall`.
    fn v_wall(&self, flow_field: &FlowField, i: usize, j: usize, x_wall: f64, dy: f64) -> f64 {
        self.wall_velocity
            .as_ref()
            .map_or(flow_field.v[(j, i)], |f| f(x_wall, j as f64 * dy).1)
    }

    /// Momentum source contribution for a u-face, already multiplied by the
    /// control volume so it is a force, matching the pressure-gradient term.
    ///
    /// On this staggered layout u-face `i` sits at `x = i dx`, mid-height of
    /// row `j`, i.e. `y = (j + 0.5) dy`.
    fn u_source_term(&self, i: usize, j: usize, dx: f64, dy: f64) -> f64 {
        self.momentum_source
            .as_ref()
            .map_or(0.0, |f| f(i as f64 * dx, (j as f64 + 0.5) * dy).0 * dx * dy)
    }

    /// Momentum source contribution for a v-face, at `x = (i + 0.5) dx`,
    /// `y = j dy`.
    fn v_source_term(&self, i: usize, j: usize, dx: f64, dy: f64) -> f64 {
        self.momentum_source
            .as_ref()
            .map_or(0.0, |f| f((i as f64 + 0.5) * dx, j as f64 * dy).1 * dx * dy)
    }

    /// Solve one SIMPLE iteration
    pub async fn solve_simple_iteration(
        &mut self,
        flow_field: &mut FlowField,
        boundary_conditions: &BoundaryConditions,
        dt: f64,
    ) -> CfdResult<(f64, f64)> {
        // Step 1: Solve momentum equations with current pressure field.
        // This begins by storing the current iterate in `u_old`, which the
        // transient term, the under-relaxation and the residual all read.
        self.momentum_prediction_step(flow_field, dt).await?;

        // The predicted field must satisfy the velocity boundary conditions
        // before its divergence is used as the pressure source.
        //
        // Boundary conditions were previously applied only at the end of the
        // iteration, so `u*` carried whatever the momentum sweep happened to
        // write on the boundary faces — values the wall then overwrote with
        // zero anyway. Their divergence entered the pressure equation as a
        // spurious mass source concentrated at the two lid corners, where the
        // moving lid meets a stationary wall. Because the swept value scales
        // with the relaxation factor, so did the spurious source, and so did
        // the converged solution: the interior momentum equations were
        // satisfied to machine precision at every relaxation factor, but each
        // one satisfied them around a different corner condition.
        flow_field.apply_boundary_conditions(boundary_conditions)?;
        flow_field.copy_to_starred();

        // Step 2: Solve pressure correction equation
        let mass_residual = self.pressure_correction_step(flow_field).await?;

        // Step 3: Correct velocities
        self.velocity_correction_step(flow_field).await?;

        // Step 4: Update pressure field
        self.pressure_update_step(flow_field).await?;

        // Step 5: Solve turbulence equations if enabled
        if self.parameters.use_turbulence {
            self.solve_turbulence(flow_field, dt).await?;
        }

        // Step 6: Apply boundary conditions
        flow_field.apply_boundary_conditions(boundary_conditions)?;

        // Step 7: Apply pressure under-relaxation.
        //
        // Velocity relaxation is *not* applied here: it is folded into the
        // momentum coefficients (see `compute_u_momentum_coefficients`).
        // Doing both would relax twice, and the explicit blend would also
        // undo part of the continuity the pressure correction just enforced,
        // since the blended field is not the divergence-free one.
        flow_field.apply_pressure_relaxation(self.parameters.pressure_relaxation)?;

        // Compute momentum residual
        let momentum_residual = self.compute_momentum_residual(flow_field, dt)?;

        Ok((mass_residual, momentum_residual))
    }

    /// Imbalance of the discretised momentum equations, normalised.
    ///
    /// For each interior velocity point this is
    /// `|a_p u_P - Σ a_nb u_nb - b|`, summed and divided by a reference
    /// momentum flux `ρ U² L`. It measures how far the current field is from
    /// satisfying the equations being solved.
    ///
    /// The previous measure was `|u - u_old|` — the change between successive
    /// iterates. That is not a residual: it reports how far the iteration
    /// *moved*, which depends on how heavily the iteration is damped, and the
    /// damping here is set by the pseudo-time step. A field far from the
    /// solution but advancing slowly registers as converged, and it does so at
    /// a different distance for every `dt`. That is why the converged answer
    /// appeared to depend on the time step.
    ///
    /// Normalising matters as much as the measure. The imbalance is divided by
    /// `Σ|a_p u_P|`, the scale of the equation's own diagonal terms, which is
    /// the convention CFD solvers report. An unnormalised sum grows with the
    /// cell count and with the coefficient magnitudes — which themselves
    /// depend on `dt` through `a_p0` — so the same numeric tolerance would
    /// mean a different thing on every grid and at every time step.
    fn compute_momentum_residual(&self, flow_field: &FlowField, dt: f64) -> CfdResult<f64> {
        let (nx, ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;
        let mu = self.config.viscosity;

        let mut residual = 0.0;
        let mut scale = 0.0;

        // Measure the *unrelaxed* momentum equation — the one actually being
        // solved for. Under-relaxation inflates the diagonal by `1/alpha` and
        // adds a matching source term; reporting the residual of that relaxed
        // system makes the same numeric tolerance correspond to a different
        // true error for every relaxation factor, so converged solutions would
        // still appear to depend on alpha. Undo both to recover the steady
        // equation before measuring it.
        let alpha = self.parameters.velocity_relaxation;

        // Over exactly the unknowns the sweeps solve for. Measuring a smaller
        // set would let the near-wall lines converge to anything at all without
        // the reported residual noticing.
        for j in 0..ny {
            for i in 1..nx {
                let cu =
                    self.compute_u_momentum_coefficients(flow_field, i, j, dt, rho, mu, dx, dy)?;
                let ap = cu.center * alpha;
                let source = cu.source - (1.0 - alpha) * cu.center * flow_field.u_old[(j, i)];
                let diagonal_u = ap * flow_field.u[(j, i)];
                let north = if j + 1 < ny {
                    cu.north * flow_field.u[(j + 1, i)]
                } else {
                    0.0
                };
                let south = if j > 0 {
                    cu.south * flow_field.u[(j - 1, i)]
                } else {
                    0.0
                };
                let imbalance_u = diagonal_u
                    - (source
                        + cu.east * flow_field.u[(j, i + 1)]
                        + cu.west * flow_field.u[(j, i - 1)]
                        + north
                        + south);
                residual += imbalance_u.abs();
                scale += diagonal_u.abs();
            }
        }

        for j in 1..ny {
            for i in 0..nx {
                let cv =
                    self.compute_v_momentum_coefficients(flow_field, i, j, dt, rho, mu, dx, dy)?;
                let ap = cv.center * alpha;
                let source = cv.source - (1.0 - alpha) * cv.center * flow_field.v_old[(j, i)];
                let diagonal_v = ap * flow_field.v[(j, i)];
                let east = if i + 1 < nx {
                    cv.east * flow_field.v[(j, i + 1)]
                } else {
                    0.0
                };
                let west = if i > 0 {
                    cv.west * flow_field.v[(j, i - 1)]
                } else {
                    0.0
                };
                let imbalance_v = diagonal_v
                    - (source
                        + east
                        + west
                        + cv.north * flow_field.v[(j + 1, i)]
                        + cv.south * flow_field.v[(j - 1, i)]);
                residual += imbalance_v.abs();
                scale += diagonal_v.abs();
            }
        }

        Ok(if scale > 1e-30 {
            residual / scale
        } else {
            residual
        })
    }

    /// Momentum prediction step: solve momentum equations with current pressure
    pub async fn momentum_prediction_step(
        &self,
        flow_field: &mut FlowField,
        dt: f64,
    ) -> CfdResult<()> {
        let (_nx, _ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;
        let mu = self.config.viscosity;

        // Copy current velocities to old values for time derivatives
        flow_field.update_old_values();

        // One Gauss-Seidel sweep of each momentum equation.
        //
        // Deliberately not more. SIMPLE lags the pressure, so driving the
        // momentum equations hard against a pressure field that is still wrong
        // converges them to the wrong intermediate state. Measured on the
        // Re=100 cavity, twenty sweeps per outer iteration left a momentum
        // residual two to three orders of magnitude *worse* than one sweep,
        // and moved the vortex further from the reference solution.
        self.solve_u_momentum(flow_field, dt, rho, mu, dx, dy)
            .await?;
        self.solve_v_momentum(flow_field, dt, rho, mu, dx, dy)
            .await?;

        // Store predicted velocities
        flow_field.copy_to_starred();

        Ok(())
    }

    /// Solve u-momentum equation
    async fn solve_u_momentum(
        &self,
        flow_field: &mut FlowField,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<()> {
        let (nx, ny, _, _) = flow_field.grid_info();

        // Every u row is an unknown.
        //
        // `u[(j, i)]` sits at `y = (j + 0.5) dy`, which is strictly interior
        // for every `j`, so there is no u row on a horizontal wall to hold a
        // boundary value. Only the faces `i = 0` and `i = nx` lie on a domain
        // boundary, which is why the `i` range stops short of them and the `j`
        // range does not stop at all.
        //
        // This previously swept `1..ny - 1`, freezing the two near-wall rows
        // and treating whatever was stored there as a boundary condition —
        // imposing the wall half a cell inside the domain.
        for j in 0..ny {
            for i in 1..nx {
                // Discretize u-momentum equation at (i, j)
                let coeffs =
                    self.compute_u_momentum_coefficients(flow_field, i, j, dt, rho, mu, dx, dy)?;

                // The north and south coefficients are zero on a near-wall row,
                // where the wall's contribution is already in `source`; the
                // guards keep the index off the end of the array.
                let north = if j + 1 < ny {
                    coeffs.north * flow_field.u[(j + 1, i)]
                } else {
                    0.0
                };
                let south = if j > 0 {
                    coeffs.south * flow_field.u[(j - 1, i)]
                } else {
                    0.0
                };

                // Solve for new u velocity using Gauss-Seidel
                let u_new = (coeffs.source
                    + coeffs.east * flow_field.u[(j, i + 1)]
                    + coeffs.west * flow_field.u[(j, i - 1)]
                    + north
                    + south)
                    / coeffs.center;

                flow_field.u[(j, i)] = u_new;
            }
        }

        Ok(())
    }

    /// Solve v-momentum equation
    async fn solve_v_momentum(
        &self,
        flow_field: &mut FlowField,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<()> {
        let (nx, ny, _, _) = flow_field.grid_info();

        // Every v column is an unknown, mirroring the u sweep: `v[(j, i)]` sits
        // at `x = (i + 0.5) dx`, strictly interior for every `i`, and only the
        // faces `j = 0` and `j = ny` lie on a domain boundary.
        for j in 1..ny {
            for i in 0..nx {
                // Discretize v-momentum equation at (i, j)
                let coeffs =
                    self.compute_v_momentum_coefficients(flow_field, i, j, dt, rho, mu, dx, dy)?;

                let east = if i + 1 < nx {
                    coeffs.east * flow_field.v[(j, i + 1)]
                } else {
                    0.0
                };
                let west = if i > 0 {
                    coeffs.west * flow_field.v[(j, i - 1)]
                } else {
                    0.0
                };

                // Solve for new v velocity using Gauss-Seidel
                let v_new = (coeffs.source
                    + east
                    + west
                    + coeffs.north * flow_field.v[(j + 1, i)]
                    + coeffs.south * flow_field.v[(j - 1, i)])
                    / coeffs.center;

                flow_field.v[(j, i)] = v_new;
            }
        }

        Ok(())
    }

    /// Pressure correction step: solve pressure Poisson equation
    pub async fn pressure_correction_step(&self, flow_field: &mut FlowField) -> CfdResult<f64> {
        let (nx, ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;

        // The pressure correction starts from zero every outer iteration.
        //
        // `p'` is a correction to the current pressure field, not a field in
        // its own right: `pressure_update_step` folds it into `p` at the end
        // of the iteration, so carrying it into the next one applies the same
        // correction twice.
        flow_field.p_prime.fill(0.0);

        // Continuity is enforced on EVERY cell.
        //
        // The domain is tiled by cells; there is no such thing as a cell that
        // does not have to conserve mass. Restricting this to `1..nx - 1` left
        // the outer ring of cells with no continuity equation at all, so
        // nothing ever removed their divergence — 3.1e-2 on the ring against
        // 2.7e-3 in the interior on a converged 32x32 manufactured solve, with
        // the resulting pressure error *growing* under refinement.
        //
        // This only became possible once the momentum sweeps stopped freezing
        // the near-wall lines. Extending continuity first leaves ring cells
        // whose faces are all prescribed — no correctable face, no solution —
        // and it duly broke convergence when tried in that order.
        let mut mass_imbalance: f64 = 0.0;
        for j in 0..ny {
            for i in 0..nx {
                // Compute mass source (continuity equation residual)
                let mass_source = self.compute_mass_source(flow_field, i, j, dx, dy, rho)?;
                flow_field.sp[(j, i)] = mass_source;
                mass_imbalance += mass_source.abs();
            }
        }

        // Neighbour coefficients, evaluated per face from the momentum
        // equation's own diagonal at that face — the same `a_p` the velocity
        // correction divides by, so the two remain each other's inverse.
        // Cached because `a_p` depends on the velocity field, which does not
        // change during the inner sweeps.
        //
        // A coefficient is zero exactly when the face it crosses is a genuine
        // domain boundary — where the velocity is prescribed and therefore not
        // correctable. Every cell keeps at least two correctable faces, so
        // every cell can be made divergence-free.
        let mut coefficients = Vec::with_capacity(nx * ny);
        for j in 0..ny {
            for i in 0..nx {
                let volume = dx * dy;
                let ae = if i + 1 == nx {
                    0.0
                } else {
                    let d = volume
                        / self.compute_u_momentum_center_coefficient(
                            flow_field,
                            i + 1,
                            j,
                            dx,
                            dy,
                            rho,
                        )?;
                    rho * d * dy / dx
                };
                let aw = if i == 0 {
                    0.0
                } else {
                    let d = volume
                        / self
                            .compute_u_momentum_center_coefficient(flow_field, i, j, dx, dy, rho)?;
                    rho * d * dy / dx
                };
                let an = if j + 1 == ny {
                    0.0
                } else {
                    let d = volume
                        / self.compute_v_momentum_center_coefficient(
                            flow_field,
                            i,
                            j + 1,
                            dx,
                            dy,
                            rho,
                        )?;
                    rho * d * dx / dy
                };
                let as_ = if j == 0 {
                    0.0
                } else {
                    let d = volume
                        / self
                            .compute_v_momentum_center_coefficient(flow_field, i, j, dx, dy, rho)?;
                    rho * d * dx / dy
                };

                coefficients.push(MomentumEquationCoeffs {
                    center: ae + aw + an + as_,
                    east: ae,
                    west: aw,
                    north: an,
                    south: as_,
                    source: 0.0,
                });
            }
        }

        // Solve pressure correction equation using Gauss-Seidel
        for _iteration in 0..200 {
            // Inner pressure correction iterations
            let mut residual = 0.0;

            for j in 0..ny {
                for i in 0..nx {
                    // Anchor one cell to fix the pressure level.
                    //
                    // With velocity prescribed on every boundary the pressure
                    // correction equation is pure Neumann and therefore
                    // singular: `p'` is determined only up to an additive
                    // constant, and Gauss-Seidel lets that constant drift.
                    // Anchoring a reference cell fixes the level without
                    // altering any pressure *gradient*, which is all the
                    // momentum equation uses.
                    //
                    // Enforcing the Neumann solvability condition instead — by
                    // subtracting the mean source — is the textbook remedy but
                    // is wrong here: this source is assembled from face fluxes
                    // that include the boundaries, so it is not required to sum
                    // to zero, and subtracting its mean injects a spurious
                    // source into every cell. Tried; it diverged.
                    //
                    // Dropping this one cell's continuity equation is legitimate
                    // now that the equation covers the whole domain: the sum of
                    // the sources over all cells telescopes to the net flux
                    // through the domain boundary, which is zero for a closed
                    // box, so the system has rank `n - 1` and exactly one
                    // equation is redundant.
                    if i == 1 && j == 1 {
                        flow_field.p_prime[(j, i)] = 0.0;
                        continue;
                    }

                    let coeffs = &coefficients[j * nx + i];

                    // A zero coefficient still guards its index: the neighbour
                    // it refers to is outside the domain.
                    let east = if i + 1 < nx {
                        coeffs.east * flow_field.p_prime[(j, i + 1)]
                    } else {
                        0.0
                    };
                    let west = if i > 0 {
                        coeffs.west * flow_field.p_prime[(j, i - 1)]
                    } else {
                        0.0
                    };
                    let north = if j + 1 < ny {
                        coeffs.north * flow_field.p_prime[(j + 1, i)]
                    } else {
                        0.0
                    };
                    let south = if j > 0 {
                        coeffs.south * flow_field.p_prime[(j - 1, i)]
                    } else {
                        0.0
                    };

                    let p_new =
                        (flow_field.sp[(j, i)] + east + west + north + south) / coeffs.center;

                    let correction = p_new - flow_field.p_prime[(j, i)];
                    residual += correction * correction;
                    flow_field.p_prime[(j, i)] = p_new;
                }
            }

            if residual.sqrt() < 1e-10 {
                break;
            }
        }

        // Report the mass imbalance, not the inner Gauss-Seidel residual.
        //
        // The outer loop treats this as its convergence measure, and the
        // inner residual only says how well the pressure-correction equation
        // was solved — it goes to zero whether or not the flow satisfies
        // continuity, so the solver could report convergence while the field
        // was still divergent.
        //
        // Normalised by a reference mass flux `ρ U L` so the same tolerance
        // means the same thing on every grid; an unnormalised sum grows with
        // the cell count.
        let reference = rho * self.config.reference_velocity * self.config.reference_length;
        Ok(if reference > 0.0 {
            mass_imbalance / reference
        } else {
            mass_imbalance
        })
    }

    /// Velocity correction step: correct velocities with pressure correction
    pub async fn velocity_correction_step(&self, flow_field: &mut FlowField) -> CfdResult<()> {
        let (nx, ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;

        // Correct every face the pressure equation treated as correctable —
        // which is every face that is not on a domain boundary. The ranges must
        // match `pressure_correction_step`'s coefficients exactly, or the
        // divergence the pressure correction was constructed to remove is not
        // the divergence that gets removed.
        for j in 0..ny {
            for i in 1..nx {
                let dp_dx = (flow_field.p_prime[(j, i)] - flow_field.p_prime[(j, i - 1)]) / dx;
                let ap_u =
                    self.compute_u_momentum_center_coefficient(flow_field, i, j, dx, dy, rho)?;
                flow_field.u[(j, i)] = flow_field.u_star[(j, i)] - (dx * dy / ap_u) * dp_dx;
            }
        }

        // Correct v-velocities
        for j in 1..ny {
            for i in 0..nx {
                let dp_dy = (flow_field.p_prime[(j, i)] - flow_field.p_prime[(j - 1, i)]) / dy;
                let ap_v =
                    self.compute_v_momentum_center_coefficient(flow_field, i, j, dx, dy, rho)?;
                flow_field.v[(j, i)] = flow_field.v_star[(j, i)] - (dx * dy / ap_v) * dp_dy;
            }
        }

        Ok(())
    }

    /// Pressure update step: add pressure correction to pressure
    pub async fn pressure_update_step(&self, flow_field: &mut FlowField) -> CfdResult<()> {
        let (nx, ny, _, _) = flow_field.grid_info();

        for j in 0..ny {
            for i in 0..nx {
                flow_field.p[(j, i)] +=
                    self.parameters.pressure_relaxation * flow_field.p_prime[(j, i)];
                flow_field.p_prime[(j, i)] = 0.0; // Reset pressure correction
            }
        }

        Ok(())
    }

    /// Solve turbulence model equations
    async fn solve_turbulence(&mut self, flow_field: &FlowField, dt: f64) -> CfdResult<()> {
        if self.parameters.use_turbulence {
            // Initialize turbulence model if not already done
            let (nx, ny, dx, dy) = flow_field.grid_info();
            let n_cells = nx * ny;

            if self.turbulence_model.is_none() {
                self.turbulence_model =
                    Some(KEpsilonModel::new(KEpsilonVariant::Standard, n_cells));
            }

            let turbulence_model = self.turbulence_model.as_mut().unwrap();

            // Convert velocity field to Vector3 format
            let mut velocity_vec = Vec::with_capacity(n_cells);
            for j in 0..ny {
                for i in 0..nx {
                    let u = if i < nx && j < ny {
                        flow_field.u[(j, i)]
                    } else {
                        0.0
                    };
                    let v = if i < nx && j < ny {
                        flow_field.v[(j, i)]
                    } else {
                        0.0
                    };
                    velocity_vec.push(Vector3::new(u, v, 0.0)); // 2D case, w=0
                }
            }

            // Simplified velocity gradients (zero for now)
            let velocity_gradients = vec![[[0.0; 3]; 3]; n_cells];

            // Create pressure vector
            let mut pressure = DVector::zeros(n_cells);
            for j in 0..ny {
                for i in 0..nx {
                    if i < nx && j < ny {
                        pressure[j * nx + i] = flow_field.p[(j, i)];
                    }
                }
            }

            let turbulence_state = TurbulenceState {
                velocity: velocity_vec,
                velocity_gradients,
                pressure,
                turbulent_ke: None, // Will be initialized by model
                epsilon: None,      // Will be initialized by model
                omega: None,
                wall_distance: DVector::from_element(n_cells, 1.0), // Simplified
                cell_volumes: DVector::from_element(n_cells, dx * dy), // 2D cell volume
                molecular_viscosity: self.config.viscosity / self.config.density, // kinematic viscosity
                density: self.config.density,
            };

            // Update turbulence model with new state
            turbulence_model.update(&turbulence_state, dt)?;
        }
        Ok(())
    }

    /// Compute effective viscosity (molecular + turbulent)
    fn compute_effective_viscosity(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        mu: f64,
    ) -> f64 {
        if let Some(ref turbulence_model) = self.turbulence_model {
            // Get turbulent viscosity from the model
            let (nx, _ny, _, _) = flow_field.grid_info();
            let cell_idx = j * nx + i;
            if cell_idx < turbulence_model.nu_t.len() {
                let nu_t = turbulence_model.nu_t[cell_idx]; // kinematic turbulent viscosity
                let mu_t = nu_t * self.config.density; // convert to dynamic viscosity
                mu + mu_t
            } else {
                mu
            }
        } else {
            mu
        }
    }

    /// Deferred-correction source for the u-momentum equation: the limited
    /// high-order convective fluxes minus their upwind counterparts, moved to
    /// the right-hand side with the sign that puts convection on the left.
    ///
    /// Face stencils run along the flow direction: for each face the upwind
    /// node `C`, downwind node `D` and far-upwind node `U` define
    /// `r = (C - U) / (D - C)`, and the correction is
    /// `psi(r) (D - C) / 2`. A face whose far-upwind node lies outside the
    /// domain falls back to pure upwind, and a wall face has zero mass flux,
    /// so its correction never enters.
    #[allow(clippy::too_many_arguments)]
    fn u_deferred_correction(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        nx: usize,
        ny: usize,
        fe: f64,
        fw: f64,
        fn_: f64,
        fs: f64,
    ) -> f64 {
        let scheme = self.parameters.convection_scheme;
        if scheme == ConvectionScheme::Upwind {
            return 0.0;
        }
        let u = &flow_field.u;

        // East face of the u control volume, between u faces `i` and `i + 1`.
        let delta_e = if fe >= 0.0 {
            // `i >= 1` for every unknown, so the far-upwind node exists.
            scheme.face_correction(Some(u[(j, i - 1)]), u[(j, i)], u[(j, i + 1)])
        } else {
            let far = (i + 2 <= nx).then(|| u[(j, i + 2)]);
            scheme.face_correction(far, u[(j, i + 1)], u[(j, i)])
        };

        // West face, between u faces `i - 1` and `i`.
        let delta_w = if fw >= 0.0 {
            let far = (i >= 2).then(|| u[(j, i - 2)]);
            scheme.face_correction(far, u[(j, i - 1)], u[(j, i)])
        } else {
            scheme.face_correction(Some(u[(j, i + 1)]), u[(j, i)], u[(j, i - 1)])
        };

        // North face, between rows `j` and `j + 1`; a wall face passes no mass.
        let delta_n = if j + 1 >= ny {
            0.0
        } else if fn_ >= 0.0 {
            let far = (j >= 1).then(|| u[(j - 1, i)]);
            scheme.face_correction(far, u[(j, i)], u[(j + 1, i)])
        } else {
            let far = (j + 2 < ny).then(|| u[(j + 2, i)]);
            scheme.face_correction(far, u[(j + 1, i)], u[(j, i)])
        };

        // South face, between rows `j - 1` and `j`.
        let delta_s = if j == 0 {
            0.0
        } else if fs >= 0.0 {
            let far = (j >= 2).then(|| u[(j - 2, i)]);
            scheme.face_correction(far, u[(j - 1, i)], u[(j, i)])
        } else {
            let far = (j + 1 < ny).then(|| u[(j + 1, i)]);
            scheme.face_correction(far, u[(j, i)], u[(j - 1, i)])
        };

        -(fe * delta_e - fw * delta_w + fn_ * delta_n - fs * delta_s)
    }

    /// Deferred-correction source for the v-momentum equation; mirrors
    /// [`Self::u_deferred_correction`] with the roles of the axes swapped.
    #[allow(clippy::too_many_arguments)]
    fn v_deferred_correction(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        nx: usize,
        ny: usize,
        fe: f64,
        fw: f64,
        fn_: f64,
        fs: f64,
    ) -> f64 {
        let scheme = self.parameters.convection_scheme;
        if scheme == ConvectionScheme::Upwind {
            return 0.0;
        }
        let v = &flow_field.v;

        // North face of the v control volume, between v faces `j` and `j + 1`.
        let delta_n = if fn_ >= 0.0 {
            scheme.face_correction(Some(v[(j - 1, i)]), v[(j, i)], v[(j + 1, i)])
        } else {
            let far = (j + 2 <= ny).then(|| v[(j + 2, i)]);
            scheme.face_correction(far, v[(j + 1, i)], v[(j, i)])
        };

        // South face, between v faces `j - 1` and `j`.
        let delta_s = if fs >= 0.0 {
            let far = (j >= 2).then(|| v[(j - 2, i)]);
            scheme.face_correction(far, v[(j - 1, i)], v[(j, i)])
        } else {
            scheme.face_correction(Some(v[(j + 1, i)]), v[(j, i)], v[(j - 1, i)])
        };

        // East face, between columns `i` and `i + 1`; a wall face passes no
        // mass.
        let delta_e = if i + 1 >= nx {
            0.0
        } else if fe >= 0.0 {
            let far = (i >= 1).then(|| v[(j, i - 1)]);
            scheme.face_correction(far, v[(j, i)], v[(j, i + 1)])
        } else {
            let far = (i + 2 < nx).then(|| v[(j, i + 2)]);
            scheme.face_correction(far, v[(j, i + 1)], v[(j, i)])
        };

        // West face, between columns `i - 1` and `i`.
        let delta_w = if i == 0 {
            0.0
        } else if fw >= 0.0 {
            let far = (i >= 2).then(|| v[(j, i - 2)]);
            scheme.face_correction(far, v[(j, i - 1)], v[(j, i)])
        } else {
            let far = (i + 1 < nx).then(|| v[(j, i + 1)]);
            scheme.face_correction(far, v[(j, i)], v[(j, i - 1)])
        };

        -(fe * delta_e - fw * delta_w + fn_ * delta_n - fs * delta_s)
    }

    /// Compute coefficients for u-momentum equation
    fn compute_u_momentum_coefficients(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<MomentumEquationCoeffs> {
        let (nx, ny, _, _) = flow_field.grid_info();

        // Compute effective viscosity (molecular + turbulent)
        let mu_eff = self.compute_effective_viscosity(flow_field, i, j, mu);

        // Diffusion coefficients using effective viscosity
        // Diffusion conductances: `Gamma * A / delta`, the face area over the
        // distance between the nodes it separates.
        //
        // These previously read `mu / dx` and `mu / dy`, omitting the face
        // area entirely. On a square grid that makes them a factor `1/h` too
        // large — 65 times too much diffusion on a 65x65 mesh — so the solver
        // ran at an effective Reynolds number far below the one requested.
        // Every other term is already a force: the pressure term is
        // `dp * dy`, the convective flux is `rho u dy`, so the mismatch was
        // confined to diffusion.
        let gamma_e = mu_eff * dy / dx;
        let gamma_w = mu_eff * dy / dx;
        let gamma_n = mu_eff * dx / dy;
        let gamma_s = mu_eff * dx / dy;

        // Convective mass fluxes through the four faces of the u control
        // volume, which on a staggered grid is centred on the u face `i` and
        // spans from cell centre `i-1` to cell centre `i`. Its east and west
        // faces therefore sit at those cell centres, where the velocity is the
        // average of the two neighbouring u values.
        //
        // All four fluxes were previously taken from a single cell-centred
        // velocity, so `fe` and `fw` were literally the same number, as were
        // `fn_` and `fs`. Upwinding then chose the same direction on opposite
        // faces of the volume, which cannot represent transport across it: the
        // scheme reduced to diffusion plus a spurious diagonal term.
        let fe = rho * 0.5 * (flow_field.u[(j, i)] + flow_field.u[(j, i + 1)]) * dy;
        let fw = rho * 0.5 * (flow_field.u[(j, i - 1)] + flow_field.u[(j, i)]) * dy;

        // Rows `j = 0` and `j = ny - 1` are *not* boundaries — every u row sits
        // at `y = (j + 0.5) dy`, strictly inside the domain. They are near-wall
        // interior lines whose control volume happens to have the wall for its
        // south (respectively north) face.
        //
        // Two things change at such a face and nothing else does. A solid wall
        // passes no mass, so the convective flux through it is zero whatever
        // the stored normal velocity happens to be. And the node on the far
        // side of the face is the wall itself, half a cell away rather than a
        // full cell, so the conductance is `mu A / (dy/2)` — twice the interior
        // value — and the value there is the wall's own tangential velocity,
        // which is data rather than an unknown. Data belongs in the source, so
        // the returned neighbour coefficient is zero while `a_p` still carries
        // the conductance.
        //
        // Freezing these rows instead, as this sweep previously did, imposes
        // the wall value half a cell inside the domain and leaves the outer
        // ring of cells with too few correctable faces.
        let south_is_wall = j == 0;
        let north_is_wall = j + 1 == ny;

        let fn_ = if north_is_wall {
            0.0
        } else {
            rho * 0.5 * (flow_field.v[(j + 1, i - 1)] + flow_field.v[(j + 1, i)]) * dx
        };
        let fs = if south_is_wall {
            0.0
        } else {
            rho * 0.5 * (flow_field.v[(j, i - 1)] + flow_field.v[(j, i)]) * dx
        };

        // Compute coefficients with upwind scheme
        let ae = gamma_e + f64::max(-fe, 0.0);
        let aw = gamma_w + f64::max(fw, 0.0);

        // `_p` enters the diagonal; `_nb` multiplies a stored neighbour and is
        // zero at a wall, where the contribution goes to `wall_source` instead.
        let gamma_wall = mu_eff * dx / (0.5 * dy);
        let mut wall_source = 0.0;
        let (an, an_nb) = if north_is_wall {
            wall_source += gamma_wall * self.u_wall(flow_field, i, j, ny as f64 * dy, dx);
            (gamma_wall, 0.0)
        } else {
            let a = gamma_n + f64::max(-fn_, 0.0);
            (a, a)
        };
        let (as_, as_nb) = if south_is_wall {
            wall_source += gamma_wall * self.u_wall(flow_field, i, j, 0.0, dx);
            (gamma_wall, 0.0)
        } else {
            let a = gamma_s + f64::max(fs, 0.0);
            (a, a)
        };

        // Transient term. Zero for a steady solve: standard SIMPLE has no
        // pseudo-time term, and keeping one makes the converged answer depend
        // on `time_step`.
        let ap0 = if self.parameters.steady {
            0.0
        } else {
            rho * dx * dy / dt
        };

        // Central coefficient.
        //
        // The net flux `(F_e - F_w) + (F_n - F_s)` is deliberately *not*
        // included. It vanishes identically once continuity holds, but during
        // the iteration it does not, and it can exceed the sum of the
        // neighbour coefficients — driving `a_p` through zero and the solve to
        // NaN. Omitting it is what guarantees `a_p = Σ a_nb (+ a_p0) > 0`, so
        // upwinding keeps the matrix diagonally dominant unconditionally.
        let ap_unrelaxed = ae + aw + an + as_ + ap0;

        // Source term (pressure gradient + old time step)
        let pressure_gradient = -(flow_field.p[(j, i)] - flow_field.p[(j, i - 1)]) * dy;
        let time_term = ap0 * flow_field.u_old[(j, i)];

        // Patankar's implicit under-relaxation: divide the diagonal by alpha
        // and add `(1-alpha)/alpha * a_p * u_prev` to the source.
        //
        // At a fixed point `u = u_prev` the two added terms cancel exactly, so
        // the converged solution is independent of alpha -- relaxation changes
        // the path, never the answer. Applying relaxation instead as a
        // post-hoc blend of the whole field, as this solver previously did,
        // has no such guarantee, and it also leaves the pressure equation
        // using an unrelaxed `a_p` while the velocities have been relaxed.
        let alpha = self.parameters.velocity_relaxation;
        let ap = ap_unrelaxed / alpha;
        // `u_source_term` was computed and then never added — the x-momentum
        // equation carried no body force at all, while the y-momentum equation
        // carried its own. Any manufactured solution was therefore imposed on
        // one component and not the other, which is why `u` came out markedly
        // further from exact than `v` on the same mesh.
        let source = pressure_gradient
            + time_term
            + wall_source
            + self.u_source_term(i, j, dx, dy)
            + self.u_deferred_correction(flow_field, i, j, nx, ny, fe, fw, fn_, fs)
            + (1.0 - alpha) / alpha * ap_unrelaxed * flow_field.u_old[(j, i)];

        Ok(MomentumEquationCoeffs {
            center: ap,
            east: ae,
            west: aw,
            north: an_nb,
            south: as_nb,
            source,
        })
    }

    /// Compute coefficients for v-momentum equation
    fn compute_v_momentum_coefficients(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<MomentumEquationCoeffs> {
        let (nx, ny, _, _) = flow_field.grid_info();

        // Compute effective viscosity (molecular + turbulent)
        let mu_eff = self.compute_effective_viscosity(flow_field, i, j, mu);

        // Similar to u-momentum but for v-component
        // See the note in the u-momentum routine: `Gamma * A / delta`, not
        // `Gamma / delta`.
        let gamma_e = mu_eff * dy / dx;
        let gamma_w = mu_eff * dy / dx;
        let gamma_n = mu_eff * dx / dy;
        let gamma_s = mu_eff * dx / dy;

        // Face fluxes for the v control volume, centred on the v face `j` and
        // spanning cell centre `j-1` to cell centre `j`. Mirrors the u case
        // above; see the note there on why a single cell-centred velocity for
        // all four faces cannot represent transport.
        let fn_ = rho * 0.5 * (flow_field.v[(j, i)] + flow_field.v[(j + 1, i)]) * dx;
        let fs = rho * 0.5 * (flow_field.v[(j - 1, i)] + flow_field.v[(j, i)]) * dx;

        // Columns `i = 0` and `i = nx - 1` are near-wall interior lines, not
        // boundaries: every v column sits at `x = (i + 0.5) dx`. The west and
        // east walls bound them at half-cell distance. See the u-momentum
        // routine for why that changes the conductance and kills the
        // convective flux, and nothing else.
        let west_is_wall = i == 0;
        let east_is_wall = i + 1 == nx;

        let fe = if east_is_wall {
            0.0
        } else {
            rho * 0.5 * (flow_field.u[(j - 1, i + 1)] + flow_field.u[(j, i + 1)]) * dy
        };
        let fw = if west_is_wall {
            0.0
        } else {
            rho * 0.5 * (flow_field.u[(j - 1, i)] + flow_field.u[(j, i)]) * dy
        };

        let an = gamma_n + f64::max(-fn_, 0.0);
        let as_ = gamma_s + f64::max(fs, 0.0);

        let gamma_wall = mu_eff * dy / (0.5 * dx);
        let mut wall_source = 0.0;
        let (ae, ae_nb) = if east_is_wall {
            wall_source += gamma_wall * self.v_wall(flow_field, i, j, nx as f64 * dx, dy);
            (gamma_wall, 0.0)
        } else {
            let a = gamma_e + f64::max(-fe, 0.0);
            (a, a)
        };
        let (aw, aw_nb) = if west_is_wall {
            wall_source += gamma_wall * self.v_wall(flow_field, i, j, 0.0, dy);
            (gamma_wall, 0.0)
        } else {
            let a = gamma_w + f64::max(fw, 0.0);
            (a, a)
        };

        let ap0 = if self.parameters.steady {
            0.0
        } else {
            rho * dx * dy / dt
        };
        // Net flux omitted, as in the u-momentum routine, to keep `a_p`
        // positive while continuity is still being established.
        let ap_unrelaxed = ae + aw + an + as_ + ap0;

        // Pressure gradient in y-direction
        let pressure_gradient = -(flow_field.p[(j, i)] - flow_field.p[(j - 1, i)]) * dx;
        let time_term = ap0 * flow_field.v_old[(j, i)];

        // Implicit under-relaxation; see the u-momentum routine.
        let alpha = self.parameters.velocity_relaxation;
        let ap = ap_unrelaxed / alpha;
        let source = pressure_gradient
            + time_term
            + wall_source
            + self.v_source_term(i, j, dx, dy)
            + self.v_deferred_correction(flow_field, i, j, nx, ny, fe, fw, fn_, fs)
            + (1.0 - alpha) / alpha * ap_unrelaxed * flow_field.v_old[(j, i)];

        Ok(MomentumEquationCoeffs {
            center: ap,
            east: ae_nb,
            west: aw_nb,
            north: an,
            south: as_,
            source,
        })
    }

    /// Compute mass source term for pressure correction equation
    fn compute_mass_source(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dx: f64,
        dy: f64,
        rho: f64,
    ) -> CfdResult<f64> {
        // Mass source = ρ * ∇·u*
        let u_e = flow_field.u_star[(j, i + 1)];
        let u_w = flow_field.u_star[(j, i)];
        let v_n = flow_field.v_star[(j + 1, i)];
        let v_s = flow_field.v_star[(j, i)];

        let mass_flux_imbalance = rho * ((u_e - u_w) * dy + (v_n - v_s) * dx);

        Ok(-mass_flux_imbalance) // Negative because we want ∇²p' = -∇·u*
    }

    /// Compute coefficients for the pressure correction equation.
    ///
    /// The coefficients are not free: SIMPLE requires that substituting the
    /// corrected velocities back into the continuity equation *reproduces*
    /// this equation. With the velocity correction
    /// `u_e = u*_e - d (p'_E - p'_P) / dx` and `d = ΔV / a_p`, continuity
    /// gives
    ///
    /// ```text
    ///     a_E = a_W = rho d dy/dx,   a_N = a_S = rho d dx/dy
    /// ```
    ///
    /// so the neighbour coefficients carry `d` — the momentum equation's own
    /// diagonal — and the mass imbalance in `compute_mass_source` is the
    /// source. Any other scaling breaks the link between the pressure the
    /// equation produces and the velocity correction it is supposed to drive.
    ///
    /// This previously used a bare Laplacian, `1/dx²` and `1/dy²`, while
    /// `velocity_correction_step` divided by `a_p = rho dx dy / dt`. The two
    /// therefore disagreed by a factor of roughly `1 / (h² dt)` — about
    /// 2 x 10^4 on a 16 x 16 cavity — so the pressure correction was that many
    /// times too small to enforce continuity. The consequence was not slow
    /// convergence but the wrong physics: with the pressure field pinned near
    /// zero, a lid-driven cavity produced a monotonic Couette profile with no
    /// recirculation at all, since the return flow in a cavity is created
    /// entirely by the pressure gradient.
    fn compute_pressure_coefficients(
        &self,
        dx: f64,
        dy: f64,
        rho: f64,
    ) -> CfdResult<MomentumEquationCoeffs> {
        // d = ΔV / a_p, with a_p = rho dx dy / dt as used by the velocity
        // correction, so d = dt / rho and `rho * d` is just the time step.
        let rho_d = rho * (dx * dy) / (rho * dx * dy / self.parameters.time_step);

        let ae = rho_d * dy / dx;
        let aw = ae;
        let an = rho_d * dx / dy;
        let as_ = an;
        let ap = ae + aw + an + as_;

        Ok(MomentumEquationCoeffs {
            center: ap,
            east: ae,
            west: aw,
            north: an,
            south: as_,
            source: 0.0, // Source is set separately
        })
    }

    /// Diagonal coefficient of the u-momentum equation, as used by the
    /// velocity correction and the pressure equation.
    ///
    /// This must be the *same* `a_p` the momentum equation was solved with —
    /// convection and diffusion included, not only the transient term — or the
    /// correction `u = u* - (ΔV / a_p) ∂p'/∂x` does not undo the momentum
    /// imbalance it is meant to.
    ///
    /// It previously returned `rho dx dy / dt`, which is only the transient
    /// contribution `a_p0`. On a Re = 100 cavity the convective and diffusive
    /// terms are of the same order as `a_p0`, so the correction was roughly
    /// twice as large as it should have been.
    fn compute_u_momentum_center_coefficient(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dx: f64,
        dy: f64,
        rho: f64,
    ) -> CfdResult<f64> {
        let coeffs = self.compute_u_momentum_coefficients(
            flow_field,
            i,
            j,
            self.parameters.time_step,
            rho,
            self.config.viscosity,
            dx,
            dy,
        )?;
        Ok(coeffs.center)
    }

    /// Diagonal coefficient of the v-momentum equation. See
    /// [`Self::compute_u_momentum_center_coefficient`].
    fn compute_v_momentum_center_coefficient(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dx: f64,
        dy: f64,
        rho: f64,
    ) -> CfdResult<f64> {
        let coeffs = self.compute_v_momentum_coefficients(
            flow_field,
            i,
            j,
            self.parameters.time_step,
            rho,
            self.config.viscosity,
            dx,
            dy,
        )?;
        Ok(coeffs.center)
    }
}

/// Coefficients for momentum equation discretization
#[derive(Debug, Clone)]
struct MomentumEquationCoeffs {
    pub center: f64,
    pub east: f64,
    pub west: f64,
    pub north: f64,
    pub south: f64,
    pub source: f64,
}

#[async_trait]
impl IncompressibleSolver for SimpleSolver {
    type Parameters = SimpleParameters;
    type Result = SimpleResult;

    fn new(config: CfdConfig, params: Self::Parameters) -> CfdResult<Self> {
        Self::new(config, params)
    }

    async fn solve_time_step(
        &mut self,
        flow_field: &mut FlowField,
        boundary_conditions: &BoundaryConditions,
        dt: f64,
    ) -> CfdResult<Self::Result> {
        let start_time = Instant::now();
        let mut residual_history = Vec::new();
        let pressure_iterations = Vec::new();
        let mut best_residual = f64::INFINITY;

        for iteration in 0..self.parameters.max_iterations {
            let (mass_residual, momentum_residual) = self
                .solve_simple_iteration(flow_field, boundary_conditions, dt)
                .await?;

            let total_residual =
                (mass_residual * mass_residual + momentum_residual * momentum_residual).sqrt();

            // Stop on divergence rather than running on to overflow.
            //
            // A residual that has grown by orders of magnitude above its best
            // value is diverging, and continuing only turns a large number
            // into an enormous one — the 8x8 cavity at a Reynolds number of a
            // million reached 1e149 before anything caught it, because
            // `is_finite` stays true right up to the moment it does not.
            if total_residual.is_finite()
                && best_residual.is_finite()
                && total_residual > best_residual * Self::DIVERGENCE_GROWTH
            {
                let solve_time = start_time.elapsed();
                return Ok(SimpleResult {
                    solver_result: SolverResult {
                        converged: false,
                        iterations: iteration + 1,
                        final_residual: total_residual,
                        residual_history,
                        solve_time,
                    },
                    pressure_iterations,
                    mass_residual,
                    momentum_residual,
                });
            }
            best_residual = best_residual.min(total_residual);

            // Stop on divergence rather than returning NaN.
            //
            // A solver asked for something it cannot do — here an 8x8 cavity
            // at a Reynolds number of a million — should say it did not
            // converge, not hand back a field of NaN that silently poisons
            // everything downstream. Reports the last finite residual so the
            // caller can see how far it got before it blew up.
            if !total_residual.is_finite() {
                let solve_time = start_time.elapsed();
                let last_finite = residual_history
                    .iter()
                    .rev()
                    .copied()
                    .find(|r: &f64| r.is_finite())
                    .unwrap_or(f64::MAX);
                return Ok(SimpleResult {
                    solver_result: SolverResult {
                        converged: false,
                        iterations: iteration + 1,
                        final_residual: last_finite,
                        residual_history,
                        solve_time,
                    },
                    pressure_iterations,
                    mass_residual: last_finite,
                    momentum_residual: last_finite,
                });
            }

            residual_history.push(total_residual);

            if total_residual < self.parameters.tolerance {
                let solve_time = start_time.elapsed();

                return Ok(SimpleResult {
                    solver_result: SolverResult {
                        converged: true,
                        iterations: iteration + 1,
                        final_residual: total_residual,
                        residual_history,
                        solve_time,
                    },
                    pressure_iterations,
                    mass_residual,
                    momentum_residual,
                });
            }
        }

        // Did not converge
        let solve_time = start_time.elapsed();
        Ok(SimpleResult {
            solver_result: SolverResult {
                converged: false,
                iterations: self.parameters.max_iterations,
                final_residual: residual_history.last().copied().unwrap_or(f64::INFINITY),
                residual_history,
                solve_time,
            },
            pressure_iterations,
            mass_residual: f64::INFINITY,
            momentum_residual: f64::INFINITY,
        })
    }

    async fn solve(
        &mut self,
        flow_field: &mut FlowField,
        boundary_conditions: &BoundaryConditions,
    ) -> CfdResult<Self::Result> {
        // For steady-state solve, use default time step
        self.solve_time_step(flow_field, boundary_conditions, self.parameters.time_step)
            .await
    }

    fn config(&self) -> &CfdConfig {
        &self.config
    }

    fn parameters(&self) -> &Self::Parameters {
        &self.parameters
    }
}