//! Unit tests of the multigrid-preconditioned CG Poisson solver — each one
//! a claim a wrong solver fails: manufactured solutions recovered to
//! rounding, anchor semantics exact, inactive cells untouched, the Galerkin
//! coarse operator identical to `R A P`, the V-cycle symmetric, the
//! iteration count grid-independent, and the timing against plain SOR.

use super::*;
use std::f64::consts::PI;

/// Deterministic pseudo-random numbers in `[-1, 1)` (no crate version
/// dependence in the tests).
struct Lcg(u64);
impl Lcg {
    fn next(&mut self) -> f64 {
        self.0 = self
            .0
            .wrapping_mul(6_364_136_223_846_793_005)
            .wrapping_add(1_442_695_040_888_963_407);
        ((self.0 >> 11) as f64 / (1u64 << 53) as f64).mul_add(2.0, -1.0)
    }
}

/// Unit-conductance five-point problem on an `nx × ny` grid: coefficient
/// 1 across every face between two active cells, zero across a domain
/// edge or a non-active face. `dirichlet_sides`: `[left, right, bottom,
/// top]` edges carry `p = 0` half a cell outside (`extra_diag += 2`).
fn assemble(
    nx: usize,
    ny: usize,
    active: impl Fn(usize, usize) -> bool,
    dirichlet_sides: [bool; 4],
) -> PoissonProblem {
    let mut pr = PoissonProblem::new(nx, ny);
    for j in 0..ny {
        for i in 0..nx {
            pr.active[j * nx + i] = active(j, i);
        }
    }
    for j in 0..ny {
        for i in 0..nx {
            let idx = j * nx + i;
            if !pr.active[idx] {
                continue;
            }
            if i + 1 < nx {
                if pr.active[idx + 1] {
                    pr.ae[idx] = 1.0;
                }
            } else if dirichlet_sides[1] {
                pr.extra_diag[idx] += 2.0;
            }
            if i > 0 {
                if pr.active[idx - 1] {
                    pr.aw[idx] = 1.0;
                }
            } else if dirichlet_sides[0] {
                pr.extra_diag[idx] += 2.0;
            }
            if j + 1 < ny {
                if pr.active[idx + nx] {
                    pr.an[idx] = 1.0;
                }
            } else if dirichlet_sides[3] {
                pr.extra_diag[idx] += 2.0;
            }
            if j > 0 {
                if pr.active[idx - nx] {
                    pr.as_[idx] = 1.0;
                }
            } else if dirichlet_sides[2] {
                pr.extra_diag[idx] += 2.0;
            }
        }
    }
    pr.validate().expect("assembled problem is valid");
    pr
}

/// `A p` written out directly from the definition (independent of the
/// solver's stencil code); zero on inactive cells.
fn apply_operator(pr: &PoissonProblem, p: &[f64]) -> Vec<f64> {
    let (nx, ny) = (pr.nx, pr.ny);
    let mut out = vec![0.0; nx * ny];
    for j in 0..ny {
        for i in 0..nx {
            let idx = j * nx + i;
            if !pr.active[idx] {
                continue;
            }
            let mut v = pr.diagonal(idx) * p[idx];
            if i + 1 < nx && pr.active[idx + 1] {
                v -= pr.ae[idx] * p[idx + 1];
            }
            if i > 0 && pr.active[idx - 1] {
                v -= pr.aw[idx] * p[idx - 1];
            }
            if j + 1 < ny && pr.active[idx + nx] {
                v -= pr.an[idx] * p[idx + nx];
            }
            if j > 0 && pr.active[idx - nx] {
                v -= pr.as_[idx] * p[idx - nx];
            }
            out[idx] = v;
        }
    }
    out
}

fn active_indices(pr: &PoissonProblem) -> Vec<usize> {
    (0..pr.nx * pr.ny).filter(|&i| pr.active[i]).collect()
}

fn remove_mean(pr: &PoissonProblem, v: &mut [f64]) {
    let cells = active_indices(pr);
    let mean = cells.iter().map(|&i| v[i]).sum::<f64>() / cells.len() as f64;
    for &i in &cells {
        v[i] -= mean;
    }
}

fn l1_active(pr: &PoissonProblem, v: &[f64]) -> f64 {
    active_indices(pr).iter().map(|&i| v[i].abs()).sum()
}

/// Smooth field on cell centres of the unit square.
fn smooth_field(nx: usize, ny: usize) -> Vec<f64> {
    let (hx, hy) = (1.0 / nx as f64, 1.0 / ny as f64);
    let mut f = vec![0.0; nx * ny];
    for j in 0..ny {
        for i in 0..nx {
            let (x, y) = ((i as f64 + 0.5) * hx, (j as f64 + 0.5) * hy);
            f[j * nx + i] = (PI * x).cos() * (2.0 * PI * y).cos() + 0.3 * (3.0 * PI * x).sin();
        }
    }
    f
}

/// Zero-mean pseudo-random rhs on the active cells.
fn random_rhs(pr: &mut PoissonProblem, seed: u64) {
    let mut g = Lcg(seed);
    for idx in active_indices(pr) {
        pr.rhs[idx] = g.next();
    }
    let mut rhs = std::mem::take(&mut pr.rhs);
    remove_mean(pr, &mut rhs);
    pr.rhs = rhs;
}

/// Point SOR at the optimal Poisson factor with the true-residual stop —
/// the reference the multigrid PCG is timed against. `anchor = Some`
/// pins that cell to zero as the production projections do (on a pure
/// Neumann problem that pins the level but slows the near-null mode);
/// `None` runs SOR on the singular consistent system, whose iterate
/// converges with a floating level.
fn sor_reference(
    pr: &PoissonProblem,
    p: &mut [f64],
    tolerance: f64,
    max_sweeps: usize,
    anchor: Option<usize>,
) -> (usize, f64) {
    let (nx, ny) = (pr.nx, pr.ny);
    let omega = 2.0 / (1.0 + (PI / nx.max(ny) as f64).sin());
    let singular = pr.is_singular();
    let anchor = anchor.filter(|_| singular);
    for sweep in 1..=max_sweeps {
        let mut residual = 0.0;
        for j in 0..ny {
            for i in 0..nx {
                let idx = j * nx + i;
                if !pr.active[idx] {
                    continue;
                }
                if anchor == Some(idx) {
                    p[idx] = 0.0;
                    continue;
                }
                let ap = pr.diagonal(idx);
                let mut nb = 0.0;
                if i + 1 < nx && pr.active[idx + 1] {
                    nb += pr.ae[idx] * p[idx + 1];
                }
                if i > 0 && pr.active[idx - 1] {
                    nb += pr.aw[idx] * p[idx - 1];
                }
                if j + 1 < ny && pr.active[idx + nx] {
                    nb += pr.an[idx] * p[idx + nx];
                }
                if j > 0 && pr.active[idx - nx] {
                    nb += pr.as_[idx] * p[idx - nx];
                }
                let rhs = pr.rhs[idx] + nb;
                let old = p[idx];
                residual += (rhs - ap * old).abs();
                p[idx] = (1.0 - omega) * old + omega * rhs / ap;
            }
        }
        // The in-sweep (half-sweep lagged) sum over-estimates the residual
        // after the sweep at omega ~ 2, so the stop is the TRUE residual,
        // checked every 8 sweeps (the count is honest to within 8).
        if residual < tolerance || sweep % 8 == 0 {
            let true_res = pr.residual_l1(p);
            if true_res < tolerance {
                return (sweep, true_res);
            }
        }
    }
    (max_sweeps, pr.residual_l1(p))
}

fn max_abs_diff(pr: &PoissonProblem, a: &[f64], b: &[f64]) -> f64 {
    active_indices(pr)
        .iter()
        .map(|&i| (a[i] - b[i]).abs())
        .fold(0.0, f64::max)
}

fn max_abs(pr: &PoissonProblem, a: &[f64]) -> f64 {
    active_indices(pr)
        .iter()
        .map(|&i| a[i].abs())
        .fold(0.0, f64::max)
}

#[test]
fn dirichlet_manufactured_solution_recovered() {
    let n = 40;
    let mut pr = assemble(n, n, |_, _| true, [true; 4]);
    let h = 1.0 / n as f64;
    let exact: Vec<f64> = (0..n * n)
        .map(|idx| {
            let (i, j) = (idx % n, idx / n);
            (PI * (i as f64 + 0.5) * h).sin() * (PI * (j as f64 + 0.5) * h).sin()
        })
        .collect();
    pr.rhs = apply_operator(&pr, &exact);
    assert!(!pr.is_singular());
    let scale = l1_active(&pr, &pr.rhs);
    let mut p = vec![0.0; n * n];
    let sol = solve_multigrid_pcg(
        &pr,
        &mut p,
        &MultigridParameters::default(),
        1e-12 * scale,
        None,
    );
    let err = max_abs_diff(&pr, &p, &exact);
    println!(
        "dirichlet {n}^2: {} iterations, residual {:.3e} (scale {:.3e}), max error {err:.3e}",
        sol.iterations, sol.residual, scale
    );
    assert!(sol.converged, "{sol:?}");
    assert!(err <= 1e-10 * max_abs(&pr, &exact), "max error {err:.3e}");
    assert!(sol.iterations < 40, "{} iterations", sol.iterations);
}

#[test]
fn neumann_box_recovered_up_to_constant_with_anchor() {
    let n = 32;
    let mut pr = assemble(n, n, |_, _| true, [false; 4]);
    assert!(pr.is_singular());
    let field = smooth_field(n, n);
    let mut rhs = apply_operator(&pr, &field);
    remove_mean(&pr, &mut rhs);
    pr.rhs = rhs;
    let scale = l1_active(&pr, &pr.rhs);
    let anchor = pr.index(1, 1);
    // Start from a deliberately shifted guess: the level must be fixed
    // by the anchor, not by the initial guess.
    let mut p = vec![3.0; n * n];
    let sol = solve_multigrid_pcg(
        &pr,
        &mut p,
        &MultigridParameters::default(),
        1e-13 * scale,
        Some(anchor),
    );
    assert!(sol.converged, "{sol:?}");
    assert_eq!(p[anchor], 0.0, "anchor semantics must be exact");
    let shifted: Vec<f64> = field.iter().map(|v| v - field[anchor]).collect();
    let err = max_abs_diff(&pr, &p, &shifted);
    println!(
        "neumann {n}^2: {} iterations, residual {:.3e} (scale {:.3e}), max error {err:.3e}",
        sol.iterations, sol.residual, scale
    );
    assert!(err <= 1e-10 * max_abs(&pr, &shifted), "max error {err:.3e}");

    // Without an anchor: mean zero over the active cells.
    let mut p2 = vec![-7.0; n * n];
    let sol2 = solve_multigrid_pcg(
        &pr,
        &mut p2,
        &MultigridParameters::default(),
        1e-13 * scale,
        None,
    );
    assert!(sol2.converged, "{sol2:?}");
    let mean = p2.iter().sum::<f64>() / (n * n) as f64;
    assert!(mean.abs() <= 1e-12, "mean {mean:e}");
}

fn circle_inactive(n: usize) -> impl Fn(usize, usize) -> bool {
    move |j, i| {
        let h = 1.0 / n as f64;
        let (x, y) = ((i as f64 + 0.5) * h, (j as f64 + 0.5) * h);
        (x - 0.5).powi(2) + (y - 0.5).powi(2) >= 0.04
    }
}

#[test]
fn masked_circle_neumann_recovered_and_inactive_cells_untouched() {
    let n = 48;
    let mut pr = assemble(n, n, circle_inactive(n), [false; 4]);
    let inactive = (0..n * n).filter(|&i| !pr.active[i]).count();
    assert!(inactive > 200, "the circle must remove cells ({inactive})");
    assert!(pr.is_singular());
    let field = smooth_field(n, n);
    pr.rhs = apply_operator(&pr, &field);
    let scale = l1_active(&pr, &pr.rhs);
    let anchor = active_indices(&pr)[0];
    let mut p: Vec<f64> = (0..n * n)
        .map(|i| if pr.active[i] { 0.0 } else { f64::NAN })
        .collect();
    let sol = solve_multigrid_pcg(
        &pr,
        &mut p,
        &MultigridParameters::default(),
        1e-13 * scale,
        Some(anchor),
    );
    println!(
        "masked {n}^2 ({} active): {} iterations, residual {:.3e} (scale {:.3e})",
        n * n - inactive,
        sol.iterations,
        sol.residual,
        scale
    );
    assert!(sol.converged, "{sol:?}");
    for i in 0..n * n {
        if pr.active[i] {
            assert!(p[i].is_finite(), "active cell {i} is {}", p[i]);
        } else {
            assert!(p[i].is_nan(), "inactive cell {i} was written: {}", p[i]);
        }
    }
    assert_eq!(p[anchor], 0.0);
    let shifted: Vec<f64> = field.iter().map(|v| v - field[anchor]).collect();
    let err = max_abs_diff(&pr, &p, &shifted);
    println!("masked max error {err:.3e}");
    assert!(err <= 1e-10 * max_abs(&pr, &shifted), "max error {err:.3e}");
}

/// Iterations to cut the initial residual by `reduction` on the Neumann
/// box with a zero-mean random rhs.
fn neumann_box_iterations(n: usize, reduction: f64) -> (usize, f64) {
    let mut pr = assemble(n, n, |_, _| true, [false; 4]);
    random_rhs(&mut pr, 17 + n as u64);
    let scale = l1_active(&pr, &pr.rhs);
    let mut p = vec![0.0; n * n];
    let sol = solve_multigrid_pcg(
        &pr,
        &mut p,
        &MultigridParameters::default(),
        reduction * scale,
        Some(pr.index(1, 1)),
    );
    assert!(sol.converged, "{n}^2: {sol:?}");
    (sol.iterations, sol.residual / scale)
}

#[test]
fn iteration_count_is_grid_independent() {
    let mut counts = Vec::new();
    for &n in &[32usize, 64, 128, 256] {
        let (it, rel) = neumann_box_iterations(n, 1e-8);
        println!("neumann {n}^2: {it} PCG iterations (final residual {rel:.2e} of rhs)");
        counts.push(it);
    }
    println!("iteration counts 32..256: {counts:?}");
    assert!(
        counts[3] <= 2 * counts[0],
        "256^2 took {} iterations vs {} at 32^2",
        counts[3],
        counts[0]
    );
    assert!(counts[3] <= 60, "256^2 took {} iterations", counts[3]);
}

/// Ragged mask: the circle plus scattered single-cell obstacles, so many
/// aggregates are partial and the over-corrected coarse correction meets
/// cells where the factor-2 argument does not hold.
fn ragged_active(n: usize) -> impl Fn(usize, usize) -> bool {
    let circle = circle_inactive(n);
    move |j, i| circle(j, i) && !((i + j) % 17 == 0 && i % 3 != 0)
}

#[test]
fn ragged_mask_iteration_count_stays_bounded() {
    let mut counts = Vec::new();
    for &n in &[48usize, 96, 192] {
        let mut pr = assemble(n, n, ragged_active(n), [false; 4]);
        random_rhs(&mut pr, 23);
        let scale = l1_active(&pr, &pr.rhs);
        let anchor = active_indices(&pr)[0];
        let mut p: Vec<f64> = (0..n * n)
            .map(|i| if pr.active[i] { 0.0 } else { f64::NAN })
            .collect();
        let sol = solve_multigrid_pcg(
            &pr,
            &mut p,
            &MultigridParameters::default(),
            1e-8 * scale,
            Some(anchor),
        );
        println!(
            "ragged {n}^2 ({} active): {} iterations, residual {:.2e} of rhs",
            active_indices(&pr).len(),
            sol.iterations,
            sol.residual / scale
        );
        assert!(sol.converged, "{n}^2: {sol:?}");
        counts.push(sol.iterations);
    }
    println!("ragged iteration counts 48..192: {counts:?}");
    assert!(counts[2] <= 2 * counts[0] && counts[2] <= 60, "{counts:?}");
}

#[test]
fn galerkin_coarse_operator_matches_r_a_p() {
    let n = 37;
    let mut pr = assemble(n, n, circle_inactive(n), [false, true, false, false]);
    // Add an uneven conductance so a wrong summation cannot hide behind
    // unit coefficients.
    for j in 0..n {
        for i in 0..n {
            let idx = j * n + i;
            let w = 1.0 + 0.5 * ((i * 7 + j * 3) % 5) as f64;
            if pr.ae[idx] != 0.0 {
                pr.ae[idx] *= w;
                pr.aw[idx + 1] *= w;
            }
            if pr.an[idx] != 0.0 {
                pr.an[idx] *= w;
                pr.as_[idx + n] *= w;
            }
        }
    }
    pr.validate().expect("weighted problem is valid");
    let hier = Hierarchy::<f64>::build(&pr, &MultigridParameters::default());
    assert!(hier.depth() >= 3, "depth {}", hier.depth());
    let mut g = Lcg(5);
    for l in 0..hier.depth() - 1 {
        let fine = hier.problem(l);
        let coarse = hier.problem(l + 1);
        coarse.validate().expect("coarse problem is valid");
        let coarse_of = hier.coarse_of(l);
        let v: Vec<f64> = (0..coarse.nx * coarse.ny)
            .map(|i| if coarse.active[i] { g.next() } else { 0.0 })
            .collect();
        let direct = apply_operator(coarse, &v);
        // P v on the fine level, A (P v), then R = summation.
        let pv: Vec<f64> = (0..fine.nx * fine.ny)
            .map(|i| {
                if coarse_of[i] != usize::MAX {
                    v[coarse_of[i]]
                } else {
                    0.0
                }
            })
            .collect();
        let apv = apply_operator(fine, &pv);
        let mut rap = vec![0.0; coarse.nx * coarse.ny];
        for &i in hier.cells(l) {
            rap[coarse_of[i]] += apv[i];
        }
        let scale = max_abs(coarse, &direct);
        let diff = max_abs_diff(coarse, &direct, &rap);
        println!(
            "level {l} -> {}: {}x{} ({} active), |A_c v - R A P v| = {diff:.3e} (scale {scale:.3e})",
            l + 1,
            coarse.nx,
            coarse.ny,
            hier.cells(l + 1).len()
        );
        assert!(scale > 0.0);
        assert!(diff <= 1e-12 * scale, "level {l}: {diff:e}");
        // The outlet column's Dirichlet contribution survives coarsening.
        assert!(!coarse.is_singular());
    }
}

#[test]
fn odd_sizes_and_one_wide_strips_converge() {
    // 37 x 23 Neumann box.
    let (nx, ny) = (37, 23);
    let mut pr = assemble(nx, ny, |_, _| true, [false; 4]);
    random_rhs(&mut pr, 3);
    let scale = l1_active(&pr, &pr.rhs);
    let mut p = vec![0.0; nx * ny];
    let sol = solve_multigrid_pcg(
        &pr,
        &mut p,
        &MultigridParameters::default(),
        1e-10 * scale,
        Some(0),
    );
    println!("37x23: {sol:?}");
    assert!(sol.converged, "{sol:?}");
    assert_eq!(p[0], 0.0);

    // A single row of 64 cells (ny = 1).
    let mut strip = assemble(64, 1, |_, _| true, [false; 4]);
    random_rhs(&mut strip, 4);
    let scale = l1_active(&strip, &strip.rhs);
    let mut p = vec![0.0; 64];
    let sol = solve_multigrid_pcg(
        &strip,
        &mut p,
        &MultigridParameters::default(),
        1e-10 * scale,
        None,
    );
    println!("64x1 strip: {sol:?}");
    assert!(sol.converged, "{sol:?}");
    assert!(strip.residual_l1(&p) < 1e-10 * scale);

    // A 1-wide column of active cells inside a 2-D grid, plus an isolated
    // active cell with no equation (left untouched).
    let (nx, ny) = (9, 50);
    let mut col = assemble(nx, ny, |j, i| i == 3 || (j == 0 && i == 7), [false; 4]);
    random_rhs(&mut col, 5);
    col.rhs[7] = 0.0;
    let scale = l1_active(&col, &col.rhs);
    let mut p: Vec<f64> = (0..nx * ny)
        .map(|i| if col.active[i] { 0.0 } else { f64::NAN })
        .collect();
    p[7] = 42.0;
    let sol = solve_multigrid_pcg(
        &col,
        &mut p,
        &MultigridParameters::default(),
        1e-10 * scale,
        Some(3),
    );
    println!("1-wide column in 9x50: {sol:?}");
    assert!(sol.converged, "{sol:?}");
    assert_eq!(p[7], 42.0, "an isolated cell has no equation");
    assert_eq!(p[3], 0.0);
    for i in 0..nx * ny {
        if col.active[i] {
            assert!(p[i].is_finite());
        } else {
            assert!(p[i].is_nan());
        }
    }
}

#[test]
fn v_cycle_preconditioner_is_symmetric() {
    let n = 30;
    let params = MultigridParameters::default();
    for (name, pr) in [
        (
            "masked neumann",
            assemble(n, n, circle_inactive(n), [false; 4]),
        ),
        (
            "outlet column",
            assemble(n, n, circle_inactive(n), [false, true, false, false]),
        ),
        ("dirichlet", assemble(n, n, |_, _| true, [true; 4])),
    ] {
        let mut hier = Hierarchy::<f64>::build(&pr, &params);
        let mut g = Lcg(11);
        let cells = active_indices(&pr);
        let mut x = vec![0.0; n * n];
        let mut y = vec![0.0; n * n];
        for &i in &cells {
            x[i] = g.next();
            y[i] = g.next();
        }
        let mut mx = vec![0.0; n * n];
        let mut my = vec![0.0; n * n];
        hier.apply_preconditioner(&x, &mut mx);
        hier.apply_preconditioner(&y, &mut my);
        let lhs: f64 = cells.iter().map(|&i| x[i] * my[i]).sum();
        let rhs: f64 = cells.iter().map(|&i| mx[i] * y[i]).sum();
        println!("{name}: <x, M^-1 y> = {lhs:.15e}, <M^-1 x, y> = {rhs:.15e}");
        assert!(lhs.abs() > 0.0);
        assert!(
            (lhs - rhs).abs() <= 1e-12 * lhs.abs().max(rhs.abs()),
            "{name}: {lhs:e} vs {rhs:e}"
        );
        // And positive on the range: <x, M^-1 x> > 0.
        let xx: f64 = cells.iter().map(|&i| x[i] * mx[i]).sum();
        assert!(xx > 0.0, "{name}: <x, M^-1 x> = {xx:e}");
    }
}

#[test]
fn validate_rejects_asymmetry_and_bad_lengths() {
    let mut pr = assemble(8, 8, |_, _| true, [false; 4]);
    pr.ae[0] = 2.0;
    assert!(pr.validate().unwrap_err().contains("asymmetric"));
    let mut pr = assemble(8, 8, |_, _| true, [false; 4]);
    pr.rhs.pop();
    assert!(pr.validate().unwrap_err().contains("length"));
    let mut pr = assemble(8, 8, |_, _| true, [false; 4]);
    pr.an[3] = -1.0;
    pr.as_[11] = -1.0;
    assert!(pr.validate().unwrap_err().contains("negative"));
}

#[test]
fn timing_against_sor_reference() {
    for &n in &[128usize, 256] {
        let mut pr = assemble(n, n, |_, _| true, [false; 4]);
        random_rhs(&mut pr, 99);
        let scale = l1_active(&pr, &pr.rhs);
        let tol = 1e-8 * scale;
        let anchor = pr.index(1, 1);

        let mut p = vec![0.0; n * n];
        let t0 = std::time::Instant::now();
        let sol = solve_multigrid_pcg(
            &pr,
            &mut p,
            &MultigridParameters::default(),
            tol,
            Some(anchor),
        );
        let t_mg = t0.elapsed().as_secs_f64();
        assert!(sol.converged, "{sol:?}");

        let mut ps = vec![0.0; n * n];
        let t0 = std::time::Instant::now();
        let (sweeps, res_sor) = sor_reference(&pr, &mut ps, tol, 20_000, None);
        let t_sor = t0.elapsed().as_secs_f64();
        println!(
            "{n}^2 Neumann, stop {tol:.2e}: MG-PCG {} it, {:.3e} res, {t_mg:.3} s | free SOR {sweeps} sweeps, {res_sor:.3e} res, {t_sor:.3} s | SOR/MG time ratio {:.1}",
            sol.iterations,
            sol.residual,
            t_sor / t_mg
        );
        assert!(
            res_sor < tol,
            "free SOR did not reach the stop in {sweeps} sweeps"
        );
        // Same residual stop, so the two solutions agree to the stop
        // level once both are shifted to the anchor.
        let shift = ps[anchor];
        for &i in &active_indices(&pr) {
            ps[i] -= shift;
        }
        let diff = max_abs_diff(&pr, &p, &ps);
        println!("{n}^2: max |p_mg - p_sor| = {diff:.3e}");

        // The production-style anchored SOR, capped: reported, not timed
        // to the stop (on 128^2 it does not reach it in 20,000 sweeps).
        let cap = 3_000;
        let mut pa = vec![0.0; n * n];
        let t0 = std::time::Instant::now();
        let (sweeps_a, res_a) = sor_reference(&pr, &mut pa, tol, cap, Some(anchor));
        let t_a = t0.elapsed().as_secs_f64();
        println!(
            "{n}^2: anchored SOR {sweeps_a} sweeps (cap {cap}), residual {res_a:.3e} vs stop {tol:.2e}, {t_a:.3} s"
        );
    }
}

// ---------------------------------------------------------------------------
// Multi-component domains (from the adversarial review)
// ---------------------------------------------------------------------------

/// Two pure-Neumann components (a 64² box split by an inactive wall column)
/// whose right-hand sides are each slightly incompatible with OPPOSITE
/// signs, so the global mean is zero: a single global mean projection
/// leaves both blocks inconsistent and CG diverged (review: max |p| 3.8e9
/// at 1e-8 relative imbalance). Per-component projection must converge and
/// recover the known field in each component up to that component's
/// constant.
#[test]
fn two_neumann_components_with_opposite_imbalances_converge() {
    let n = 64;
    let wall = n / 2;
    let mut pr = assemble(n, n, |_, i| i != wall, [false; 4]);
    let exact = smooth_field(n, n);
    let rhs = apply_operator(&pr, &exact);
    // Per-component imbalance ±eps × scale, zero overall.
    let scale = l1_active(&pr, &rhs) / active_indices(&pr).len() as f64;
    for j in 0..n {
        for i in 0..n {
            let idx = j * n + i;
            if !pr.active[idx] {
                continue;
            }
            let sign = if i < wall { 1.0 } else { -1.0 };
            pr.rhs[idx] = rhs[idx] + sign * 1e-6 * scale;
        }
    }
    let mut p = vec![0.0; n * n];
    let tolerance = 1e-10 * l1_active(&pr, &pr.rhs);
    let sol = solve_multigrid_pcg(
        &pr,
        &mut p,
        &MultigridParameters::default(),
        tolerance,
        Some(n + 1),
    );
    assert!(
        sol.converged,
        "PCG did not converge on two imbalanced Neumann components: {sol:?}"
    );
    assert!(sol.iterations <= 20, "iterations {}", sol.iterations);
    // Each component matches the field up to its own constant.
    for half in 0..2 {
        let members: Vec<usize> = active_indices(&pr)
            .into_iter()
            .filter(|&idx| (idx % n < wall) == (half == 0))
            .collect();
        let shift =
            members.iter().map(|&idx| p[idx] - exact[idx]).sum::<f64>() / members.len() as f64;
        let err = members
            .iter()
            .map(|&idx| (p[idx] - exact[idx] - shift).abs())
            .fold(0.0, f64::max);
        let amp = members
            .iter()
            .map(|&idx| exact[idx].abs())
            .fold(0.0, f64::max);
        assert!(
            err < 1e-5 * amp,
            "component {half}: max error {err:.3e} vs amplitude {amp:.3e}"
        );
    }
    // Anchor semantics hold in the anchor's component; the other is mean zero.
    assert_eq!(p[n + 1], 0.0);
    let right: Vec<usize> = active_indices(&pr)
        .into_iter()
        .filter(|&idx| idx % n > wall)
        .collect();
    let right_mean = right.iter().map(|&idx| p[idx]).sum::<f64>() / right.len() as f64;
    assert!(
        right_mean.abs() < 1e-12,
        "right component mean {right_mean:.3e}"
    );
}

/// A Dirichlet component next to a singular one: an imbalance in the
/// Neumann half must not corrupt the Dirichlet half (review: the Dirichlet
/// half read 3.3e3 against an exact 10.0 under a global treatment).
#[test]
fn dirichlet_component_is_untouched_by_an_imbalanced_neumann_neighbour() {
    let n = 48;
    let wall = n / 2;
    // Left side Dirichlet (p = 0 half a cell outside the left edge), wall
    // column inactive, right half pure Neumann.
    let mut pr = assemble(n, n, |_, i| i != wall, [true, false, false, false]);
    // Exact: left half p = 10 + the discrete solution of ap p = rhs with
    // rhs chosen from a known field; simplest: take a known field on both
    // halves and build rhs = A field, then perturb the right half only.
    let exact = smooth_field(n, n);
    let rhs = apply_operator(&pr, &exact);
    let scale = l1_active(&pr, &rhs) / active_indices(&pr).len() as f64;
    for j in 0..n {
        for i in 0..n {
            let idx = j * n + i;
            if pr.active[idx] {
                pr.rhs[idx] = rhs[idx] + if i > wall { 1e-6 * scale } else { 0.0 };
            }
        }
    }
    let mut p = vec![0.0; n * n];
    let tolerance = 1e-10 * l1_active(&pr, &pr.rhs);
    let sol = solve_multigrid_pcg(
        &pr,
        &mut p,
        &MultigridParameters::default(),
        tolerance,
        None,
    );
    assert!(sol.converged, "{sol:?}");
    // Left (Dirichlet) component exact — no constant freedom there.
    let left: Vec<usize> = active_indices(&pr)
        .into_iter()
        .filter(|&idx| idx % n < wall)
        .collect();
    let err = left
        .iter()
        .map(|&idx| (p[idx] - exact[idx]).abs())
        .fold(0.0, f64::max);
    let amp = left.iter().map(|&idx| exact[idx].abs()).fold(0.0, f64::max);
    assert!(
        err < 1e-8 * amp,
        "Dirichlet half corrupted: max error {err:.3e} vs {amp:.3e}"
    );
}