//! The added-mass piston: partitioned FSI on a real discretised fluid,
//! against a closed form.
//!
//! An elastic piston (mass `M`, spring `K`, one degree of freedom) closes
//! the left end of a fluid channel `[s(t), L] x [0, H]`; slip walls top and
//! bottom, pressure outlet on the right. The incompressible column moves
//! rigidly with the piston, so the fluid acts on it as pure added mass and
//! the coupled dynamics have a closed form:
//!
//! ```text
//!     (M + rho (L - s) H) s'' + K s = 0,
//!     omega ~ sqrt( K / (M + rho L H) )        for s0 << L.
//! ```
//!
//! `rtx-cfd`'s `tests/ale_piston_channel.rs` verifies the fluid half alone
//! (the discrete column is exact in space); what this test adds is the
//! coupling: `rtx-fsi`'s `Subiterated` driver iterating a genuine
//! fluid-solve/structure-solve pass to the interface fixed point each step.
//!
//! The mass ratio is deliberately heavy — added mass `rho L H = 0.25`
//! against `M = 0.04`, ratio 6.25 — putting the staggered exchange beyond
//! its divergence threshold (Causin, Gerbeau & Nobile 2005: the interface
//! fixed point turns repulsive once the added mass outweighs what the
//! structure presents, at any time step). One discrete subtlety this test
//! itself surfaced: the *continuous* threshold `m_a > M` is not the
//! discrete one. Newmark average acceleration weights the interface force
//! by `beta dt^2`, so the staggered iteration gain here is
//! `beta m_a / (M + K beta dt^2)` — a mass ratio of 2.5 at `beta = 1/4`
//! gives gain 0.625 and *converges* (measured: ~17 passes/step). The
//! benchmark therefore runs at ratio 6.25, gain ~1.56, which is genuinely
//! repulsive. `rtx-fsi` reproduced the instability on a linear model map;
//! here it must reproduce it against the real solver, and Aitken
//! relaxation must recover it — in the *same* configuration. Three claims:
//!
//! 1. Plain staggered (unit relaxation) coupling **diverges**.
//! 2. Aitken converges in a handful of subiterations per step.
//! 3. The converged oscillation period matches the added-mass closed form
//!    — and is nowhere near the dry-structure period `2 pi sqrt(M/K)`,
//!    so the agreement could not have happened without the fluid.
//!
//! Measured: staggered diverges after 7 subiterations; Aitken runs at 3.0
//! subiterations/step and lands on T = 1.07009 vs the closed form's
//! 1.06999 (9.8e-5 relative, halving to 4.8e-5 at dt/2 — first order,
//! from the half-step centring between the fluid's backward-difference
//! acceleration and Newmark's); outlet flux matches the piston sweep to
//! ~1e-9 every step. Finding on the way in: the absolute epsilon guard in
//! `aitken_factor` disabled Aitken below residual ~1e-8 — see
//! `coupling.rs` and the `aitken_is_scale_invariant` unit test.

use rtx_cfd::CfdConfig;
use rtx_cfd::solvers::incompressible::ale::{
    AleBoundaries, AleField, AleParameters, AlePisoSolver, SideBoundary,
};
use rtx_fsi::{FsiError, Subiterated};
use std::cell::RefCell;
use std::f64::consts::PI;

const RHO: f64 = 1.0;
const L: f64 = 1.0;
const H: f64 = 0.25;
const NX: usize = 32;
const NY: usize = 4;
const M: f64 = 0.04;
const K: f64 = 10.0;
const S0: f64 = 0.02;

const ADDED_MASS: f64 = RHO * L * H;

fn lines_x(s: f64) -> Vec<f64> {
    (0..=NX)
        .map(|i| s + (L - s) * i as f64 / NX as f64)
        .collect()
}

fn lines_y() -> Vec<f64> {
    (0..=NY).map(|j| H * j as f64 / NY as f64).collect()
}

/// One-DOF structure: Newmark average acceleration (beta 1/4, gamma 1/2),
/// unconditionally stable and second order.
#[derive(Debug, Clone, Copy)]
struct Piston {
    s: f64,
    v: f64,
    a: f64,
}

impl Piston {
    fn step(self, dt: f64, force: f64) -> Piston {
        const BETA: f64 = 0.25;
        const GAMMA: f64 = 0.5;
        let s_pred = self.s + dt * self.v + 0.5 * dt * dt * (1.0 - 2.0 * BETA) * self.a;
        let a_new = (force - K * s_pred) / (M + K * BETA * dt * dt);
        Piston {
            s: s_pred + BETA * dt * dt * a_new,
            v: self.v + dt * ((1.0 - GAMMA) * self.a + GAMMA * a_new),
            a: a_new,
        }
    }
}

/// The pressure on the piston face: the discrete field is exactly linear
/// in x, so two cell columns extrapolate it to the wall exactly.
fn wall_pressure(field: &AleField) -> f64 {
    let xc0 = 0.5 * (field.x[0] + field.x[1]);
    let xc1 = 0.5 * (field.x[1] + field.x[2]);
    let mean =
        |column: usize| -> f64 { (0..NY).map(|j| field.p[(j, column)]).sum::<f64>() / NY as f64 };
    let p0 = mean(0);
    let p1 = mean(1);
    p0 + (p0 - p1) * (xc0 - field.x[0]) / (xc1 - xc0)
}

fn fluid_solver() -> AlePisoSolver {
    let config = CfdConfig::new()
        .with_density(RHO)
        .with_viscosity(1e-3)
        .with_reference_velocity(1.0)
        .with_reference_length(L);
    let params = AleParameters {
        corrector_steps: 30,
        tolerance: 1e-10,
        boundaries: AleBoundaries {
            left: SideBoundary::Velocity,
            right: SideBoundary::PressureOutlet,
            bottom: SideBoundary::SlipWall,
            top: SideBoundary::SlipWall,
        },
        ..AleParameters::default()
    };
    AlePisoSolver::new(config, params).expect("valid solver configuration")
}

/// Consistent rest start: piston displaced to `s0` and released; the fluid
/// is at rest and the initial pressure field is the added-mass reaction to
/// the initial coupled acceleration.
fn initial_state() -> (AleField, Piston) {
    let a0 = -K * S0 / (M + RHO * (L - S0) * H);
    let mut field = AleField::new(lines_x(S0), lines_y()).expect("valid grid");
    let xc: Vec<f64> = field.x.windows(2).map(|w| 0.5 * (w[0] + w[1])).collect();
    for j in 0..NY {
        for i in 0..NX {
            field.p[(j, i)] = RHO * a0 * (L - xc[i]);
        }
    }
    let piston = Piston {
        s: S0,
        v: 0.0,
        a: a0,
    };
    (field, piston)
}

struct CoupledRun {
    /// Downward zero-crossing times of s(t).
    crossings: Vec<f64>,
    mean_subiterations: f64,
    /// Worst |outlet volume flux - piston sweep rate| over the run.
    worst_flux_mismatch: f64,
}

/// March the coupled system, driving each step's interface (the end-of-step
/// piston position) to a fixed point with the given scheme.
fn run_coupled(dt: f64, t_end: f64, scheme: &mut Subiterated) -> Result<CoupledRun, FsiError> {
    let (field0, piston0) = initial_state();
    let solver = RefCell::new(fluid_solver());
    let base = RefCell::new(field0);
    let piston_committed = RefCell::new(piston0);
    // The state the most recent coupling pass produced, committed after
    // the interface converges.
    let latest: RefCell<Option<(AleField, Piston)>> = RefCell::new(None);

    let steps = (t_end / dt).round() as usize;
    let mut crossings = Vec::new();
    let mut total_iterations = 0usize;
    let mut worst_flux: f64 = 0.0;
    let mut previous_s = piston0.s;

    for step in 0..steps {
        let t0 = step as f64 * dt;
        let piston_n = *piston_committed.borrow();
        let s_n = piston_n.s;

        let pass = |state: &[f64]| -> Vec<f64> {
            let s_candidate = state[0];
            let mut field = base.borrow().clone();
            let mut fluid = solver.borrow_mut();
            fluid.set_time(t0);
            let wall = (s_candidate - s_n) / dt;
            fluid.set_boundary_velocity(
                move |x, _y, _t| {
                    if x < 0.5 * L { (wall, 0.0) } else { (0.0, 0.0) }
                },
            );
            let result = futures::executor::block_on(fluid.advance(
                &mut field,
                &lines_x(s_candidate),
                &lines_y(),
                dt,
            ))
            .expect("fluid step");
            assert!(
                result.solver_result.converged,
                "fluid mass residual {:.3e}",
                result.solver_result.final_residual
            );
            // Pressure pushes the piston out of the fluid (toward -x).
            let force = -wall_pressure(&field) * H;
            let candidate = piston_n.step(dt, force);
            *latest.borrow_mut() = Some((field, candidate));
            vec![candidate.s]
        };

        // Predict the interface with the structure alone, then iterate.
        let s_predicted = piston_n.step(dt, -wall_pressure(&base.borrow()) * H).s;
        let converged = scheme.solve(&[s_predicted], pass)?;
        // One final pass at the agreed interface leaves fluid and structure
        // consistent with it.
        pass(&converged.state);
        let (field_new, piston_new) = latest.borrow_mut().take().expect("pass ran");

        // Discrete mass bookkeeping: what leaves the outlet must equal what
        // the piston sweeps, every step.
        let outlet_flux: f64 = (0..NY)
            .map(|j| field_new.u[(j, NX)] * (field_new.y[j + 1] - field_new.y[j]))
            .sum();
        let sweep_rate = (piston_new.s - s_n) / dt * H;
        worst_flux = worst_flux.max((outlet_flux - sweep_rate).abs());

        total_iterations += converged.iterations;
        base.replace(field_new);
        piston_committed.replace(piston_new);

        let t1 = (step + 1) as f64 * dt;
        if previous_s > 0.0 && piston_new.s <= 0.0 {
            crossings.push(t1 - dt * piston_new.s / (piston_new.s - previous_s));
        }
        previous_s = piston_new.s;
    }

    Ok(CoupledRun {
        crossings,
        mean_subiterations: total_iterations as f64 / steps as f64,
        worst_flux_mismatch: worst_flux,
    })
}

fn mean_period(crossings: &[f64]) -> f64 {
    assert!(
        crossings.len() >= 3,
        "need at least three crossings, got {}",
        crossings.len()
    );
    let periods: Vec<f64> = crossings.windows(2).map(|w| w[1] - w[0]).collect();
    periods.iter().sum::<f64>() / periods.len() as f64
}

#[test]
fn aitken_coupling_lands_on_the_added_mass_frequency() {
    let t_exact = 2.0 * PI * ((M + ADDED_MASS) / K).sqrt();
    let t_dry = 2.0 * PI * (M / K).sqrt();

    let mut scheme = Subiterated::aitken(50, 1e-11).expect("valid scheme");
    let coarse = run_coupled(2e-3, 4.0, &mut scheme).expect("coupled run");
    let fine = run_coupled(1e-3, 4.0, &mut scheme).expect("coupled run");

    let t_coarse = mean_period(&coarse.crossings);
    let t_fine = mean_period(&fine.crossings);
    let err_coarse = (t_coarse - t_exact).abs() / t_exact;
    let err_fine = (t_fine - t_exact).abs() / t_exact;
    println!(
        "  closed form T = {t_exact:.5} (dry {t_dry:.5})\n  dt 2e-3: T = {t_coarse:.5} \
         (err {err_coarse:.2e}), {:.1} subiterations/step, flux mismatch {:.2e}\n  dt 1e-3: \
         T = {t_fine:.5} (err {err_fine:.2e}), {:.1} subiterations/step, flux mismatch {:.2e}",
        coarse.mean_subiterations,
        coarse.worst_flux_mismatch,
        fine.mean_subiterations,
        fine.worst_flux_mismatch
    );

    // The coupled period matches the closed form and refines toward it.
    assert!(
        err_coarse < 0.01,
        "period {t_coarse:.5} vs closed form {t_exact:.5}: error {err_coarse:.3e}"
    );
    assert!(
        err_fine < err_coarse,
        "period error did not fall with dt: {err_coarse:.3e} -> {err_fine:.3e}"
    );
    // The added mass is what it matched: the dry period is 47% shorter. If
    // the fluid force were wrong or missing, the measurement would sit
    // near t_dry, not t_exact.
    assert!(
        (t_coarse - t_dry).abs() > 0.4 * t_dry,
        "measured period {t_coarse:.5} is suspiciously near the dry period {t_dry:.5}"
    );
    // Aitken earns its keep: a handful of passes per step, not the budget.
    assert!(
        coarse.mean_subiterations < 10.0,
        "mean subiterations {:.1}",
        coarse.mean_subiterations
    );
    // Conservation across the coupling: outlet flux equals piston sweep to
    // solver tolerance, every step.
    assert!(
        coarse.worst_flux_mismatch < 1e-8,
        "outlet flux vs piston sweep mismatch {:.3e}",
        coarse.worst_flux_mismatch
    );
}

#[test]
fn plain_staggered_coupling_diverges_under_heavy_added_mass() {
    // Unit relaxation, generous budget: the added-mass ratio of 2.5 makes
    // the interface fixed point repulsive, so this must report divergence
    // - the same failure rtx-fsi reproduces on its linear model map, now
    // on the real solver. If this ever starts converging, either the mass
    // ratio changed or the fluid stopped pushing back; both are findings.
    let mut scheme = Subiterated::relaxed(1.0, 50, 1e-11).expect("valid scheme");
    match run_coupled(2e-3, 0.5, &mut scheme) {
        Err(FsiError::CouplingDiverged {
            iterations,
            residual,
        }) => {
            println!("  diverged after {iterations} subiterations, residual {residual:.3e}");
        }
        Err(other) => panic!("expected CouplingDiverged, got {other:?}"),
        Ok(run) => panic!(
            "expected divergence, but the staggered scheme converged \
             ({:.1} subiterations/step on average)",
            run.mean_subiterations
        ),
    }
}