//! Geometry for embedded (immersed) bodies on the fixed staggered grid:
//! signed-distance descriptions, the fluid/ghost/solid classification of
//! cells and faces, the ghost-face velocity imposition, and the two load
//! routes the embedded solver's verification leans on.
//!
//! # The method, in one paragraph
//!
//! A body is a signed-distance function `phi(x, y, t)`, negative inside the
//! solid, with a surface-velocity function for the no-slip value on its
//! boundary. A pressure cell is *fluid* when `phi` at its centre is
//! positive. A velocity face is *fluid* — an unknown of the momentum and
//! continuity equations — only when both cells it separates are fluid.
//! Every other interior face is prescribed: a *ghost* face, within reach of
//! the fluid, takes the value a linear velocity profile along the surface
//! normal would have there (the profile pinned to the surface velocity at
//! the foot of the normal and to the interpolated fluid velocity at a probe
//! point one cell beyond the ghost's mirror image); a deep *solid* face
//! takes the surface velocity itself and is never read by a fluid stencil.
//! The fluid stencils need no special casing at all: diffusion and
//! convection at a fluid face read whatever its neighbours hold, and the
//! ghost values encode where the wall actually is. This is the sharp-
//! interface ghost-cell family (Tseng–Ferziger 2003, Mittal et al. 2008)
//! transcribed onto the MAC layout; it is the incompressible counterpart of
//! Farhat's embedded-boundary line, not FIVER itself.
//!
//! # Compatibility
//!
//! Prescribed faces carry the mass flux continuity sees. With every domain
//! side prescribed the projection is a pure Neumann problem, solvable only
//! if the net flux through all prescribed faces vanishes; extrapolated
//! ghost values do not sum to zero around a body on their own. [`EmbeddedMask::impose`]
//! therefore removes the net ghost flux of each step uniformly from the
//! ghost faces that bound the fluid (the target net flux of a rigid body is
//! exactly zero). The correction is reported so a test can watch its size
//! fall with the mesh; it is not a hidden mean-source subtraction on the
//! Poisson equation.
//!
//! # Limits (stated before results)
//!
//! Bodies must not touch the domain boundary (the outer ring of cells must
//! be fluid — checked). A body thinner than about two cells has ghost faces
//! on both sides whose probes reach across it; the classification still
//! works, the reconstruction degrades.

use crate::{CfdError, CfdResult};
use nalgebra::DMatrix;

type ScalarFn = Box<dyn Fn(f64, f64, f64) -> f64 + Send + Sync>;
type VectorFn = Box<dyn Fn(f64, f64, f64) -> (f64, f64) + Send + Sync>;

/// One sample of a body's surface, for load integration.
#[derive(Debug, Clone, Copy)]
pub struct SurfaceSample {
    /// Point on the surface.
    pub x: f64,
    /// Point on the surface.
    pub y: f64,
    /// Unit normal pointing out of the solid into the fluid.
    pub nx: f64,
    /// Unit normal pointing out of the solid into the fluid.
    pub ny: f64,
    /// Arc length the sample represents.
    pub ds: f64,
}

/// A rigid or prescribed-motion body embedded in the grid.
pub struct EmbeddedBody {
    sdf: ScalarFn,
    surface_velocity: VectorFn,
    /// Surface sampler at a requested spacing; `None` for a bare SDF body
    /// (loads by surface reconstruction are then unavailable).
    sampler: Option<Box<dyn Fn(f64) -> Vec<SurfaceSample> + Send + Sync>>,
}

impl EmbeddedBody {
    /// A body from its signed distance function (negative inside), at rest.
    pub fn from_sdf<F>(sdf: F) -> Self
    where
        F: Fn(f64, f64, f64) -> f64 + Send + Sync + 'static,
    {
        Self {
            sdf: Box::new(sdf),
            surface_velocity: Box::new(|_, _, _| (0.0, 0.0)),
            sampler: None,
        }
    }

    /// Prescribe the surface velocity `(x, y, t) -> (u, v)`; the default is
    /// rest. For a manufactured solution this is the exact field, evaluated
    /// wherever the mask asks.
    #[must_use]
    pub fn with_surface_velocity<F>(mut self, f: F) -> Self
    where
        F: Fn(f64, f64, f64) -> (f64, f64) + Send + Sync + 'static,
    {
        self.surface_velocity = Box::new(f);
        self
    }

    /// A fixed circle.
    pub fn circle(cx: f64, cy: f64, r: f64) -> Self {
        let mut body =
            Self::from_sdf(move |x, y, _| ((x - cx).powi(2) + (y - cy).powi(2)).sqrt() - r);
        body.sampler = Some(Box::new(move |ds| {
            let n = ((2.0 * std::f64::consts::PI * r / ds).ceil() as usize).max(8);
            let dtheta = 2.0 * std::f64::consts::PI / n as f64;
            (0..n)
                .map(|k| {
                    let theta = (k as f64 + 0.5) * dtheta;
                    let (s, c) = theta.sin_cos();
                    SurfaceSample {
                        x: cx + r * c,
                        y: cy + r * s,
                        nx: c,
                        ny: s,
                        ds: r * dtheta,
                    }
                })
                .collect()
        }));
        body
    }

    /// A fixed axis-aligned rectangle `[x0, x1] x [y0, y1]` (exact SDF).
    pub fn rectangle(x0: f64, y0: f64, x1: f64, y1: f64) -> Self {
        let (cx, cy) = (0.5 * (x0 + x1), 0.5 * (y0 + y1));
        let (hx, hy) = (0.5 * (x1 - x0), 0.5 * (y1 - y0));
        let mut body = Self::from_sdf(move |x, y, _| {
            let qx = (x - cx).abs() - hx;
            let qy = (y - cy).abs() - hy;
            let outside = (qx.max(0.0).powi(2) + qy.max(0.0).powi(2)).sqrt();
            outside + qx.max(qy).min(0.0)
        });
        body.sampler = Some(Box::new(move |ds| {
            let mut out = Vec::new();
            let mut edge = |ax: f64, ay: f64, bx: f64, by: f64, nx: f64, ny: f64| {
                let len = ((bx - ax).powi(2) + (by - ay).powi(2)).sqrt();
                let n = ((len / ds).ceil() as usize).max(1);
                for k in 0..n {
                    let s = (k as f64 + 0.5) / n as f64;
                    out.push(SurfaceSample {
                        x: ax + s * (bx - ax),
                        y: ay + s * (by - ay),
                        nx,
                        ny,
                        ds: len / n as f64,
                    });
                }
            };
            edge(x0, y0, x1, y0, 0.0, -1.0);
            edge(x1, y0, x1, y1, 1.0, 0.0);
            edge(x1, y1, x0, y1, 0.0, 1.0);
            edge(x0, y1, x0, y0, -1.0, 0.0);
            out
        }));
        body
    }

    /// Union of two bodies: the SDF is the minimum; the surface velocity and
    /// samples come from whichever body a point is closer to. Samples of one
    /// body lying inside the other are dropped.
    pub fn union(a: EmbeddedBody, b: EmbeddedBody) -> Self {
        let a = std::sync::Arc::new(a);
        let b = std::sync::Arc::new(b);
        let (a1, b1) = (a.clone(), b.clone());
        let (a2, b2) = (a.clone(), b.clone());
        let (a3, b3) = (a, b);
        let sampler: Box<dyn Fn(f64) -> Vec<SurfaceSample> + Send + Sync> = Box::new(move |ds| {
            let mut out: Vec<SurfaceSample> = a3
                .surface_samples(ds)
                .into_iter()
                .filter(|s| b3.phi(s.x, s.y, 0.0) > -1e-12)
                .collect();
            out.extend(
                b3.surface_samples(ds)
                    .into_iter()
                    .filter(|s| a3.phi(s.x, s.y, 0.0) > -1e-12),
            );
            out
        });
        Self {
            sdf: Box::new(move |x, y, t| a1.phi(x, y, t).min(b1.phi(x, y, t))),
            surface_velocity: Box::new(move |x, y, t| {
                if a2.phi(x, y, t) <= b2.phi(x, y, t) {
                    a2.surface_velocity(x, y, t)
                } else {
                    b2.surface_velocity(x, y, t)
                }
            }),
            sampler: Some(sampler),
        }
    }

    /// Signed distance, negative inside the solid.
    pub fn phi(&self, x: f64, y: f64, t: f64) -> f64 {
        (self.sdf)(x, y, t)
    }

    /// Surface (no-slip) velocity.
    pub fn surface_velocity(&self, x: f64, y: f64, t: f64) -> (f64, f64) {
        (self.surface_velocity)(x, y, t)
    }

    /// Unit normal out of the solid (the SDF gradient), by central
    /// differences at spacing `eps`.
    pub fn normal(&self, x: f64, y: f64, t: f64, eps: f64) -> (f64, f64) {
        let gx = (self.phi(x + eps, y, t) - self.phi(x - eps, y, t)) / (2.0 * eps);
        let gy = (self.phi(x, y + eps, t) - self.phi(x, y - eps, t)) / (2.0 * eps);
        let norm = (gx * gx + gy * gy).sqrt();
        if norm > 0.0 {
            (gx / norm, gy / norm)
        } else {
            (1.0, 0.0)
        }
    }

    /// Surface samples at roughly spacing `ds`; empty for a bare SDF body.
    pub fn surface_samples(&self, ds: f64) -> Vec<SurfaceSample> {
        self.sampler.as_ref().map_or_else(Vec::new, |s| s(ds))
    }
}

/// Force on a body by surface-stress reconstruction, with its bookkeeping.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SurfaceForce {
    /// Force per unit depth.
    pub fx: f64,
    /// Force per unit depth.
    pub fy: f64,
    /// Surface samples integrated.
    pub samples: usize,
    /// Samples whose probes could not be reconstructed (no fluid cells to
    /// fit, e.g. deep in a concave corner) and were skipped. Non-zero means
    /// the load is missing a piece of surface; the caller decides whether
    /// that piece matters.
    pub skipped: usize,
}

/// What a velocity face is.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FaceKind {
    /// Unknown: both adjacent cells are fluid.
    Fluid,
    /// Prescribed by the near-wall reconstruction.
    Ghost,
    /// Deep inside the body: prescribed the surface velocity, never read.
    Solid,
}

/// One lattice node entering a bilinear interpolation, with the value to
/// use instead when the node is not a fluid face: the surface velocity at
/// that node's own nearest surface point.
#[derive(Debug, Clone, Copy)]
struct StencilNode {
    j: usize,
    i: usize,
    x: f64,
    y: f64,
    weight: f64,
    fallback: Option<f64>,
}

/// The reconstruction data of one ghost face.
#[derive(Debug, Clone)]
struct Ghost {
    j: usize,
    i: usize,
    x: f64,
    y: f64,
    /// Foot of the normal on the surface.
    foot: (f64, f64),
    /// Surface velocity component at the foot of the normal.
    u_surface: f64,
    /// Signed distance of the face and of the probe (probe is positive).
    s_face: f64,
    s_probe: f64,
    nodes: [StencilNode; 4],
    /// The face bounds a fluid cell (it carries mass flux continuity sees)
    /// and its outward-from-fluid orientation sign, for the compatibility
    /// correction. `0.0` when no fluid cell is adjacent.
    flux_sign: f64,
}

/// Classification of a grid against a body at one instant.
pub struct EmbeddedMask {
    nx: usize,
    ny: usize,
    dx: f64,
    dy: f64,
    /// Row-major `(j, i)`.
    cell_fluid: Vec<bool>,
    u_kind: Vec<FaceKind>,
    v_kind: Vec<FaceKind>,
    u_ghosts: Vec<Ghost>,
    v_ghosts: Vec<Ghost>,
    /// First fluid cell, for anchoring a pure-Neumann projection.
    anchor: (usize, usize),
    fluid_cells: usize,
}

impl EmbeddedMask {
    /// Classify the `nx` by `ny` grid of spacing `dx`, `dy` against `body`
    /// at time `t`. Errors if the body touches the domain boundary.
    pub fn build(
        body: &EmbeddedBody,
        nx: usize,
        ny: usize,
        dx: f64,
        dy: f64,
        t: f64,
    ) -> CfdResult<Self> {
        let xc = |i: usize| (i as f64 + 0.5) * dx;
        let yc = |j: usize| (j as f64 + 0.5) * dy;
        let mut cell_fluid = vec![true; nx * ny];
        let mut fluid_cells = 0;
        let mut anchor = None;
        for j in 0..ny {
            for i in 0..nx {
                let fluid = body.phi(xc(i), yc(j), t) > 0.0;
                cell_fluid[j * nx + i] = fluid;
                if fluid {
                    fluid_cells += 1;
                    if anchor.is_none() {
                        anchor = Some((j, i));
                    }
                } else if i == 0 || j == 0 || i + 1 == nx || j + 1 == ny {
                    return Err(CfdError::invalid_parameter(format!(
                        "embedded body reaches the domain boundary at cell ({j}, {i}); \
                         bodies must be surrounded by fluid"
                    )));
                }
            }
        }
        let Some(anchor) = anchor else {
            return Err(CfdError::invalid_parameter(
                "embedded body covers the whole domain",
            ));
        };

        let h_min = dx.min(dy);
        // Faces further inside than this are never read by a fluid stencil
        // (a fluid face's neighbours lie within one cell; SDFs are
        // 1-Lipschitz).
        let reach = 1.5 * h_min;
        let eps = 1e-6 * h_min;

        let is_fluid = |j: usize, i: usize| cell_fluid[j * nx + i];

        let mut u_kind = vec![FaceKind::Fluid; ny * (nx + 1)];
        let mut v_kind = vec![FaceKind::Fluid; (ny + 1) * nx];
        // Interior u faces: i = 1..nx separate cells (j, i-1) and (j, i).
        for j in 0..ny {
            for i in 1..nx {
                if !(is_fluid(j, i - 1) && is_fluid(j, i)) {
                    let phi = body.phi(i as f64 * dx, yc(j), t);
                    u_kind[j * (nx + 1) + i] = if phi > -reach {
                        FaceKind::Ghost
                    } else {
                        FaceKind::Solid
                    };
                }
            }
        }
        for j in 1..ny {
            for i in 0..nx {
                if !(is_fluid(j - 1, i) && is_fluid(j, i)) {
                    let phi = body.phi(xc(i), j as f64 * dy, t);
                    v_kind[j * nx + i] = if phi > -reach {
                        FaceKind::Ghost
                    } else {
                        FaceKind::Solid
                    };
                }
            }
        }

        // Ghost reconstruction data. The probe sits one cell beyond the
        // face's mirror image: `s_probe = |s_face| + h_min`.
        let u_pos = |j: usize, i: usize| (i as f64 * dx, yc(j));
        let v_pos = |j: usize, i: usize| (xc(i), j as f64 * dy);
        let build_ghost = |j: usize,
                           i: usize,
                           x: f64,
                           y: f64,
                           component: usize,
                           kinds: &Vec<FaceKind>,
                           flux_sign: f64|
         -> Ghost {
            let s_face = body.phi(x, y, t);
            let (nrm_x, nrm_y) = body.normal(x, y, t, eps);
            let foot = (x - s_face * nrm_x, y - s_face * nrm_y);
            let s_probe = s_face.abs() + h_min;
            let probe = (foot.0 + s_probe * nrm_x, foot.1 + s_probe * nrm_y);
            let vel = body.surface_velocity(foot.0, foot.1, t);
            let u_surface = if component == 0 { vel.0 } else { vel.1 };
            let nodes = bilinear_nodes(probe, component, nx, ny, dx, dy, |jj, ii| {
                let k = if component == 0 {
                    kinds[jj * (nx + 1) + ii]
                } else {
                    kinds[jj * nx + ii]
                };
                if k == FaceKind::Fluid {
                    None
                } else {
                    let (px, py) = if component == 0 {
                        u_pos(jj, ii)
                    } else {
                        v_pos(jj, ii)
                    };
                    let s = body.phi(px, py, t);
                    let (nx_, ny_) = body.normal(px, py, t, eps);
                    let f = body.surface_velocity(px - s * nx_, py - s * ny_, t);
                    Some(if component == 0 { f.0 } else { f.1 })
                }
            });
            Ghost {
                j,
                i,
                x,
                y,
                foot,
                u_surface,
                s_face,
                s_probe,
                nodes,
                flux_sign,
            }
        };

        let mut u_ghosts = Vec::new();
        for j in 0..ny {
            for i in 1..nx {
                if u_kind[j * (nx + 1) + i] == FaceKind::Ghost {
                    let (x, y) = u_pos(j, i);
                    // Outward from the fluid cell: +1 if the fluid cell is
                    // west of the face, -1 if east.
                    let flux_sign = if is_fluid(j, i - 1) {
                        1.0
                    } else if is_fluid(j, i) {
                        -1.0
                    } else {
                        0.0
                    };
                    u_ghosts.push(build_ghost(j, i, x, y, 0, &u_kind, flux_sign));
                }
            }
        }
        let mut v_ghosts = Vec::new();
        for j in 1..ny {
            for i in 0..nx {
                if v_kind[j * nx + i] == FaceKind::Ghost {
                    let (x, y) = v_pos(j, i);
                    let flux_sign = if is_fluid(j - 1, i) {
                        1.0
                    } else if is_fluid(j, i) {
                        -1.0
                    } else {
                        0.0
                    };
                    v_ghosts.push(build_ghost(j, i, x, y, 1, &v_kind, flux_sign));
                }
            }
        }

        Ok(Self {
            nx,
            ny,
            dx,
            dy,
            cell_fluid,
            u_kind,
            v_kind,
            u_ghosts,
            v_ghosts,
            anchor,
            fluid_cells,
        })
    }

    /// Is pressure cell `(j, i)` fluid?
    #[inline]
    pub fn is_fluid_cell(&self, j: usize, i: usize) -> bool {
        self.cell_fluid[j * self.nx + i]
    }

    /// Kind of the u face `(j, i)`, `i = 0..=nx` (boundary faces read Fluid;
    /// the solver prescribes them from the side boundary).
    #[inline]
    pub fn u_kind(&self, j: usize, i: usize) -> FaceKind {
        self.u_kind[j * (self.nx + 1) + i]
    }

    /// Kind of the v face `(j, i)`, `j = 0..=ny`.
    #[inline]
    pub fn v_kind(&self, j: usize, i: usize) -> FaceKind {
        self.v_kind[j * self.nx + i]
    }

    /// The first fluid cell — the projection's anchor when it is pure Neumann.
    pub fn anchor(&self) -> (usize, usize) {
        self.anchor
    }

    /// Number of fluid cells.
    pub fn fluid_cells(&self) -> usize {
        self.fluid_cells
    }

    /// Number of ghost faces (u + v).
    pub fn ghost_faces(&self) -> usize {
        self.u_ghosts.len() + self.v_ghosts.len()
    }

    /// Write the prescribed values onto every non-fluid interior face of
    /// `u`, `v`: ghost faces by the linear normal reconstruction from the
    /// current fluid values, deep solid faces the surface velocity. Then
    /// remove the net ghost mass flux around the body uniformly from the
    /// flux-carrying ghost faces so the projection stays compatible.
    ///
    /// Returns the per-face compatibility correction applied (velocity
    /// units) — a diagnostic a test can watch shrink with the mesh.
    pub fn impose(
        &self,
        body: &EmbeddedBody,
        u: &mut DMatrix<f64>,
        v: &mut DMatrix<f64>,
        t: f64,
    ) -> f64 {
        let (u_source, v_source) = (u.clone(), v.clone());
        self.impose_from(body, &u_source, &v_source, u, v, t)
    }

    /// [`Self::impose`] with the fluid values read from a *different* field
    /// than the one written: the moving-body step reconstructs the new
    /// mask's ghost values from the previous step's corrected field (the
    /// boundary-history principle — ghost data, like domain-boundary data,
    /// is carried by what the previous step left, not by the uncorrected
    /// predictor state). With `source == target` values this is `impose`.
    #[allow(clippy::too_many_arguments)]
    pub fn impose_from(
        &self,
        body: &EmbeddedBody,
        u_source: &DMatrix<f64>,
        v_source: &DMatrix<f64>,
        u: &mut DMatrix<f64>,
        v: &mut DMatrix<f64>,
        t: f64,
    ) -> f64 {
        let (nx, ny, dx, dy) = (self.nx, self.ny, self.dx, self.dy);

        // Deep solid faces first (cheap, and ghost probes never read them
        // — fallbacks replace non-fluid nodes).
        for j in 0..ny {
            for i in 1..nx {
                if self.u_kind(j, i) == FaceKind::Solid {
                    u[(j, i)] = body
                        .surface_velocity(i as f64 * dx, (j as f64 + 0.5) * dy, t)
                        .0;
                }
            }
        }
        for j in 1..ny {
            for i in 0..nx {
                if self.v_kind(j, i) == FaceKind::Solid {
                    v[(j, i)] = body
                        .surface_velocity((i as f64 + 0.5) * dx, j as f64 * dy, t)
                        .1;
                }
            }
        }

        // Ghost values from the source fluid field (reads only fluid faces
        // and fallbacks, so order does not matter).
        let u_vals: Vec<f64> = self
            .u_ghosts
            .iter()
            .map(|g| g.reconstruct(u_source))
            .collect();
        let v_vals: Vec<f64> = self
            .v_ghosts
            .iter()
            .map(|g| g.reconstruct(v_source))
            .collect();

        // Net outward (from fluid) flux through flux-carrying ghost faces.
        let mut net = 0.0;
        let mut area = 0.0;
        for (g, &val) in self.u_ghosts.iter().zip(&u_vals) {
            if g.flux_sign != 0.0 {
                net += g.flux_sign * val * dy;
                area += dy;
            }
        }
        for (g, &val) in self.v_ghosts.iter().zip(&v_vals) {
            if g.flux_sign != 0.0 {
                net += g.flux_sign * val * dx;
                area += dx;
            }
        }
        let correction = if area > 0.0 { net / area } else { 0.0 };

        for (g, &val) in self.u_ghosts.iter().zip(&u_vals) {
            u[(g.j, g.i)] = val - g.flux_sign * correction;
        }
        for (g, &val) in self.v_ghosts.iter().zip(&v_vals) {
            v[(g.j, g.i)] = val - g.flux_sign * correction;
        }
        correction
    }

    /// Force on the body by surface-stress reconstruction. At each surface
    /// sample the pressure is linearly extrapolated to the wall from probes
    /// at `1 h` and `2 h` along the normal; the normal derivatives of the
    /// normal and tangential velocity at the wall come from the same probes
    /// with the surface velocity at the wall (quadratic fit `a s + b s^2`);
    /// the tangential derivative of the normal velocity along the surface
    /// comes from the surface-velocity function. The viscous traction is
    /// then the full `mu (grad u + grad u^T) n`, written in the sample's
    /// normal/tangent frame with `n` held fixed:
    ///
    /// ```text
    ///   tau.n = mu [ 2 n (n.grad)(u.n) + t ( (n.grad)(u.t) + (t.grad)(u.n) ) ]
    /// ```
    ///
    /// which reduces to the wall shear `mu (n.grad)(u.t) t` on a rigid
    /// no-slip wall and is exact in principle for a manufactured surface
    /// velocity, so a manufactured solution can test this route against the
    /// exact surface integral. Returns `(F_x, F_y)` per unit depth.
    ///
    /// `ds` is the surface sampling spacing. A sample whose probes cannot
    /// be reconstructed is skipped and counted in `skipped`; the force of a
    /// body with no sampler is zero with zero samples.
    #[allow(clippy::too_many_arguments)]
    pub fn surface_force(
        &self,
        body: &EmbeddedBody,
        u: &DMatrix<f64>,
        v: &DMatrix<f64>,
        p: &DMatrix<f64>,
        mu: f64,
        t: f64,
        ds: f64,
    ) -> SurfaceForce {
        let h = self.dx.min(self.dy);
        let samples = body.surface_samples(ds);
        let (mut fx, mut fy) = (0.0, 0.0);
        let mut skipped = 0;
        for s in &samples {
            let d1 = h;
            let d2 = 2.0 * h;
            let probes = (|| {
                let p1 = self.pressure_at(p, s.x + d1 * s.nx, s.y + d1 * s.ny)?;
                let p2 = self.pressure_at(p, s.x + d2 * s.nx, s.y + d2 * s.ny)?;
                let v1 = self.velocity_at(body, u, v, s.x + d1 * s.nx, s.y + d1 * s.ny, t)?;
                let v2 = self.velocity_at(body, u, v, s.x + d2 * s.nx, s.y + d2 * s.ny, t)?;
                Some((p1, p2, v1, v2))
            })();
            let Some((p1, p2, (u1, v1), (u2, v2))) = probes else {
                skipped += 1;
                continue;
            };
            let p_wall = p1 + (p1 - p2) * d1 / (d2 - d1);

            let (tx, ty) = (-s.ny, s.nx);
            let (u_s, v_s) = body.surface_velocity(s.x, s.y, t);
            let ut_wall = u_s * tx + v_s * ty;
            let un_wall = u_s * s.nx + v_s * s.ny;
            // Wall gradient of a quadratic `a s + b s^2` through the two
            // probes (values relative to the wall).
            let wall_gradient =
                |f1: f64, f2: f64| (f1 * d2 * d2 - f2 * d1 * d1) / (d1 * d2 * (d2 - d1));
            let dn_ut = wall_gradient(u1 * tx + v1 * ty - ut_wall, u2 * tx + v2 * ty - ut_wall);
            let dn_un = wall_gradient(
                u1 * s.nx + v1 * s.ny - un_wall,
                u2 * s.nx + v2 * s.ny - un_wall,
            );
            // Tangential derivative of (u . n) along the surface, n fixed.
            let eps = 1e-6 * h;
            let (up, vp) = body.surface_velocity(s.x + eps * tx, s.y + eps * ty, t);
            let (um, vm) = body.surface_velocity(s.x - eps * tx, s.y - eps * ty, t);
            let dt_un = ((up - um) * s.nx + (vp - vm) * s.ny) / (2.0 * eps);

            let traction_n = -p_wall + 2.0 * mu * dn_un;
            let traction_t = mu * (dn_ut + dt_un);
            fx += (traction_n * s.nx + traction_t * tx) * s.ds;
            fy += (traction_n * s.ny + traction_t * ty) * s.ds;
        }
        SurfaceForce {
            fx,
            fy,
            samples: samples.len(),
            skipped,
        }
    }

    /// Force on the body by a momentum balance over the rectangle of whole
    /// cells `[i0, i1) x [j0, j1)` (which must enclose the body and lie in
    /// the fluid on its boundary):
    /// `F = sum_outer (sigma.n - rho u (u.n)) A - d/dt int rho u dV + int f dV`,
    /// with the unsteady term from `(u - u_old)/dt` over the fluid cells of
    /// the box and `f` an optional volumetric momentum source evaluated at
    /// cell centres. Independent of the surface reconstruction: it reads no
    /// near-wall value at all.
    #[allow(clippy::too_many_arguments)]
    pub fn control_volume_force(
        &self,
        u: &DMatrix<f64>,
        v: &DMatrix<f64>,
        p: &DMatrix<f64>,
        u_old: &DMatrix<f64>,
        v_old: &DMatrix<f64>,
        dt: f64,
        rho: f64,
        mu: f64,
        source: Option<&dyn Fn(f64, f64) -> (f64, f64)>,
        (i0, i1, j0, j1): (usize, usize, usize, usize),
    ) -> (f64, f64) {
        let (dx, dy) = (self.dx, self.dy);
        let (mut fx, mut fy) = (0.0, 0.0);

        // East (i = i1, n = +x) and west (i = i0, n = -x) faces: u lives
        // on them; p and v are averaged across.
        for j in j0..j1 {
            for (i, sign) in [(i1, 1.0), (i0, -1.0)] {
                let un = u[(j, i)];
                let p_f = 0.5 * (p[(j, i - 1)] + p[(j, i)]);
                let dudx = (u[(j, i + 1)] - u[(j, i - 1)]) / (2.0 * dx);
                let v_c = |jj: usize, ii: usize| 0.5 * (v[(jj, ii)] + v[(jj + 1, ii)]);
                let dvdx = (v_c(j, i) - v_c(j, i - 1)) / dx;
                let dudy = if j == 0 {
                    (u[(j + 1, i)] - u[(j, i)]) / dy
                } else if j + 1 == self.ny {
                    (u[(j, i)] - u[(j - 1, i)]) / dy
                } else {
                    (u[(j + 1, i)] - u[(j - 1, i)]) / (2.0 * dy)
                };
                let v_f = 0.5 * (v_c(j, i - 1) + v_c(j, i));
                let sxx = -p_f + 2.0 * mu * dudx;
                let sxy = mu * (dudy + dvdx);
                // sigma.n - rho u (u.n), n = sign * e_x
                fx += sign * (sxx - rho * un * un) * dy;
                fy += sign * (sxy - rho * v_f * un) * dy;
            }
        }
        // North (j = j1, n = +y) and south (j = j0, n = -y) faces.
        for i in i0..i1 {
            for (j, sign) in [(j1, 1.0), (j0, -1.0)] {
                let vn = v[(j, i)];
                let p_f = 0.5 * (p[(j - 1, i)] + p[(j, i)]);
                let dvdy = (v[(j + 1, i)] - v[(j - 1, i)]) / (2.0 * dy);
                let u_c = |jj: usize, ii: usize| 0.5 * (u[(jj, ii)] + u[(jj, ii + 1)]);
                let dudy = (u_c(j, i) - u_c(j - 1, i)) / dy;
                let dvdx = if i == 0 {
                    (v[(j, i + 1)] - v[(j, i)]) / dx
                } else if i + 1 == self.nx {
                    (v[(j, i)] - v[(j, i - 1)]) / dx
                } else {
                    (v[(j, i + 1)] - v[(j, i - 1)]) / (2.0 * dx)
                };
                let u_f = 0.5 * (u_c(j - 1, i) + u_c(j, i));
                let syy = -p_f + 2.0 * mu * dvdy;
                let sxy = mu * (dudy + dvdx);
                fx += sign * (sxy - rho * u_f * vn) * dx;
                fy += sign * (syy - rho * vn * vn) * dx;
            }
        }
        // Unsteady term and source over the fluid cells of the box.
        for j in j0..j1 {
            for i in i0..i1 {
                if !self.is_fluid_cell(j, i) {
                    continue;
                }
                let du = 0.5 * ((u[(j, i)] - u_old[(j, i)]) + (u[(j, i + 1)] - u_old[(j, i + 1)]));
                let dv = 0.5 * ((v[(j, i)] - v_old[(j, i)]) + (v[(j + 1, i)] - v_old[(j + 1, i)]));
                fx -= rho * du / dt * dx * dy;
                fy -= rho * dv / dt * dx * dy;
                if let Some(f) = source {
                    let (sx, sy) = f((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy);
                    fx += sx * dx * dy;
                    fy += sy * dx * dy;
                }
            }
        }
        (fx, fy)
    }

    /// Pressure at a point from cell centres: bilinear when all four
    /// surrounding cells are fluid, otherwise the least-squares linear fit
    /// through the fluid ones (pressure has no wall value to add — it is a
    /// Neumann quantity there); `None` if fewer than three usable cells.
    fn pressure_at(&self, p: &DMatrix<f64>, x: f64, y: f64) -> Option<f64> {
        let (nx, ny, dx, dy) = (self.nx, self.ny, self.dx, self.dy);
        let gx = x / dx - 0.5;
        let gy = y / dy - 0.5;
        let i0 = gx.floor().clamp(0.0, (nx - 2) as f64) as usize;
        let j0 = gy.floor().clamp(0.0, (ny - 2) as f64) as usize;
        let fx = (gx - i0 as f64).clamp(0.0, 1.0);
        let fy = (gy - j0 as f64).clamp(0.0, 1.0);
        let nodes = [
            (j0, i0, (1.0 - fx) * (1.0 - fy)),
            (j0, i0 + 1, fx * (1.0 - fy)),
            (j0 + 1, i0, (1.0 - fx) * fy),
            (j0 + 1, i0 + 1, fx * fy),
        ];
        if nodes.iter().all(|&(jj, ii, _)| self.is_fluid_cell(jj, ii)) {
            return Some(nodes.iter().map(|&(jj, ii, w)| w * p[(jj, ii)]).sum());
        }
        let pts: Vec<(f64, f64, f64)> = nodes
            .iter()
            .filter(|&&(jj, ii, _)| self.is_fluid_cell(jj, ii))
            .map(|&(jj, ii, _)| ((ii as f64 + 0.5) * dx, (jj as f64 + 0.5) * dy, p[(jj, ii)]))
            .collect();
        linear_fit(&pts, (x, y))
    }

    /// Velocity at a point from the staggered lattices: bilinear when all
    /// four surrounding nodes of a component are fluid faces, otherwise the
    /// least-squares linear fit through the fluid nodes and the point's own
    /// boundary intercept (nearest surface point, surface velocity). `None`
    /// if even that is degenerate.
    fn velocity_at(
        &self,
        body: &EmbeddedBody,
        u: &DMatrix<f64>,
        v: &DMatrix<f64>,
        x: f64,
        y: f64,
        t: f64,
    ) -> Option<(f64, f64)> {
        let (nx, ny, dx, dy) = (self.nx, self.ny, self.dx, self.dy);
        let eps = 1e-6 * dx.min(dy);
        let s = body.phi(x, y, t);
        let (nrm_x, nrm_y) = body.normal(x, y, t, eps);
        let foot = (x - s * nrm_x, y - s * nrm_y);
        let vel_foot = body.surface_velocity(foot.0, foot.1, t);
        let mut out = [0.0; 2];
        for component in 0..2 {
            let nodes = bilinear_nodes((x, y), component, nx, ny, dx, dy, |_, _| None);
            let values = if component == 0 { u } else { v };
            let fluid = |n: &StencilNode| {
                if component == 0 {
                    self.u_kind(n.j, n.i) == FaceKind::Fluid
                } else {
                    self.v_kind(n.j, n.i) == FaceKind::Fluid
                }
            };
            if nodes.iter().all(fluid) {
                out[component] = nodes.iter().map(|n| n.weight * values[(n.j, n.i)]).sum();
            } else {
                let mut pts: Vec<(f64, f64, f64)> = nodes
                    .iter()
                    .filter(|n| fluid(n))
                    .map(|n| (n.x, n.y, values[(n.j, n.i)]))
                    .collect();
                let foot_val = if component == 0 {
                    vel_foot.0
                } else {
                    vel_foot.1
                };
                pts.push((foot.0, foot.1, foot_val));
                out[component] = linear_fit(&pts, (x, y))?;
            }
        }
        Some((out[0], out[1]))
    }
}

impl Ghost {
    /// The ghost value: a linear function fitted (least squares) through the
    /// fluid nodes around the probe and the boundary-intercept point with
    /// its surface velocity, evaluated at the ghost face — exact for linear
    /// fields. If the points are degenerate (fewer than three, or
    /// collinear), fall back to the linear profile along the normal with the
    /// non-fluid nodes replaced by the surface velocity at their own
    /// projections.
    fn reconstruct(&self, values: &DMatrix<f64>) -> f64 {
        let mut pts: Vec<(f64, f64, f64)> = self
            .nodes
            .iter()
            .filter(|n| n.fallback.is_none())
            .map(|n| (n.x, n.y, values[(n.j, n.i)]))
            .collect();
        pts.push((self.foot.0, self.foot.1, self.u_surface));
        if let Some(val) = linear_fit(&pts, (self.x, self.y)) {
            return val;
        }
        let mut probe = 0.0;
        for n in &self.nodes {
            probe += n.weight * n.fallback.unwrap_or_else(|| values[(n.j, n.i)]);
        }
        self.u_surface + (probe - self.u_surface) * (self.s_face / self.s_probe)
    }
}

/// Least-squares fit of `a + b (x - x0) + c (y - y0)` through `pts`,
/// evaluated at `at = (x0, y0)`; `None` if the normal matrix is singular
/// (fewer than three points or collinear ones), judged relative to its
/// scale.
fn linear_fit(pts: &[(f64, f64, f64)], at: (f64, f64)) -> Option<f64> {
    if pts.len() < 3 {
        return None;
    }
    // Normal equations for [1, dx, dy].
    let mut m = [[0.0f64; 3]; 3];
    let mut rhs = [0.0f64; 3];
    for &(x, y, val) in pts {
        let r = [1.0, x - at.0, y - at.1];
        for a in 0..3 {
            for b in 0..3 {
                m[a][b] += r[a] * r[b];
            }
            rhs[a] += r[a] * val;
        }
    }
    let det = m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
        - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
        + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]);
    let scale = m[0][0] * m[1][1] * m[2][2];
    if scale <= 0.0 || det.abs() <= 1e-10 * scale {
        return None;
    }
    // Cramer's rule for the constant term only (the value at `at`).
    let det_a = rhs[0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
        - m[0][1] * (rhs[1] * m[2][2] - m[1][2] * rhs[2])
        + m[0][2] * (rhs[1] * m[2][1] - m[1][1] * rhs[2]);
    Some(det_a / det)
}

/// The four lattice nodes of velocity `component` (0 = u on `x = i dx,
/// y = (j + 1/2) dy`; 1 = v on `x = (i + 1/2) dx, y = j dy`) surrounding
/// `point`, with bilinear weights, clamped to the lattice. `fallback(j, i)`
/// supplies the value to use for a node that is not a fluid face.
fn bilinear_nodes(
    point: (f64, f64),
    component: usize,
    nx: usize,
    ny: usize,
    dx: f64,
    dy: f64,
    fallback: impl Fn(usize, usize) -> Option<f64>,
) -> [StencilNode; 4] {
    let (gx, gy, max_i, max_j) = if component == 0 {
        (point.0 / dx, point.1 / dy - 0.5, nx, ny - 1)
    } else {
        (point.0 / dx - 0.5, point.1 / dy, nx - 1, ny)
    };
    let i0 = gx.floor().clamp(0.0, (max_i - 1) as f64) as usize;
    let j0 = gy.floor().clamp(0.0, (max_j - 1) as f64) as usize;
    let fx = (gx - i0 as f64).clamp(0.0, 1.0);
    let fy = (gy - j0 as f64).clamp(0.0, 1.0);
    let mk = |j: usize, i: usize, weight: f64| {
        let (x, y) = if component == 0 {
            (i as f64 * dx, (j as f64 + 0.5) * dy)
        } else {
            ((i as f64 + 0.5) * dx, j as f64 * dy)
        };
        StencilNode {
            j,
            i,
            x,
            y,
            weight,
            fallback: fallback(j, i),
        }
    };
    [
        mk(j0, i0, (1.0 - fx) * (1.0 - fy)),
        mk(j0, i0 + 1, fx * (1.0 - fy)),
        mk(j0 + 1, i0, (1.0 - fx) * fy),
        mk(j0 + 1, i0 + 1, fx * fy),
    ]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn circle_sdf_and_normal_are_exact() {
        let body = EmbeddedBody::circle(0.5, 0.5, 0.2);
        assert!((body.phi(0.9, 0.5, 0.0) - 0.2).abs() < 1e-14);
        assert!((body.phi(0.5, 0.5, 0.0) + 0.2).abs() < 1e-14);
        let (nx, ny) = body.normal(0.5 + 0.2 * 0.6, 0.5 + 0.2 * 0.8, 0.0, 1e-7);
        assert!((nx - 0.6).abs() < 1e-6 && (ny - 0.8).abs() < 1e-6);
        let samples = body.surface_samples(0.01);
        let len: f64 = samples.iter().map(|s| s.ds).sum();
        assert!((len - 2.0 * std::f64::consts::PI * 0.2).abs() < 1e-12);
    }

    #[test]
    fn rectangle_sdf_is_a_distance() {
        let body = EmbeddedBody::rectangle(0.25, 0.19, 0.6, 0.21);
        assert!((body.phi(0.7, 0.2, 0.0) - 0.1).abs() < 1e-14);
        assert!((body.phi(0.4, 0.2, 0.0) + 0.01).abs() < 1e-14);
        // Corner: Euclidean distance.
        assert!((body.phi(0.63, 0.25, 0.0) - 0.05).abs() < 1e-14);
    }

    #[test]
    fn classification_of_a_circle_on_a_box() {
        let n = 32;
        let h = 1.0 / n as f64;
        let body = EmbeddedBody::circle(0.5, 0.5, 0.2);
        let mask = EmbeddedMask::build(&body, n, n, h, h, 0.0).unwrap();
        let solid_cells = n * n - mask.fluid_cells();
        let expected = std::f64::consts::PI * 0.04 / (h * h);
        assert!(
            (solid_cells as f64 - expected).abs() < 0.1 * expected,
            "solid cells {solid_cells} vs area/h^2 {expected:.1}"
        );
        // A face deep inside is Solid, a face just inside is Ghost, the
        // anchor is a fluid cell, the outer ring is fluid.
        assert_eq!(mask.u_kind(n / 2, n / 2), FaceKind::Solid);
        assert!(mask.ghost_faces() > 0);
        let (aj, ai) = mask.anchor();
        assert!(mask.is_fluid_cell(aj, ai));
        for i in 0..n {
            assert!(mask.is_fluid_cell(0, i) && mask.is_fluid_cell(n - 1, i));
        }
    }

    #[test]
    fn body_touching_the_boundary_is_refused() {
        let body = EmbeddedBody::circle(0.0, 0.5, 0.2);
        assert!(EmbeddedMask::build(&body, 16, 16, 1.0 / 16.0, 1.0 / 16.0, 0.0).is_err());
    }

    /// A linear velocity field is reproduced exactly by the ghost
    /// reconstruction: surface value + linear profile + bilinear probe are
    /// all exact for linear fields, so every ghost face must read the field
    /// itself, and the compatibility correction must be zero for a
    /// divergence-free one.
    #[test]
    fn ghost_reconstruction_is_exact_for_linear_fields() {
        let n = 24;
        let h = 1.0 / n as f64;
        let lin_u = |x: f64, y: f64| 0.3 + 0.7 * x - 0.2 * y;
        let lin_v = |x: f64, y: f64| -0.1 + 0.4 * x - 0.7 * y; // div = 0
        let body = EmbeddedBody::circle(0.5, 0.5, 0.2)
            .with_surface_velocity(move |x, y, _| (lin_u(x, y), lin_v(x, y)));
        let mask = EmbeddedMask::build(&body, n, n, h, h, 0.0).unwrap();
        let mut u = DMatrix::zeros(n, n + 1);
        let mut v = DMatrix::zeros(n + 1, n);
        for j in 0..n {
            for i in 0..=n {
                u[(j, i)] = lin_u(i as f64 * h, (j as f64 + 0.5) * h);
            }
        }
        for j in 0..=n {
            for i in 0..n {
                v[(j, i)] = lin_v((i as f64 + 0.5) * h, j as f64 * h);
            }
        }
        // Scramble the non-fluid faces so the test is not vacuous.
        for j in 0..n {
            for i in 1..n {
                if mask.u_kind(j, i) != FaceKind::Fluid {
                    u[(j, i)] = 99.0;
                }
            }
        }
        for j in 1..n {
            for i in 0..n {
                if mask.v_kind(j, i) != FaceKind::Fluid {
                    v[(j, i)] = 99.0;
                }
            }
        }
        let correction = mask.impose(&body, &mut u, &mut v, 0.0);
        assert!(correction.abs() < 1e-10, "correction {correction:.3e}");
        let mut max_err: f64 = 0.0;
        for j in 0..n {
            for i in 1..n {
                if mask.u_kind(j, i) == FaceKind::Ghost {
                    max_err =
                        max_err.max((u[(j, i)] - lin_u(i as f64 * h, (j as f64 + 0.5) * h)).abs());
                }
            }
        }
        for j in 1..n {
            for i in 0..n {
                if mask.v_kind(j, i) == FaceKind::Ghost {
                    max_err =
                        max_err.max((v[(j, i)] - lin_v((i as f64 + 0.5) * h, j as f64 * h)).abs());
                }
            }
        }
        assert!(
            max_err < 1e-9,
            "ghost reconstruction error on a linear field: {max_err:.3e}"
        );
    }
}