//! The overset (chimera) hybrid: a curvilinear patch around the body over
//! the fixed background grid (`docs/overset_metal_campaign.md` §2, A-P2).
//!
//! [`OversetPisoSolver`] drives an [`EmbeddedPisoSolver`] (the background,
//! its mask from the overlap classification: holes, fringe) and a
//! [`CurvilinearPisoSolver`] (the patch, its outer row an acceptor ring)
//! through one PISO step together:
//!
//! 1. patch predictor (`begin_step`: mesh swap if the patch moves, fluxes,
//!    `u*`, matrix); if the mesh moved, the overlap is rebuilt and the
//!    fringe re-stamped from the patch's previous corrected field (values
//!    identical when the mesh did not move — the moving path with a
//!    stationary patch is the static path to the bit);
//! 2. background predictor (`begin_step`);
//! 3. per corrector, alternating Schwarz on `p'` with Dirichlet transmission
//!    both ways: fringe `p'` ← patch `p'` (zero in the first round),
//!    background solve; acceptor `p'` ← background `p'`, patch solve; stop
//!    when the exchanged values change by less than `schwarz_tolerance ×
//!    max |p'|`; then ONE application of the correction on each mesh;
//! 4. the exchange for the next step, from the corrected fields (the
//!    boundary-history principle): prescribed background faces and fringe
//!    `p` from the patch, acceptor `u, v, p` from the background.
//!
//! The overlap mass defect — continuity is enforced on neither the fringe
//! nor the acceptor cells — is measured every step on both sides.

pub mod overlap;

pub use overlap::{Acceptor, CellClass, DualDonor, FringeEntry, LatticeDonor, OverlapMap};

use crate::error::{CfdError, CfdResult};
use crate::mesh::PatchMesh;
use crate::solvers::incompressible::curvilinear::{
    CurvilinearPisoSolver, CurvilinearSolverState, PatchField,
};
use crate::solvers::incompressible::embedded::{EmbeddedPisoSolver, EmbeddedSolverState};
use crate::solvers::incompressible::flow_field::FlowField;

/// Parameters of the composite step.
#[derive(Debug, Clone)]
pub struct OversetParameters {
    /// Projections per step (both meshes).
    pub corrector_steps: usize,
    /// Schwarz stop: the largest change of an exchanged `p'` value between
    /// rounds, relative to the largest exchanged `|p'|`.
    pub schwarz_tolerance: f64,
    /// Round cap per corrector.
    pub max_rounds: usize,
    /// Anderson-acceleration depth on the acceptor `p'` vector (0 = plain
    /// alternating Schwarz). Plain Schwarz converges at ≈ 0.82 per round
    /// here — the patch's wall is Neumann for `p'`, so its pressure level
    /// is pinned only through the fringe and decays weakly across the
    /// overlap (measured: 4.5e-3 relative after 20 rounds).
    pub anderson_depth: usize,
    /// Stall detection: stop a corrector's rounds when the exchanged change
    /// has not fallen by 30% over this many rounds (0 = off, the default).
    /// Right for a steady march, where the step's `p'` sits at the inner
    /// solvers' noise floor and rounds cannot help (at n = 64/128 the cap
    /// was burned on 4% of the steps; a 7.5 h n = 128 march); WRONG for a
    /// transient — on the falsifier plate it stopped rounds that were still
    /// converging and raised the max force spike from 594 to 4044 N/m.
    pub stall_rounds: usize,
    /// Patch rows below the acceptor row kept non-hole
    /// (`overlap::DEFAULT_OVERLAP_ROWS`).
    pub overlap_rows: usize,
}

impl Default for OversetParameters {
    fn default() -> Self {
        Self {
            corrector_steps: 2,
            schwarz_tolerance: 1e-3,
            max_rounds: 20,
            anderson_depth: 3,
            stall_rounds: 0,
            overlap_rows: overlap::DEFAULT_OVERLAP_ROWS,
        }
    }
}

/// The two fields.
#[derive(Debug, Clone)]
pub struct OversetField {
    /// Background staggered field.
    pub background: FlowField,
    /// Patch collocated field.
    pub patch: PatchField,
}

/// What one composite step reports.
#[derive(Debug, Clone)]
pub struct OversetResult {
    /// Schwarz rounds per corrector.
    pub rounds: Vec<usize>,
    /// Whether every corrector's Schwarz iteration met its stop.
    pub schwarz_converged: bool,
    /// Correctors whose rounds stalled at the inner solvers' noise floor
    /// (no progress over three rounds) and were stopped there.
    pub schwarz_stalled: usize,
    /// Background mass residual after the last corrector (its own measure).
    pub background_residual: f64,
    /// Largest patch cell imbalance after the last corrector.
    pub patch_max_divergence: f64,
    /// Patch BiCGSTAB iterations, summed.
    pub patch_poisson_iterations: usize,
    /// Whether every patch pressure solve converged.
    pub patch_converged: bool,
    /// `Σ_fringe |div|` on the background (volume flux).
    pub background_mass_defect: f64,
    /// `Σ_acceptors |div|` on the patch (volume flux).
    pub patch_mass_defect: f64,
    /// `Σ |F|` over the acceptor–interior faces: the flux scale through the overlap.
    pub overlap_flux_scale: f64,
    /// Background cells that changed class in this step's overlap rebuild.
    pub reclassified_cells: usize,
    /// Background cells that jumped hole → active (patch moved > 1 cell).
    pub fresh_cells: usize,
}

/// Restorable state of the composite (the coupling re-runs a step).
#[derive(Clone)]
pub struct OversetSolverState {
    background: EmbeddedSolverState,
    patch: CurvilinearSolverState,
    overlap: OverlapMap,
    pending: Option<PatchMesh>,
    acceptor_warm: Vec<f64>,
}

/// The composite solver.
pub struct OversetPisoSolver {
    background: EmbeddedPisoSolver,
    patch: CurvilinearPisoSolver,
    overlap: OverlapMap,
    params: OversetParameters,
    grid: (usize, usize, f64, f64),
    pending: Option<PatchMesh>,
    /// The first corrector's converged acceptor `p'` of the previous step:
    /// the temporal warm start (the correction is correlated step to step).
    acceptor_warm: Vec<f64>,
}

/// Anderson acceleration of a fixed-point iteration `x ← G(x)` with residual
/// `r = G(x) − x`, depth `m` (least squares by normal equations, tiny).
struct Anderson {
    m: usize,
    xs: Vec<Vec<f64>>,
    rs: Vec<Vec<f64>>,
}

impl Anderson {
    fn new(m: usize) -> Self {
        Self {
            m,
            xs: Vec::new(),
            rs: Vec::new(),
        }
    }

    /// Next iterate from the current `x` and its residual `r`.
    fn next(&mut self, x: &[f64], r: &[f64]) -> Vec<f64> {
        self.xs.push(x.to_vec());
        self.rs.push(r.to_vec());
        while self.xs.len() > self.m + 1 {
            self.xs.remove(0);
            self.rs.remove(0);
        }
        let k = self.xs.len() - 1;
        if self.m == 0 || k == 0 {
            return x.iter().zip(r).map(|(a, b)| a + b).collect();
        }
        // Differences ΔR_i = r_{i+1} − r_i, ΔX_i = x_{i+1} − x_i, i = 0..k.
        let n = x.len();
        let dr: Vec<Vec<f64>> = (0..k)
            .map(|i| (0..n).map(|j| self.rs[i + 1][j] - self.rs[i][j]).collect())
            .collect();
        let dx: Vec<Vec<f64>> = (0..k)
            .map(|i| (0..n).map(|j| self.xs[i + 1][j] - self.xs[i][j]).collect())
            .collect();
        // Normal equations (ΔRᵀΔR) γ = ΔRᵀ r, Tikhonov-regularised.
        let mut a = vec![vec![0.0; k]; k];
        let mut b = vec![0.0; k];
        let mut trace = 0.0;
        for i in 0..k {
            for l in 0..k {
                a[i][l] = dr[i].iter().zip(&dr[l]).map(|(p, q)| p * q).sum();
            }
            trace += a[i][i];
            b[i] = dr[i].iter().zip(r).map(|(p, q)| p * q).sum();
        }
        for i in 0..k {
            a[i][i] += 1e-10 * trace.max(1e-300);
        }
        let gamma = solve_small(a, b);
        (0..n)
            .map(|j| {
                let mut v = x[j] + r[j];
                for i in 0..k {
                    v -= gamma[i] * (dx[i][j] + dr[i][j]);
                }
                v
            })
            .collect()
    }
}

/// Gaussian elimination with partial pivoting on a tiny dense system.
fn solve_small(mut a: Vec<Vec<f64>>, mut b: Vec<f64>) -> Vec<f64> {
    let k = b.len();
    for col in 0..k {
        let piv = (col..k)
            .max_by(|&i, &j| a[i][col].abs().partial_cmp(&a[j][col].abs()).unwrap())
            .unwrap();
        a.swap(col, piv);
        b.swap(col, piv);
        let d = a[col][col];
        if d.abs() <= 1e-300 {
            continue;
        }
        for row in col + 1..k {
            let f = a[row][col] / d;
            for c in col..k {
                a[row][c] -= f * a[col][c];
            }
            b[row] -= f * b[col];
        }
    }
    let mut x = vec![0.0; k];
    for row in (0..k).rev() {
        let mut v = b[row];
        for c in row + 1..k {
            v -= a[row][c] * x[c];
        }
        x[row] = if a[row][row].abs() > 1e-300 {
            v / a[row][row]
        } else {
            0.0
        };
    }
    x
}

impl OversetPisoSolver {
    /// Compose a configured background (sides, sources, boundary data) and
    /// a configured patch (its `Inner` side velocity via
    /// `set_side_velocity`) on a background of `nx × ny` cells, spacing
    /// `dx, dy`. Builds the overlap, installs the hole/fringe mask on the
    /// background and the acceptor ring on the patch.
    pub fn new(
        mut background: EmbeddedPisoSolver,
        mut patch: CurvilinearPisoSolver,
        grid: (usize, usize, f64, f64),
        params: OversetParameters,
    ) -> CfdResult<Self> {
        let (nx, ny, dx, dy) = grid;
        let overlap = OverlapMap::build(patch.mesh(), nx, ny, dx, dy, params.overlap_rows)?;
        background.set_overlap(overlap.background_mask(), overlap.fringe_flags());
        // Each background solve must be accurate below the Schwarz stop, or
        // the exchanged values never settle (measured at the default 1e-2).
        background.set_inner_stop_factor(0.1 * params.schwarz_tolerance);
        patch.set_acceptor_ring(true);
        Ok(Self {
            background,
            patch,
            overlap,
            params,
            grid,
            pending: None,
            acceptor_warm: Vec::new(),
        })
    }

    /// The background solver.
    pub fn background(&self) -> &EmbeddedPisoSolver {
        &self.background
    }
    /// The patch solver.
    pub fn patch(&self) -> &CurvilinearPisoSolver {
        &self.patch
    }
    /// The current overlap map.
    pub fn overlap(&self) -> &OverlapMap {
        &self.overlap
    }
    /// Parameters.
    pub fn parameters(&self) -> &OversetParameters {
        &self.params
    }
    /// The patch's time (the background's agrees).
    pub fn time(&self) -> f64 {
        self.patch.time()
    }

    /// Name the patch geometry the next step ends on (see
    /// `CurvilinearPisoSolver::set_mesh`; same topology; a fraction of a
    /// background cell per step).
    pub fn set_patch_mesh(&mut self, next: PatchMesh) -> CfdResult<()> {
        self.patch.set_mesh(next.clone())?;
        self.pending = Some(next);
        Ok(())
    }

    /// Stamp both exchanges from the current fields and initialise the
    /// background (boundary data at the current time). Call once after the
    /// fields are set.
    pub fn initialize(&mut self, field: &mut OversetField) -> CfdResult<()> {
        self.exchange(field);
        self.background.initialize(&mut field.background)
    }

    /// Subtract the mean of `p'` over the active background cells (the
    /// composite level pin; the fringe cells' Dirichlet values shift with
    /// it so the face corrections across active–fringe faces are unchanged).
    fn remove_background_mean(&self, field: &mut FlowField) {
        let (nx, ny, _, _) = self.grid;
        let (mut sum, mut count) = (0.0, 0usize);
        for j in 0..ny {
            for i in 0..nx {
                if self.overlap.class(j, i) == CellClass::Active {
                    sum += field.p_prime[(j, i)];
                    count += 1;
                }
            }
        }
        if count == 0 {
            return;
        }
        let mean = sum / count as f64;
        for j in 0..ny {
            for i in 0..nx {
                if self.overlap.class(j, i) != CellClass::Hole {
                    field.p_prime[(j, i)] -= mean;
                }
            }
        }
    }

    /// The exchange for the next step: prescribed background faces and
    /// fringe `p` from the patch's cell field; acceptor `u, v, p` from the
    /// background.
    fn exchange(&self, field: &mut OversetField) {
        self.overlap
            .stamp_fringe_faces(&mut field.background, &field.patch.u, &field.patch.v);
        let p = self.overlap.fringe_cell_values(&field.patch.p);
        self.overlap.stamp_fringe_cells(&mut field.background.p, &p);
        let acc = self
            .overlap
            .acceptor_values(&field.background, &field.background.p);
        self.patch.stamp_acceptors(&mut field.patch, &acc);
    }

    /// Capture the state.
    pub fn snapshot(&self) -> OversetSolverState {
        OversetSolverState {
            background: self.background.snapshot(),
            patch: self.patch.snapshot(),
            overlap: self.overlap.clone(),
            pending: self.pending.clone(),
            acceptor_warm: self.acceptor_warm.clone(),
        }
    }

    /// Restore a captured state.
    pub fn restore(&mut self, state: &OversetSolverState) {
        self.background.restore(&state.background);
        self.patch.restore(&state.patch);
        self.overlap = state.overlap.clone();
        self.pending = state.pending.clone();
        self.acceptor_warm = state.acceptor_warm.clone();
        // The background's mask and fringe flags come back with its own
        // state; its fringe Dirichlet data are transient (set every round).
    }

    /// Advance both meshes one step of `dt`.
    pub async fn advance(&mut self, field: &mut OversetField, dt: f64) -> CfdResult<OversetResult> {
        let (nx, ny, dx, dy) = self.grid;

        // 1. If the patch moves, the overlap follows the NEXT mesh before any
        //    predictor runs: the background is reclassified (fresh cells
        //    filled from neighbours), and the fringe is re-stamped from the
        //    patch's previous CORRECTED cell field through the new donors —
        //    the same numbers as the static path's exchange when the mesh did
        //    not move, so a stationary patch through this path is the static
        //    path to the bit.
        let mut reclassified = 0usize;
        let mut fresh = 0usize;
        if self.pending.take().is_some() {
            let next = self
                .patch
                .next_mesh()
                .expect("set_mesh named the next mesh");
            let new = OverlapMap::build(next, nx, ny, dx, dy, self.params.overlap_rows)?;
            for j in 0..ny {
                for i in 0..nx {
                    let (was, now) = (self.overlap.class(j, i), new.class(j, i));
                    if was != now {
                        reclassified += 1;
                        if was == CellClass::Hole && now == CellClass::Active {
                            fresh += 1;
                            // Fill from the neighbours that already carry a
                            // pressure (the embedded solver's fresh-cell rule).
                            let (mut sum, mut count) = (0.0, 0usize);
                            for (jj, ii) in [
                                (j, i + 1),
                                (j, i.wrapping_sub(1)),
                                (j + 1, i),
                                (j.wrapping_sub(1), i),
                            ] {
                                if jj < ny
                                    && ii < nx
                                    && self.overlap.class(jj, ii) != CellClass::Hole
                                {
                                    sum += field.background.p[(jj, ii)];
                                    count += 1;
                                }
                            }
                            if count > 0 {
                                field.background.p[(j, i)] = sum / count as f64;
                            }
                        }
                    }
                }
            }
            self.overlap = new;
            self.background
                .set_overlap(self.overlap.background_mask(), self.overlap.fringe_flags());
            self.overlap
                .stamp_fringe_faces(&mut field.background, &field.patch.u, &field.patch.v);
            let p = self.overlap.fringe_cell_values(&field.patch.p);
            self.overlap.stamp_fringe_cells(&mut field.background.p, &p);
        }

        // 1b. Patch predictor (swaps in the pending mesh).
        let patch_start = self.patch.begin_step(&mut field.patch, dt)?;

        // 2. Background predictor.
        let bg_start = self.background.begin_step(&mut field.background, dt)?;

        // 3. Correctors: alternating Schwarz on the acceptor p' vector `a`
        //    (patch solve with Dirichlet a → fringe p' → background solve →
        //    G(a)), Anderson-accelerated, warm-started in the first
        //    corrector from the previous step's converged `a`.
        let fringe_cells: Vec<(usize, usize)> = self
            .overlap
            .fringe_cells
            .iter()
            .map(|e| (e.j, e.i))
            .collect();
        let n_acc = self.overlap.acceptors.len();
        let mut rounds = Vec::with_capacity(self.params.corrector_steps);
        let mut schwarz_converged = true;
        let mut schwarz_stalled = 0usize;
        let mut residual_history = Vec::new();
        let mut final_residual = f64::INFINITY;
        let mut patch_iterations = 0usize;
        let mut patch_converged = true;
        let trace_rounds = std::env::var("RTX_OVERSET_ROUNDS").is_ok();
        // The stop is relative to the STEP's pressure-correction scale (the
        // first corrector's): a later corrector's own p' is a mop-up of the
        // size of the inner solver's absolute-stop noise (≈ 1e-9 in flux
        // units is ≈ 1e-9/dt² in pressure — 0.02 here, the whole second
        // correction), so a stop relative to its own magnitude cannot be
        // met and would only feed noise to the acceleration.
        let mut step_scale = 0.0_f64;
        for corrector in 0..self.params.corrector_steps.max(1) {
            let mut a: Vec<f64> = if corrector == 0 && self.acceptor_warm.len() == n_acc {
                self.acceptor_warm.clone()
            } else {
                vec![0.0; n_acc]
            };
            let mut anderson = Anderson::new(self.params.anderson_depth);
            let mut patch_pc = vec![0.0; self.patch.mesh().cell_count()];
            let mut done = false;
            let mut used = 0usize;
            let mut history: Vec<f64> = Vec::new();
            for round in 0..self.params.max_rounds.max(1) {
                used = round + 1;
                // Patch with Dirichlet a.
                self.patch.set_acceptor_correction(&a);
                match self.patch.solve_correction(&field.patch, dt) {
                    Some((pc, out)) => {
                        patch_iterations += out.iterations;
                        patch_converged &= out.converged;
                        patch_pc = pc;
                    }
                    None => patch_pc.iter_mut().for_each(|v| *v = 0.0),
                }
                // Background with the fringe p' from the patch.
                let fringe_vals = self.overlap.fringe_cell_values(&patch_pc);
                self.background
                    .set_fringe_correction(&fringe_cells, &fringe_vals);
                self.background.solve_correction(
                    &mut field.background,
                    dt,
                    corrector == 0 || round > 0,
                )?;
                // Pin the composite level: the coupled p' problem is pure
                // Neumann (walls everywhere), so its constant mode is
                // undamped by the rounds and the warm start hands each
                // step's level to the next — measured as a background
                // pressure of 1e7 growing 5e4 per step on the falsifier.
                // A constant in p' moves no velocity; zero mean over the
                // active cells pins it.
                self.remove_background_mean(&mut field.background);
                let g = self.overlap.acceptor_scalar(&field.background.p_prime);
                let r: Vec<f64> = g.iter().zip(&a).map(|(x, y)| x - y).collect();
                let g_max = g.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
                step_scale = step_scale.max(g_max);
                let scale = step_scale.max(1e-300);
                let change = r.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
                if trace_rounds {
                    println!(
                        "      corrector {corrector} round {round}: |G(a) − a| {change:.3e} / step scale {scale:.3e} = {:.3e} (max|G(a)| {g_max:.3e})",
                        change / scale
                    );
                }
                if change <= self.params.schwarz_tolerance * scale {
                    done = true;
                    break;
                }
                // Stall: no progress over two rounds means the exchanged
                // values sit at the inner solvers' noise floor (the step's
                // own p' is that small near a steady state); more rounds
                // cannot help — measured as 6560 of 150k steps burning the
                // 20-round cap at n = 64, and a 7.5 h n = 128 march.
                history.push(change);
                let k = self.params.stall_rounds;
                if k > 0 && history.len() > k && change > 0.7 * history[history.len() - 1 - k] {
                    schwarz_stalled += 1;
                    done = true;
                    break;
                }
                let next = anderson.next(&a, &r);
                // Guard: an extrapolation far beyond the data is noise-driven;
                // take the plain Schwarz step instead.
                let next_max = next.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
                a = if next_max > 4.0 * g_max.max(scale) {
                    g.clone()
                } else {
                    next
                };
            }
            schwarz_converged &= done;
            rounds.push(used);
            if corrector == 0 {
                self.acceptor_warm = a.clone();
            }
            // One application on each mesh (the patch with the a it solved).
            let residual = self.background.apply_correction(&mut field.background, dt);
            residual_history.push(residual);
            final_residual = residual;
            self.patch
                .apply_correction_pub(&mut field.patch, &patch_pc, dt);
            if corrector + 1 < self.params.corrector_steps.max(1) {
                field.background.copy_to_starred();
            }
        }
        let patch_max_div = self.patch.max_divergence_pub(&field.patch.flux);

        // 4. The exchange for the next step, then the clocks.
        self.exchange(field);
        let outer = self.overlap.acceptor_outer_velocity(&field.background);
        let (patch_defect, scale) = self.patch.acceptor_mass_defect(&field.patch, &outer);
        let bg_defect = self.overlap.background_mass_defect(&field.background);
        let bg = self.background.end_step(
            &mut field.background,
            &bg_start,
            residual_history,
            rounds.len(),
            final_residual,
        );
        self.patch.end_step(
            &patch_start,
            rounds.len(),
            patch_max_div,
            patch_iterations,
            patch_converged,
        );
        if (self.patch.time() - self.background.time()).abs()
            > 1e-12 * self.patch.time().abs().max(1.0)
        {
            return Err(CfdError::invalid_parameter(
                "overset: background and patch clocks disagree".to_string(),
            ));
        }
        Ok(OversetResult {
            rounds,
            schwarz_converged,
            schwarz_stalled,
            background_residual: bg.solver_result.final_residual,
            patch_max_divergence: patch_max_div,
            patch_poisson_iterations: patch_iterations,
            patch_converged,
            background_mass_defect: bg_defect,
            patch_mass_defect: patch_defect,
            overlap_flux_scale: scale,
            reclassified_cells: reclassified,
            fresh_cells: fresh,
        })
    }
}