//! 3D Stage 1, gate 1 (omni-cortex `docs/three_d_stage1_campaign.md`): the
//! 3D Poisson solver at `nz = 1` is the 2D solver bit for bit (solution and
//! iteration count; lexicographic and red-black; cached and uncached), and
//! an extrusion in z (decoupled planes, and periodic z with a z-invariant
//! right-hand side) is bit-identical across planes.

use rtx_cfd::solvers::incompressible::three_d::poisson::{
    PcgCache3, PoissonProblem3D, solve_multigrid_pcg3, solve_multigrid_pcg3_cached,
};
use rtx_cfd::solvers::incompressible::{
    MgSmoother, MultigridParameters, PcgCache, PoissonProblem, solve_multigrid_pcg,
    solve_multigrid_pcg_cached,
};

/// The `poisson_redblack.rs` masked channel (a hole, an outlet Dirichlet).
fn problem_2d(nx: usize, ny: usize, seed: u64) -> PoissonProblem {
    let mut p = PoissonProblem::new(nx, ny);
    let (dx, dy, dt) = (1.0 / nx as f64, 0.41 / ny as f64, 1e-3);
    let (ae, an) = (dt * dy / dx, dt * dx / dy);
    let hole = |i: usize, j: usize| {
        let (x, y) = ((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy);
        (x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.05 * 0.05
    };
    for j in 0..ny {
        for i in 0..nx {
            let idx = j * nx + i;
            if hole(i, j) {
                p.active[idx] = false;
                continue;
            }
            if i + 1 < nx && !hole(i + 1, j) {
                p.ae[idx] = ae;
            }
            if i > 0 && !hole(i - 1, j) {
                p.aw[idx] = ae;
            }
            if j + 1 < ny && !hole(i, j + 1) {
                p.an[idx] = an;
            }
            if j > 0 && !hole(i, j - 1) {
                p.as_[idx] = an;
            }
            if i + 1 == nx {
                p.extra_diag[idx] = 2.0 * ae;
            }
        }
    }
    let mut state = seed | 1;
    for idx in 0..nx * ny {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        p.rhs[idx] = if p.active[idx] {
            1e-6 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5)
        } else {
            0.0
        };
    }
    p
}

/// The 2D problem stacked `nz` times; `az` couples the planes (0 = decoupled).
fn extrude(p2: &PoissonProblem, nz: usize, az: f64, periodic_z: bool) -> PoissonProblem3D {
    let (nx, ny) = (p2.nx, p2.ny);
    let mut p = PoissonProblem3D::new(nx, ny, nz);
    p.periodic_z = periodic_z;
    for k in 0..nz {
        for idx2 in 0..nx * ny {
            let idx = k * nx * ny + idx2;
            p.active[idx] = p2.active[idx2];
            p.ae[idx] = p2.ae[idx2];
            p.aw[idx] = p2.aw[idx2];
            p.an[idx] = p2.an[idx2];
            p.as_[idx] = p2.as_[idx2];
            p.extra_diag[idx] = p2.extra_diag[idx2];
            p.rhs[idx] = p2.rhs[idx2];
            if p2.active[idx2] && az != 0.0 {
                let up = k + 1 < nz || periodic_z;
                let down = k > 0 || periodic_z;
                if up {
                    p.at[idx] = az;
                }
                if down {
                    p.ab[idx] = az;
                }
            }
        }
    }
    p
}

fn bits(v: &[f64]) -> Vec<u64> {
    v.iter().map(|x| x.to_bits()).collect()
}

#[test]
fn nz_one_is_the_two_d_solver_bit_for_bit() {
    let (nx, ny) = (96, 40);
    let tol = 1e-12;
    for (name, params) in [
        ("lexicographic", MultigridParameters::default()),
        (
            "red-black",
            MultigridParameters {
                smoother: MgSmoother::RedBlack,
                ..MultigridParameters::default()
            },
        ),
    ] {
        let mut cache2 = PcgCache::default();
        let mut cache3 = PcgCache3::default();
        for seed in [5u64, 20, 21] {
            let p2 = problem_2d(nx, ny, seed);
            let p3 = extrude(&p2, 1, 0.0, false);
            assert!(p3.validate().is_ok(), "{:?}", p3.validate());
            let (mut a, mut b) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
            let sa = solve_multigrid_pcg(&p2, &mut a, &params, tol, None);
            let sb = solve_multigrid_pcg3(&p3, &mut b, &params, tol, None);
            assert!(sa.converged && sb.converged, "{name} seed {seed} converged");
            assert_eq!(
                sa.iterations, sb.iterations,
                "{name} seed {seed} iterations"
            );
            assert_eq!(bits(&a), bits(&b), "{name} seed {seed}: 3D differs from 2D");
            let (mut c, mut d) = (vec![0.0; nx * ny], vec![0.0; nx * ny]);
            let sc = solve_multigrid_pcg_cached(&p2, &mut c, &params, tol, None, &mut cache2);
            let sd = solve_multigrid_pcg3_cached(&p3, &mut d, &params, tol, None, &mut cache3);
            assert_eq!(sc.iterations, sd.iterations);
            assert_eq!(bits(&c), bits(&d), "{name} seed {seed}: cached 3D differs");
            assert_eq!(
                bits(&a),
                bits(&c),
                "{name} seed {seed}: cached 2D differs from uncached"
            );
            assert!(p3.residual_l1(&b) < tol);
            println!(
                "  {name} seed {seed}: {} iterations, bit-identical to the 2D solver",
                sa.iterations
            );
        }
    }
}

/// Plane-to-plane identity of a z-invariant solve. Bit identity across
/// planes holds only where every plane's arithmetic path is the same:
/// decoupled planes under the lexicographic smoother. The seven-point
/// red-black colouring `(i + j + k) % 2` swaps the colours between
/// neighbouring planes on every level (the 2×2×2 aggregation merges plane
/// pairs, so the coarse levels swap again), and a lexicographic sweep of
/// coupled planes reads updated values below and old values above; in
/// both cases the planes agree to the solve's own accuracy (`1e-6 · scale`,
/// the red-black-vs-lexicographic pin's standard), not in bits.
#[test]
fn an_extrusion_in_z_is_z_invariant() {
    let (nx, ny, nz) = (48, 20, 8);
    let tol = 1e-12;
    let p2 = problem_2d(nx, ny, 7);
    let az = 1e-3 * (1.0 / 48.0) * (0.41 / 20.0) / 0.05;
    for (name, az, periodic, decoupled) in [
        ("decoupled planes", 0.0, false, true),
        ("periodic z", az, true, false),
        ("closed z (walls)", az, false, false),
    ] {
        for smoother in [MgSmoother::Lexicographic, MgSmoother::RedBlack] {
            let params = MultigridParameters {
                smoother,
                ..MultigridParameters::default()
            };
            let p3 = extrude(&p2, nz, az, periodic);
            assert!(p3.validate().is_ok(), "{name}: {:?}", p3.validate());
            let mut sol = vec![0.0; nx * ny * nz];
            let s = solve_multigrid_pcg3(&p3, &mut sol, &params, tol, None);
            assert!(
                s.converged,
                "{name} {smoother:?}: not converged ({} it)",
                s.iterations
            );
            assert!(p3.residual_l1(&sol) < tol, "{name} {smoother:?}: residual");
            let plane = |k: usize| &sol[k * nx * ny..(k + 1) * nx * ny];
            let scale = sol.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
            let mut worst = 0.0_f64;
            for k in 1..nz {
                let d = plane(k)
                    .iter()
                    .zip(plane(0))
                    .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
                worst = worst.max(d);
                if decoupled && smoother == MgSmoother::Lexicographic {
                    assert_eq!(
                        bits(plane(k)),
                        bits(plane(0)),
                        "{name} {smoother:?}: plane {k} differs from plane 0 in bits"
                    );
                }
            }
            assert!(
                worst <= 1e-6 * scale,
                "{name} {smoother:?}: planes differ by {worst:.3e} on a scale of {scale:.3e}"
            );
            println!(
                "  {name} {smoother:?}: {} iterations, planes within {worst:.2e} of {scale:.2e}{}",
                s.iterations,
                if decoupled && smoother == MgSmoother::Lexicographic {
                    " (planes bit-identical)"
                } else {
                    ""
                }
            );
        }
    }
}

#[test]
#[ignore = "probe: prints CG iteration counts across nz, coupling and depth"]
fn probe_iteration_counts() {
    let (nx, ny) = (48, 20);
    let p2 = problem_2d(nx, ny, 7);
    let ae = 1e-3 * (0.41 / 20.0) / (1.0 / 48.0);
    for (nz, az) in [(1usize, 0.0), (8, 0.0), (8, ae), (16, ae)] {
        for smoother in [MgSmoother::Lexicographic, MgSmoother::RedBlack] {
            for coarsest in [32usize, usize::MAX / 2] {
                let params = MultigridParameters {
                    smoother,
                    coarsest_cells: coarsest,
                    ..MultigridParameters::default()
                };
                let p3 = extrude(&p2, nz, az, true);
                let mut sol = vec![0.0; nx * ny * nz];
                let s = solve_multigrid_pcg3(&p3, &mut sol, &params, 1e-12, None);
                println!(
                    "  nz {nz} az/ae {:.0} {smoother:?} coarsest {}: {} iterations",
                    az / ae,
                    if coarsest == 32 { "32" } else { "single level" },
                    s.iterations
                );
            }
        }
    }
}