//! The nonlinear Newmark stepper under violent sudden loads: the rescue
//! path (backtracking line search, then step subdivision) behind
//! [`NonlinearDynamicStepper::step`].
//!
//! Why this exists: both FSI3 study-march deaths (2026-08-24) were the
//! flag's SVK Newton returning `ConvergenceFailed { iterations: 60 }`
//! inside a coupling pass at a violent mid-cycle load — the structural
//! solver, not the coupling. A full Newton step from the Newmark
//! predictor under a load far from the current configuration can leave
//! SVK's region of convergence; the plain loop had no line search and no
//! subdivision, so the first such step killed a five-hour march.
//!
//! The failure shape, measured by a probe before the rescue was written
//! (2026-08-24, this mesh, 60-iteration budget):
//! - A static tip load from a quiescent state NEVER failed — up to
//!   1e6 N (tip deflection past the flag's own length) plain Newton
//!   converged in ≤ 19 iterations. From rest the predictor IS the
//!   current configuration and `M/(β Δt²)` regularizes the walk.
//!   Pinned below (`static_loads_from_rest_never_need_rescue`).
//! - The kill is MID-SWING: three steps of swing-up under a 1e4 N tip
//!   load at dt = 5e-3 (tip at −0.31 m, −31 m/s), then the load
//!   REVERSED — plain Newton dead in 60 iterations. A turning point, the
//!   same shape as the FSI3 deaths. Pinned below as the rescue's test.
//!
//! Contract:
//! 1. The plain Newton path is UNTOUCHED — a step it converges reports
//!    zero rescues (the FSI2/FSI3 committed defaults are re-verified
//!    bit-identical separately).
//! 2. The measured killer step must be rescued, deterministically, to a
//!    state consistent with a fine-dt reference march of the same total
//!    interval.

use nalgebra::Vector3;
use rtx_fea::analysis::{
    AnalysisConfig, ConvergenceCriteria, DynamicState, NonlinearDynamicAnalysis,
    NonlinearDynamicStepper,
};
use rtx_fea::assembly::dof_mapping::DofComponent;
use rtx_fea::boundary::dirichlet::{DirichletBC, DirichletType};
use rtx_fea::boundary::{BoundaryCondition, BoundaryConditionSet, SpatialFunction};
use rtx_fea::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};

const E_MOD: f64 = 1.4e6;
const NU: f64 = 0.4;
const RHO: f64 = 1000.0;

/// The measured plain-Newton killer (see the module docs): swing up for
/// three steps under this tip load, then reverse it.
const SWING_LOAD: f64 = 1e4;
const SWING_DT: f64 = 5e-3;
const SWING_STEPS: usize = 3;

/// `nx` by `ny` Quad8 mesh of `[x0, x1] x [y0, y1]` (serendipity
/// lattice), as in `tests/nonlinear_newmark_csm3.rs`.
fn quad8_rect_mesh(x0: f64, x1: f64, y0: f64, y1: f64, nx: usize, ny: usize) -> Mesh {
    let mut mesh = Mesh::new(2).unwrap();
    let (lx, ly) = (2 * nx + 1, 2 * ny + 1);
    let mut grid = vec![vec![None; ly]; lx];
    for (i, column) in grid.iter_mut().enumerate() {
        for (j, slot) in column.iter_mut().enumerate() {
            if i % 2 == 1 && j % 2 == 1 {
                continue;
            }
            let x = x0 + (x1 - x0) * i as f64 / (2 * nx) as f64;
            let y = y0 + (y1 - y0) * j as f64 / (2 * ny) as f64;
            *slot = Some(mesh.add_node(Node::new_2d(x, y)));
        }
    }
    for i in 0..nx {
        for j in 0..ny {
            let (a, b) = (2 * i, 2 * j);
            let nodes = vec![
                grid[a][b].unwrap(),
                grid[a + 2][b].unwrap(),
                grid[a + 2][b + 2].unwrap(),
                grid[a][b + 2].unwrap(),
                grid[a + 1][b].unwrap(),
                grid[a + 2][b + 1].unwrap(),
                grid[a + 1][b + 2].unwrap(),
                grid[a][b + 1].unwrap(),
            ];
            mesh.add_element(Element::new(ElementType::Quad8, nodes, MaterialId(0)).unwrap())
                .unwrap();
        }
    }
    mesh
}

fn materials() -> MaterialDatabase {
    let mut db = MaterialDatabase::new();
    db.add_material(
        MaterialId(0),
        LinearElastic::new(E_MOD, NU).with_density(RHO),
        None,
    );
    db
}

fn clamp_left(mesh: &Mesh, x_left: f64) -> BoundaryConditionSet {
    let clamped: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| (node.position().x - x_left).abs() < 1e-12)
        .map(|(&id, _)| id)
        .collect();
    let mut set = BoundaryConditionSet::new();
    for component in [DofComponent::DisplacementX, DofComponent::DisplacementY] {
        set.add_condition(BoundaryCondition::Dirichlet(DirichletBC {
            nodes: clamped.clone(),
            components: vec![component],
            condition_type: DirichletType::Spatial(SpatialFunction(Box::new(|_| 0.0))),
            time_range: None,
            ramping_factor: 1.0,
            gradual_enforcement: false,
        }));
    }
    set
}

fn point_a(mesh: &Mesh, x: f64, y: f64) -> NodeId {
    mesh.nodes
        .iter()
        .find(|(_, node)| {
            (node.position().x - x).abs() < 1e-12 && (node.position().y - y).abs() < 1e-12
        })
        .map(|(&id, _)| id)
        .expect("tracking point must be a mesh node")
}

/// The FSI2/FSI3 flag geometry at 10x2 Quad8 with the harness's
/// 60-iteration Newton budget.
fn flag_analysis(dt: f64) -> (NonlinearDynamicAnalysis, NodeId) {
    let mesh = quad8_rect_mesh(0.25, 0.6, 0.19, 0.21, 10, 2);
    let a_node = point_a(&mesh, 0.6, 0.2);
    let analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, 0.25),
        dt,
        1,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian()
    .with_convergence_criteria(ConvergenceCriteria {
        max_iterations: 60,
        ..ConvergenceCriteria::default()
    });
    (analysis, a_node)
}

/// Swing the flag up for [`SWING_STEPS`] steps under `(0, -SWING_LOAD)`
/// from a quiescent start, through the PLAIN path (asserted).
fn swing_up(stepper: &mut NonlinearDynamicStepper<'_>, a_node: NodeId) -> DynamicState {
    stepper.set_nodal_forces(&[(a_node, Vector3::new(0.0, -SWING_LOAD, 0.0))]);
    let mut state = stepper.rest_state().unwrap();
    state.acceleration.fill(0.0);
    for _ in 0..SWING_STEPS {
        let (next, _) = stepper.step(&state).expect("swing-up step");
        state = next;
    }
    assert_eq!(
        stepper.rescue_counts(),
        (0, 0),
        "the swing-up must not need rescuing — it is the reference plain path"
    );
    state
}

/// Measured negative result, pinned: a static tip load from a quiescent
/// state does not defeat plain Newton even at 1e6 N (tip deflection
/// beyond the flag's own length, 19 iterations). The rescue must stay
/// out of the way.
#[test]
fn static_loads_from_rest_never_need_rescue() {
    for load in [1e4, 1e6] {
        let (analysis, a_node) = flag_analysis(SWING_DT);
        let mut stepper = analysis.stepper().unwrap();
        stepper.set_nodal_forces(&[(a_node, Vector3::new(0.0, -load, 0.0))]);
        let mut state = stepper.rest_state().unwrap();
        state.acceleration.fill(0.0);
        let (next, iterations) = stepper
            .step(&state)
            .expect("static load from rest must converge on the plain path");
        assert!(next.displacement.iter().all(|v| v.is_finite()));
        assert!(
            iterations <= 25,
            "static {load:.0e} N took {iterations} iterations"
        );
        assert_eq!(
            stepper.rescue_counts(),
            (0, 0),
            "static load from rest engaged the rescue"
        );
    }
}

/// A benign step must go through the untouched plain path: zero rescues,
/// few iterations.
#[test]
fn benign_step_is_never_rescued() {
    let (analysis, a_node) = flag_analysis(SWING_DT);
    let mut stepper = analysis.stepper().unwrap();
    stepper.set_nodal_forces(&[(a_node, Vector3::new(0.0, -0.1, 0.0))]);
    let mut state = stepper.rest_state().unwrap();
    state.acceleration.fill(0.0);
    let (next, iterations) = stepper.step(&state).expect("benign step must converge");
    assert!(next.displacement.iter().all(|v| v.is_finite()));
    assert!(
        iterations <= 5,
        "benign step took {iterations} Newton iterations"
    );
    assert_eq!(stepper.rescue_counts(), (0, 0));
}

/// The measured killer step (mid-swing load reversal — plain Newton dies
/// in 60 iterations here, the FSI3 death shape) must be rescued:
/// deterministically, and to a state the fine-dt reference march
/// corroborates.
#[test]
fn mid_swing_load_reversal_is_rescued_and_matches_fine_dt_reference() {
    let (analysis, a_node) = flag_analysis(SWING_DT);
    let mut stepper = analysis.stepper().unwrap();
    let a_dofs = stepper.node_dofs(a_node);
    let state = swing_up(&mut stepper, a_node);

    // The reversal.
    stepper.set_nodal_forces(&[(a_node, Vector3::new(0.0, SWING_LOAD, 0.0))]);
    let (rescued, _) = stepper
        .step(&state)
        .expect("the rescue path must carry the mid-swing load reversal");
    let rescues = stepper.rescue_counts();
    assert!(
        rescues.0 + rescues.1 > 0,
        "this step was measured to defeat plain Newton (60 iterations); zero \
         rescues means the scenario no longer bites and this test is vacuous"
    );
    assert!(rescued.displacement.iter().all(|v| v.is_finite()));
    let uy = rescued.displacement[a_dofs[1]];

    // Determinism: the coupling subiterates by re-running the same step
    // from the same state — the rescue must be a pure function of
    // (state, forces) too.
    let (again, _) = stepper.step(&state).unwrap();
    assert_eq!(
        rescued.displacement, again.displacement,
        "rescued step is not deterministic"
    );
    assert_eq!(rescued.velocity, again.velocity);

    // Fine-dt reference: the same interval marched at dt/32 under the
    // same constant reversed load from the same mid-swing state (the
    // DOF numbering is the same mesh's). Newmark at two different steps
    // agrees to O(dt²) — but at a violent reversal the one-step coarse
    // answer legitimately differs in detail, so the band is generous; it
    // still catches a wrong-branch answer (a different deformation
    // scale) or a sign error.
    let (fine_analysis, fine_a_node) = flag_analysis(SWING_DT / 32.0);
    let mut fine = fine_analysis.stepper().unwrap();
    fine.set_nodal_forces(&[(fine_a_node, Vector3::new(0.0, SWING_LOAD, 0.0))]);
    let mut ref_state = state.clone();
    for _ in 0..32 {
        let (next, _) = fine.step(&ref_state).expect("fine-dt reference step");
        ref_state = next;
    }
    let uy_ref = ref_state.displacement[a_dofs[1]];
    println!(
        "  reversal from uy {:.4e} (v_tip {:.3e}): rescued uy {uy:.4e}, fine-dt \
         (dt/32) reference {uy_ref:.4e}; rescues (line-search, subdivision) = {rescues:?}",
        state.displacement[a_dofs[1]], state.velocity[a_dofs[1]],
    );
    // Measured 2026-08-24: rescued −0.341 vs reference −0.195 (0.47x of
    // the scale) — a one-step coarse Newmark answer at a violent reversal
    // legitimately differs in detail; the band only has to catch a wrong
    // branch or a sign error, both of which sit at a different scale.
    assert!(
        (uy - uy_ref).abs() < 0.75 * uy_ref.abs().max(state.displacement[a_dofs[1]].abs()),
        "rescued step uy {uy:.4e} inconsistent with the fine-dt reference {uy_ref:.4e}"
    );
}

/// After a rescued step, the march must be able to CONTINUE — the
/// coupling re-steps and then keeps marching from whatever the rescue
/// returned. Ten further steps under the reversed load must all
/// converge (plain or rescued) and stay finite.
#[test]
fn march_continues_after_a_rescued_step() {
    let (analysis, a_node) = flag_analysis(SWING_DT);
    let mut stepper = analysis.stepper().unwrap();
    let state = swing_up(&mut stepper, a_node);
    stepper.set_nodal_forces(&[(a_node, Vector3::new(0.0, SWING_LOAD, 0.0))]);
    let (mut state, _) = stepper.step(&state).expect("rescued step");
    for k in 0..10 {
        let (next, _) = stepper
            .step(&state)
            .unwrap_or_else(|e| panic!("step {k} after the rescue failed: {e:?}"));
        assert!(
            next.displacement.iter().all(|v| v.is_finite()),
            "step {k} after the rescue went non-finite"
        );
        state = next;
    }
    let (line_search, subdivision) = stepper.rescue_counts();
    println!("  post-rescue march: rescues line-search {line_search}, subdivision {subdivision}");
}