//! P0 gates 1, 2, 4 (`docs/overset_metal_campaign.md` §5.3): the
//! curvilinear collocated PISO on a Cartesian patch reproduces the
//! staggered PISO's manufactured-solution order and error to within 2×
//! (not bit-identical — a different discretisation); on a skewed,
//! stretched, periodic annulus it reaches order ≥ 1.8 in the Stokes limit
//! and ≈ 1 with upwind; every step is divergence-free to the solver's
//! tolerance.

use rtx_cfd::mesh::PatchMesh;
use rtx_cfd::mesh::patch_gen::{annulus_skewed, cartesian};
use rtx_cfd::solvers::incompressible::{
    CurvilinearParameters, CurvilinearPisoSolver, NormalDiffusion, PatchConvection, PatchField,
};
use rtx_cfd::{CfdConfig, CfdResult};
use std::f64::consts::PI;

const RHO: f64 = 1.0;
const MU: f64 = 0.05;

fn u_exact(x: f64, y: f64) -> f64 {
    (PI * x).sin() * (PI * y).cos()
}
fn v_exact(x: f64, y: f64) -> f64 {
    -(PI * x).cos() * (PI * y).sin()
}
/// `f = ρ u·∇u (if convecting) − μ ∇²u + ∇p`, `p = sin(πx) sin(πy)`
/// (the `mms_piso.rs` forcing).
fn source(x: f64, y: f64, convecting: bool) -> (f64, f64) {
    let conv = if convecting { RHO * 0.5 * PI } else { 0.0 };
    let fx = conv * (2.0 * PI * x).sin()
        + 2.0 * PI * PI * MU * u_exact(x, y)
        + PI * (PI * x).cos() * (PI * y).sin();
    let fy = conv * (2.0 * PI * y).sin()
        + 2.0 * PI * PI * MU * v_exact(x, y)
        + PI * (PI * x).sin() * (PI * y).cos();
    (fx, fy)
}

struct Measurement {
    l2_velocity: f64,
    max_div_rel: f64,
    steps: usize,
}

/// Smallest across-patch cell size (the explicit diffusion limit).
fn min_spacing(mesh: &PatchMesh) -> f64 {
    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());
        }
    }
    h
}

async fn march(
    mesh: PatchMesh,
    convection: PatchConvection,
    diffusion: NormalDiffusion,
    steady_tol: f64,
) -> CfdResult<Measurement> {
    let nu = MU / RHO;
    let h = min_spacing(&mesh);
    let env = |k: &str, d: f64| {
        std::env::var(k)
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(d)
    };
    let dt = env("RTX_CURV_DTFRAC", 0.4) * (h * h / (4.0 * nu)).min(h);
    let tol = env("RTX_CURV_TOL", 1e-5);
    let config = CfdConfig::new()
        .with_density(RHO)
        .with_viscosity(MU)
        .with_reference_velocity(1.0)
        .with_reference_length(1.0);
    let params = CurvilinearParameters {
        tolerance: tol,
        convection,
        normal_diffusion: diffusion,
        ..CurvilinearParameters::default()
    };
    let convecting = convection != PatchConvection::None;
    let mut solver = CurvilinearPisoSolver::new(config, params, mesh)?;
    solver.set_boundary_velocity(|x, y, _t| (u_exact(x, y), v_exact(x, y)));
    solver.set_momentum_source(move |x, y, _t| source(x, y, convecting));
    let mut field = PatchField::new(solver.mesh());
    solver.initialize(&mut field, |_, _| (0.0, 0.0));

    let mut steady = f64::INFINITY;
    let mut max_div_rel = 0.0_f64;
    let mut steps = 0;
    for step in 0..400_000 {
        let before = (field.u.clone(), field.v.clone());
        let r = solver.advance(&mut field, dt).await?;
        assert!(
            r.poisson_converged,
            "pressure solve did not converge at step {step}: {r:?}"
        );
        let flux_scale: f64 = field.flux.iter().map(|f| f.abs()).sum();
        max_div_rel = max_div_rel.max(r.max_divergence / flux_scale.max(1e-300));
        steps = step + 1;
        let change = field
            .u
            .iter()
            .zip(&before.0)
            .chain(field.v.iter().zip(&before.1))
            .map(|(a, b)| (a - b).abs())
            .fold(0.0, f64::max);
        steady = change / dt;
        if step % 20_000 == 0 && std::env::var("RTX_CURV_TRACE").is_ok() {
            println!(
                "  step {step} t={:.2} |du/dt| {steady:.3e} poisson iters {} div {:.2e}",
                solver.time(),
                r.poisson_iterations,
                r.max_divergence
            );
        }
        if steady < steady_tol {
            break;
        }
    }
    assert!(
        steady < steady_tol,
        "no steady state: |du/dt| = {steady:.3e}"
    );

    let mesh = solver.mesh();
    let (mut sq, mut vol) = (0.0, 0.0);
    for c in 0..mesh.cell_count() {
        let xy = mesh.centre(c);
        let eu = field.u[c] - u_exact(xy[0], xy[1]);
        let ev = field.v[c] - v_exact(xy[0], xy[1]);
        sq += (eu * eu + ev * ev) * mesh.area(c);
        vol += mesh.area(c);
    }
    Ok(Measurement {
        l2_velocity: (sq / vol).sqrt(),
        max_div_rel,
        steps,
    })
}

fn orders(errs: &[f64]) -> Vec<f64> {
    errs.windows(2).map(|p| (p[0] / p[1]).log2()).collect()
}

#[tokio::test]
async fn cartesian_patch_reproduces_the_staggered_piso_order_and_error() -> CfdResult<()> {
    // mms_piso.rs (staggered PISO, upwind): 3.516214e-2 / 1.953750e-2 /
    // 1.037512e-2 at 16/32/64, orders 0.85 / 0.91.
    let reference = [3.516214e-2, 1.953750e-2, 1.037512e-2];
    let only: Option<usize> = std::env::var("RTX_CURV_N")
        .ok()
        .and_then(|v| v.parse().ok());
    let mut errs = Vec::new();
    for (&n, &r) in [16usize, 32, 64].iter().zip(&reference) {
        if only.is_some_and(|o| o != n) {
            continue;
        }
        let m = march(
            cartesian(n, n, 1.0, 1.0, false)?,
            PatchConvection::Upwind,
            NormalDiffusion::Explicit,
            1e-6,
        )
        .await?;
        println!(
            "cartesian n={n}: L2 {:.6e} (staggered {r:.6e}, ratio {:.2}), max div {:.2e}, {} steps",
            m.l2_velocity,
            m.l2_velocity / r,
            m.max_div_rel,
            m.steps
        );
        assert!(m.max_div_rel < 1e-9, "divergence {:.3e}", m.max_div_rel);
        assert!(
            m.l2_velocity < 2.0 * r,
            "L2 {:.3e} > 2x staggered {r:.3e}",
            m.l2_velocity
        );
        errs.push(m.l2_velocity);
    }
    let o = orders(&errs);
    println!("cartesian orders {o:?}");
    if only.is_none() {
        assert!(o.iter().all(|&x| x > 0.75 && x < 2.3), "orders {o:?}");
    }
    Ok(())
}

#[tokio::test]
async fn skewed_annulus_stokes_limit_is_second_order() -> CfdResult<()> {
    for diffusion in [NormalDiffusion::Explicit, NormalDiffusion::LineImplicit] {
        let mut errs = Vec::new();
        for ns in [32usize, 64, 128] {
            let m = march(
                annulus_skewed([0.0, 0.0], 0.5, 1.5, ns, ns / 4, 0.3, 3.0)?,
                PatchConvection::None,
                diffusion,
                1e-7,
            )
            .await?;
            println!(
                "annulus {diffusion:?} ns={ns}: L2 {:.6e}, max div {:.2e}, {} steps",
                m.l2_velocity, m.max_div_rel, m.steps
            );
            assert!(m.max_div_rel < 1e-9, "divergence {:.3e}", m.max_div_rel);
            errs.push(m.l2_velocity);
        }
        let o = orders(&errs);
        println!("annulus {diffusion:?} Stokes orders {o:?}");
        assert!(
            o.iter().all(|&x| x >= 1.8),
            "Stokes-limit orders {o:?} (gate >= 1.8)"
        );
    }
    Ok(())
}

#[tokio::test]
async fn skewed_annulus_with_upwind_is_first_order() -> CfdResult<()> {
    let mut errs = Vec::new();
    for ns in [32usize, 64, 128] {
        let m = march(
            annulus_skewed([0.0, 0.0], 0.5, 1.5, ns, ns / 4, 0.3, 3.0)?,
            PatchConvection::Upwind,
            NormalDiffusion::Explicit,
            1e-6,
        )
        .await?;
        println!(
            "annulus upwind ns={ns}: L2 {:.6e}, max div {:.2e}, {} steps",
            m.l2_velocity, m.max_div_rel, m.steps
        );
        assert!(m.max_div_rel < 1e-9, "divergence {:.3e}", m.max_div_rel);
        errs.push(m.l2_velocity);
    }
    let o = orders(&errs);
    println!("annulus upwind orders {o:?}");
    assert!(o.iter().all(|&x| x > 0.7 && x < 1.6), "upwind orders {o:?}");
    Ok(())
}

/// P4 step 2 gate (ii): the van Albada deferred correction on the skewed
/// annulus beats upwind at every rung (upwind: 1.151502e-1 / 5.445770e-2 /
/// 3.119486e-2) with orders >= 1.4.
#[tokio::test]
async fn skewed_annulus_with_tvd_beats_upwind() -> CfdResult<()> {
    let upwind = [1.151502e-1, 5.445770e-2, 3.119486e-2];
    let mut errs = Vec::new();
    for (&ns, &u) in [32usize, 64, 128].iter().zip(&upwind) {
        let m = march(
            annulus_skewed([0.0, 0.0], 0.5, 1.5, ns, ns / 4, 0.3, 3.0)?,
            PatchConvection::TvdVanAlbada,
            NormalDiffusion::Explicit,
            1e-6,
        )
        .await?;
        println!(
            "annulus tvd ns={ns}: L2 {:.6e} (upwind {u:.6e}, ratio {:.2}), max div {:.2e}, {} steps",
            m.l2_velocity,
            m.l2_velocity / u,
            m.max_div_rel,
            m.steps
        );
        assert!(m.max_div_rel < 1e-9, "divergence {:.3e}", m.max_div_rel);
        assert!(
            m.l2_velocity < u,
            "TVD {:.4e} not below upwind {u:.4e}",
            m.l2_velocity
        );
        errs.push(m.l2_velocity);
    }
    let o = orders(&errs);
    println!("annulus tvd orders {o:?}");
    assert!(
        o.iter().all(|&x| x >= 1.4),
        "tvd orders {o:?} (gate >= 1.4)"
    );
    Ok(())
}

#[tokio::test]
async fn snapshot_restore_rerun_is_bit_identical() -> CfdResult<()> {
    let mesh = annulus_skewed([0.0, 0.0], 0.5, 1.5, 24, 6, 0.3, 2.0)?;
    let config = CfdConfig::new().with_density(RHO).with_viscosity(MU);
    let mut solver = CurvilinearPisoSolver::new(config, CurvilinearParameters::default(), mesh)?;
    solver.set_boundary_velocity(|x, y, t| (u_exact(x, y) * (1.0 + 0.1 * t), v_exact(x, y)));
    solver.set_momentum_source(|x, y, _| source(x, y, true));
    let mut field = PatchField::new(solver.mesh());
    solver.initialize(&mut field, |x, y| (u_exact(x, y), v_exact(x, y)));
    let dt = 1e-3;
    for _ in 0..5 {
        solver.advance(&mut field, dt).await?;
    }
    let saved = solver.snapshot();
    let field_saved = field.clone();
    for _ in 0..10 {
        solver.advance(&mut field, dt).await?;
    }
    let reference = (field.clone(), solver.time());
    solver.restore(&saved);
    field = field_saved;
    for _ in 0..10 {
        solver.advance(&mut field, dt).await?;
    }
    assert_eq!(solver.time().to_bits(), reference.1.to_bits());
    let mut max_diff = 0.0_f64;
    for (a, b) in field
        .u
        .iter()
        .zip(&reference.0.u)
        .chain(field.v.iter().zip(&reference.0.v))
        .chain(field.p.iter().zip(&reference.0.p))
        .chain(field.flux.iter().zip(&reference.0.flux))
    {
        max_diff = max_diff.max((a - b).abs());
    }
    assert!(max_diff == 0.0, "re-run differs by {max_diff:.3e}");
    Ok(())
}