//! PISO (Pressure-Implicit with Splitting of Operators) algorithm
//!
//! A transient pressure-velocity coupling method: one explicit momentum
//! predictor per time step, followed by pressure-correction (projection)
//! steps that make the velocity field divergence-free. Marching it in time
//! with a steady forcing converges to the steady discrete solution, which is
//! how `tests/mms_piso.rs` verifies it against a manufactured solution.
//!
//! # Grid convention
//!
//! The staggered layout is the one `FlowField` and the SIMPLE solver define:
//! `u[(j, i)]` lives at `(i dx, (j + 0.5) dy)` for `i = 0..=nx`, `v[(j, i)]`
//! at `((i + 0.5) dx, j dy)` for `j = 0..=ny`, `p[(j, i)]` at cell centres.
//! The only velocity components on a domain boundary are the normal ones —
//! u faces `i = 0`, `i = nx` and v faces `j = 0`, `j = ny`. Everything else,
//! including the near-wall lines, is an unknown and is updated every step.
//!
//! # History
//!
//! The previous implementation had never had a test of any kind, and
//! inspection plus the manufactured-solution harness found the same defect
//! species the SIMPLE census recorded:
//!
//! - **The pressure correction had its sign inverted.** It solved
//!   `-lap(p') = +rho div(u*) / dt` and then corrected with
//!   `u = u* - (dt/rho) grad(p')`, so each projection *doubled* the
//!   divergence instead of removing it.
//! - The momentum sweeps froze the near-wall lines (`1..ny-1`), imposing the
//!   wall half a cell inside the domain, and the pressure correction skipped
//!   the outer ring of cells (`1..nx-1`), so ring cells had no continuity
//!   equation — both exactly as in SIMPLE before its repair.
//! - The predictor read neighbours that the same sweep had already
//!   overwritten, so the "explicit" step mixed old and new values in sweep
//!   order.
//! - The pressure gradient was dropped entirely on the last interior face
//!   (`if i < nx - 1 { ... } else { 0.0 }`).
//! - Convective face fluxes fell back to the centre value at the sweep edges
//!   instead of using the prescribed boundary faces that exist there.

use super::poisson::{
    MgPrecision, MultigridParameters, PoissonProblem, PoissonSolverKind, solve_multigrid_pcg,
};
use super::{BoundaryConditions, FlowField, IncompressibleSolver, SolverResult};
use crate::{CfdConfig, CfdResult};
use async_trait::async_trait;

/// Parameters for PISO algorithm
#[derive(Debug, Clone)]
pub struct PisoParameters {
    /// Number of corrector steps (typically 2-3). With an explicit predictor
    /// the first projection already removes the divergence; the extra
    /// correctors are cheap no-ops kept for the algorithm's shape.
    pub corrector_steps: usize,
    /// Time step size. The predictor is explicit, so stability requires
    /// `dt < dx^2 / (4 nu)` and `dt < dx / |u|_max`.
    pub time_step: f64,
    /// Convergence tolerance on the normalised mass imbalance after
    /// correction.
    pub tolerance: f64,
    /// Inner solver of the pressure-correction system (default
    /// [`PoissonSolverKind::Sor`]). Both solve the same system to the same
    /// true-residual stop; multigrid's cost is mesh-independent.
    pub poisson_solver: PoissonSolverKind,
    /// Precision of the multigrid V-cycle (default [`MgPrecision::F64`] =
    /// bit-identical; `F32` = the M1 precision probe, no effect with SOR).
    pub poisson_precision: MgPrecision,
}

impl Default for PisoParameters {
    fn default() -> Self {
        Self {
            corrector_steps: 2,
            time_step: 0.001,
            tolerance: 1e-6,
            poisson_solver: PoissonSolverKind::Sor,
            poisson_precision: MgPrecision::F64,
        }
    }
}

/// Result of PISO algorithm execution
#[derive(Debug, Clone)]
pub struct PisoResult {
    /// Base solver result information
    pub solver_result: SolverResult,
    /// Number of corrector steps performed
    pub corrector_steps_performed: usize,
}

/// PISO algorithm implementation
pub struct PisoSolver {
    config: CfdConfig,
    parameters: PisoParameters,
    /// Optional volumetric momentum source `f(x, y) -> (f_x, f_y)`, per unit
    /// volume — the hook a manufactured solution enters through, exactly as
    /// on [`super::SimpleSolver`].
    #[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 control volumes need the tangential
    /// wall velocity for their half-cell diffusion term, and on this
    /// staggered layout there is nowhere to store it. Falls back to the value
    /// on the near-wall line itself when unset.
    #[allow(clippy::type_complexity)]
    wall_velocity: Option<Box<dyn Fn(f64, f64) -> (f64, f64) + Send + Sync>>,
}

impl PisoSolver {
    /// Create new PISO solver
    pub fn new(config: CfdConfig, parameters: PisoParameters) -> CfdResult<Self> {
        config.validate()?;

        Ok(Self {
            config,
            parameters,
            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));
    }

    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_old[(j, i)], |f| f(i as f64 * dx, y_wall).0)
    }

    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_old[(j, i)], |f| f(x_wall, j as f64 * dy).1)
    }

    /// Upwind face value: the value carried across the face is the one from
    /// the side the flow comes from.
    fn upwind(face_velocity: f64, upstream: f64, downstream: f64) -> f64 {
        if face_velocity >= 0.0 {
            upstream
        } else {
            downstream
        }
    }

    /// Explicit momentum predictor:
    /// `u* = u_old + dt (-conv + nu lap(u) - grad(p)/rho + f/rho)`,
    /// every term evaluated from `u_old`/`v_old`, so the step is genuinely
    /// explicit and independent of sweep order.
    fn momentum_predictor(&self, flow_field: &mut FlowField, dt: f64) -> CfdResult<()> {
        let (nx, ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;
        let nu = self.config.viscosity / rho;

        // u faces: every row is an unknown; only i = 0 and i = nx are
        // boundary data.
        for j in 0..ny {
            for i in 1..nx {
                let uo = &flow_field.u_old;
                let vo = &flow_field.v_old;
                let u_p = uo[(j, i)];

                // Cell-centre velocities on the east/west faces of the u
                // control volume. The neighbours i-1 and i+1 always exist:
                // they are boundary faces at the sweep edges, which hold
                // prescribed data rather than needing a fallback.
                let ue_face = 0.5 * (uo[(j, i)] + uo[(j, i + 1)]);
                let uw_face = 0.5 * (uo[(j, i - 1)] + uo[(j, i)]);

                let south_is_wall = j == 0;
                let north_is_wall = j + 1 == ny;

                // Transverse face velocities; a solid wall passes no mass.
                let vn_face = if north_is_wall {
                    0.0
                } else {
                    0.5 * (vo[(j + 1, i - 1)] + vo[(j + 1, i)])
                };
                let vs_face = if south_is_wall {
                    0.0
                } else {
                    0.5 * (vo[(j, i - 1)] + vo[(j, i)])
                };

                let conv_x = (ue_face * Self::upwind(ue_face, uo[(j, i)], uo[(j, i + 1)])
                    - uw_face * Self::upwind(uw_face, uo[(j, i - 1)], uo[(j, i)]))
                    / dx;
                let conv_y = (vn_face
                    * if north_is_wall {
                        0.0
                    } else {
                        Self::upwind(vn_face, uo[(j, i)], uo[(j + 1, i)])
                    }
                    - vs_face
                        * if south_is_wall {
                            0.0
                        } else {
                            Self::upwind(vs_face, uo[(j - 1, i)], uo[(j, i)])
                        })
                    / dy;

                let diff_x = nu * (uo[(j, i + 1)] - 2.0 * u_p + uo[(j, i - 1)]) / (dx * dx);

                // Wall-adjacent diffusive fluxes act over half a cell: the
                // node beyond the wall face is the wall itself, dy/2 away.
                let flux_north = if north_is_wall {
                    nu * (self.u_wall(flow_field, i, j, ny as f64 * dy, dx) - u_p) / (0.5 * dy)
                } else {
                    nu * (uo[(j + 1, i)] - u_p) / dy
                };
                let flux_south = if south_is_wall {
                    nu * (u_p - self.u_wall(flow_field, i, j, 0.0, dx)) / (0.5 * dy)
                } else {
                    nu * (u_p - uo[(j - 1, i)]) / dy
                };
                let diff_y = (flux_north - flux_south) / dy;

                // The pressure gradient acts on every unknown face — dropping
                // it anywhere solves a different equation there.
                let pressure_gradient =
                    -(flow_field.p[(j, i)] - flow_field.p[(j, i - 1)]) / (rho * dx);

                let body_force = self
                    .momentum_source
                    .as_ref()
                    .map_or(0.0, |f| f(i as f64 * dx, (j as f64 + 0.5) * dy).0 / rho);

                flow_field.u[(j, i)] = u_p
                    + dt * (-conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force);
            }
        }

        // v faces, mirrored.
        for j in 1..ny {
            for i in 0..nx {
                let uo = &flow_field.u_old;
                let vo = &flow_field.v_old;
                let v_p = vo[(j, i)];

                let vn_face = 0.5 * (vo[(j, i)] + vo[(j + 1, i)]);
                let vs_face = 0.5 * (vo[(j - 1, i)] + vo[(j, i)]);

                let west_is_wall = i == 0;
                let east_is_wall = i + 1 == nx;

                let ue_face = if east_is_wall {
                    0.0
                } else {
                    0.5 * (uo[(j - 1, i + 1)] + uo[(j, i + 1)])
                };
                let uw_face = if west_is_wall {
                    0.0
                } else {
                    0.5 * (uo[(j - 1, i)] + uo[(j, i)])
                };

                let conv_y = (vn_face * Self::upwind(vn_face, vo[(j, i)], vo[(j + 1, i)])
                    - vs_face * Self::upwind(vs_face, vo[(j - 1, i)], vo[(j, i)]))
                    / dy;
                let conv_x = (ue_face
                    * if east_is_wall {
                        0.0
                    } else {
                        Self::upwind(ue_face, vo[(j, i)], vo[(j, i + 1)])
                    }
                    - uw_face
                        * if west_is_wall {
                            0.0
                        } else {
                            Self::upwind(uw_face, vo[(j, i - 1)], vo[(j, i)])
                        })
                    / dx;

                let diff_y = nu * (vo[(j + 1, i)] - 2.0 * v_p + vo[(j - 1, i)]) / (dy * dy);

                let flux_east = if east_is_wall {
                    nu * (self.v_wall(flow_field, i, j, nx as f64 * dx, dy) - v_p) / (0.5 * dx)
                } else {
                    nu * (vo[(j, i + 1)] - v_p) / dx
                };
                let flux_west = if west_is_wall {
                    nu * (v_p - self.v_wall(flow_field, i, j, 0.0, dy)) / (0.5 * dx)
                } else {
                    nu * (v_p - vo[(j, i - 1)]) / dx
                };
                let diff_x = (flux_east - flux_west) / dx;

                let pressure_gradient =
                    -(flow_field.p[(j, i)] - flow_field.p[(j - 1, i)]) / (rho * dy);

                let body_force = self
                    .momentum_source
                    .as_ref()
                    .map_or(0.0, |f| f((i as f64 + 0.5) * dx, j as f64 * dy).1 / rho);

                flow_field.v[(j, i)] = v_p
                    + dt * (-conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force);
            }
        }

        flow_field.copy_to_starred();
        Ok(())
    }

    /// One projection: solve the pressure-correction Poisson equation and
    /// subtract `(dt/rho) grad(p')` from the predicted velocities, so the
    /// corrected field is discretely divergence-free.
    ///
    /// Continuity is enforced on every cell. A coefficient is zero exactly
    /// when its face is a domain boundary, where the normal velocity is
    /// prescribed and not correctable. With velocity prescribed on the whole
    /// boundary the system is pure Neumann; one cell is anchored to fix the
    /// level, which is legitimate because the source telescopes to the net
    /// boundary flux — zero for a closed box — so exactly one equation is
    /// redundant.
    ///
    /// Returns the normalised mass imbalance of the *corrected* field — what
    /// the projection failed to remove, which is the inner solver's
    /// truncation and is the step's honest convergence measure.
    fn project(&self, flow_field: &mut FlowField, dt: f64) -> CfdResult<f64> {
        let (nx, ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;

        flow_field.p_prime.fill(0.0);

        // Mass imbalance of the predicted field, per cell, as a flux. Its
        // absolute sum is the scale the inner solve converges relative to.
        let mut source_scale = 0.0;
        for j in 0..ny {
            for i in 0..nx {
                let divergence_flux = rho
                    * ((flow_field.u_star[(j, i + 1)] - flow_field.u_star[(j, i)]) * dy
                        + (flow_field.v_star[(j + 1, i)] - flow_field.v_star[(j, i)]) * dx);
                flow_field.sp[(j, i)] = -divergence_flux;
                source_scale += divergence_flux.abs();
            }
        }

        // With the correction `u = u* - (dt/rho) (p'_P - p'_W)/dx`, continuity
        // of the corrected field gives neighbour coefficients
        // `rho (dt/rho) A / delta = dt A / delta`.
        let ae_interior = dt * dy / dx;
        let an_interior = dt * dx / dy;

        // Successive over-relaxation at the optimal Poisson factor
        // `omega = 2 / (1 + sin(pi h))`. Plain Gauss-Seidel contracts the
        // smooth modes by only ~(1 - O(h^2)) per sweep, so on a 64^2 grid a
        // 400-sweep cap left a divergence of ~1e-2 that *grew* with mesh
        // size; SOR brings the contraction to ~(1 - O(h)) and the same
        // tolerance costs tens of sweeps instead of thousands.
        //
        // The inner stop measures the TRUE residual of the pressure-correction
        // equation, `|b + sum(a_nb p'_nb) - a_p p'_P|` summed over cells (one
        // half-sweep lagged). An earlier version summed the per-sweep iterate
        // CHANGE instead — the same movement-not-residual pseudo-criterion the
        // SIMPLE census flagged: slow modes move little per sweep while their
        // residual is still large, so the loop declared victory with an
        // unremoved divergence. The Taylor-Green benchmark caught both.
        // Converge relative to this projection's own source, floored at a
        // tenth of the divergence level the outer corrector loop is checking
        // for: the corrector loop measures the true post-correction
        // divergence and re-projects (compounding the reduction), so the
        // inner solve only needs a solid contraction per pass, not machine zero — which on
        // a long steady march would spend a hundred sweeps per step
        // polishing a correction that is already far below tolerance.
        let reference_flux = rho * self.config.reference_velocity * self.config.reference_length;
        let inner_stop =
            (1e-2 * 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 five-point system the SOR loop below sweeps — the
            // same coefficients, right-hand side, anchor cell and stop —
            // handed to the multigrid-preconditioned CG solver. The SOR
            // loop pins `p'(1, 1) = 0` and solves the remaining equations;
            // on the compatible (closed-box) source that is the singular
            // system's solution shifted to `p'(1, 1) = 0`, which is what
            // `anchor` requests.
            let mut problem = PoissonProblem::new(nx, ny);
            for j in 0..ny {
                for i in 0..nx {
                    let idx = j * nx + i;
                    problem.ae[idx] = if i + 1 == nx { 0.0 } else { ae_interior };
                    problem.aw[idx] = if i == 0 { 0.0 } else { ae_interior };
                    problem.an[idx] = if j + 1 == ny { 0.0 } else { an_interior };
                    problem.as_[idx] = if j == 0 { 0.0 } else { an_interior };
                    problem.rhs[idx] = flow_field.sp[(j, i)];
                }
            }
            let mut p_prime = vec![0.0; nx * ny];
            let solution = solve_multigrid_pcg(
                &problem,
                &mut p_prime,
                &MultigridParameters {
                    precision: self.parameters.poisson_precision,
                    ..MultigridParameters::default()
                },
                inner_stop,
                Some(nx + 1),
            );
            // An unconverged multigrid solve (iteration cap, rounding floor,
            // inconsistent system) is not applied: the SOR sweeps below take
            // over for this projection, so the worst case is the old cost,
            // never a silently wrong correction.
            multigrid_converged = solution.converged;
            if multigrid_converged {
                for j in 0..ny {
                    for i in 0..nx {
                        flow_field.p_prime[(j, i)] = p_prime[j * nx + i];
                    }
                }
            }
        }
        if !multigrid_converged {
            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 i == 1 && j == 1 {
                            flow_field.p_prime[(j, i)] = 0.0;
                            continue;
                        }

                        let ae = if i + 1 == nx { 0.0 } else { ae_interior };
                        let aw = if i == 0 { 0.0 } else { ae_interior };
                        let an = if j + 1 == ny { 0.0 } else { an_interior };
                        let as_ = if j == 0 { 0.0 } else { an_interior };
                        let ap = ae + aw + an + as_;

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

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

        // Correct exactly the faces the equation above treated as
        // correctable: every interior face.
        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;
                flow_field.u[(j, i)] = flow_field.u_star[(j, i)] - (dt / rho) * dp_dx;
            }
        }
        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;
                flow_field.v[(j, i)] = flow_field.v_star[(j, i)] - (dt / rho) * dp_dy;
            }
        }

        // Fold the correction into the pressure. No under-relaxation: PISO
        // corrects rather than iterates within the step.
        for j in 0..ny {
            for i in 0..nx {
                flow_field.p[(j, i)] += flow_field.p_prime[(j, i)];
            }
        }

        // What is left after the correction.
        let mut mass_imbalance = 0.0;
        for j in 0..ny {
            for i in 0..nx {
                let divergence_flux = rho
                    * ((flow_field.u[(j, i + 1)] - flow_field.u[(j, i)]) * dy
                        + (flow_field.v[(j + 1, i)] - flow_field.v[(j, i)]) * dx);
                mass_imbalance += divergence_flux.abs();
            }
        }

        let reference = rho * self.config.reference_velocity * self.config.reference_length;
        Ok(if reference > 0.0 {
            mass_imbalance / reference
        } else {
            mass_imbalance
        })
    }
}

#[async_trait]
impl IncompressibleSolver for PisoSolver {
    type Parameters = PisoParameters;
    type Result = PisoResult;

    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 = std::time::Instant::now();
        let mut residual_history = Vec::new();

        // The state at the start of the step is what the explicit predictor
        // differentiates.
        flow_field.apply_boundary_conditions(boundary_conditions)?;
        flow_field.update_old_values();

        self.momentum_predictor(flow_field, dt)?;

        let mut total_corrector_steps = 0;
        let mut final_residual = f64::INFINITY;
        for _corrector in 0..self.parameters.corrector_steps.max(1) {
            let mass_residual = self.project(flow_field, dt)?;
            residual_history.push(mass_residual);
            final_residual = mass_residual;
            total_corrector_steps += 1;

            if mass_residual < self.parameters.tolerance {
                break;
            }

            // Re-project from the corrected field: with an explicit predictor
            // the second pass mops up the inner solver's truncation.
            flow_field.copy_to_starred();
        }

        let solve_time = start_time.elapsed();
        Ok(PisoResult {
            solver_result: SolverResult {
                converged: final_residual < self.parameters.tolerance,
                iterations: total_corrector_steps,
                final_residual,
                residual_history,
                solve_time,
            },
            corrector_steps_performed: total_corrector_steps,
        })
    }

    async fn solve(
        &mut self,
        flow_field: &mut FlowField,
        boundary_conditions: &BoundaryConditions,
    ) -> CfdResult<Self::Result> {
        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
    }
}