//! 3D Stage 1, gate 3: the device-resident CG against the host PCG on the
//! same operator — converged, true residual below the tolerance, solutions
//! agreeing to the solve's accuracy, run-to-run bit-identical; the singular
//! closed box honours the anchor / the zero mean.
//!
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test three_d_cg_device -- --nocapture`
#![cfg(feature = "cuda")]

use rtx_cfd::solvers::incompressible::three_d::poisson::device_cg::{
    DevicePcgCache3, solve_multigrid_pcg3_device_cached,
};
use rtx_cfd::solvers::incompressible::three_d::poisson::{PoissonProblem3D, solve_multigrid_pcg3};
use rtx_cfd::solvers::incompressible::{MgSmoother, MultigridParameters};

/// A channel on cubic cells with a z-cylinder hole and an outlet Dirichlet
/// (`dirichlet = true`), or a closed Neumann box (`false`), seeded rhs.
fn channel(
    nx: usize,
    ny: usize,
    nz: usize,
    periodic_z: bool,
    dirichlet: bool,
    seed: u64,
) -> PoissonProblem3D {
    let mut p = PoissonProblem3D::new(nx, ny, nz);
    p.periodic_z = periodic_z;
    let h = 0.41 / ny as f64;
    let a = 3.24e-4 * h;
    let hole = |i: usize, j: usize| {
        let (x, y) = ((i as f64 + 0.5) * h, (j as f64 + 0.5) * h);
        (x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.05 * 0.05
    };
    for k in 0..nz {
        for j in 0..ny {
            for i in 0..nx {
                let idx = p.index(k, j, i);
                if hole(i, j) {
                    p.active[idx] = false;
                    continue;
                }
                if i + 1 < nx && !hole(i + 1, j) {
                    p.ae[idx] = a;
                }
                if i > 0 && !hole(i - 1, j) {
                    p.aw[idx] = a;
                }
                if j + 1 < ny && !hole(i, j + 1) {
                    p.an[idx] = a;
                }
                if j > 0 && !hole(i, j - 1) {
                    p.as_[idx] = a;
                }
                if k + 1 < nz || periodic_z {
                    p.at[idx] = a;
                }
                if k > 0 || periodic_z {
                    p.ab[idx] = a;
                }
                if dirichlet && i + 1 == nx {
                    p.extra_diag[idx] = 2.0 * a;
                }
            }
        }
    }
    let mut state = seed | 1;
    for idx in 0..nx * ny * nz {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        p.rhs[idx] = if p.active[idx] {
            1e-3 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5)
        } else {
            0.0
        };
    }
    if !dirichlet {
        // A compatible right-hand side: zero mean over the active cells.
        let cells: Vec<usize> = (0..nx * ny * nz).filter(|&i| p.active[i]).collect();
        let mean = cells.iter().map(|&i| p.rhs[i]).sum::<f64>() / cells.len() as f64;
        for &i in &cells {
            p.rhs[i] -= mean;
        }
    }
    p
}

fn bits(v: &[f64]) -> Vec<u64> {
    v.iter().map(|x| x.to_bits()).collect()
}

#[test]
fn device_cg_matches_the_host_pcg() {
    let params = MultigridParameters {
        smoother: MgSmoother::RedBlack,
        ..MultigridParameters::default()
    };
    let tol = 1e-12;
    let mut cache = DevicePcgCache3::default();
    for (nz, periodic, seed) in [(1usize, false, 3u64), (8, true, 5), (8, false, 7)] {
        let p = channel(96, 40, nz, periodic, true, seed);
        let n = p.nx * p.ny * p.nz;
        let mut host = vec![0.0; n];
        let sh = solve_multigrid_pcg3(&p, &mut host, &params, tol, None);
        let mut dev = vec![0.0; n];
        let sd = solve_multigrid_pcg3_device_cached(&p, &mut dev, &params, tol, None, &mut cache);
        let mut dev2 = vec![0.0; n];
        let sd2 = solve_multigrid_pcg3_device_cached(&p, &mut dev2, &params, tol, None, &mut cache);
        assert!(
            sh.converged && sd.converged,
            "nz {nz}: host {} / device {}",
            sh.converged,
            sd.converged
        );
        let res_dev = p.residual_l1(&dev);
        assert!(res_dev < tol, "nz {nz}: device true residual {res_dev:.3e}");
        let scale = host.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
        let worst = host
            .iter()
            .zip(&dev)
            .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
        println!(
            "  96×40×{nz} periodic {periodic}: host {} it / device {} it (cache hit second solve: {} it); |Δp| {worst:.3e} of {scale:.3e}; device residual {res_dev:.3e}",
            sh.iterations, sd.iterations, sd2.iterations
        );
        assert!(
            worst <= 1e-10 * scale,
            "nz {nz}: device and host differ by {worst:.3e} of {scale:.3e}"
        );
        assert_eq!(
            bits(&dev),
            bits(&dev2),
            "nz {nz}: the device solve is not run-to-run identical"
        );
        assert_eq!(sd.iterations, sd2.iterations);
    }
}

#[test]
fn the_singular_box_honours_the_anchor_and_the_mean() {
    let params = MultigridParameters {
        smoother: MgSmoother::RedBlack,
        ..MultigridParameters::default()
    };
    let tol = 1e-12;
    let p = channel(48, 20, 8, false, false, 9);
    let n = p.nx * p.ny * p.nz;
    assert!(p.is_singular());
    let cells: Vec<usize> = (0..n).filter(|&i| p.active[i]).collect();
    let anchor = cells[cells.len() / 3];
    let mut cache = DevicePcgCache3::default();
    let mut host = vec![0.0; n];
    let sh = solve_multigrid_pcg3(&p, &mut host, &params, tol, Some(anchor));
    let mut dev = vec![0.0; n];
    let sd =
        solve_multigrid_pcg3_device_cached(&p, &mut dev, &params, tol, Some(anchor), &mut cache);
    assert!(sh.converged && sd.converged);
    assert_eq!(
        dev[anchor].to_bits(),
        0.0f64.to_bits(),
        "anchor not at zero: {}",
        dev[anchor]
    );
    let scale = host.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
    let worst = host
        .iter()
        .zip(&dev)
        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
    println!(
        "  singular box with anchor: host {} it / device {} it; |Δp| {worst:.3e} of {scale:.3e}",
        sh.iterations, sd.iterations
    );
    assert!(worst <= 1e-10 * scale);
    let mut dev0 = vec![0.0; n];
    let s0 = solve_multigrid_pcg3_device_cached(&p, &mut dev0, &params, tol, None, &mut cache);
    assert!(s0.converged);
    let mean = cells.iter().map(|&i| dev0[i]).sum::<f64>() / cells.len() as f64;
    println!("  singular box without anchor: mean {mean:.3e} of {scale:.3e}");
    assert!(mean.abs() <= 1e-13 * scale, "mean {mean:.3e}");
}