//! embedded3 item 9b: the device step carrying a static cut-cell mask —
//! host vs device on the manufactured sphere (CutCell, merging in) and on
//! the CFD1 cylinder (nz 4 periodic), fields agreeing to the registered
//! `1e-9·scale` over 100 steps under TIGHT tolerances (inner stop 1e-6,
//! mass 1e-12 — item 7's identity rule; under the default tolerances the
//! two CGs differ at the inner-stop level and the agreement is reported),
//! equal corrector counts; the device step time recorded.
//!
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test embedded3_device_cut -- --nocapture`
#![cfg(feature = "cuda")]

mod embedded3_sphere;

use embedded3_sphere::{C, MU, R, RHO, boundary3, source3, u3, v3, w3};
use rtx_cfd::solvers::incompressible::ConvectionScheme;
use rtx_cfd::solvers::incompressible::embedded3::step::device::DeviceStep;
use rtx_cfd::solvers::incompressible::embedded3::{
    Body, Boundaries, Field, Fluid, Grid, Parameters, Side, Solver, WallScheme,
};

fn max_diff(a: &[f64], b: &[f64]) -> f64 {
    a.iter()
        .zip(b)
        .fold(0.0_f64, |m, (&x, &y)| m.max((x - y).abs()))
}

fn scale(a: &[f64]) -> f64 {
    a.iter().fold(0.0_f64, |m, &x| m.max(x.abs()))
}

/// Builds two identical solvers (host and device) and marches both.
fn march(
    make: &dyn Fn(bool) -> (Solver, Grid),
    dt: f64,
    steps: usize,
    label: &str,
    bound: Option<f64>,
    tight: bool,
) {
    let make = |t: bool| make(t);
    let (mut host, g) = make(tight);
    let mut fh = Field::new(g);
    host.initialize(&mut fh);
    let (mut dev_solver, _) = make(tight);
    let mut fd = Field::new(g);
    dev_solver.initialize(&mut fd);
    let mut device = DeviceStep::new(dev_solver, g);
    device.upload(&fd);
    let merged = device.merged_cells();
    let start = std::time::Instant::now();
    let mut differ = 0;
    for _ in 0..steps {
        let rh = host.advance(&mut fh, dt);
        let rd = device.advance(dt);
        if rh.corrector_steps_performed != rd.corrector_steps_performed {
            differ += 1;
        }
    }
    let device_seconds = start.elapsed().as_secs_f64();
    device.download(&mut fd);
    let du = max_diff(&fh.u, &fd.u)
        .max(max_diff(&fh.v, &fd.v))
        .max(max_diff(&fh.w, &fd.w));
    let su = scale(&fh.u).max(scale(&fh.v)).max(scale(&fh.w));
    let dp = max_diff(&fh.p, &fd.p);
    let sp = scale(&fh.p).max(RHO);
    println!(
        "  {label} (tight {tight}): {steps} steps, {merged} merged cells; host vs device max |Δu| {du:.3e} on {su:.3e}, max |Δp| {dp:.3e} on {sp:.3e}; corrector counts differ on {differ} steps; {:.1} ms per step (host + device)",
        1e3 * device_seconds / steps as f64
    );
    if let Some(bound) = bound {
        assert!(du < bound * su, "velocity differs: {du:.3e} on {su:.3e}");
        assert!(dp < bound * sp, "pressure differs: {dp:.3e} on {sp:.3e}");
        // The corrector counts are reported, not gated: at a mass
        // tolerance of 1e-12 the two residuals differ in their last digits
        // (the fields agree to rounding regardless).
    }
}

fn tolerances(tight: bool) -> (f64, f64) {
    if tight { (1e-12, 1e-6) } else { (1e-8, 1e-2) }
}

#[test]
fn sphere_cut_cell_host_equals_device() {
    let n = 12;
    let h = 1.0 / n as f64;
    let dt = 0.4 * (h * h / (4.0 * MU / RHO)).min(h);
    let make = |tight: bool| {
        let (tolerance, inner_stop_factor) = tolerances(tight);
        let mut solver = Solver::new(
            Fluid {
                density: RHO,
                viscosity: MU,
                reference_velocity: 1.0,
                reference_length: 1.0,
            },
            Parameters {
                corrector_steps: 2,
                tolerance,
                inner_stop_factor,
                convection_scheme: ConvectionScheme::Upwind,
                wall_scheme: WallScheme::CutCell,
                ..Parameters::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));
        solver.set_body(
            Body::sphere(move |_t| C, R)
                .with_surface_velocity(|x, y, z, _t| (u3(x, y, z), v3(x, y, z), w3(x, y, z))),
        );
        (solver, Grid::cubic(n, n, n, h))
    };
    march(&make, dt, 100, "sphere MMS n 12 CutCell", None, false);
    march(&make, dt, 100, "sphere MMS n 12 CutCell", Some(1e-9), true);
}

#[test]
fn cylinder_cut_cell_host_equals_device() {
    use rtx_cfd::solvers::incompressible::EmbeddedBody;
    let ny = 41;
    let nz = 4;
    let h = 0.41 / ny as f64;
    let nx = (2.2 / h).round() as usize;
    let dt = 2e-3;
    let make = |tight: bool| {
        let (tolerance, inner_stop_factor) = tolerances(tight);
        let mut solver = Solver::new(
            Fluid {
                density: 1.0,
                viscosity: 1e-3,
                reference_velocity: 0.2,
                reference_length: 0.1,
            },
            Parameters {
                corrector_steps: 2,
                tolerance,
                inner_stop_factor,
                convection_scheme: ConvectionScheme::Upwind,
                wall_scheme: WallScheme::CutCell,
                boundaries: Boundaries {
                    x1: Side::PressureOutlet,
                    z0: Side::Periodic,
                    z1: Side::Periodic,
                    ..Boundaries::default()
                },
                ..Parameters::default()
            },
        );
        solver.set_boundary_velocity(move |x, y, _z, _t| {
            if x <= 0.0 {
                (1.5 * 0.2 * 4.0 * y * (0.41 - y) / (0.41 * 0.41), 0.0, 0.0)
            } else {
                (0.0, 0.0, 0.0)
            }
        });
        let circle = EmbeddedBody::from_sdf(|x, y, _t| {
            ((x - 0.2_f64).powi(2) + (y - 0.2_f64).powi(2)).sqrt() - 0.05
        });
        solver.set_body(Body::extruded(circle, nz as f64 * h));
        (solver, Grid::cubic(nx, ny, nz, h))
    };
    march(
        &make,
        dt,
        100,
        "CFD1 cylinder ny 41 nz 4 periodic CutCell",
        None,
        false,
    );
    march(
        &make,
        dt,
        100,
        "CFD1 cylinder ny 41 nz 4 periodic CutCell",
        Some(1e-9),
        true,
    );
}