//! S2-7b instrument: the cut wall on a CURVED wall with an exact velocity
//! and an exact pressure that has a wall-normal gradient. Taylor–Couette
//! flow between two embedded concentric cylinders (inner `R1` rotating at
//! `OMEGA`, outer `R2` at rest), z periodic; the domain sides lie in the
//! solid. Exact: `u_θ = A r + B/r`, `p = ρ (A² r²/2 + 2AB ln r − B²/(2r²))`.
//! Reads: both walls' effective radii from the profile fit on full faces
//! (the S2-7 form), and the pressure error by cell class against the
//! exact p(r) (the S2-7b form). `RTX_E3_CURVED_MODE=rigid` is the
//! linear-exactness mode: solid-body rotation of both cylinders
//! (`u = Ω r e_θ`, `p = ρ Ω² r²/2`).
use rtx_cfd::solvers::incompressible::ConvectionScheme;
use rtx_cfd::solvers::incompressible::embedded3::{
    Body, Boundaries, FaceKind, Field, Fluid, Grid, Parameters, Side, Solver, WallScheme,
};

const MU: f64 = 0.1;
const RHO: f64 = 1.0;
const LX: f64 = 2.0;
const R1: f64 = 0.3;
const R2: f64 = 0.8;
const OMEGA: f64 = 1.0;
const CENTRE: (f64, f64) = (1.013, 1.017);

struct Exact {
    a: f64,
    b: f64,
}

impl Exact {
    fn new(rigid: bool) -> Self {
        let d = R2 * R2 - R1 * R1;
        if rigid {
            Exact { a: OMEGA, b: 0.0 }
        } else if outer_drives() {
            // The OUTER cylinder rotates at Ω, the inner is at rest: the
            // static wall is the convex one (the DFG's kind).
            Exact {
                a: OMEGA * R2 * R2 / d,
                b: -OMEGA * R1 * R1 * R2 * R2 / d,
            }
        } else {
            Exact {
                a: -OMEGA * R1 * R1 / d,
                b: OMEGA * R1 * R1 * R2 * R2 / d,
            }
        }
    }
    fn u_theta(&self, r: f64) -> f64 {
        self.a * r + self.b / r
    }
    fn p(&self, r: f64) -> f64 {
        RHO * (0.5 * self.a * self.a * r * r + 2.0 * self.a * self.b * r.ln()
            - 0.5 * self.b * self.b / (r * r))
    }
    /// The velocity at (x, y): the exact profile in the gap, the walls' own
    /// motion outside it (the inner body rotates, the outer is at rest).
    fn velocity(&self, x: f64, y: f64, rigid: bool) -> (f64, f64) {
        let (dx, dy) = (x - CENTRE.0, y - CENTRE.1);
        let r = (dx * dx + dy * dy).sqrt().max(1e-12);
        let ut = if r < R1 {
            if rigid || !outer_drives() { OMEGA * r } else { 0.0 }
        } else if r > R2 {
            if rigid || outer_drives() { OMEGA * r } else { 0.0 }
        } else {
            self.u_theta(r)
        };
        (-ut * dy / r, ut * dx / r)
    }
}

/// `RTX_E3_CURVED_MODE=outer`: the outer cylinder drives, the inner is at rest.
fn outer_drives() -> bool {
    std::env::var("RTX_E3_CURVED_MODE").is_ok_and(|v| v == "outer")
}

fn parameters() -> Parameters {
    Parameters {
        corrector_steps: 2,
        tolerance: 1e-10,
        convection_scheme: if std::env::var("RTX_E3_CURVED_SCHEME").is_ok_and(|v| v == "upwind") {
            ConvectionScheme::Upwind
        } else {
            ConvectionScheme::TvdVanAlbada
        },
        wall_scheme: WallScheme::CutCell,
        boundaries: Boundaries {
            z0: Side::Periodic,
            z1: Side::Periodic,
            ..Boundaries::default()
        },
        ..Parameters::default()
    }
}

/// Least squares of `u_θ = a r + b / r` through `(r, u_θ)` points.
fn fit_ab(points: &[(f64, f64)]) -> (f64, f64) {
    let (mut s11, mut s12, mut s22, mut t1, mut t2) = (0.0, 0.0, 0.0, 0.0, 0.0);
    for &(r, u) in points {
        let (f1, f2) = (r, 1.0 / r);
        s11 += f1 * f1;
        s12 += f1 * f2;
        s22 += f2 * f2;
        t1 += f1 * u;
        t2 += f2 * u;
    }
    let det = s11 * s22 - s12 * s12;
    ((t1 * s22 - t2 * s12) / det, (s11 * t2 - s12 * t1) / det)
}

fn reading(n: usize, rigid: bool) {
    let h = 1.0 / n as f64;
    let (nx, ny, nz) = ((LX * n as f64) as usize, (LX * n as f64) as usize, 2);
    let ex = Exact::new(rigid);
    let mut solver = Solver::new(
        Fluid {
            density: RHO,
            viscosity: MU,
            reference_velocity: OMEGA * R1,
            reference_length: R2 - R1,
        },
        parameters(),
    );
    let exb = Exact::new(rigid);
    solver.set_boundary_velocity(move |x, y, _z, _t| {
        let (u, v) = exb.velocity(x, y, rigid);
        (u, v, 0.0)
    });
    // The fluid is the gap: φ > 0 there.
    solver.set_body(
        Body::from_sdf(move |x, y, _z, _t| {
            let r = ((x - CENTRE.0).powi(2) + (y - CENTRE.1).powi(2)).sqrt();
            (r - R1).min(R2 - r)
        })
        .with_surface_velocity(move |x, y, _z, _t| {
            let (dx, dy) = (x - CENTRE.0, y - CENTRE.1);
            let r = (dx * dx + dy * dy).sqrt().max(1e-12);
            let inner_side = r < 0.5 * (R1 + R2);
            let moving = rigid || (inner_side != outer_drives());
            let ut = if moving { OMEGA * r } else { 0.0 };
            (-ut * dy / r, ut * dx / r, 0.0)
        }),
    );
    let g = Grid::cubic(nx, ny, nz, h);
    let mut field = Field::new(g);
    for k in 0..nz {
        for j in 0..ny {
            for i in 0..=nx {
                field.u[g.uface(k, j, i)] = ex.velocity(i as f64 * h, (j as f64 + 0.5) * h, rigid).0;
            }
        }
        for j in 0..=ny {
            for i in 0..nx {
                field.v[g.vface(k, j, i)] = ex.velocity((i as f64 + 0.5) * h, j as f64 * h, rigid).1;
            }
        }
    }
    solver.initialize(&mut field);
    let dt = 0.5 * h * h / (6.0 * MU);
    let t_end: f64 = std::env::var("RTX_E3_CURVED_T")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(4.0);
    let steps = (t_end / dt).ceil() as usize;
    let mut last_res = 0.0;
    for _ in 0..steps {
        last_res = solver.advance(&mut field, dt).final_residual;
    }
    let mask = solver.mask().expect("mask");
    let r_of = |x: f64, y: f64| ((x - CENTRE.0).powi(2) + (y - CENTRE.1).powi(2)).sqrt();
    // The profile on full faces two to N cells off both walls.
    let mut pts = Vec::new();
    let inner = |r: f64| r > R1 + 2.0 * h && r < R2 - 2.0 * h;
    for j in 0..ny {
        for i in 0..=nx {
            let f = g.uface(0, j, i);
            let (x, y) = (i as f64 * h, (j as f64 + 0.5) * h);
            let r = r_of(x, y);
            if inner(r) && mask.a_u(f) >= 1.0 {
                // u = −u_θ dy/r
                let dy = y - CENTRE.1;
                if dy.abs() > 0.3 * r {
                    pts.push((r, -field.u[f] * r / dy));
                }
            }
        }
    }
    for j in 0..=ny {
        for i in 0..nx {
            let f = g.vface(0, j, i);
            let (x, y) = ((i as f64 + 0.5) * h, j as f64 * h);
            let r = r_of(x, y);
            if inner(r) && mask.a_v(f) >= 1.0 {
                let dx = x - CENTRE.0;
                if dx.abs() > 0.3 * r {
                    pts.push((r, field.v[f] * r / dx));
                }
            }
        }
    }
    let (a, b) = fit_ab(&pts);
    // The walls: outer where u_θ = 0 (Couette) or the fit's own (rigid: the
    // inner/outer are not separable, report A and B); inner where u_θ = Ω r.
    let (off_in, off_out) = if rigid {
        (f64::NAN, f64::NAN)
    } else if outer_drives() {
        // inner: u_θ = 0 → r² = −b/a; outer: u_θ = Ω r → r² = b/(Ω − a)
        let r_in = (-b / a).sqrt();
        let r_out = (b / (OMEGA - a)).sqrt();
        ((r_in - R1) / h, (R2 - r_out) / h)
    } else {
        let r_out = (-b / a).sqrt();
        let r_in = (b / (OMEGA - a)).sqrt();
        ((r_in - R1) / h, (R2 - r_out) / h)
    };
    // The pressure against the exact p(r): mean-free over full cells in the gap.
    let scale = RHO * (OMEGA * R1).powi(2);
    let is_cut = |c: usize| mask.vol(c) < 1.0 - 1e-9;
    let (mut sum, mut cnt) = (0.0, 0usize);
    let cell_r = |c: usize| {
        let (_, j, i) = g.kji(c);
        r_of((i as f64 + 0.5) * h, (j as f64 + 0.5) * h)
    };
    for j in 0..ny {
        for i in 0..nx {
            let c = g.cell(0, j, i);
            if mask.cell_active(c) && !is_cut(c) && mask.master(c).is_none() {
                sum += field.p[c] - ex.p(cell_r(c));
                cnt += 1;
            }
        }
    }
    let level = sum / cnt.max(1) as f64;
    let (mut sq_full, mut sq_cut, mut n_cut, mut sum_cut) = (0.0, 0.0, 0usize, 0.0);
    let (mut sq_small, mut n_small, mut sq_large, mut n_large) = (0.0, 0usize, 0.0, 0usize);
    let (mut sq_in, mut n_in, mut sq_out, mut n_out) = (0.0, 0usize, 0.0, 0usize);
    for j in 0..ny {
        for i in 0..nx {
            let c = g.cell(0, j, i);
            if !mask.cell_active(c) || mask.master(c).is_some() {
                continue;
            }
            let e = (field.p[c] - level - ex.p(cell_r(c))) / scale;
            if is_cut(c) {
                sq_cut += e * e;
                sum_cut += e;
                n_cut += 1;
                if mask.vol(c) < 0.5 {
                    sq_small += e * e;
                    n_small += 1;
                } else {
                    sq_large += e * e;
                    n_large += 1;
                }
                if cell_r(c) < 0.5 * (R1 + R2) {
                    sq_in += e * e;
                    n_in += 1;
                } else {
                    sq_out += e * e;
                    n_out += 1;
                }
            } else {
                sq_full += e * e;
            }
        }
    }
    let rms = |sq: f64, n: usize| (sq / n.max(1) as f64).sqrt();
    // S2-7b diagnostics: the wall's mass flux per cut cell (the body's
    // velocity is tangential: any flux is the facet normal's), in units of
    // Ω R1 h², and the ghost faces' error against the exact field (u faces
    // of kind Ghost), in units of Ω R1.
    let body = solver.body().expect("body");
    let (wf, _) = mask.wall_flux_table(body, solver.time());
    let (mut wsum, mut wmax, mut wn) = (0.0f64, 0.0f64, 0usize);
    for j in 0..ny {
        for i in 0..nx {
            let c = g.cell(0, j, i);
            if mask.cell_active(c) && is_cut(c) {
                let q = wf[c].abs() / (OMEGA * R1 * h * h);
                wsum += q;
                wmax = wmax.max(q);
                wn += 1;
            }
        }
    }
    let (mut gsq, mut gn, mut gmax) = (0.0f64, 0usize, 0.0f64);
    for j in 0..ny {
        for i in 0..=nx {
            let f = g.uface(0, j, i);
            if mask.u_kind(f) == rtx_cfd::solvers::incompressible::embedded3::FaceKind::Ghost {
                let e = (field.u[f] - ex.velocity(i as f64 * h, (j as f64 + 0.5) * h, rigid).0) / (OMEGA * R1);
                gsq += e * e;
                gn += 1;
                gmax = gmax.max(e.abs());
            }
        }
    }
    println!(
        "    wall flux per cut cell (of ΩR1 h²): mean {:.3e} max {:.3e} ({wn}); ghost u faces vs exact (of ΩR1): rms {:.3e} max {:.3e} ({gn})",
        wsum / wn.max(1) as f64,
        wmax,
        (gsq / gn.max(1) as f64).sqrt(),
        gmax
    );
    println!(
        "  {} n {n}: walls' offsets {off_in:+.4} h (inner) {off_out:+.4} h (outer), positive = inside the fluid; fit A {a:.5} B {b:.5} (exact {:.5} {:.5}, {} points); pressure error of ρ(ΩR1)²: full cells {:.3e} ({cnt}), cut cells {:.3e} mean {:+.3e} ({n_cut}), fraction < 0.5 {:.3e} ({n_small}), ≥ 0.5 {:.3e} ({n_large}), inner wall {:.3e} ({n_in}), outer wall {:.3e} ({n_out}); merged {}; residual {last_res:.1e}",
        if rigid { "rigid" } else if outer_drives() { "outer-driven" } else { "couette" },
        ex.a,
        ex.b,
        pts.len(),
        rms(sq_full, cnt),
        rms(sq_cut, n_cut),
        sum_cut / n_cut.max(1) as f64,
        rms(sq_small, n_small),
        rms(sq_large, n_large),
        rms(sq_in, n_in),
        rms(sq_out, n_out),
        mask.merged_cells()
    );
}

#[test]
#[ignore = "S2-7b instrument: Taylor–Couette between embedded cylinders (minutes per rung on the host)"]
fn curved_wall_effective_position_and_pressure() {
    let ns: Vec<usize> = std::env::var("RTX_E3_CURVED_NS")
        .ok()
        .map(|v| v.split(',').filter_map(|t| t.trim().parse().ok()).collect())
        .unwrap_or_else(|| vec![16, 32]);
    let rigid = std::env::var("RTX_E3_CURVED_MODE").is_ok_and(|v| v == "rigid");
    for n in ns {
        reading(n, rigid);
    }
}

/// S2-7b B1: the operator probe on this instrument — one predictor and one
/// corrector from the exact field at the faces' open-part centroids (exact
/// p at the cells); the predictor's residual per face, recovered from the
/// one correction, in units of the exact field's largest acceleration
/// (Ω² R2), by aperture band and by WALL (inner / outer); the correction's
/// pressure at each wall's cut cells (of ρ(ΩR1)²).
fn probe(n: usize, rigid: bool) {
    let h = 1.0 / n as f64;
    let (nx, ny, nz) = ((LX * n as f64) as usize, (LX * n as f64) as usize, 2);
    let ex = Exact::new(rigid);
    let mut params = parameters();
    params.corrector_steps = 1;
    let mut solver = Solver::new(
        Fluid {
            density: RHO,
            viscosity: MU,
            reference_velocity: OMEGA * R1,
            reference_length: R2 - R1,
        },
        params,
    );
    let exb = Exact::new(rigid);
    solver.set_boundary_velocity(move |x, y, _z, _t| {
        let (u, v) = exb.velocity(x, y, rigid);
        (u, v, 0.0)
    });
    solver.set_body(
        Body::from_sdf(move |x, y, _z, _t| {
            let r = ((x - CENTRE.0).powi(2) + (y - CENTRE.1).powi(2)).sqrt();
            (r - R1).min(R2 - r)
        })
        .with_surface_velocity(move |x, y, _z, _t| {
            let (dx, dy) = (x - CENTRE.0, y - CENTRE.1);
            let r = (dx * dx + dy * dy).sqrt().max(1e-12);
            let inner_side = r < 0.5 * (R1 + R2);
            let moving = rigid || (inner_side != outer_drives());
            let ut = if moving { OMEGA * r } else { 0.0 };
            (-ut * dy / r, ut * dx / r, 0.0)
        }),
    );
    let g = Grid::cubic(nx, ny, nz, h);
    let mut field = Field::new(g);
    solver.initialize(&mut field);
    let tables = solver.mask().expect("mask").face_shift_tables().expect("shift tables").clone();
    for k in 0..nz {
        for j in 0..ny {
            for i in 0..=nx {
                let f = g.uface(k, j, i);
                let t = &tables[0][3 * f..3 * f + 3];
                field.u[f] = ex.velocity(i as f64 * h + t[0], (j as f64 + 0.5) * h + t[1], rigid).0;
            }
        }
        for j in 0..=ny {
            for i in 0..nx {
                let f = g.vface(k, j, i);
                let t = &tables[1][3 * f..3 * f + 3];
                field.v[f] = ex.velocity((i as f64 + 0.5) * h + t[0], j as f64 * h + t[1], rigid).1;
            }
        }
    }
    let r_of = |x: f64, y: f64| ((x - CENTRE.0).powi(2) + (y - CENTRE.1).powi(2)).sqrt();
    for idx in 0..g.cells() {
        let (_, j, i) = g.kji(idx);
        let r = r_of((i as f64 + 0.5) * h, (j as f64 + 0.5) * h).clamp(R1, R2);
        field.p[idx] = ex.p(r);
    }
    {
        let (body, mask) = (solver.body().expect("body"), solver.mask().expect("mask"));
        mask.impose(body, &mut field.u, &mut field.v, &mut field.w, solver.time());
    }
    let dt = 0.5 * h * h / (6.0 * MU);
    solver.advance(&mut field, dt);
    let mask = solver.mask().expect("mask");
    let pp = &field.p_prime;
    let a_max = OMEGA * OMEGA * R2;
    let is_cut = |c: usize| mask.vol(c) < 1.0 - 1e-9;
    // [wall][band]: bands α<¼, ¼–½, ½–¾, ¾–1, full next to cut.
    let mut acc: [[Vec<f64>; 5]; 2] = Default::default();
    let bin_of = |a: f64, near: bool| -> Option<usize> {
        if a < 1.0 {
            Some(((a * 4.0).floor() as usize).min(3))
        } else if near {
            Some(4)
        } else {
            None
        }
    };
    let wall_of = |x: f64, y: f64| usize::from(r_of(x, y) >= 0.5 * (R1 + R2));
    for j in 0..ny {
        for i in 1..nx {
            let f = g.uface(0, j, i);
            if mask.u_kind(f) != FaceKind::Fluid {
                continue;
            }
            let (cm, cp) = (g.cell(0, j, i - 1), g.cell(0, j, i));
            let Some(b) = bin_of(mask.a_u(f), is_cut(cm) || is_cut(cp)) else { continue };
            let star = field.u[f] + (dt / RHO) * mask.grad_weight(0, f) * (pp[cp] - pp[cm]) / h;
            acc[wall_of(i as f64 * h, (j as f64 + 0.5) * h)][b].push((star - field.u_old[f]) / dt / a_max);
        }
    }
    for j in 1..ny {
        for i in 0..nx {
            let f = g.vface(0, j, i);
            if mask.v_kind(f) != FaceKind::Fluid {
                continue;
            }
            let (cm, cp) = (g.cell(0, j - 1, i), g.cell(0, j, i));
            let Some(b) = bin_of(mask.a_v(f), is_cut(cm) || is_cut(cp)) else { continue };
            let star = field.v[f] + (dt / RHO) * mask.grad_weight(1, f) * (pp[cp] - pp[cm]) / h;
            acc[wall_of((i as f64 + 0.5) * h, j as f64 * h)][b].push((star - field.v_old[f]) / dt / a_max);
        }
    }
    let scale = RHO * (OMEGA * R1).powi(2);
    let (mut sum, mut cnt) = (0.0, 0usize);
    for j in 0..ny {
        for i in 0..nx {
            let c = g.cell(0, j, i);
            if mask.cell_active(c) && !is_cut(c) && mask.master(c).is_none() {
                sum += pp[c];
                cnt += 1;
            }
        }
    }
    let lvl = sum / cnt.max(1) as f64;
    let mut psq = [0.0f64; 3];
    let mut pn = [0usize; 3];
    for j in 0..ny {
        for i in 0..nx {
            let c = g.cell(0, j, i);
            if !mask.cell_active(c) || mask.master(c).is_some() {
                continue;
            }
            let e = (pp[c] - lvl) / scale;
            let w = if is_cut(c) { wall_of((i as f64 + 0.5) * h, (j as f64 + 0.5) * h) } else { 2 };
            psq[w] += e * e;
            pn[w] += 1;
        }
    }
    let rms = |v: &[f64]| (v.iter().map(|x| x * x).sum::<f64>() / v.len().max(1) as f64).sqrt();
    let names = ["α<¼", "¼–½", "½–¾", "¾–1", "full next to cut"];
    let mode = if rigid { "rigid" } else if outer_drives() { "outer-driven" } else { "couette" };
    for (w, wname) in ["inner wall", "outer wall"].iter().enumerate() {
        let mut line = format!("  probe {mode} n {n} {wname}:");
        for (b, name) in names.iter().enumerate() {
            line += &format!(" {name} {} rms {:.3e};", acc[w][b].len(), rms(&acc[w][b]));
        }
        line += &format!(" p' at its cut cells {:.3e} ({})", (psq[w] / pn[w].max(1) as f64).sqrt(), pn[w]);
        println!("{line}");
    }
    println!("  probe {mode} n {n} interior p' {:.3e} ({})", (psq[2] / pn[2].max(1) as f64).sqrt(), pn[2]);
}

#[test]
#[ignore = "S2-7b B1 probe: the discrete operator on the exact Taylor–Couette field (seconds per rung)"]
fn curved_operator_probe() {
    let ns: Vec<usize> = std::env::var("RTX_E3_CURVED_NS")
        .ok()
        .map(|v| v.split(',').filter_map(|t| t.trim().parse().ok()).collect())
        .unwrap_or_else(|| vec![16, 32, 64]);
    let rigid = std::env::var("RTX_E3_CURVED_MODE").is_ok_and(|v| v == "rigid");
    for n in ns {
        probe(n, rigid);
    }
}