//! 3D Stage 1, gate 4: (a) z-invariance identity — the 2D manufactured
//! problem (`embedded_mms.rs`, no body) on the 2D embedded solver with the
//! multigrid Poisson vs the 3D solver at nz = 1 (dz = 1, z slip): the same
//! values over 200 steps; (b) a three-dimensional manufactured solution
//! marched to steady state on `n³` cubes: errors monotone under
//! refinement, orders in [0.75, 2.3], TVD's error below upwind's, and the
//! field divergence-free on every cell.

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

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

// ---- the 2D manufactured problem (embedded_mms.rs) ----
fn u2(x: f64, y: f64) -> f64 {
    (PI * x).sin() * (PI * y).cos()
}
fn v2(x: f64, y: f64) -> f64 {
    -(PI * x).cos() * (PI * y).sin()
}
fn source2(x: f64, y: f64) -> (f64, f64) {
    let fx = RHO * 0.5 * PI * (2.0 * PI * x).sin()
        + 2.0 * PI * PI * MU * u2(x, y)
        + PI * (PI * x).cos() * (PI * y).sin();
    let fy = RHO * 0.5 * PI * (2.0 * PI * y).sin()
        + 2.0 * PI * PI * MU * v2(x, y)
        + PI * (PI * x).sin() * (PI * y).cos();
    (fx, fy)
}
fn boundary2(x: f64, y: f64) -> (f64, f64) {
    let u = if x <= 0.0 || x >= 1.0 { 0.0 } else { u2(x, y) };
    let v = if y <= 0.0 || y >= 1.0 { 0.0 } else { v2(x, y) };
    (u, v)
}

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

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

#[tokio::test]
async fn nz_one_reproduces_the_two_d_embedded_mms_march() {
    let n = 16;
    let h = 1.0 / n as f64;
    let dt = time_step(n);
    let config = CfdConfig::new()
        .with_density(RHO)
        .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");
    two.set_momentum_source(|x, y, _| source2(x, y));
    two.set_boundary_velocity(|x, y, _| boundary2(x, y));
    let mut three = Piso3Solver::new(
        fluid(),
        Piso3Parameters {
            corrector_steps: 2,
            tolerance: 1e-8,
            boundaries: Boundaries3 {
                z0: SideBoundary3::SlipWall,
                z1: SideBoundary3::SlipWall,
                ..Boundaries3::default()
            },
            ..Piso3Parameters::default()
        },
    );
    three.set_momentum_source(|x, y, _z, _t| {
        let (fx, fy) = source2(x, y);
        (fx, fy, 0.0)
    });
    three.set_boundary_velocity(|x, y, _z, _t| {
        let (u, v) = boundary2(x, y);
        (u, v, 0.0)
    });
    let mut a = FlowField::new(n, n, h, h).expect("field");
    let g = Grid3 {
        nx: n,
        ny: n,
        nz: 1,
        dx: h,
        dy: h,
        dz: 1.0,
    };
    let mut b = FlowField3D::new(g);
    for j in 0..n {
        let y = (j as f64 + 0.5) * h;
        a.u[(j, 0)] = boundary2(0.0, y).0;
        a.u[(j, n)] = boundary2(1.0, y).0;
        b.u[g.uface(0, j, 0)] = boundary2(0.0, y).0;
        b.u[g.uface(0, j, n)] = boundary2(1.0, y).0;
    }
    for i in 0..n {
        let x = (i as f64 + 0.5) * h;
        a.v[(0, i)] = boundary2(x, 0.0).1;
        a.v[(n, i)] = boundary2(x, 1.0).1;
        b.v[g.vface(0, 0, i)] = boundary2(x, 0.0).1;
        b.v[g.vface(0, n, i)] = boundary2(x, 1.0).1;
    }
    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 {
                worst = worst.max((a.u[(j, i)] - b.u[g.uface(0, j, i)]).abs());
            }
        }
        for j in 0..=n {
            for i in 0..n {
                worst = worst.max((a.v[(j, i)] - b.v[g.vface(0, j, i)]).abs());
            }
        }
        for j in 0..n {
            for i in 0..n {
                worst = worst.max((a.p[(j, i)] - b.p[g.cell(0, j, i)]).abs());
            }
        }
        assert!(
            worst == 0.0,
            "step {step}: 3D departs from the 2D embedded MMS march by {worst:.3e}"
        );
    }
    println!("  200 steps of the manufactured problem value-identical to the 2D embedded solver");
}

// ---- the 3D manufactured solution ----
// u = sin πx cos πy cos πz, v = cos πx sin πy cos πz, w = −2 cos πx cos πy sin πz
// (divergence-free), p = sin πx sin πy sin πz; source = ρ(u·∇)u + ∇p − μ∇²u.
fn u3(x: f64, y: f64, z: f64) -> f64 {
    (PI * x).sin() * (PI * y).cos() * (PI * z).cos()
}
fn v3(x: f64, y: f64, z: f64) -> f64 {
    (PI * x).cos() * (PI * y).sin() * (PI * z).cos()
}
fn w3(x: f64, y: f64, z: f64) -> f64 {
    -2.0 * (PI * x).cos() * (PI * y).cos() * (PI * z).sin()
}
fn p3(x: f64, y: f64, z: f64) -> f64 {
    (PI * x).sin() * (PI * y).sin() * (PI * z).sin()
}
fn source3(x: f64, y: f64, z: f64) -> (f64, f64, f64) {
    let (sx, cx) = (PI * x).sin_cos();
    let (sy, cy) = (PI * y).sin_cos();
    let (sz, cz) = (PI * z).sin_cos();
    let (u, v, w) = (u3(x, y, z), v3(x, y, z), w3(x, y, z));
    // Gradients.
    let (ux, uy, uz) = (PI * cx * cy * cz, -PI * sx * sy * cz, -PI * sx * cy * sz);
    let (vx, vy, vz) = (-PI * sx * sy * cz, PI * cx * cy * cz, -PI * cx * sy * sz);
    let (wx, wy, wz) = (
        2.0 * PI * sx * cy * sz,
        2.0 * PI * cx * sy * sz,
        -2.0 * PI * cx * cy * cz,
    );
    let (px, py, pz) = (PI * cx * sy * sz, PI * sx * cy * sz, PI * sx * sy * cz);
    // ∇²(product of three π-trig functions) = −3π² (itself).
    let lap = -3.0 * PI * PI;
    let fx = RHO * (u * ux + v * uy + w * uz) + px - MU * lap * u;
    let fy = RHO * (u * vx + v * vy + w * vz) + py - MU * lap * v;
    let fz = RHO * (u * wx + v * wy + w * wz) + pz - MU * lap * w;
    (fx, fy, fz)
}
/// The exact field on the cube's boundary with the normal components
/// snapped to their analytic zero.
fn boundary3(x: f64, y: f64, z: f64) -> (f64, f64, f64) {
    let u = if x <= 0.0 || x >= 1.0 {
        0.0
    } else {
        u3(x, y, z)
    };
    let v = if y <= 0.0 || y >= 1.0 {
        0.0
    } else {
        v3(x, y, z)
    };
    let w = if z <= 0.0 || z >= 1.0 {
        0.0
    } else {
        w3(x, y, z)
    };
    (u, v, w)
}

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

fn measure3(n: usize, scheme: ConvectionScheme) -> Measurement {
    let h = 1.0 / n as f64;
    let dt = time_step(n);
    let mut solver = Piso3Solver::new(
        fluid(),
        Piso3Parameters {
            corrector_steps: 2,
            tolerance: 1e-8,
            convection_scheme: scheme,
            ..Piso3Parameters::default()
        },
    );
    solver.set_momentum_source(|x, y, z, _t| source3(x, y, z));
    solver.set_boundary_velocity(|x, y, z, _t| boundary3(x, y, z));
    let g = Grid3 {
        nx: n,
        ny: n,
        nz: n,
        dx: h,
        dy: h,
        dz: h,
    };
    let mut f = FlowField3D::new(g);
    solver.initialize(&mut f);
    let mut steps = 0;
    for step in 0..200_000 {
        let (bu, bv, bw) = (f.u.clone(), f.v.clone(), f.w.clone());
        solver.advance(&mut f, dt);
        steps = step + 1;
        let mut change = 0.0_f64;
        for (a, b) in
            f.u.iter()
                .zip(&bu)
                .chain(f.v.iter().zip(&bv))
                .chain(f.w.iter().zip(&bw))
        {
            change = change.max((a - b).abs());
        }
        if change / dt < 1e-6 {
            break;
        }
    }
    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)]
                    - u3(i as f64 * h, (j as f64 + 0.5) * h, (k as f64 + 0.5) * h);
                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)]
                    - v3((i as f64 + 0.5) * h, j as f64 * h, (k as f64 + 0.5) * h);
                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)]
                    - w3((i as f64 + 0.5) * h, (j as f64 + 0.5) * h, k as f64 * h);
                sq += e * e * dv;
                vol += dv;
            }
        }
    }
    Measurement {
        l2_velocity: (sq / vol).sqrt(),
        max_div: f.max_divergence(),
        steps,
    }
}

fn ladder(resolutions: &[usize], scheme: ConvectionScheme) -> Vec<Measurement> {
    let ms: Vec<Measurement> = resolutions.iter().map(|&n| measure3(n, scheme)).collect();
    let errors: Vec<f64> = ms.iter().map(|m| m.l2_velocity).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!(
            "  {scheme:?} n = {n:3} ({:5} steps)  L2 velocity {:.6e}  order {rate}  max |div| {:.3e}",
            ms[i].steps, errors[i], ms[i].max_div
        );
    }
    assert!(
        errors.windows(2).all(|w| w[1] < w[0]),
        "{scheme:?}: errors not monotone {errors:?}"
    );
    for w in errors.windows(2) {
        let rate = (w[0] / w[1]).log2();
        assert!(
            rate > 0.75 && rate < 2.3,
            "{scheme:?}: observed order {rate:.3} outside [0.75, 2.3]; errors {errors:?}"
        );
    }
    for m in &ms {
        assert!(m.max_div < 1e-5, "{scheme:?}: max |div| {:.3e}", m.max_div);
    }
    ms
}

#[test]
fn three_d_mms_orders() {
    let resolutions = [12usize, 24];
    let up = ladder(&resolutions, ConvectionScheme::Upwind);
    let tvd = ladder(&resolutions, ConvectionScheme::TvdVanAlbada);
    let tvd_rate = (tvd[0].l2_velocity / tvd[1].l2_velocity).log2();
    assert!(tvd_rate > 1.1, "TVD order {tvd_rate:.3} not above 1.1");
    for (a, b) in up.iter().zip(&tvd) {
        assert!(
            b.l2_velocity < a.l2_velocity,
            "TVD error not below upwind's"
        );
    }
}

#[test]
#[ignore = "the three-rung ladder to n = 48 (minutes on the host)"]
fn three_d_mms_orders_three_rungs() {
    let resolutions = [12usize, 24, 48];
    ladder(&resolutions, ConvectionScheme::Upwind);
    ladder(&resolutions, ConvectionScheme::TvdVanAlbada);
}