//! 3D Stage 1, gates 4–6 on the DEVICE: the device-resident step against
//! the host step from the same start — the manufactured problem (upwind
//! and TVD), Beltrami with time-dependent boundary data, and Poiseuille
//! at nz = 1 and on the periodic extrusion (plane agreement). The
//! predictors are compiled without FMA contraction, so the only difference
//! from the host is the CG's reduction order: agreement to 1e-12 relative.
//!
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test three_d_device_step -- --nocapture`
#![cfg(feature = "cuda")]

use rtx_cfd::solvers::incompressible::three_d::piso_device::Piso3Device;
use rtx_cfd::solvers::incompressible::three_d::{
    Boundaries3, FlowField3D, Fluid3, Grid3, Piso3Parameters, Piso3Solver, SideBoundary3,
};
use rtx_cfd::solvers::incompressible::{ConvectionScheme, MgSmoother};
use std::f64::consts::PI;

const RHO: f64 = 1.0;

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

/// `tight` = the identity configuration: inner stop 1e-6 of the source
/// scale AND mass tolerance 1e-12 (its 0.1× floor on the inner stop is
/// what dominates near a steady state); otherwise the defaults.
fn params(scheme: ConvectionScheme, z: SideBoundary3, tight: bool) -> Piso3Parameters {
    Piso3Parameters {
        corrector_steps: 2,
        tolerance: if tight { 1e-12 } else { 1e-8 },
        inner_stop_factor: if tight { 1e-6 } else { 1e-2 },
        boundaries: Boundaries3 {
            z0: z,
            z1: z,
            ..Boundaries3::default()
        },
        // The device V-cycle is red-black; the host uses the same so the
        // preconditioners match.
        poisson_smoother: MgSmoother::RedBlack,
        convection_scheme: scheme,
        ..Piso3Parameters::default()
    }
}

/// Max |Δ| over u, v, w between two fields on the velocity scale, and max
/// |Δp| on the pressure scale floored at the dynamic pressure `ρ U²` (a
/// flat pressure field must not inflate a rounding-level difference).
fn compare(a: &FlowField3D, b: &FlowField3D) -> (f64, f64) {
    let mut worst = 0.0_f64;
    let mut scale = 0.0_f64;
    for (x, y) in
        a.u.iter()
            .zip(&b.u)
            .chain(a.v.iter().zip(&b.v))
            .chain(a.w.iter().zip(&b.w))
    {
        worst = worst.max((x - y).abs());
        scale = scale.max(x.abs());
    }
    let mut worst_p = 0.0_f64;
    let mut scale_p = 0.0_f64;
    for (x, y) in a.p.iter().zip(&b.p) {
        worst_p = worst_p.max((x - y).abs());
        scale_p = scale_p.max(x.abs());
    }
    let scale_p = scale_p.max(RHO * scale * scale);
    // One number: the larger of the two relative differences, on the velocity scale.
    (worst.max(worst_p / scale_p * scale), scale)
}

// ---- the manufactured solution of three_d_mms.rs ----
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 source3(mu: f64, 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));
    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);
    let lap = -3.0 * PI * PI;
    (
        RHO * (u * ux + v * uy + w * uz) + px - mu * lap * u,
        RHO * (u * vx + v * vy + w * vz) + py - mu * lap * v,
        RHO * (u * wx + v * wy + w * wz) + pz - mu * lap * w,
    )
}
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)
}

fn mms_pair(
    n: usize,
    scheme: ConvectionScheme,
    isf: bool,
) -> (Piso3Solver, Piso3Solver, Grid3, f64) {
    let mu = 0.05;
    let h = 1.0 / n as f64;
    let dt = 0.4 * (h * h / (4.0 * mu / RHO)).min(h);
    let mk = || {
        let mut s = Piso3Solver::new(fluid(mu), params(scheme, SideBoundary3::Velocity, isf));
        s.set_momentum_source(move |x, y, z, _t| source3(mu, x, y, z));
        s.set_boundary_velocity(|x, y, z, _t| boundary3(x, y, z));
        s
    };
    (
        mk(),
        mk(),
        Grid3 {
            nx: n,
            ny: n,
            nz: n,
            dx: h,
            dy: h,
            dz: h,
        },
        dt,
    )
}

fn run_pair(
    host: Piso3Solver,
    dev: Piso3Solver,
    g: Grid3,
    dt: f64,
    steps: usize,
    label: &str,
) -> (f64, f64) {
    let mut host = host;
    let mut fh = FlowField3D::new(g);
    let mut device = Piso3Device::new(dev, g);
    device.upload(&fh);
    let mut mismatched = 0usize;
    for _ in 0..steps {
        let rh = host.advance(&mut fh, dt);
        let rd = device.advance(dt);
        if rh.poisson_iterations != rd.poisson_iterations {
            mismatched += 1;
        }
    }
    let mut fd = FlowField3D::new(g);
    device.download(&mut fd);
    let (worst, scale) = compare(&fh, &fd);
    println!(
        "  {label}: {steps} steps, host vs device max |Δ| {worst:.3e} on a scale of {scale:.3e}; CG iteration counts differ on {mismatched} steps"
    );
    (worst, scale)
}

/// The CG's stop is a threshold on a reduction; host and device reduce in
/// different orders, so on a marginal step one side takes one more
/// iteration and the answers differ by the inner tolerance. The gate is
/// therefore taken at a TIGHT inner stop (1e-6 of the source scale), where
/// that flip cannot show above 1e-12; the default stop's difference is
/// reported alongside.
#[test]
fn device_step_matches_the_host_step_on_the_manufactured_problem() {
    for scheme in [ConvectionScheme::Upwind, ConvectionScheme::TvdVanAlbada] {
        let (h, d, g, dt) = mms_pair(12, scheme, false);
        run_pair(
            h,
            d,
            g,
            dt,
            100,
            &format!("MMS n 12 {scheme:?} (default tolerances)"),
        );
        let (h, d, g, dt) = mms_pair(12, scheme, true);
        let (worst, scale) = run_pair(h, d, g, dt, 100, &format!("MMS n 12 {scheme:?} (tight)"));
        assert!(
            worst <= 1e-11 * scale,
            "{scheme:?}: {worst:.3e} of {scale:.3e}"
        );
    }
}

// ---- Beltrami (three_d_beltrami.rs) ----
const NU: f64 = 0.02;
const A: f64 = PI / 4.0;
const D: f64 = PI / 2.0;
fn exact(x: f64, y: f64, z: f64, t: f64) -> (f64, f64, f64) {
    let decay = (-D * D * NU * t).exp();
    (
        -A * ((A * x).exp() * (A * y + D * z).sin() + (A * z).exp() * (A * x + D * y).cos())
            * decay,
        -A * ((A * y).exp() * (A * z + D * x).sin() + (A * x).exp() * (A * y + D * z).cos())
            * decay,
        -A * ((A * z).exp() * (A * x + D * y).sin() + (A * y).exp() * (A * z + D * x).cos())
            * decay,
    )
}

#[test]
fn device_step_matches_the_host_step_on_beltrami() {
    let n = 16;
    let h = 1.0 / n as f64;
    let dt = 0.25 * h * h / (4.0 * NU);
    let g = Grid3 {
        nx: n,
        ny: n,
        nz: n,
        dx: h,
        dy: h,
        dz: h,
    };
    for tight in [false, true] {
        let mk = || {
            let mut s = Piso3Solver::new(
                fluid(NU * RHO),
                params(ConvectionScheme::Upwind, SideBoundary3::Velocity, tight),
            );
            s.set_boundary_velocity(|x, y, z, t| exact(x, y, z, t));
            s
        };
        let mut host = mk();
        let mut fh = FlowField3D::new(g);
        for k in 0..n {
            for j in 0..n {
                for i in 0..=n {
                    fh.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 {
                    fh.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 {
                    fh.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 mut device = Piso3Device::new(mk(), g);
        device.upload(&fh);
        let mut mismatched = 0;
        for _ in 0..40 {
            let rh = host.advance(&mut fh, dt);
            let rd = device.advance(dt);
            if rh.poisson_iterations != rd.poisson_iterations {
                mismatched += 1;
            }
        }
        let mut fd = FlowField3D::new(g);
        device.download(&mut fd);
        let (worst, scale) = compare(&fh, &fd);
        println!(
            "  Beltrami n 16 (tight {tight}): 40 steps (time-dependent tables), host vs device max |Δ| {worst:.3e} on {scale:.3e}; CG counts differ on {mismatched} steps"
        );
        if tight {
            assert!(worst <= 1e-11 * scale, "{worst:.3e} of {scale:.3e}");
        }
    }
}

// ---- Poiseuille (three_d_poiseuille_identity.rs) ----
const MU_P: 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_P;
    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
}

#[test]
fn device_step_matches_the_host_step_on_poiseuille_and_is_z_invariant() {
    let n = 16;
    let h = 1.0 / n as f64;
    let dt = 0.4 * (h * h / (4.0 * MU_P)).min(h);
    for (nz, dz, z, tight) in [
        (1usize, 1.0, SideBoundary3::SlipWall, false),
        (1, 1.0, SideBoundary3::SlipWall, true),
        (4, h, SideBoundary3::Periodic, true),
    ] {
        let g = Grid3 {
            nx: n,
            ny: n,
            nz,
            dx: h,
            dy: h,
            dz,
        };
        let mk = || {
            let mut s = Piso3Solver::new(fluid(MU_P), params(ConvectionScheme::Upwind, z, tight));
            s.set_momentum_source(|_x, _y, _z, _t| (G, 0.0, 0.0));
            let u_hat = discrete_profile(n);
            s.set_boundary_velocity(move |x, y, _z, _t| {
                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], 0.0, 0.0)
                } else {
                    (0.0, 0.0, 0.0)
                }
            });
            s
        };
        let mut host = mk();
        let mut fh = FlowField3D::new(g);
        let u_hat = discrete_profile(n);
        for k in 0..nz {
            for (j, &uj) in u_hat.iter().enumerate() {
                fh.u[g.uface(k, j, 0)] = uj;
                fh.u[g.uface(k, j, n)] = uj;
            }
        }
        let mut device = Piso3Device::new(mk(), g);
        device.upload(&fh);
        let mut mismatched = 0;
        for _ in 0..300 {
            let rh = host.advance(&mut fh, dt);
            let rd = device.advance(dt);
            if rh.poisson_iterations != rd.poisson_iterations {
                mismatched += 1;
            }
        }
        let mut fd = FlowField3D::new(g);
        device.download(&mut fd);
        let (worst, scale) = compare(&fh, &fd);
        println!(
            "  Poiseuille n 16 nz {nz} (tight {tight}): 300 steps, host vs device max |Δ| {worst:.3e} on {scale:.3e}; CG counts differ on {mismatched} steps"
        );
        if tight {
            assert!(
                worst <= 1e-11 * scale,
                "nz {nz}: {worst:.3e} of {scale:.3e}"
            );
        }
        if nz > 1 {
            let plane =
                |f: &FlowField3D, k: usize| f.u[k * n * (n + 1)..(k + 1) * n * (n + 1)].to_vec();
            let p0 = plane(&fd, 0);
            let mut worst_plane = 0.0_f64;
            for k in 1..nz {
                for (a, b) in plane(&fd, k).iter().zip(&p0) {
                    worst_plane = worst_plane.max((a - b).abs());
                }
            }
            println!("  Poiseuille nz {nz}: device planes within {worst_plane:.3e} of {scale:.3e}");
            assert!(worst_plane <= 1e-11 * scale);
        }
    }
}

/// Diagnostic: where and when the Poiseuille host/device difference enters.
#[test]
#[ignore = "diagnostic: per-component host/device differences on Poiseuille variants"]
fn poiseuille_difference_diagnostic() {
    let n = 16;
    let h = 1.0 / n as f64;
    let dt = 0.4 * (h * h / (4.0 * MU_P)).min(h);
    for (label, dz, inlet, source) in [
        ("as is (dz 1, inlet profile, source G)", 1.0, true, true),
        ("dz = h", h, true, true),
        ("closed box (no inlet), source G", 1.0, false, true),
        ("inlet profile, no source", 1.0, true, false),
    ] {
        let g = Grid3 {
            nx: n,
            ny: n,
            nz: 1,
            dx: h,
            dy: h,
            dz,
        };
        let mk = || {
            let mut s = Piso3Solver::new(
                fluid(MU_P),
                params(ConvectionScheme::Upwind, SideBoundary3::SlipWall, true),
            );
            if source {
                s.set_momentum_source(|_x, _y, _z, _t| (G, 0.0, 0.0));
            }
            let u_hat = discrete_profile(n);
            s.set_boundary_velocity(move |x, y, _z, _t| {
                if inlet && (x <= 0.0 || x >= 1.0) {
                    let j = ((y / h - 0.5).round().max(0.0) as usize).min(n - 1);
                    (u_hat[j], 0.0, 0.0)
                } else {
                    (0.0, 0.0, 0.0)
                }
            });
            s
        };
        let mut host = mk();
        let mut fh = FlowField3D::new(g);
        if inlet {
            let u_hat = discrete_profile(n);
            for (j, &uj) in u_hat.iter().enumerate() {
                fh.u[g.uface(0, j, 0)] = uj;
                fh.u[g.uface(0, j, n)] = uj;
            }
        }
        let mut device = Piso3Device::new(mk(), g);
        device.upload(&fh);
        let mut fd = FlowField3D::new(g);
        for step in 1..=300 {
            let rh = host.advance(&mut fh, dt);
            let rd = device.advance(dt);
            if [1, 2, 10, 100, 300].contains(&step) {
                device.download(&mut fd);
                let du =
                    fh.u.iter()
                        .zip(&fd.u)
                        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
                let dv =
                    fh.v.iter()
                        .zip(&fd.v)
                        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
                let dp =
                    fh.p.iter()
                        .zip(&fd.p)
                        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
                let dpp = fh
                    .p_prime
                    .iter()
                    .zip(&fd.p_prime)
                    .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
                let dsp = fh
                    .sp
                    .iter()
                    .zip(&fd.sp)
                    .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
                println!(
                    "  {label} step {step}: Δu {du:.2e} Δv {dv:.2e} Δp {dp:.2e} Δp' {dpp:.2e} Δsp {dsp:.2e}; CG it host {} / dev {}; correctors {} / {}; residual {:.2e} / {:.2e}",
                    rh.poisson_iterations,
                    rd.poisson_iterations,
                    rh.corrector_steps_performed,
                    rd.corrector_steps_performed,
                    rh.final_residual,
                    rd.final_residual
                );
            }
        }
    }
}

/// The device step's cost at the anchor size (378 × 62 × 62, cubic cells,
/// the manufactured source, red-black + device CG), `RTX_PROFILE=1` for the
/// phase split. Recorded, not asserted.
#[test]
#[ignore = "bench: ms per device step at 378×62×62"]
fn bench_device_step_anchor_size() {
    let (nx, ny, nz) = (378usize, 62usize, 62usize);
    let h = 0.41 / ny as f64;
    let mu = 1.0e-3;
    let g = Grid3 {
        nx,
        ny,
        nz,
        dx: h,
        dy: h,
        dz: h,
    };
    let dt = 0.4 * (h * h / (4.0 * mu / RHO)).min(h / 2.0);
    let mut s = Piso3Solver::new(
        fluid(mu),
        params(
            ConvectionScheme::TvdVanAlbada,
            SideBoundary3::Velocity,
            false,
        ),
    );
    s.set_momentum_source(move |x, y, z, _t| source3(mu, x / 2.5, y / 0.41, z / 0.41));
    s.set_boundary_velocity(|x, y, z, _t| boundary3(x / 2.5, y / 0.41, z / 0.41));
    let mut device = Piso3Device::new(s, g);
    let f = FlowField3D::new(g);
    device.upload(&f);
    device.advance(dt);
    let t0 = std::time::Instant::now();
    let steps = 20;
    let mut it = 0;
    for _ in 0..steps {
        it += device.advance(dt).poisson_iterations;
    }
    let ms = t0.elapsed().as_secs_f64() * 1e3 / steps as f64;
    println!(
        "  device step at {nx}×{ny}×{nz} ({} cells): {ms:.1} ms per step, {:.1} CG iterations per step",
        g.cells(),
        it as f64 / steps as f64
    );
    if let Some(t) = device.timers() {
        let n = t.steps.max(1) as f64;
        println!(
            "  split per step: predictor {:.1} ms, poisson {:.1} ms, apply {:.1} ms ({} steps timed)",
            t.predictor_ns as f64 / n / 1e6,
            t.poisson_ns as f64 / n / 1e6,
            t.apply_ns as f64 / n / 1e6,
            t.steps
        );
    }
}