//! 3D Stage 1, gate 6: plane Poiseuille flow on the 3D solver. (a) At
//! `nz = 1` (`dz = 1`, z sides slip) the 3D host step reproduces the 2D
//! embedded solver (no body, multigrid Poisson) to the value over 200
//! steps; (b) the 2D `poiseuille.rs` gates on the 3D field at nz = 1 and on
//! a periodic-z extrusion: `|u − û| < 1e-7`, `max |v|, |w| < 1e-7`,
//! `p spread < 1e-6`, with û the discrete channel profile.

use rtx_cfd::CfdConfig;
use rtx_cfd::solvers::incompressible::three_d::{
    Boundaries3, FlowField3D, Fluid3, Grid3, Piso3Parameters, Piso3Solver, SideBoundary3,
};
use rtx_cfd::solvers::incompressible::{
    EmbeddedParameters, EmbeddedPisoSolver, FlowField, PoissonSolverKind,
};

const MU: f64 = 0.1;
const G: f64 = 0.8;

fn discrete_profile(n: usize) -> Vec<f64> {
    let h = 1.0 / n as f64;
    let rhs_value = -G * h * h / MU;
    let mut diag = vec![-2.0; n];
    diag[0] = -3.0;
    diag[n - 1] = -3.0;
    let mut rhs = vec![rhs_value; n];
    let upper = vec![1.0; n];
    for j in 1..n {
        let factor = 1.0 / diag[j - 1];
        diag[j] -= factor * upper[j - 1];
        rhs[j] -= factor * rhs[j - 1];
    }
    let mut u = vec![0.0; n];
    u[n - 1] = rhs[n - 1] / diag[n - 1];
    for j in (0..n - 1).rev() {
        u[j] = (rhs[j] - upper[j] * u[j + 1]) / diag[j];
    }
    u
}

fn fluid() -> Fluid3 {
    Fluid3 {
        density: 1.0,
        viscosity: MU,
        reference_velocity: 1.0,
        reference_length: 1.0,
    }
}

/// The inlet/outlet carry the discrete profile (the embedded solvers
/// re-stamp every Velocity side from the boundary function each step);
/// the walls are no-slip.
fn profile_boundary(n: usize) -> impl Fn(f64, f64) -> f64 + Clone {
    let u_hat = discrete_profile(n);
    let h = 1.0 / n as f64;
    move |x: f64, y: f64| {
        if x <= 0.0 || x >= 1.0 {
            let j = ((y / h - 0.5).round().max(0.0) as usize).min(n - 1);
            u_hat[j]
        } else {
            0.0
        }
    }
}

fn solver3(n: usize, nz_periodic: bool) -> Piso3Solver {
    let z = if nz_periodic {
        SideBoundary3::Periodic
    } else {
        SideBoundary3::SlipWall
    };
    let mut s = Piso3Solver::new(
        fluid(),
        Piso3Parameters {
            corrector_steps: 2,
            tolerance: 1e-8,
            boundaries: Boundaries3 {
                z0: z,
                z1: z,
                ..Boundaries3::default()
            },
            ..Piso3Parameters::default()
        },
    );
    s.set_momentum_source(|_x, _y, _z, _t| (G, 0.0, 0.0));
    let ub = profile_boundary(n);
    s.set_boundary_velocity(move |x, y, _z, _t| (ub(x, y), 0.0, 0.0));
    s
}

fn field3(n: usize, nz: usize, dz: f64) -> FlowField3D {
    let h = 1.0 / n as f64;
    let g = Grid3 {
        nx: n,
        ny: n,
        nz,
        dx: h,
        dy: h,
        dz,
    };
    let mut f = FlowField3D::new(g);
    let u_hat = discrete_profile(n);
    for k in 0..nz {
        for (j, &uj) in u_hat.iter().enumerate() {
            f.u[g.uface(k, j, 0)] = uj;
            f.u[g.uface(k, j, n)] = uj;
        }
    }
    f
}

fn field2(n: usize) -> FlowField {
    let h = 1.0 / n as f64;
    let mut f = FlowField::new(n, n, h, h).expect("field");
    let u_hat = discrete_profile(n);
    for (j, &uj) in u_hat.iter().enumerate() {
        f.u[(j, 0)] = uj;
        f.u[(j, n)] = uj;
    }
    f
}

fn dt_for(n: usize) -> f64 {
    let h = 1.0 / n as f64;
    0.4 * (h * h / (4.0 * MU)).min(h)
}

/// (a) the value identity at nz = 1.
#[tokio::test]
async fn nz_one_is_the_two_d_embedded_solver() {
    let n = 16;
    let dt = dt_for(n);
    let config = CfdConfig::new()
        .with_density(1.0)
        .with_viscosity(MU)
        .with_reference_velocity(1.0)
        .with_reference_length(1.0);
    let mut two = EmbeddedPisoSolver::new(
        config,
        EmbeddedParameters {
            corrector_steps: 2,
            tolerance: 1e-8,
            poisson_solver: PoissonSolverKind::Multigrid,
            ..EmbeddedParameters::default()
        },
    )
    .expect("2D solver");
    two.set_momentum_source(|_x, _y, _t| (G, 0.0));
    let ub = profile_boundary(n);
    two.set_boundary_velocity(move |x, y, _t| (ub(x, y), 0.0));
    let mut three = solver3(n, false);
    let mut a = field2(n);
    let mut b = field3(n, 1, 1.0);
    let g = b.grid;
    let mut signed_zero = 0usize;
    for step in 0..200 {
        two.advance(&mut a, dt).await.expect("2D step");
        three.advance(&mut b, dt);
        let mut worst = 0.0_f64;
        for j in 0..n {
            for i in 0..=n {
                let (x, y) = (a.u[(j, i)], b.u[g.uface(0, j, i)]);
                if x != y {
                    worst = worst.max((x - y).abs());
                } else if x.to_bits() != y.to_bits() {
                    signed_zero += 1;
                }
            }
        }
        for j in 0..=n {
            for i in 0..n {
                let (x, y) = (a.v[(j, i)], b.v[g.vface(0, j, i)]);
                if x != y {
                    worst = worst.max((x - y).abs());
                } else if x.to_bits() != y.to_bits() {
                    signed_zero += 1;
                }
            }
        }
        for j in 0..n {
            for i in 0..n {
                let (x, y) = (a.p[(j, i)], b.p[g.cell(0, j, i)]);
                if x != y {
                    worst = worst.max((x - y).abs());
                }
            }
        }
        assert!(
            worst == 0.0,
            "step {step}: the 3D field departs from the 2D embedded solver by {worst:.3e}"
        );
    }
    println!(
        "  200 steps value-identical to the 2D embedded solver ({signed_zero} ±0 sign differences)"
    );
}

struct Measurement {
    max_u_vs_discrete: f64,
    max_v: f64,
    max_w: f64,
    p_spread: f64,
    steps: usize,
}

fn measure(n: usize, nz: usize, dz: f64, periodic: bool) -> Measurement {
    let dt = dt_for(n);
    let mut solver = solver3(n, periodic);
    let mut f = field3(n, nz, dz);
    let g = f.grid;
    let u_hat = discrete_profile(n);
    let mut steps = 0;
    for step in 0..200_000 {
        let before = f.u.clone();
        solver.advance(&mut f, dt);
        steps = step + 1;
        let change =
            f.u.iter()
                .zip(&before)
                .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()))
                / dt;
        if change < 1e-8 {
            break;
        }
    }
    let mut max_u_vs_discrete = 0.0_f64;
    for k in 0..nz {
        for (j, &uj) in u_hat.iter().enumerate() {
            for i in 1..n {
                max_u_vs_discrete = max_u_vs_discrete.max((f.u[g.uface(k, j, i)] - uj).abs());
            }
        }
    }
    let max_v = f.v.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
    let max_w = f.w.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
    let (mut p_min, mut p_max) = (f64::INFINITY, f64::NEG_INFINITY);
    for &p in &f.p {
        p_min = p_min.min(p);
        p_max = p_max.max(p);
    }
    Measurement {
        max_u_vs_discrete,
        max_v,
        max_w,
        p_spread: p_max - p_min,
        steps,
    }
}

/// (b) the 2D gates on the 3D field.
#[test]
fn poiseuille_is_the_discrete_profile_in_three_d() {
    for (n, nz, dz, periodic) in [
        (16usize, 1usize, 1.0, false),
        (16, 4, 1.0 / 16.0, true),
        (32, 1, 1.0, false),
    ] {
        let m = measure(n, nz, dz, periodic);
        println!(
            "  n {n} nz {nz} periodic {periodic}: {} steps; |u − û| {:.3e}, max |v| {:.3e}, max |w| {:.3e}, p spread {:.3e}",
            m.steps, m.max_u_vs_discrete, m.max_v, m.max_w, m.p_spread
        );
        assert!(
            m.max_u_vs_discrete < 1e-7,
            "u departs from the discrete profile by {:.3e}",
            m.max_u_vs_discrete
        );
        assert!(m.max_v < 1e-7, "spurious transverse flow {:.3e}", m.max_v);
        assert!(m.max_w < 1e-7, "spurious spanwise flow {:.3e}", m.max_w);
        assert!(m.p_spread < 1e-6, "spurious pressure {:.3e}", m.p_spread);
    }
}