//! The pressure step of the embedded-boundary PISO: the masked five-point
//! problem and the projection, split into its solve and apply halves so
//! the overset coupling (`overset/`) can iterate the solve across two
//! meshes before applying the correction once.

use super::EmbeddedPisoSolver;
use crate::CfdResult;
use crate::solvers::incompressible::ale::SideBoundary;
use crate::solvers::incompressible::poisson::{
    MultigridParameters, PoissonProblem, PoissonSolverKind, solve_multigrid_pcg_cached,
};
use crate::solvers::incompressible::{EmbeddedMask, FlowField};

impl EmbeddedPisoSolver {
    /// The pressure-correction system of one projection as a
    /// [`PoissonProblem`]: active = fluid cells, coefficient `dt A / delta`
    /// across every fluid interior face and zero across every prescribed
    /// one (domain Velocity / SlipWall sides, non-fluid interior faces),
    /// the outlet's Dirichlet `p' = 0` half a cell away as a diagonal-only
    /// `extra_diag`, right-hand side the mass imbalance `sp`. Arm for arm
    /// the coefficients the SOR loop of [`Self::project`] forms in place.
    pub(super) fn poisson_problem(&self, field: &FlowField, dt: f64) -> PoissonProblem {
        let (nx, ny, dx, dy) = field.grid_info();
        let b = self.parameters.boundaries;
        let outlet = SideBoundary::PressureOutlet;
        let ae_interior = dt * dy / dx;
        let an_interior = dt * dx / dy;
        let ae_outlet = dt * dy / (0.5 * dx);
        let an_outlet = dt * dx / (0.5 * dy);

        let mut problem = PoissonProblem::new(nx, ny);
        for j in 0..ny {
            for i in 0..nx {
                let idx = j * nx + i;
                if !self.cell_is_fluid(j, i) {
                    problem.active[idx] = false;
                    continue;
                }
                let mut extra = 0.0;
                // An overset fringe neighbour is a Dirichlet cell: its
                // coefficient moves to the diagonal and its known p' to
                // the right-hand side (the outlet's construction).
                let mut rhs_extra = 0.0;
                if i + 1 == nx {
                    if b.right == outlet {
                        extra += ae_outlet;
                    }
                } else if self.u_is_fluid(j, i + 1) {
                    if self.is_fringe(j, i + 1) {
                        extra += ae_interior;
                        rhs_extra += ae_interior * self.fringe_correction(j, i + 1);
                    } else {
                        problem.ae[idx] = ae_interior;
                    }
                }
                if i == 0 {
                    if b.left == outlet {
                        extra += ae_outlet;
                    }
                } else if self.u_is_fluid(j, i) {
                    if self.is_fringe(j, i - 1) {
                        extra += ae_interior;
                        rhs_extra += ae_interior * self.fringe_correction(j, i - 1);
                    } else {
                        problem.aw[idx] = ae_interior;
                    }
                }
                if j + 1 == ny {
                    if b.top == outlet {
                        extra += an_outlet;
                    }
                } else if self.v_is_fluid(j + 1, i) {
                    if self.is_fringe(j + 1, i) {
                        extra += an_interior;
                        rhs_extra += an_interior * self.fringe_correction(j + 1, i);
                    } else {
                        problem.an[idx] = an_interior;
                    }
                }
                if j == 0 {
                    if b.bottom == outlet {
                        extra += an_outlet;
                    }
                } else if self.v_is_fluid(j, i) {
                    if self.is_fringe(j - 1, i) {
                        extra += an_interior;
                        rhs_extra += an_interior * self.fringe_correction(j - 1, i);
                    } else {
                        problem.as_[idx] = an_interior;
                    }
                }
                problem.extra_diag[idx] = extra;
                problem.rhs[idx] = field.sp[(j, i)] + rhs_extra;
            }
        }
        problem
    }

    /// One projection on the fluid cells: the fixed-grid PISO's, with a
    /// zero coefficient across every prescribed face (domain Velocity /
    /// SlipWall sides and every non-fluid interior face), a Dirichlet `p' =
    /// 0` half a cell beyond an outlet face, and the anchor on the first
    /// fluid cell when no outlet exists. Returns the normalised mass
    /// imbalance of the corrected field over the fluid cells.
    #[allow(clippy::too_many_lines)]
    /// One projection on the fluid cells: [`Self::solve_correction`] then
    /// [`Self::apply_correction`] — the two halves the overset coupling
    /// calls separately (several solves, one application).
    pub(super) fn project(
        &self,
        field: &mut FlowField,
        dt: f64,
        warm_start: bool,
    ) -> CfdResult<f64> {
        self.solve_correction(field, dt, warm_start)?;
        Ok(self.apply_correction(field, dt))
    }

    /// The solve half of a projection: the continuity source from `u*`
    /// into `field.sp`, then `p'` into `field.p_prime` (multigrid PCG with
    /// the SOR fallback). Nothing else in `field` is touched.
    #[allow(clippy::too_many_lines)]
    pub(crate) fn solve_correction(
        &self,
        field: &mut FlowField,
        dt: f64,
        warm_start: bool,
    ) -> CfdResult<()> {
        let (nx, ny, dx, dy) = field.grid_info();
        let rho = self.config.density;
        let b = self.parameters.boundaries;
        let outlet = SideBoundary::PressureOutlet;
        let any_outlet = [b.left, b.right, b.bottom, b.top].contains(&outlet);
        let anchor = self.mask.as_ref().map_or((1, 1), EmbeddedMask::anchor);

        // With `warm_start` (the FIRST corrector only), the previous
        // step's correction is the multigrid initial guess — the
        // correction field is temporally correlated step to step
        // (measured 2026-08-30 on the FSI2 rigid phase: 2.96 PCG
        // iterations/solve from zero, 1.27 warm). Later correctors
        // solve for a much SMALLER correction, and the first
        // corrector's p' is a WORSE guess than zero there (measured:
        // the all-correctors draft cost 3.9 iters/solve in the coupled
        // phase). The SOR fallback below still starts from zero,
        // exactly as before.

        let mut source_scale = 0.0;
        for j in 0..ny {
            for i in 0..nx {
                if !self.cell_is_fluid(j, i) {
                    field.sp[(j, i)] = 0.0;
                    continue;
                }
                let divergence_flux = rho
                    * ((field.u_star[(j, i + 1)] - field.u_star[(j, i)]) * dy
                        + (field.v_star[(j + 1, i)] - field.v_star[(j, i)]) * dx);
                field.sp[(j, i)] = -divergence_flux;
                source_scale += divergence_flux.abs();
            }
        }
        // Swept-volume source (knob; see `swept_volume`): the corrected
        // field must satisfy Σ u·n A = −dV_f/dt in every interface cell.
        if let (true, Some(a_new), Some(a_old)) =
            (self.swept_volume != 0.0, &self.alpha_new, &self.alpha_old)
        {
            for j in 0..ny {
                for i in 0..nx {
                    if !self.cell_is_fluid(j, i) {
                        continue;
                    }
                    let k = j * nx + i;
                    let da = a_new[k] - a_old[k];
                    if da != 0.0 {
                        field.sp[(j, i)] -= self.swept_volume * rho * da * dx * dy / dt;
                    }
                }
            }
        }
        // Reporting-only divergence trace (RTX_EMBEDDED_TRACE_SP): where the
        // projection's source sits relative to the step's fresh cells, in
        // units of one whole cell volume per step (rho dx dy / dt).
        if warm_start && !self.fresh_trace.is_empty() {
            let unit = rho * dx * dy / dt;
            let is_fresh =
                |j: usize, i: usize| self.fresh_trace.iter().any(|&(a, b)| a == j && b == i);
            let is_nbr = |j: usize, i: usize| {
                self.fresh_trace.iter().any(|&(a, b)| {
                    (a == j && (b + 1 == i || i + 1 == b)) || (b == i && (a + 1 == j || j + 1 == a))
                })
            };
            let (mut mf, mut mn, mut mo) = (0.0f64, 0.0f64, 0.0f64);
            let (mut arg, mut argv) = ((0usize, 0usize), 0.0f64);
            let mut sum_fresh = 0.0f64;
            for j in 0..ny {
                for i in 0..nx {
                    if !self.cell_is_fluid(j, i) {
                        continue;
                    }
                    let v = field.sp[(j, i)] / unit;
                    if is_fresh(j, i) {
                        mf = mf.max(v.abs());
                        sum_fresh += v;
                    } else if is_nbr(j, i) {
                        mn = mn.max(v.abs());
                    } else {
                        mo = mo.max(v.abs());
                    }
                    if v.abs() > argv.abs() {
                        argv = v;
                        arg = (j, i);
                    }
                }
            }
            let class = if is_fresh(arg.0, arg.1) {
                "FRESH"
            } else if is_nbr(arg.0, arg.1) {
                "NEIGHBOUR"
            } else {
                "other"
            };
            // The argmax cell's 3x3 neighbourhood: F = fluid, S = solid,
            // * = fresh this step (row above first).
            let mut hood = String::new();
            for dj in [1i64, 0, -1] {
                for di in [-1i64, 0, 1] {
                    let (jj, ii) = (arg.0 as i64 + dj, arg.1 as i64 + di);
                    let c = if jj < 0 || ii < 0 || jj >= ny as i64 || ii >= nx as i64 {
                        '#'
                    } else if is_fresh(jj as usize, ii as usize) {
                        '*'
                    } else if self.cell_is_fluid(jj as usize, ii as usize) {
                        'F'
                    } else {
                        'S'
                    };
                    hood.push(c);
                }
                hood.push('/');
            }
            let fj = self.fresh_trace.iter().map(|c| c.0);
            let fi = self.fresh_trace.iter().map(|c| c.1);
            println!(
                "  SP-TRACE fresh rows {:?}..{:?} cols {:?}..{:?}; argmax hood {hood}",
                fj.clone().min(),
                fj.max(),
                fi.clone().min(),
                fi.max()
            );
            println!(
                "  SP-TRACE t = {:.6}: {} fresh cells; max |sp| {:+.3} cell-volumes/step at ({}, {}) [{class}]; \
                 max over fresh {:.3}, neighbours {:.3}, others {:.3}; sum over fresh {:+.3}",
                self.time + dt,
                self.fresh_trace.len(),
                argv,
                arg.0,
                arg.1,
                mf,
                mn,
                mo,
                sum_fresh
            );
        }

        let ae_interior = dt * dy / dx;
        let an_interior = dt * dx / dy;
        // Outlet face: p' = 0 half a cell away.
        let ae_outlet = dt * dy / (0.5 * dx);
        let an_outlet = dt * dx / (0.5 * dy);

        let reference_flux = rho * self.config.reference_velocity * self.config.reference_length;
        let inner_stop = (self.inner_stop_factor * source_scale)
            .max(0.1 * self.parameters.tolerance * reference_flux)
            + 1e-14;
        let mut multigrid_converged = false;
        if self.parameters.poisson_solver == PoissonSolverKind::Multigrid {
            // The same system the SOR loop below sweeps, handed to the
            // multigrid-preconditioned CG solver: anchored on the first
            // fluid cell when there is no outlet (the SOR loop pins it to
            // zero), level-free otherwise. Non-fluid cells and isolated
            // fluid cells are never written and keep `p' = 0`.
            let problem = self.poisson_problem(field, dt);
            // Warm start from the previous correction on the CURRENT
            // fluid cells; everything else stays zero, preserving the
            // p' = 0 invariant on non-fluid cells through the copy-back.
            let mut p_prime = vec![0.0; nx * ny];
            if warm_start {
                for j in 0..ny {
                    for i in 0..nx {
                        if self.cell_is_fluid(j, i) {
                            p_prime[j * nx + i] = field.p_prime[(j, i)];
                        }
                    }
                }
            }
            let anchor_cell =
                (!any_outlet && !self.has_fringe()).then_some(anchor.0 * nx + anchor.1);
            let solution = solve_multigrid_pcg_cached(
                &problem,
                &mut p_prime,
                &MultigridParameters {
                    precision: self.parameters.poisson_precision,
                    smoother: self.parameters.poisson_smoother,
                    threads: self.poisson_threads,
                    ..MultigridParameters::default()
                },
                inner_stop,
                anchor_cell,
                &mut self.pcg_cache.borrow_mut(),
            );
            let (s0, i0, c0, k0) = self.poisson_profile.get();
            self.poisson_profile.set((
                s0 + solution.setup_ns,
                i0 + solution.iterate_ns,
                c0 + 1,
                k0 + solution.iterations as u64,
            ));
            // Unconverged: fall back to the SOR sweeps for this projection
            // rather than apply a correction that did not reach the stop.
            multigrid_converged = solution.converged;
            if multigrid_converged {
                for j in 0..ny {
                    for i in 0..nx {
                        field.p_prime[(j, i)] = p_prime[j * nx + i];
                    }
                }
                self.stamp_fringe_correction(&mut field.p_prime);
            }
        }
        if !multigrid_converged {
            // The fallback is unchanged: SOR from zero, as it always ran.
            field.p_prime.fill(0.0);
            self.stamp_fringe_correction(&mut field.p_prime);
            let omega = 2.0 / (1.0 + (std::f64::consts::PI / nx.max(ny) as f64).sin());
            for _sweep in 0..2000 {
                let mut residual = 0.0;
                for j in 0..ny {
                    for i in 0..nx {
                        if !self.cell_is_fluid(j, i) {
                            continue;
                        }
                        if !any_outlet && !self.has_fringe() && (j, i) == anchor {
                            field.p_prime[(j, i)] = 0.0;
                            continue;
                        }

                        // A coefficient is zero exactly when the face is
                        // prescribed: a domain side with velocity data, or a
                        // non-fluid interior face.
                        let ae = if i + 1 == nx {
                            if b.right == outlet { ae_outlet } else { 0.0 }
                        } else if self.u_is_fluid(j, i + 1) {
                            ae_interior
                        } else {
                            0.0
                        };
                        let aw = if i == 0 {
                            if b.left == outlet { ae_outlet } else { 0.0 }
                        } else if self.u_is_fluid(j, i) {
                            ae_interior
                        } else {
                            0.0
                        };
                        let an = if j + 1 == ny {
                            if b.top == outlet { an_outlet } else { 0.0 }
                        } else if self.v_is_fluid(j + 1, i) {
                            an_interior
                        } else {
                            0.0
                        };
                        let as_ = if j == 0 {
                            if b.bottom == outlet { an_outlet } else { 0.0 }
                        } else if self.v_is_fluid(j, i) {
                            an_interior
                        } else {
                            0.0
                        };
                        let ap = ae + aw + an + as_;
                        if ap == 0.0 {
                            // An isolated fluid cell enclosed by prescribed
                            // faces has no equation; leave p' = 0 there.
                            continue;
                        }

                        let east = if i + 1 < nx {
                            ae * field.p_prime[(j, i + 1)]
                        } else {
                            0.0
                        };
                        let west = if i > 0 {
                            aw * field.p_prime[(j, i - 1)]
                        } else {
                            0.0
                        };
                        let north = if j + 1 < ny {
                            an * field.p_prime[(j + 1, i)]
                        } else {
                            0.0
                        };
                        let south = if j > 0 {
                            as_ * field.p_prime[(j - 1, i)]
                        } else {
                            0.0
                        };

                        let rhs = field.sp[(j, i)] + east + west + north + south;
                        let p_old = field.p_prime[(j, i)];
                        residual += (rhs - ap * p_old).abs();
                        field.p_prime[(j, i)] = (1.0 - omega) * p_old + omega * rhs / ap;
                    }
                }
                if residual < inner_stop {
                    break;
                }
            }
        }
        Ok(())
    }

    /// The apply half of a projection: correct the fluid faces from `u*`
    /// with the `p'` in `field.p_prime` (outlet faces against `p' = 0`
    /// outside), add `p'` to `p` on the fluid cells, and return the
    /// normalised mass imbalance of the corrected field.
    pub(crate) fn apply_correction(&self, field: &mut FlowField, dt: f64) -> f64 {
        let (nx, ny, dx, dy) = field.grid_info();
        let rho = self.config.density;
        let b = self.parameters.boundaries;
        let outlet = SideBoundary::PressureOutlet;
        let reference_flux = rho * self.config.reference_velocity * self.config.reference_length;

        // Correct exactly the faces the equations treated as correctable:
        // fluid interior faces, and outlet faces against p' = 0 outside.
        for j in 0..ny {
            for i in 1..nx {
                if self.u_is_fluid(j, i) {
                    let dp_dx = (field.p_prime[(j, i)] - field.p_prime[(j, i - 1)]) / dx;
                    field.u[(j, i)] = field.u_star[(j, i)] - (dt / rho) * dp_dx;
                }
            }
            if b.left == outlet {
                let dp_dx = (field.p_prime[(j, 0)] - 0.0) / (0.5 * dx);
                field.u[(j, 0)] = field.u_star[(j, 0)] - (dt / rho) * dp_dx;
            }
            if b.right == outlet {
                let dp_dx = (0.0 - field.p_prime[(j, nx - 1)]) / (0.5 * dx);
                field.u[(j, nx)] = field.u_star[(j, nx)] - (dt / rho) * dp_dx;
            }
        }
        for i in 0..nx {
            for j in 1..ny {
                if self.v_is_fluid(j, i) {
                    let dp_dy = (field.p_prime[(j, i)] - field.p_prime[(j - 1, i)]) / dy;
                    field.v[(j, i)] = field.v_star[(j, i)] - (dt / rho) * dp_dy;
                }
            }
            if b.bottom == outlet {
                let dp_dy = (field.p_prime[(0, i)] - 0.0) / (0.5 * dy);
                field.v[(0, i)] = field.v_star[(0, i)] - (dt / rho) * dp_dy;
            }
            if b.top == outlet {
                let dp_dy = (0.0 - field.p_prime[(ny - 1, i)]) / (0.5 * dy);
                field.v[(ny, i)] = field.v_star[(ny, i)] - (dt / rho) * dp_dy;
            }
        }

        for j in 0..ny {
            for i in 0..nx {
                if self.cell_is_fluid(j, i) {
                    field.p[(j, i)] += field.p_prime[(j, i)];
                }
            }
        }

        let mut mass_imbalance = 0.0;
        for j in 0..ny {
            for i in 0..nx {
                if !self.cell_is_fluid(j, i) {
                    continue;
                }
                let divergence_flux = rho
                    * ((field.u[(j, i + 1)] - field.u[(j, i)]) * dy
                        + (field.v[(j + 1, i)] - field.v[(j, i)]) * dx);
                mass_imbalance += divergence_flux.abs();
            }
        }

        if reference_flux > 0.0 {
            mass_imbalance / reference_flux
        } else {
            mass_imbalance
        }
    }
}