//! P0 gate 3 (`docs/overset_metal_campaign.md` §5.3): plane Poiseuille
//! flow on the patch. On the Cartesian AND the affinely sheared periodic
//! channel the discrete fixed point is exactly the `poiseuille.rs`
//! profile (congruent parallelogram cells: the non-orthogonal corrections
//! cancel by translation invariance) — recovered to 1e-8 with `v` and the
//! pressure spread at the same level. On a channel of smoothly varying
//! skew and stretch the error against the parabola falls at second
//! order. Cell mass is conserved to 1e-12 everywhere.

use rtx_cfd::mesh::PatchMesh;
use rtx_cfd::mesh::patch_gen::{cartesian, channel_sheared, channel_varying_skew};
use rtx_cfd::solvers::incompressible::{
    CurvilinearParameters, CurvilinearPisoSolver, NormalDiffusion, PatchField,
};
use rtx_cfd::{CfdConfig, CfdResult};

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

fn u_exact(y: f64) -> f64 {
    G / (2.0 * MU) * y * (1.0 - y)
}

/// The 1-D discrete channel profile with half-cell wall closures
/// (`tests/poiseuille.rs`).
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
}

struct Steady {
    field: PatchField,
    mesh: PatchMesh,
    worst_mass: f64,
}

async fn steady(mesh: PatchMesh, diffusion: NormalDiffusion, dt_factor: f64) -> CfdResult<Steady> {
    let mut h = f64::INFINITY;
    for c in 0..mesh.cell_count() {
        for (f, _) in mesh.cell_faces(c) {
            let d = mesh.faces()[f].d;
            h = h.min((d[0] * d[0] + d[1] * d[1]).sqrt());
        }
    }
    let dt = dt_factor * 0.4 * h * h / (4.0 * MU);
    let config = CfdConfig::new().with_density(1.0).with_viscosity(MU);
    let params = CurvilinearParameters {
        tolerance: 1e-6,
        normal_diffusion: diffusion,
        ..CurvilinearParameters::default()
    };
    let mut solver = CurvilinearPisoSolver::new(config, params, mesh)?;
    solver.set_boundary_velocity(|_, _, _| (0.0, 0.0));
    solver.set_momentum_source(|_, _, _| (G, 0.0));
    let mut field = PatchField::new(solver.mesh());
    solver.initialize(&mut field, |_, _| (0.0, 0.0));
    // Mass defect at the steady state, relative to the largest face flux
    // (during the transient it is the pressure solver's residual).
    let env = |k: &str, d: f64| {
        std::env::var(k)
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(d)
    };
    // 1e-9: the rounding floor of |du/dt| sits at ~2e-11 on the 32² channels
    // (measured: 0 pressure iterations, divergence 1e-16) and the gates are
    // at 1e-8.
    let steady_tol = env("RTX_CURV_STEADY", 1e-9);
    let trace = std::env::var("RTX_CURV_TRACE").is_ok();
    let mut worst_mass = 0.0_f64;
    let mut converged = false;
    for step in 0..2_000_000 {
        let before = field.u.clone();
        let r = solver.advance(&mut field, dt).await?;
        assert!(r.poisson_converged, "{r:?}");
        let scale: f64 = field.flux.iter().map(|f| f.abs()).fold(0.0, f64::max);
        worst_mass = r.max_divergence / scale.max(1e-300);
        let change = field
            .u
            .iter()
            .zip(&before)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0, f64::max);
        if trace && step % 50_000 == 0 {
            println!(
                "  step {step} t={:.2} |du/dt| {:.3e} iters {} div {:.2e}",
                solver.time(),
                change / dt,
                r.poisson_iterations,
                r.max_divergence
            );
        }
        if change / dt < steady_tol {
            converged = true;
            break;
        }
    }
    assert!(converged, "no steady state");
    let mesh = solver.mesh().clone();
    Ok(Steady {
        field,
        mesh,
        worst_mass,
    })
}

fn max_vs_discrete(s: &Steady, n: usize) -> (f64, f64, f64) {
    let u_hat = discrete_profile(n);
    let (mut du, mut v, mut pmin, mut pmax) = (0.0_f64, 0.0_f64, f64::INFINITY, f64::NEG_INFINITY);
    for c in 0..s.mesh.cell_count() {
        let (k, _) = s.mesh.cell_ki(c);
        du = du.max((s.field.u[c] - u_hat[k]).abs());
        v = v.max(s.field.v[c].abs());
        pmin = pmin.min(s.field.p[c]);
        pmax = pmax.max(s.field.p[c]);
    }
    (du, v, pmax - pmin)
}

#[tokio::test]
async fn cartesian_and_sheared_channels_hit_the_discrete_profile_exactly() -> CfdResult<()> {
    let n = 16;
    let shapes: [(&str, Box<dyn Fn() -> CfdResult<PatchMesh>>); 3] = [
        (
            "cartesian",
            Box::new(move || cartesian(n, n, 1.0, 1.0, true)),
        ),
        (
            "sheared 0.4",
            Box::new(move || channel_sheared(1.0, 1.0, n, n, 0.4, true)),
        ),
        (
            "sheared -0.7",
            Box::new(move || channel_sheared(1.0, 1.0, n, n, -0.7, true)),
        ),
    ];
    for (name, mesh) in &shapes {
        for diffusion in [NormalDiffusion::Explicit, NormalDiffusion::LineImplicit] {
            let s = steady(mesh()?, diffusion, 1.0).await?;
            let (du, v, dp) = max_vs_discrete(&s, n);
            println!(
                "{name} {diffusion:?}: |u - u_hat| {du:.3e}, |v| {v:.3e}, p spread {dp:.3e}, mass {:.3e}",
                s.worst_mass
            );
            assert!(du < 1e-8, "{name} {diffusion:?}: |u - u_hat| = {du:.3e}");
            assert!(v < 1e-8, "{name} {diffusion:?}: |v| = {v:.3e}");
            assert!(dp < 1e-8, "{name} {diffusion:?}: p spread {dp:.3e}");
            assert!(
                s.worst_mass < 1e-12,
                "{name} {diffusion:?}: mass defect {:.3e}",
                s.worst_mass
            );
        }
    }
    Ok(())
}

#[tokio::test]
async fn varying_skew_channel_converges_to_the_parabola_at_second_order() -> CfdResult<()> {
    let mut errs = Vec::new();
    let mut vs = Vec::new();
    let only: Option<usize> = std::env::var("RTX_CURV_N")
        .ok()
        .and_then(|v| v.parse().ok());
    for n in [16usize, 32, 64] {
        if only.is_some_and(|o| o != n) {
            continue;
        }
        let s = steady(
            channel_varying_skew(1.0, 1.0, n, n, 0.1, 2.0, true)?,
            NormalDiffusion::Explicit,
            1.0,
        )
        .await?;
        let (mut e, mut v) = (0.0_f64, 0.0_f64);
        for c in 0..s.mesh.cell_count() {
            let y = s.mesh.centre(c)[1];
            e = e.max((s.field.u[c] - u_exact(y)).abs());
            v = v.max(s.field.v[c].abs());
        }
        println!(
            "varying skew n={n}: |u - parabola| {e:.3e}, |v| {v:.3e}, mass {:.3e}",
            s.worst_mass
        );
        assert!(s.worst_mass < 1e-12);
        errs.push(e);
        vs.push(v);
    }
    let o: Vec<f64> = errs.windows(2).map(|p| (p[0] / p[1]).log2()).collect();
    let ov: Vec<f64> = vs.windows(2).map(|p| (p[0] / p[1]).log2()).collect();
    println!("varying skew orders u {o:?}, v {ov:?}");
    if only.is_none() {
        assert!(o.iter().all(|&x| x >= 1.8), "orders {o:?}");
    }
    Ok(())
}