//! Collocated finite-volume PISO on a structured curvilinear patch — the
//! body-fitted half of the overset hybrid (`docs/overset_metal_campaign.md`
//! §2.1, §5.3). Phase P0: a STATIC patch, manufactured-solution gated.
//!
//! The step is the Zang–Street–Koseff (1994) fractional step in the
//! incremental form the background PISO uses:
//!
//! ```text
//! û  = u^n + dt (−C(F^n, u^n) + ν D(u^n) + f/ρ)
//! u* = û − (dt/ρ) G_c p^n                 cell gradient (least squares)
//! F* = interp(û)·S_f − (dt/ρ) L_f(p^n)    compact face operator
//! Σ_f (dt/ρ) L_f(p') = Σ_f F*             pressure correction
//! F  = F* − (dt/ρ) L_f(p'),  u = u* − (dt/ρ) G_c p',  p += p'
//! ```
//!
//! so the face flux — the primary variable for convection, divergence and
//! the next step — sees the whole pressure through one compact operator
//! (no checkerboard), and the cell velocity is a slave. Further correctors
//! re-project the STORED fluxes. `L_f` is the orthogonal difference plus a
//! node-based tangential correction (a 9-point stencil), assembled and
//! applied by the same code.

mod operators;
mod predictor;
mod projection;

pub use operators::Operators;

use crate::mesh::{PatchMesh, PatchSide};
use crate::solvers::incompressible::sparse_bicgstab::CsrMatrix;
use crate::{CfdConfig, CfdResult};

/// What one side of the patch is.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SideBc {
    /// Prescribed velocity (a wall or an inflow): the boundary function
    /// gives the velocity; mass flux `u_b · S_f`; pressure Neumann.
    #[default]
    Velocity,
    /// Open boundary at gauge pressure zero: zero-gradient velocity,
    /// flux from the predictor, corrected by the projection.
    Outlet,
}

/// The four sides. `s_start`/`s_end` are ignored on a periodic patch.
#[derive(Debug, Clone, Copy, Default)]
pub struct PatchBoundaries {
    /// `k = 0` (the wall of an O-grid).
    pub inner: SideBc,
    /// `k = nn`.
    pub outer: SideBc,
    /// `i = 0`.
    pub s_start: SideBc,
    /// `i = ns`.
    pub s_end: SideBc,
}

impl PatchBoundaries {
    /// The condition on a side.
    pub fn get(&self, side: PatchSide) -> SideBc {
        match side {
            PatchSide::Inner => self.inner,
            PatchSide::Outer => self.outer,
            PatchSide::SStart => self.s_start,
            PatchSide::SEnd => self.s_end,
        }
    }
}

/// Convection treatment.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PatchConvection {
    /// First-order upwind on the face fluxes (the background's scheme).
    #[default]
    Upwind,
    /// No convection: the Stokes limit, for the second-order MMS gate.
    None,
}

/// How the across-patch diffusion is time-stepped.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum NormalDiffusion {
    /// Forward Euler, like the background (limit `~Δn²/(4ν)`).
    #[default]
    Explicit,
    /// The orthogonal part of the n-face diffusion implicit along each
    /// s-line (tridiagonal); the tangential part and everything else
    /// explicit.
    LineImplicit,
}

/// Solver parameters.
#[derive(Debug, Clone)]
pub struct CurvilinearParameters {
    /// Projections per step (the first removes the divergence; the rest
    /// mop up inner-solver truncation).
    pub corrector_steps: usize,
    /// Pressure-correction stop: each projection reduces the L1 cell mass
    /// imbalance by this factor (plus a rounding floor of 1e-15 × Σ|F|).
    /// Relative to the incoming divergence, not to the flux scale, so the
    /// stop tightens as the flow settles — an absolute stop leaves a
    /// velocity-noise floor that grows with the grid (measured: |du/dt|
    /// floored at 2e-4 on the 64² Cartesian MMS with `1e-10 × Σ|F|`).
    pub tolerance: f64,
    /// Convection scheme.
    pub convection: PatchConvection,
    /// Across-patch diffusion treatment.
    pub normal_diffusion: NormalDiffusion,
    /// Side conditions.
    pub boundaries: PatchBoundaries,
    /// BiCGSTAB iteration cap.
    pub max_poisson_iterations: usize,
}

impl Default for CurvilinearParameters {
    fn default() -> Self {
        Self {
            corrector_steps: 2,
            tolerance: 1e-4,
            convection: PatchConvection::Upwind,
            normal_diffusion: NormalDiffusion::Explicit,
            boundaries: PatchBoundaries::default(),
            max_poisson_iterations: 5000,
        }
    }
}

/// The patch state: cell-centred velocity and pressure, face mass fluxes.
#[derive(Debug, Clone)]
pub struct PatchField {
    /// Cell x-velocity.
    pub u: Vec<f64>,
    /// Cell y-velocity.
    pub v: Vec<f64>,
    /// Cell pressure.
    pub p: Vec<f64>,
    /// Volume flux through every face, oriented +s / +n.
    pub flux: Vec<f64>,
}

impl PatchField {
    /// Zero field on `mesh`.
    pub fn new(mesh: &PatchMesh) -> Self {
        let n = mesh.cell_count();
        Self {
            u: vec![0.0; n],
            v: vec![0.0; n],
            p: vec![0.0; n],
            flux: vec![0.0; mesh.faces().len()],
        }
    }
}

/// What one step reports.
#[derive(Debug, Clone, Copy)]
pub struct CurvilinearResult {
    /// Correctors performed.
    pub corrector_steps_performed: usize,
    /// Largest cell mass imbalance after the last corrector.
    pub max_divergence: f64,
    /// Pressure-solver iterations, summed over the correctors.
    pub poisson_iterations: usize,
    /// Whether every pressure solve reached its tolerance.
    pub poisson_converged: bool,
    /// Boundary-flux defect removed on a closed patch (zero if an outlet exists).
    pub boundary_flux_adjustment: f64,
}

/// Restorable solver state (the coupling re-runs a step).
#[derive(Debug, Clone)]
pub struct CurvilinearSolverState {
    time: f64,
}

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

/// PISO on a curvilinear patch.
pub struct CurvilinearPisoSolver {
    config: CfdConfig,
    params: CurvilinearParameters,
    mesh: PatchMesh,
    ops: Operators,
    boundary_velocity: Option<VelocityFn>,
    momentum_source: Option<VelocityFn>,
    time: f64,
    matrix: Option<(f64, CsrMatrix, Option<usize>)>,
}

impl CurvilinearPisoSolver {
    /// Build on `mesh`.
    pub fn new(
        config: CfdConfig,
        params: CurvilinearParameters,
        mesh: PatchMesh,
    ) -> CfdResult<Self> {
        let ops = Operators::new(&mesh, &params.boundaries);
        Ok(Self {
            config,
            params,
            mesh,
            ops,
            boundary_velocity: None,
            momentum_source: None,
            time: 0.0,
            matrix: None,
        })
    }

    /// Velocity on every `Velocity` side, `(x, y, t) -> (u, v)`.
    pub fn set_boundary_velocity<F>(&mut self, f: F)
    where
        F: Fn(f64, f64, f64) -> (f64, f64) + Send + Sync + 'static,
    {
        self.boundary_velocity = Some(Box::new(f));
    }
    /// Body force per unit volume, `(x, y, t) -> (fx, fy)`.
    pub fn set_momentum_source<F>(&mut self, f: F)
    where
        F: Fn(f64, f64, f64) -> (f64, f64) + Send + Sync + 'static,
    {
        self.momentum_source = Some(Box::new(f));
    }
    /// The mesh.
    pub fn mesh(&self) -> &PatchMesh {
        &self.mesh
    }
    /// The operators.
    pub fn operators(&self) -> &Operators {
        &self.ops
    }
    /// Parameters.
    pub fn parameters(&self) -> &CurvilinearParameters {
        &self.params
    }
    /// Configuration.
    pub fn config(&self) -> &CfdConfig {
        &self.config
    }
    /// Current time.
    pub fn time(&self) -> f64 {
        self.time
    }
    /// Set the time.
    pub fn set_time(&mut self, t: f64) {
        self.time = t;
    }
    /// Capture the state.
    pub fn snapshot(&self) -> CurvilinearSolverState {
        CurvilinearSolverState { time: self.time }
    }
    /// Restore a captured state.
    pub fn restore(&mut self, state: &CurvilinearSolverState) {
        self.time = state.time;
    }

    pub(crate) fn boundary_velocity(&self, x: f64, y: f64, t: f64) -> (f64, f64) {
        self.boundary_velocity
            .as_ref()
            .map_or((0.0, 0.0), |f| f(x, y, t))
    }
    pub(crate) fn source_at(&self, xy: [f64; 2], t: f64) -> (f64, f64) {
        self.momentum_source
            .as_ref()
            .map_or((0.0, 0.0), |f| f(xy[0], xy[1], t))
    }

    /// Set the cell velocities from a function and make the face fluxes
    /// consistent (interpolated; prescribed on velocity sides).
    pub fn initialize<F>(&self, field: &mut PatchField, velocity: F)
    where
        F: Fn(f64, f64) -> (f64, f64),
    {
        let mesh = &self.mesh;
        for c in 0..mesh.cell_count() {
            let xy = mesh.centre(c);
            let (u, v) = velocity(xy[0], xy[1]);
            field.u[c] = u;
            field.v[c] = v;
        }
        let zero = vec![0.0; mesh.cell_count()];
        field.flux =
            self.predicted_fluxes(&field.u.clone(), &field.v.clone(), &zero, 0.0, self.time);
    }

    /// Advance one step of `dt`.
    pub async fn advance(
        &mut self,
        field: &mut PatchField,
        dt: f64,
    ) -> CfdResult<CurvilinearResult> {
        let t_old = self.time;
        let t_new = t_old + dt;
        let rho = self.config.density;
        let mesh = &self.mesh;

        let (uh, vh) = self.predict(field, dt, t_old);
        let mut flux = self.predicted_fluxes(&uh, &vh, &field.p, dt, t_new);
        let adjustment = self.adjust_boundary_flux(&mut flux);
        for c in 0..mesh.cell_count() {
            let g = self.pressure_gradient(&field.p, c);
            field.u[c] = uh[c] - dt / rho * g[0];
            field.v[c] = vh[c] - dt / rho * g[1];
        }
        field.flux = flux;

        if self.matrix.as_ref().is_none_or(|(d, _, _)| *d != dt) {
            let (m, anchor) = self.assemble_pressure_matrix(dt);
            self.matrix = Some((dt, m, anchor));
        }
        let (_, matrix, anchor) = self.matrix.as_ref().expect("assembled");
        let flux_scale: f64 = field.flux.iter().map(|f| f.abs()).sum::<f64>().max(1e-300);
        let floor = 1e-15 * flux_scale;

        let mut iterations = 0;
        let mut converged = true;
        let mut performed = 0;
        let mut max_div = self.max_divergence(&field.flux);
        for _ in 0..self.params.corrector_steps {
            let incoming = self.divergence_l1(&field.flux);
            if incoming <= floor {
                break;
            }
            let tolerance = self.params.tolerance * incoming + floor;
            let (pc, out) = self.solve_pressure_correction(matrix, *anchor, &field.flux, tolerance);
            iterations += out.iterations;
            converged &= out.converged;
            self.apply_correction(field, &pc, dt);
            performed += 1;
            max_div = self.max_divergence(&field.flux);
        }
        self.time = t_new;
        Ok(CurvilinearResult {
            corrector_steps_performed: performed,
            max_divergence: max_div,
            poisson_iterations: iterations,
            poisson_converged: converged,
            boundary_flux_adjustment: adjustment,
        })
    }
}