//! 3D Stage 1, gate 5: the Ethier–Steinman (Beltrami) exact unsteady
//! solution on the unit cube with time-dependent Dirichlet data — the
//! transient machinery with no source term. (1) The L2 velocity error at
//! `T` falls under `dt ~ h²` refinement at order ≥ 0.75; (2) the kinetic
//! energy decay follows the closed form within the discretisation error.

use rtx_cfd::solvers::incompressible::three_d::{
    FlowField3D, Fluid3, Grid3, Piso3Parameters, Piso3Solver,
};
use std::f64::consts::PI;

const RHO: f64 = 1.0;
const NU: f64 = 0.02;
const A: f64 = PI / 4.0;
const D: f64 = PI / 2.0;
const T_END: f64 = 0.25;

fn exact(x: f64, y: f64, z: f64, t: f64) -> (f64, f64, f64) {
    let decay = (-D * D * NU * t).exp();
    let u = -A
        * ((A * x).exp() * (A * y + D * z).sin() + (A * z).exp() * (A * x + D * y).cos())
        * decay;
    let v = -A
        * ((A * y).exp() * (A * z + D * x).sin() + (A * x).exp() * (A * y + D * z).cos())
        * decay;
    let w = -A
        * ((A * z).exp() * (A * x + D * y).sin() + (A * y).exp() * (A * z + D * x).cos())
        * decay;
    (u, v, w)
}

/// The boundary data as FACE AVERAGES (3 × 3 Gauss over the face): the
/// exact field has a non-zero normal velocity on the closed box, and its
/// face-centre samples leave an O(h²) net inflow a pure-Neumann projection
/// can only spread uniformly; the face-averaged fluxes of a divergence-free
/// field sum to zero to quadrature accuracy, so the box is compatible.
fn face_averaged(h: f64) -> impl Fn(f64, f64, f64, f64) -> (f64, f64, f64) {
    const G: [f64; 3] = [-0.774_596_669_241_483_4, 0.0, 0.774_596_669_241_483_4];
    const W: [f64; 3] = [5.0 / 9.0, 8.0 / 9.0, 5.0 / 9.0];
    move |x: f64, y: f64, z: f64, t: f64| {
        let on_x = x <= 0.0 || x >= 1.0;
        let on_y = y <= 0.0 || y >= 1.0;
        let on_z = z <= 0.0 || z >= 1.0;
        if !(on_x || on_y || on_z) {
            return exact(x, y, z, t);
        }
        let (mut u, mut v, mut w) = (0.0, 0.0, 0.0);
        for (a, wa) in G.iter().zip(&W) {
            for (b, wb) in G.iter().zip(&W) {
                let (xx, yy, zz) = if on_x {
                    (x, y + 0.5 * h * a, z + 0.5 * h * b)
                } else if on_y {
                    (x + 0.5 * h * a, y, z + 0.5 * h * b)
                } else {
                    (x + 0.5 * h * a, y + 0.5 * h * b, z)
                };
                let e = exact(xx, yy, zz, t);
                u += 0.25 * wa * wb * e.0;
                v += 0.25 * wa * wb * e.1;
                w += 0.25 * wa * wb * e.2;
            }
        }
        (u, v, w)
    }
}

struct Measurement {
    l2: f64,
    /// `max |div − mean(div)|`, with the mean reported (the residual
    /// incompatibility of the face-averaged data: quadrature level).
    max_div: f64,
    mean_div: f64,
    energy_ratio: f64,
    steps: usize,
}

fn measure(n: usize) -> Measurement {
    let h = 1.0 / n as f64;
    // dt ~ h²: the diffusion limit with a margin.
    let dt = 0.25 * h * h / (4.0 * NU);
    let steps = (T_END / dt).ceil() as usize;
    let dt = T_END / steps as f64;
    let mut solver = Piso3Solver::new(
        Fluid3 {
            density: RHO,
            viscosity: NU * RHO,
            reference_velocity: 1.0,
            reference_length: 1.0,
        },
        Piso3Parameters {
            corrector_steps: 2,
            tolerance: 1e-8,
            ..Piso3Parameters::default()
        },
    );
    solver.set_boundary_velocity(face_averaged(h));
    let g = Grid3 {
        nx: n,
        ny: n,
        nz: n,
        dx: h,
        dy: h,
        dz: h,
    };
    let mut f = FlowField3D::new(g);
    for k in 0..n {
        for j in 0..n {
            for i in 0..=n {
                f.u[g.uface(k, j, i)] = exact(
                    i as f64 * h,
                    (j as f64 + 0.5) * h,
                    (k as f64 + 0.5) * h,
                    0.0,
                )
                .0;
            }
        }
    }
    for k in 0..n {
        for j in 0..=n {
            for i in 0..n {
                f.v[g.vface(k, j, i)] = exact(
                    (i as f64 + 0.5) * h,
                    j as f64 * h,
                    (k as f64 + 0.5) * h,
                    0.0,
                )
                .1;
            }
        }
    }
    for k in 0..=n {
        for j in 0..n {
            for i in 0..n {
                f.w[g.wface(k, j, i)] = exact(
                    (i as f64 + 0.5) * h,
                    (j as f64 + 0.5) * h,
                    k as f64 * h,
                    0.0,
                )
                .2;
            }
        }
    }
    let e0 = f.kinetic_energy(RHO);
    for _ in 0..steps {
        solver.advance(&mut f, dt);
    }
    let e1 = f.kinetic_energy(RHO);
    let (mut sq, mut vol) = (0.0, 0.0);
    let dv = h * h * h;
    for k in 0..n {
        for j in 0..n {
            for i in 1..n {
                let e = f.u[g.uface(k, j, i)]
                    - exact(
                        i as f64 * h,
                        (j as f64 + 0.5) * h,
                        (k as f64 + 0.5) * h,
                        T_END,
                    )
                    .0;
                sq += e * e * dv;
                vol += dv;
            }
        }
    }
    for k in 0..n {
        for j in 1..n {
            for i in 0..n {
                let e = f.v[g.vface(k, j, i)]
                    - exact(
                        (i as f64 + 0.5) * h,
                        j as f64 * h,
                        (k as f64 + 0.5) * h,
                        T_END,
                    )
                    .1;
                sq += e * e * dv;
                vol += dv;
            }
        }
    }
    for k in 1..n {
        for j in 0..n {
            for i in 0..n {
                let e = f.w[g.wface(k, j, i)]
                    - exact(
                        (i as f64 + 0.5) * h,
                        (j as f64 + 0.5) * h,
                        k as f64 * h,
                        T_END,
                    )
                    .2;
                sq += e * e * dv;
                vol += dv;
            }
        }
    }
    let mut divs = Vec::with_capacity(n * n * n);
    for k in 0..n {
        for j in 0..n {
            for i in 0..n {
                divs.push(
                    (f.u[g.uface(k, j, i + 1)] - f.u[g.uface(k, j, i)]) / h
                        + (f.v[g.vface(k, j + 1, i)] - f.v[g.vface(k, j, i)]) / h
                        + (f.w[g.wface(k + 1, j, i)] - f.w[g.wface(k, j, i)]) / h,
                );
            }
        }
    }
    let mean_div = divs.iter().sum::<f64>() / divs.len() as f64;
    let max_div = divs
        .iter()
        .fold(0.0_f64, |m, d| m.max((d - mean_div).abs()));
    Measurement {
        l2: (sq / vol).sqrt(),
        max_div,
        mean_div,
        energy_ratio: e1 / e0,
        steps,
    }
}

#[test]
fn beltrami_error_falls_under_space_time_refinement() {
    let resolutions = [8usize, 16, 32];
    let ms: Vec<Measurement> = resolutions.iter().map(|&n| measure(n)).collect();
    let exact_ratio = (-2.0 * D * D * NU * T_END).exp();
    let errors: Vec<f64> = ms.iter().map(|m| m.l2).collect();
    for (i, &n) in resolutions.iter().enumerate() {
        let rate = if i == 0 {
            "    -".to_string()
        } else {
            format!("{:5.2}", (errors[i - 1] / errors[i]).log2())
        };
        println!(
            "  n = {n:2} ({:5} steps)  L2 = {:.6e}  order {rate}  E(T)/E(0) = {:.5} (exact {exact_ratio:.5}, error {:.2e})  max |div − mean| {:.2e} (mean {:.2e})",
            ms[i].steps,
            ms[i].l2,
            ms[i].energy_ratio,
            (ms[i].energy_ratio - exact_ratio).abs(),
            ms[i].max_div,
            ms[i].mean_div
        );
    }
    assert!(
        errors.windows(2).all(|w| w[1] < w[0]),
        "errors not monotone: {errors:?}"
    );
    for w in errors.windows(2) {
        let rate = (w[0] / w[1]).log2();
        assert!(
            rate >= 0.75,
            "observed order {rate:.3} below 0.75; errors {errors:?}"
        );
    }
    for m in &ms {
        assert!(m.max_div < 1e-6, "max |div − mean| {:.3e}", m.max_div);
    }
    // The energy decay: the discretisation error at each rung bounds it.
    let mut e_err: Vec<f64> = ms
        .iter()
        .map(|m| (m.energy_ratio - exact_ratio).abs())
        .collect();
    assert!(
        e_err.windows(2).all(|w| w[1] < w[0]),
        "energy error not falling: {e_err:?}"
    );
    e_err.clear();
}