//! embedded3 gate 9b: Turek–Hron CFD1 (the cylinder with the rigid flag,
//! Re 20) on the 3D solver at ny = 41 — the 2D geometry extruded, at nz = 1
//! (dz = 1, z slip) and nz = 4 periodic: the settled control-volume drag
//! 15.6156 and surface drag 15.7126 of the 2D embedded record to
//! `rel < 5e-4` (printed-digit identity across the regimes).

use rtx_cfd::solvers::incompressible::embedded3::{
    Body, Boundaries, Field, Fluid, Grid, Parameters, Side, Solver,
};
use rtx_cfd::solvers::incompressible::{EmbeddedBody, MgSmoother};

const L: f64 = 2.5;
const H: f64 = 0.41;
const RHO: f64 = 1000.0;
const NU: f64 = 1e-3;
const U_MEAN: f64 = 0.2;
const SOR_DRAG_CV: f64 = 15.6156;
const SOR_DRAG_SURFACE: f64 = 15.7126;

fn inflow(y: f64) -> f64 {
    1.5 * U_MEAN * y * (H - y) / (0.5 * H).powi(2)
}

fn body2() -> EmbeddedBody {
    EmbeddedBody::union(
        EmbeddedBody::circle(0.2, 0.2, 0.05),
        EmbeddedBody::rectangle(0.20, 0.19, 0.6, 0.21),
    )
}

fn run(ny: usize, nz: usize, dz: f64, periodic: bool) -> (f64, f64, usize, usize) {
    let h = H / ny as f64;
    let nx = (L / h).round() as usize;
    let mu = RHO * NU;
    let u_peak = 1.5 * 1.5 * U_MEAN;
    let dt = 0.25 / (2.0 * u_peak / h + 4.0 * NU / (h * h));
    let z = if periodic {
        Side::Periodic
    } else {
        Side::SlipWall
    };
    let mut solver = Solver::new(
        Fluid {
            density: RHO,
            viscosity: mu,
            reference_velocity: U_MEAN,
            reference_length: 0.1,
        },
        Parameters {
            corrector_steps: 2,
            tolerance: 1e-7,
            boundaries: Boundaries {
                x1: Side::PressureOutlet,
                z0: z,
                z1: z,
                ..Boundaries::default()
            },
            poisson_smoother: MgSmoother::Lexicographic,
            ..Parameters::default()
        },
    );
    solver.set_boundary_velocity(|x, y, _z, _t| {
        if x <= 0.0 {
            (inflow(y), 0.0, 0.0)
        } else {
            (0.0, 0.0, 0.0)
        }
    });
    let lz = nz as f64 * dz;
    solver.set_body(Body::extruded(body2(), lz));
    let g = Grid {
        nx,
        ny,
        nz,
        dx: h,
        dy: h,
        dz,
    };
    let mut f = Field::new(g);
    for k in 0..nz {
        for j in 0..ny {
            let u0 = inflow((j as f64 + 0.5) * h);
            for i in 0..=nx {
                f.u[g.uface(k, j, i)] = u0;
            }
        }
    }
    solver.initialize(&mut f);
    let cv = (
        (0.10 / h).round() as usize,
        (0.75 / h).round() as usize,
        (0.05 / h).round() as usize,
        (0.36 / h).round() as usize,
        0,
        nz,
    );
    let flow_through = L / U_MEAN;
    let min_steps = (flow_through / dt).ceil() as usize;
    let mut history: Vec<f64> = Vec::new();
    let mut steps = 0;
    loop {
        solver.advance(&mut f, dt);
        steps += 1;
        if steps % 50 == 0 {
            let fx = solver
                .mask()
                .unwrap()
                .control_volume_force(&f, dt, RHO, mu, None, cv)[0]
                / lz;
            history.push(fx);
            let umax = f.u.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
            assert!(umax.is_finite(), "non-finite at step {steps}");
            if steps >= min_steps && history.len() > 4 {
                let now = history[history.len() - 1];
                let then = history[history.len() - 5];
                if ((now - then) / now).abs() < 1e-4 {
                    break;
                }
            }
        }
        assert!(steps < 400_000, "did not settle");
    }
    let mask = solver.mask().unwrap();
    let surface = mask.surface_force(solver.body().unwrap(), &f, mu, solver.time(), 0.5 * h);
    let drag_cv = mask.control_volume_force(&f, dt, RHO, mu, None, cv)[0] / lz;
    (drag_cv, surface.f[0] / lz, surface.skipped, steps)
}

#[test]
fn cfd1_at_ny_41_reproduces_the_two_d_record() {
    let ny = 41;
    let h = H / ny as f64;
    for (nz, dz, periodic) in [(1usize, 1.0, false), (4, h, true)] {
        let (cv, surface, skipped, steps) = run(ny, nz, dz, periodic);
        let rel_cv = ((cv - SOR_DRAG_CV) / SOR_DRAG_CV).abs();
        let rel_s = ((surface - SOR_DRAG_SURFACE) / SOR_DRAG_SURFACE).abs();
        println!(
            "  ny 41 nz {nz} periodic {periodic}: {steps} steps; CV drag {cv:.4} (record 15.6156, rel {rel_cv:.2e}); surface drag {surface:.4} (record 15.7126, rel {rel_s:.2e}, skipped {skipped})"
        );
        assert!(rel_cv < 5e-4, "CV drag {cv:.4} vs the record 15.6156");
        assert!(
            rel_s < 5e-4,
            "surface drag {surface:.4} vs the record 15.7126"
        );
    }
}

/// Diagnostic: which probes fail on the skipped surface samples, and the
/// surface force per z level, at nz 4 periodic after 200 steps.
#[test]
#[ignore = "diagnostic: skipped surface samples and per-level force on CFD1 at nz 4"]
fn skipped_samples_diagnostic() {
    let ny = 41;
    let h = H / ny as f64;
    let nx = (L / h).round() as usize;
    let mu = RHO * NU;
    let u_peak = 1.5 * 1.5 * U_MEAN;
    let dt = 0.25 / (2.0 * u_peak / h + 4.0 * NU / (h * h));
    let (nz, dz) = (4usize, h);
    let lz = nz as f64 * dz;
    let mut solver = Solver::new(
        Fluid {
            density: RHO,
            viscosity: mu,
            reference_velocity: U_MEAN,
            reference_length: 0.1,
        },
        Parameters {
            corrector_steps: 2,
            tolerance: 1e-7,
            boundaries: Boundaries {
                x1: Side::PressureOutlet,
                z0: Side::Periodic,
                z1: Side::Periodic,
                ..Boundaries::default()
            },
            poisson_smoother: MgSmoother::Lexicographic,
            ..Parameters::default()
        },
    );
    solver.set_boundary_velocity(|x, y, _z, _t| {
        if x <= 0.0 {
            (inflow(y), 0.0, 0.0)
        } else {
            (0.0, 0.0, 0.0)
        }
    });
    solver.set_body(Body::extruded(body2(), lz));
    let g = Grid {
        nx,
        ny,
        nz,
        dx: h,
        dy: h,
        dz,
    };
    let mut f = Field::new(g);
    for k in 0..nz {
        for j in 0..ny {
            let u0 = inflow((j as f64 + 0.5) * h);
            for i in 0..=nx {
                f.u[g.uface(k, j, i)] = u0;
            }
        }
    }
    solver.initialize(&mut f);
    for _ in 0..200 {
        solver.advance(&mut f, dt);
    }
    let mask = solver.mask().unwrap();
    let body = solver.body().unwrap();
    let samples = body.surface_samples(0.5 * h);
    let mut by_z: std::collections::BTreeMap<i64, (usize, usize, f64)> =
        std::collections::BTreeMap::new();
    let mut shown = 0;
    for s in &samples {
        let n = [s.nx, s.ny, s.nz];
        let key = (s.z * 1e4).round() as i64;
        let e = by_z.entry(key).or_insert((0, 0, 0.0));
        e.0 += 1;
        match mask.traction_at(body, &f, mu, solver.time(), [s.x, s.y, s.z], n) {
            Some(tr) => e.2 += tr[0] * s.area,
            None => {
                e.1 += 1;
                if shown < 6 {
                    shown += 1;
                    let at = |d: f64| [s.x + d * n[0], s.y + d * n[1], s.z + d * n[2]];
                    let (x1, x2) = (at(h), at(2.0 * h));
                    println!(
                        "  skipped ({:.4}, {:.4}, {:.4}) n ({:.2}, {:.2}): p1 {} p2 {} u1 {} u2 {}",
                        s.x,
                        s.y,
                        s.z,
                        s.nx,
                        s.ny,
                        mask.pressure_at(&f.p, x1[0], x1[1], x1[2]).is_some(),
                        mask.pressure_at(&f.p, x2[0], x2[1], x2[2]).is_some(),
                        mask.velocity_at(body, &f, x1[0], x1[1], x1[2], 0.0)
                            .is_some(),
                        mask.velocity_at(body, &f, x2[0], x2[1], x2[2], 0.0)
                            .is_some()
                    );
                }
            }
        }
    }
    for (z, (n, sk, fx)) in &by_z {
        println!(
            "  z {:.4}: {n} samples, {sk} skipped, drag contribution per unit depth {:.4}",
            *z as f64 / 1e4,
            fx / (lz / by_z.len() as f64)
        );
    }
}

/// Diagnostic: is the periodic nz 4 solution z-invariant, and does its
/// plane 0 equal the nz 1 solution, after 200 steps from the same start?
#[test]
#[ignore = "diagnostic: plane symmetry of CFD1 at nz 4 periodic"]
fn plane_symmetry_diagnostic() {
    let ny = 41;
    let h = H / ny as f64;
    let nx = (L / h).round() as usize;
    let mu = RHO * NU;
    let u_peak = 1.5 * 1.5 * U_MEAN;
    let dt = 0.25 / (2.0 * u_peak / h + 4.0 * NU / (h * h));
    let mk = |nz: usize, dz: f64, z: Side| {
        let mut s = Solver::new(
            Fluid {
                density: RHO,
                viscosity: mu,
                reference_velocity: U_MEAN,
                reference_length: 0.1,
            },
            Parameters {
                corrector_steps: 2,
                tolerance: 1e-7,
                boundaries: Boundaries {
                    x1: Side::PressureOutlet,
                    z0: z,
                    z1: z,
                    ..Boundaries::default()
                },
                poisson_smoother: MgSmoother::Lexicographic,
                ..Parameters::default()
            },
        );
        s.set_boundary_velocity(|x, y, _z, _t| {
            if x <= 0.0 {
                (inflow(y), 0.0, 0.0)
            } else {
                (0.0, 0.0, 0.0)
            }
        });
        s.set_body(Body::extruded(body2(), nz as f64 * dz));
        let g = Grid {
            nx,
            ny,
            nz,
            dx: h,
            dy: h,
            dz,
        };
        let mut f = Field::new(g);
        for k in 0..nz {
            for j in 0..ny {
                let u0 = inflow((j as f64 + 0.5) * h);
                for i in 0..=nx {
                    f.u[g.uface(k, j, i)] = u0;
                }
            }
        }
        s.initialize(&mut f);
        (s, f, g)
    };
    let (mut s1, mut f1, g1) = mk(1, 1.0, Side::SlipWall);
    let (mut s4, mut f4, g4) = mk(4, h, Side::Periodic);
    println!(
        "  ghost faces: nz 1 {} / nz 4 {} (per plane {})",
        s1.mask().unwrap().ghost_faces(),
        s4.mask().unwrap().ghost_faces(),
        s4.mask().unwrap().ghost_faces() / 4
    );
    for step in 1..=200 {
        s1.advance(&mut f1, dt);
        s4.advance(&mut f4, dt);
        if [1, 2, 10, 50, 200].contains(&step) {
            let plane = |f: &Field, g: &Grid, k: usize| {
                f.u[k * g.ny * (g.nx + 1)..(k + 1) * g.ny * (g.nx + 1)].to_vec()
            };
            let p0 = plane(&f4, &g4, 0);
            let mut zinv = 0.0_f64;
            for k in 1..4 {
                for (a, b) in plane(&f4, &g4, k).iter().zip(&p0) {
                    zinv = zinv.max((a - b).abs());
                }
            }
            let p1 = plane(&f1, &g1, 0);
            let vs1 = p0
                .iter()
                .zip(&p1)
                .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
            let wmax = f4.w.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
            println!(
                "  step {step}: nz 4 planes within {zinv:.3e}; plane 0 vs nz 1 {vs1:.3e}; max |w| {wmax:.3e}; ghost corr nz1 {:.3e} / nz4 {:.3e}",
                s1.ghost_correction(),
                s4.ghost_correction()
            );
        }
    }
}