//! The pressure step: face fluxes from the predictor with the compact
//! pressure term, boundary-flux adjustment for closed patches, the
//! pressure-correction equation on the 9-point operator, and the flux and
//! velocity corrections.

use super::{CurvilinearPisoSolver, PatchField, PressureSystem, SideBc, StepGeometry};
use crate::mesh::PatchSide;
use crate::solvers::incompressible::sparse_bicgstab::{
    BicgstabResult, CsrMatrix, bicgstab_jacobi, project_mean,
};

impl CurvilinearPisoSolver {
    /// Node values of a pressure-like cell field: zero on outlet sides,
    /// extrapolated elsewhere.
    pub(super) fn pressure_nodes(&self, p: &[f64]) -> Vec<f64> {
        let b = &self.params.boundaries;
        self.ops
            .node_values(&self.mesh, p, &|side: PatchSide, _| match b.get(side) {
                SideBc::Outlet => Some(0.0),
                SideBc::Velocity => None,
            })
    }

    /// `L_f(p)` on every face (zero on Neumann faces, outlet value zero).
    pub(super) fn pressure_face_gradients(&self, p: &[f64]) -> Vec<f64> {
        let mesh = &self.mesh;
        let pn = self.pressure_nodes(p);
        (0..mesh.faces().len())
            .map(|f| {
                let bval = mesh
                    .side(f)
                    .and_then(|s| match self.params.boundaries.get(s) {
                        SideBc::Outlet => Some(0.0),
                        SideBc::Velocity => None,
                    });
                self.ops.face_gradient_flux(mesh, f, p, &pn, bval)
            })
            .collect()
    }

    /// Least-squares cell gradient of a pressure-like field (outlet faces
    /// at zero, Neumann faces left out).
    pub(super) fn pressure_gradient(&self, p: &[f64], c: usize) -> [f64; 2] {
        let mesh = &self.mesh;
        self.ops.gradient(mesh, c, p, &|f| match mesh.side(f) {
            Some(s) if self.params.boundaries.get(s) == SideBc::Outlet => Some(0.0),
            _ => None,
        })
    }

    /// `F* = interp(û)·S̄ − (dt/ρ) L_f(p^n)` on interior and outlet faces,
    /// the prescribed flux `u_b · S̄` on velocity faces (at `t_new`, on the
    /// end-of-step face centres). `S̄` is the step's face vector (`geo`).
    pub(super) fn predicted_fluxes(
        &self,
        uh: &[f64],
        vh: &[f64],
        p: &[f64],
        dt: f64,
        t_new: f64,
        geo: &StepGeometry,
    ) -> Vec<f64> {
        let mesh = &self.mesh;
        let rho = self.config.density;
        let lp = self.pressure_face_gradients(p);
        mesh.faces()
            .iter()
            .enumerate()
            .map(|(f, face)| {
                let s = geo.s_bar[f];
                match (face.owner, face.neigh) {
                    (Some(o), Some(n)) => {
                        let w = face.w;
                        let uf = w * uh[o] + (1.0 - w) * uh[n];
                        let vf = w * vh[o] + (1.0 - w) * vh[n];
                        uf * s[0] + vf * s[1] - dt / rho * lp[f]
                    }
                    _ => {
                        let c = mesh.boundary_cell(f);
                        let side = mesh.side(f).expect("boundary");
                        if side == PatchSide::Outer && self.acceptors.is_some() {
                            // The acceptor ring's outer faces: the stamped
                            // velocity's own flux (read only by the overlap
                            // mass-defect measure).
                            return uh[c] * s[0] + vh[c] * s[1];
                        }
                        match self.params.boundaries.get(side) {
                            SideBc::Velocity => {
                                let (ub, vb) = self.boundary_velocity(
                                    side,
                                    face.centre[0],
                                    face.centre[1],
                                    t_new,
                                );
                                ub * s[0] + vb * s[1]
                            }
                            SideBc::Outlet => uh[c] * s[0] + vh[c] * s[1] - dt / rho * lp[f],
                        }
                    }
                }
            })
            .collect()
    }

    /// On a patch with no outlet the prescribed boundary fluxes must sum
    /// to zero for the projection to be solvable; the O(h²) defect of
    /// face-centre sampling is spread over the velocity faces by area
    /// (OpenFOAM's `adjustPhi`). Returns the defect removed.
    pub(super) fn adjust_boundary_flux(&self, flux: &mut [f64]) -> f64 {
        let mesh = &self.mesh;
        let has_outlet = [
            PatchSide::Inner,
            PatchSide::Outer,
            PatchSide::SStart,
            PatchSide::SEnd,
        ]
        .iter()
        .any(|&s| self.params.boundaries.get(s) == SideBc::Outlet);
        // A Robin wall absorbs the net flux through its compliance (the
        // pressure system is then not pure Neumann).
        if has_outlet || self.acceptors.is_some() || self.robin.is_some() {
            return 0.0;
        }
        let (mut net, mut total_len) = (0.0, 0.0);
        for (f, face) in mesh.faces().iter().enumerate() {
            if mesh.side(f).is_some() {
                let out_sign = if face.owner.is_some() { 1.0 } else { -1.0 };
                net += out_sign * flux[f];
                total_len += (face.s[0] * face.s[0] + face.s[1] * face.s[1]).sqrt();
            }
        }
        if total_len == 0.0 {
            return net;
        }
        for (f, face) in mesh.faces().iter().enumerate() {
            if mesh.side(f).is_some() {
                let out_sign = if face.owner.is_some() { 1.0 } else { -1.0 };
                let len = (face.s[0] * face.s[0] + face.s[1] * face.s[1]).sqrt();
                flux[f] -= out_sign * net * len / total_len;
            }
        }
        net
    }

    /// Assemble `−Σ_f sign (dt/ρ) L_f` (positive diagonal) and pick the
    /// anchor for the pure-Neumann case. Acceptor cells get identity rows
    /// (their `p'` is Dirichlet) and the interior rows' couplings to them
    /// are recorded as links, eliminated at solve time.
    pub(super) fn assemble_pressure_matrix(&self, dt: f64) -> PressureSystem {
        let mesh = &self.mesh;
        let rho = self.config.density;
        let n = mesh.cell_count();
        let mut tri = Vec::with_capacity(n * 12);
        let mut links = Vec::new();
        let mut coefs = Vec::new();
        let mut any_dirichlet = self.acceptors.is_some();
        for c in 0..n {
            if self.is_acceptor(c) {
                tri.push((c, c, 1.0));
                continue;
            }
            for (f, sign) in mesh.cell_faces(c) {
                if let Some(w) = &self.robin {
                    if mesh.side(f) == Some(PatchSide::Inner) {
                        // The compliant wall: outward flux `+|S| p'_c / alpha`.
                        let sv = mesh.faces()[f].s;
                        let len = (sv[0] * sv[0] + sv[1] * sv[1]).sqrt();
                        tri.push((c, c, len / w.alpha));
                        any_dirichlet = true;
                    }
                }
                self.ops
                    .face_gradient_coeffs(mesh, &self.params.boundaries, f, &mut coefs);
                if mesh.side(f).is_some() && !coefs.is_empty() {
                    any_dirichlet = true;
                }
                for &(col, v) in &coefs {
                    let coef = -sign * dt / rho * v;
                    if self.is_acceptor(col) {
                        links.push((c, col, coef));
                    } else {
                        tri.push((c, col, coef));
                    }
                }
            }
            tri.push((c, c, 0.0)); // guarantee a diagonal entry
        }
        let mut matrix = CsrMatrix::from_triplets(n, &tri);
        let anchor = if any_dirichlet {
            None
        } else {
            // An interior cell away from the seam: (1, 1).
            let a_cell = mesh.cell(1.min(mesh.nn() - 1), 1.min(mesh.ns() - 1));
            matrix.set_row_identity(a_cell);
            Some(a_cell)
        };
        PressureSystem {
            dt,
            matrix,
            anchor,
            links,
        }
    }

    /// The right-hand side `−Σ sign F` on the equation-carrying cells, with
    /// the acceptor couplings eliminated (`rhs −= coef · p'_acceptor`) and
    /// zero on acceptor rows; mean-projected and anchored when pure Neumann.
    /// Also returns its L1 norm BEFORE the projection (the incoming
    /// imbalance the stop is relative to — the static path's
    /// `divergence_l1`, unchanged to the bit).
    pub(super) fn pressure_rhs(&self, system: &PressureSystem, flux: &[f64]) -> (Vec<f64>, f64) {
        let mesh = &self.mesh;
        let n = mesh.cell_count();
        let mut rhs = vec![0.0; n];
        for c in 0..n {
            if self.is_acceptor(c) {
                continue;
            }
            let mut div = 0.0;
            for (f, sign) in mesh.cell_faces(c) {
                div += sign * flux[f];
            }
            rhs[c] = -div;
        }
        for &(row, acc, coef) in &system.links {
            rhs[row] -= coef * self.acceptor_correction(acc).unwrap_or(0.0);
        }
        let incoming: f64 = rhs.iter().map(|r| r.abs()).sum();
        if let Some(a) = system.anchor {
            project_mean(&mut rhs);
            rhs[a] = 0.0;
        }
        (rhs, incoming)
    }

    /// Solve the assembled system for `p'` (zero start); acceptor entries
    /// are then set to their Dirichlet values.
    pub(super) fn solve_pressure_correction(
        &self,
        system: &PressureSystem,
        rhs: Vec<f64>,
        tolerance: f64,
    ) -> (Vec<f64>, BicgstabResult) {
        let n = self.mesh.cell_count();
        let mut pc = vec![0.0; n];
        let out = bicgstab_jacobi(
            &system.matrix,
            &rhs,
            &mut pc,
            tolerance,
            self.params.max_poisson_iterations,
        );
        if self.acceptors.is_some() {
            for c in 0..n {
                if let Some(v) = self.acceptor_correction(c) {
                    pc[c] = v;
                }
            }
        }
        (pc, out)
    }

    /// `F −= (dt/ρ) L_f(p')`, `u −= (dt/ρ) ∇p'`, `p += p'`.
    pub(super) fn apply_correction(&self, field: &mut PatchField, pc: &[f64], dt: f64) {
        let mesh = &self.mesh;
        let rho = self.config.density;
        let lp = self.pressure_face_gradients(pc);
        for f in 0..mesh.faces().len() {
            field.flux[f] -= dt / rho * lp[f];
        }
        if let Some(w) = &self.robin {
            // The compliant wall's flux answer: `δu_b = −p' S / (alpha |S|)`,
            // `δF = δu_b · S = −p' |S| / alpha` in the face's own orientation.
            for &f in &self.robin_faces {
                let c = mesh.boundary_cell(f);
                let sv = mesh.faces()[f].s;
                let len = (sv[0] * sv[0] + sv[1] * sv[1]).sqrt();
                field.flux[f] -= pc[c] * len / w.alpha;
            }
        }
        for c in 0..mesh.cell_count() {
            if self.is_acceptor(c) {
                continue;
            }
            let g = self.pressure_gradient(pc, c);
            field.u[c] -= dt / rho * g[0];
            field.v[c] -= dt / rho * g[1];
            field.p[c] += pc[c];
        }
    }

    /// Total cell mass imbalance `Σ_c |Σ_f sign F_f|` over the cells that
    /// carry continuity (acceptors excluded).
    pub(super) fn divergence_l1(&self, flux: &[f64]) -> f64 {
        let mesh = &self.mesh;
        (0..mesh.cell_count())
            .filter(|&c| !self.is_acceptor(c))
            .map(|c| {
                mesh.cell_faces(c)
                    .iter()
                    .map(|&(f, sign)| sign * flux[f])
                    .sum::<f64>()
                    .abs()
            })
            .sum()
    }

    /// Largest cell mass imbalance `|Σ sign F_f|` over the cells that
    /// carry continuity (acceptors excluded).
    pub(super) fn max_divergence(&self, flux: &[f64]) -> f64 {
        let mesh = &self.mesh;
        (0..mesh.cell_count())
            .filter(|&c| !self.is_acceptor(c))
            .map(|c| {
                mesh.cell_faces(c)
                    .iter()
                    .map(|&(f, sign)| sign * flux[f])
                    .sum::<f64>()
                    .abs()
            })
            .fold(0.0, f64::max)
    }
}