//! Nonlinear transient analysis: Newmark-β time integration with a full
//! Newton solve on the internal force inside every step.
//!
//! The existing [`super::dynamic_analysis`] stepper is linear by
//! construction — it factorises `M + γΔt C + βΔt² K` once and reuses it,
//! which is exactly right for constant matrices and exactly wrong for
//! finite deformation. This analysis solves, at every step,
//!
//! ```text
//!   M a_{n+1} + f_int(u_{n+1}) = F_ext(t_{n+1})
//!   a_{n+1} = (u_{n+1} - u_pred) / (β Δt²),
//!   u_pred  = u_n + Δt v_n + Δt² (1/2 - β) a_n,
//!   v_{n+1} = v_n + Δt ((1 - γ) a_n + γ a_{n+1})
//! ```
//!
//! by Newton on `R(u) = F_ext - f_int(u) - M a(u)` with the consistent
//! Jacobian `K_T(u) + M / (β Δt²)`, where `f_int` and `K_T` come from the
//! total-Lagrangian St. Venant–Kirchhoff path
//! ([`crate::elements::total_lagrangian`]) or the small-strain path, the
//! same seam the nonlinear static analysis uses. The consistent mass is
//! assembled once (element mass matrices are configuration-independent in
//! a total-Lagrangian setting); no damping (Rayleigh damping can be added
//! when something needs it — the Turek–Hron CSM3 benchmark is undamped).
//!
//! # Two ways to drive it
//!
//! [`NonlinearDynamicAnalysis::run`] marches `num_steps` steps from rest —
//! the benchmark shape (CSM3: gravity switched on at rest).
//!
//! [`NonlinearDynamicAnalysis::stepper`] hands out the same machinery one
//! step at a time, for a partitioned coupling loop: the caller owns the
//! state ([`DynamicState`]), sets the interface load with
//! [`NonlinearDynamicStepper::set_nodal_forces`], and calls
//! [`NonlinearDynamicStepper::step`] — which reads the start-of-step state
//! and *does not commit anything*, so a subiteration can re-run the same
//! step from the same state under an updated load as many times as the
//! interface fixed point takes (the semantics the coupled piston benchmark
//! established). `run` is implemented on the stepper, so the benchmark
//! tests pin both.
//!
//! Limits, stated up front: Dirichlet conditions must be homogeneous
//! (`u = 0` — a clamped edge); the body force is constant in time, applied
//! fully from `t = 0` (CSM3's definition: gravity switched on at rest, the
//! structure oscillates about its static deflection). Nodal forces may
//! change between steps (and between subiterations of one step) through
//! the stepper.

use super::{AnalysisConfig, ConvergenceCriteria};
use crate::assembly::dof_mapping::{AdvancedDofNumbering, DofComponent, DofMappingStrategy};
use crate::assembly::SparseMatrix;
use crate::boundary::{BoundaryCondition, BoundaryConditionSet};
use crate::elements::total_lagrangian::{self, saint_venant_kirchhoff};
use crate::elements::{ElementMatrixComputer, StandardFiniteElement};
use crate::error::{AnalysisError, FeaResult};
use crate::materials::{reduced_constitutive, MaterialDatabase};
use crate::mesh::{Mesh, NodeId};
use crate::solvers::{BandedLu, LinearSolver, SolverOptions};
use nalgebra::{DMatrix, DVector, Vector3};

/// Time histories and final state of a nonlinear transient run.
#[derive(Debug, Clone)]
pub struct NonlinearDynamicResults {
    /// Sample times `t_1..t_N` (end of each step).
    pub times: Vec<f64>,
    /// Per tracked node: its displacement components at every sample time,
    /// in the order the nodes were passed to `track_node`.
    pub tracked: Vec<Vec<DVector<f64>>>,
    /// Full displacement, velocity, acceleration at the final time.
    pub displacement: DVector<f64>,
    /// Final velocity.
    pub velocity: DVector<f64>,
    /// Final acceleration.
    pub acceleration: DVector<f64>,
    /// Newton iterations summed over all steps.
    pub total_iterations: usize,
    /// Largest Newton iteration count of any step.
    pub max_iterations_per_step: usize,
}

/// The full kinematic state at one instant: displacement, velocity and
/// acceleration as full-length vectors under the analysis's DOF numbering
/// (constrained entries zero). The caller owns it; a coupling loop clones
/// the committed state and re-steps from it freely.
#[derive(Debug, Clone)]
pub struct DynamicState {
    /// Displacement.
    pub displacement: DVector<f64>,
    /// Velocity.
    pub velocity: DVector<f64>,
    /// Acceleration.
    pub acceleration: DVector<f64>,
}

/// Per-element setup computed once: coordinates, DOFs, the DOF-expanded
/// consistent mass (configuration-independent), and the material.
struct ElementCache {
    coords: Vec<Vector3<f64>>,
    dofs: Vec<usize>,
    element_type: crate::mesh::ElementType,
    mass: DMatrix<f64>,
    material_id: crate::mesh::MaterialId,
}

/// Nonlinear Newmark transient analysis. See the module docs.
pub struct NonlinearDynamicAnalysis {
    mesh: Mesh,
    materials: MaterialDatabase,
    boundary_conditions: BoundaryConditionSet,
    #[allow(dead_code)]
    config: AnalysisConfig,
    criteria: ConvergenceCriteria,
    dt: f64,
    num_steps: usize,
    gamma: f64,
    beta: f64,
    total_lagrangian: bool,
    #[allow(clippy::type_complexity)]
    body_force: Option<Box<dyn Fn(Vector3<f64>) -> Vector3<f64> + Send + Sync>>,
    nodal_forces: Vec<(NodeId, Vector3<f64>)>,
    tracked_nodes: Vec<NodeId>,
}

impl NonlinearDynamicAnalysis {
    /// Average-acceleration Newmark (γ = 1/2, β = 1/4), the benchmark's
    /// scheme and the unconditionally stable one for linear problems.
    pub fn new(
        mesh: Mesh,
        materials: MaterialDatabase,
        boundary_conditions: BoundaryConditionSet,
        dt: f64,
        num_steps: usize,
        config: AnalysisConfig,
    ) -> Self {
        Self {
            mesh,
            materials,
            boundary_conditions,
            config,
            criteria: ConvergenceCriteria::default(),
            dt,
            num_steps,
            gamma: 0.5,
            beta: 0.25,
            total_lagrangian: false,
            body_force: None,
            nodal_forces: Vec::new(),
            tracked_nodes: Vec::new(),
        }
    }

    /// Switch to the total-Lagrangian St. Venant–Kirchhoff formulation
    /// (plane strain in 2-D), as on the nonlinear static analysis.
    #[must_use]
    pub fn with_total_lagrangian(mut self) -> Self {
        self.total_lagrangian = true;
        self
    }

    /// Newmark parameters (default γ = 1/2, β = 1/4).
    #[must_use]
    pub fn with_newmark_parameters(mut self, gamma: f64, beta: f64) -> Self {
        self.gamma = gamma;
        self.beta = beta;
        self
    }

    /// Convergence criteria for the per-step Newton loop.
    #[must_use]
    pub fn with_convergence_criteria(mut self, criteria: ConvergenceCriteria) -> Self {
        self.criteria = criteria;
        self
    }

    /// Constant body force per unit (reference) volume, applied from t = 0.
    pub fn set_body_force<F>(&mut self, f: F)
    where
        F: Fn(Vector3<f64>) -> Vector3<f64> + Send + Sync + 'static,
    {
        self.body_force = Some(Box::new(f));
    }

    /// Concentrated nodal forces, added to the external force (as on
    /// [`super::NonlinearStaticAnalysis`]). For a load that changes in
    /// time, use [`Self::stepper`] and set the forces before each step.
    pub fn set_nodal_forces(&mut self, forces: Vec<(NodeId, Vector3<f64>)>) {
        self.nodal_forces = forces;
    }

    /// Record this node's displacement at every step of [`Self::run`].
    pub fn track_node(&mut self, node: NodeId) {
        self.tracked_nodes.push(node);
    }

    /// Build the single-step driver: DOF numbering, element caches, the
    /// consistent mass, and the constant external force, assembled once.
    pub fn stepper(&self) -> FeaResult<NonlinearDynamicStepper<'_>> {
        NonlinearDynamicStepper::build(self)
    }

    /// March `num_steps` steps of `dt` from rest, via the same stepper a
    /// coupling loop would drive.
    pub fn run(&mut self) -> FeaResult<NonlinearDynamicResults> {
        let mut stepper = self.stepper()?;
        let mut state = stepper.rest_state()?;

        let mut times = Vec::with_capacity(self.num_steps);
        let mut tracked: Vec<Vec<DVector<f64>>> =
            vec![Vec::with_capacity(self.num_steps); self.tracked_nodes.len()];
        let mut total_iterations = 0usize;
        let mut max_iterations_per_step = 0usize;

        for step in 1..=self.num_steps {
            let (new_state, iterations) = stepper.step(&state)?;
            state = new_state;
            total_iterations += iterations;
            max_iterations_per_step = max_iterations_per_step.max(iterations);

            times.push(step as f64 * self.dt);
            for (slot, node) in self.tracked_nodes.iter().enumerate() {
                let dofs = stepper.node_dofs(*node);
                let mut value = DVector::zeros(dofs.len());
                for (c, &dof) in dofs.iter().enumerate() {
                    value[c] = state.displacement[dof];
                }
                tracked[slot].push(value);
            }
        }

        Ok(NonlinearDynamicResults {
            times,
            tracked,
            displacement: state.displacement,
            velocity: state.velocity,
            acceleration: state.acceleration,
            total_iterations,
            max_iterations_per_step,
        })
    }

    /// The DOF indices of a node under the analysis's own numbering, for
    /// reading the returned full-length vectors.
    pub fn node_dofs(&self, node: NodeId) -> FeaResult<Vec<usize>> {
        let numbering =
            AdvancedDofNumbering::displacement_only(&self.mesh, DofMappingStrategy::Sequential)?;
        Ok(numbering.get_node_dofs(node))
    }
}

/// The single-step Newmark–Newton driver behind
/// [`NonlinearDynamicAnalysis`]. Holds everything assembled once (DOF
/// numbering, element caches, consistent mass, the constant body-force
/// vector); the mutable pieces are the nodal forces and the linear solver.
///
/// [`Self::step`] is a pure function of the start-of-step [`DynamicState`]
/// and the current forces: nothing is committed, so a partitioned coupling
/// can re-run one step under updated interface loads until the interface
/// converges, then keep the accepted state.
pub struct NonlinearDynamicStepper<'a> {
    analysis: &'a NonlinearDynamicAnalysis,
    dof_numbering: AdvancedDofNumbering,
    free_dofs: Vec<usize>,
    free_index: Vec<Option<usize>>,
    total_dofs: usize,
    caches: Vec<ElementCache>,
    /// Free-free consistent mass, for consistent initial accelerations.
    mass_free: SparseMatrix,
    /// A lumped mass added per DOF (global numbering) — the partitioned
    /// coupling's fictitious added mass (`set_added_lumped_mass`): it enters
    /// the Newmark inertial residual and the effective tangent, never the
    /// consistent mass or the rest state. Zero by default.
    added_mass: DVector<f64>,
    /// The body-force part of the external force (constant).
    external_body: DVector<f64>,
    /// Body force plus the current nodal forces.
    external: DVector<f64>,
    /// Banded LU on the Newton tangent: the tangent's bandwidth is set
    /// by the mesh numbering (~26 on the 35×2 flag), and the dense
    /// factorization it replaces was 98% of the structural step.
    solver: BandedLu,
    solver_options: SolverOptions,
    /// Steps carried by the line-search rescue after the plain Newton
    /// loop failed (bookkeeping only — never touches the step result).
    rescued_line_search: usize,
    /// Steps carried by step subdivision after both the plain loop and
    /// the line search failed.
    rescued_subdivision: usize,
}

impl<'a> NonlinearDynamicStepper<'a> {
    #[allow(clippy::too_many_lines)]
    fn build(analysis: &'a NonlinearDynamicAnalysis) -> FeaResult<Self> {
        let dim = analysis.mesh.spatial_dimension;
        let mut dof_numbering = AdvancedDofNumbering::displacement_only(
            &analysis.mesh,
            DofMappingStrategy::Sequential,
        )?;

        // Homogeneous Dirichlet only (see the module docs).
        for condition in analysis.boundary_conditions.conditions() {
            if let BoundaryCondition::Dirichlet(dirichlet) = condition {
                for &node in &dirichlet.nodes {
                    let position = analysis
                        .mesh
                        .get_node(node)
                        .ok_or_else(|| {
                            AnalysisError::InvalidConfiguration(format!(
                                "Dirichlet condition names missing node {node:?}"
                            ))
                        })?
                        .position();
                    for component in &dirichlet.components {
                        if component.canonical_index() >= dim {
                            continue;
                        }
                        let value = dirichlet.get_value(0.0, &position);
                        if value.abs() > 1e-14 {
                            return Err(AnalysisError::InvalidConfiguration(
                                "NonlinearDynamicAnalysis supports homogeneous Dirichlet \
                                 conditions only"
                                    .to_string(),
                            )
                            .into());
                        }
                        let Some(dof) = dof_numbering.get_dof(node, *component) else {
                            continue;
                        };
                        dof_numbering.constrain_dof(dof)?;
                    }
                }
            }
        }
        let total_dofs = dof_numbering.total_dofs;
        let free_dofs = dof_numbering.free_dofs.clone();
        let num_free = free_dofs.len();
        let mut free_index = vec![None; total_dofs];
        for (k, &dof) in free_dofs.iter().enumerate() {
            free_index[dof] = Some(k);
        }

        let mut caches = Vec::with_capacity(analysis.mesh.elements.len());
        for element in analysis.mesh.elements.values() {
            let coords: Vec<Vector3<f64>> = element
                .nodes
                .iter()
                .map(|id| analysis.mesh.get_node(*id).unwrap().position())
                .collect();
            let dofs: Vec<usize> = element
                .nodes
                .iter()
                .flat_map(|node| dof_numbering.get_node_dofs(*node))
                .collect();
            let material = analysis
                .materials
                .get_material(element.material_id)
                .ok_or_else(|| {
                    AnalysisError::InvalidConfiguration(format!(
                        "Material {} not found",
                        element.material_id.0
                    ))
                })?;
            let density = material.properties().density;
            let fe = StandardFiniteElement::new(element.element_type, coords.clone());
            let scalar =
                ElementMatrixComputer::compute_consistent_mass_matrix(&fe, &coords, density, None)?;
            let nodes = element.nodes.len();
            let mut mass = DMatrix::zeros(nodes * dim, nodes * dim);
            for a in 0..nodes {
                for b in 0..nodes {
                    let m = scalar.matrix[(a, b)];
                    for d in 0..dim {
                        mass[(a * dim + d, b * dim + d)] = m;
                    }
                }
            }
            caches.push(ElementCache {
                coords,
                dofs,
                element_type: element.element_type,
                mass,
                material_id: element.material_id,
            });
        }

        // Free-free consistent mass (for initial accelerations).
        let mut mass_free = SparseMatrix::new(num_free, num_free);
        for cache in &caches {
            for (local_row, &dof_row) in cache.dofs.iter().enumerate() {
                let Some(free_row) = free_index[dof_row] else {
                    continue;
                };
                for (local_col, &dof_col) in cache.dofs.iter().enumerate() {
                    if let Some(free_col) = free_index[dof_col] {
                        let value = cache.mass[(local_row, local_col)];
                        if value != 0.0 {
                            mass_free.add_entry(free_row, free_col, value)?;
                        }
                    }
                }
            }
        }
        mass_free.finalize()?;

        // Constant consistent external force from the body-force field.
        let mut external_body: DVector<f64> = DVector::zeros(num_free);
        if let Some(force) = &analysis.body_force {
            for cache in &caches {
                let fe = StandardFiniteElement::new(cache.element_type, cache.coords.clone());
                let f_e = ElementMatrixComputer::compute_body_force_vector(
                    &fe,
                    &cache.coords,
                    force.as_ref(),
                    None,
                )?;
                for (local, &dof) in cache.dofs.iter().enumerate() {
                    if let Some(free) = free_index[dof] {
                        external_body[free] += f_e[local];
                    }
                }
            }
        }

        let added_mass = DVector::zeros(total_dofs);
        let mut stepper = Self {
            added_mass,
            analysis,
            dof_numbering,
            free_dofs,
            free_index,
            total_dofs,
            caches,
            mass_free,
            external_body: external_body.clone(),
            external: external_body,
            solver: BandedLu::new(),
            solver_options: SolverOptions::default(),
            rescued_line_search: 0,
            rescued_subdivision: 0,
        };
        stepper.set_nodal_forces(&analysis.nodal_forces);
        Ok(stepper)
    }

    /// Replace the concentrated nodal forces (the interface load of a
    /// coupling subiteration). The body-force part is unaffected.
    pub fn set_nodal_forces(&mut self, forces: &[(NodeId, Vector3<f64>)]) {
        self.external.copy_from(&self.external_body);
        for (node, force) in forces {
            let dofs = self.dof_numbering.get_node_dofs(*node);
            for (component, &dof) in dofs.iter().enumerate() {
                if let Some(free) = self.free_index[dof] {
                    self.external[free] += force[component];
                }
            }
        }
    }

    /// Set the lumped mass added to every DOF of each node (the coupling
    /// loop's fictitious added mass, `docs/overset_metal_campaign.md` §5.17
    /// in omni-cortex): `(M + M_f) ü = F + M_f ü_k` contracts at any mass
    /// ratio and leaves the fixed point unchanged when the caller adds the
    /// load `M_f ü_k` of the previous subiterate. Entries not listed keep
    /// their value; `0.0` restores the plain stepper bit for bit.
    pub fn set_added_lumped_mass(&mut self, entries: &[(NodeId, f64)]) {
        for (node, m) in entries {
            for dof in self.dof_numbering.get_node_dofs(*node) {
                self.added_mass[dof] = *m;
            }
        }
    }

    /// The state at rest under the *current* external force: `u = v = 0`,
    /// the acceleration consistent with `M a0 = F_ext - f_int(0)`.
    pub fn rest_state(&mut self) -> FeaResult<DynamicState> {
        let u = DVector::zeros(self.total_dofs);
        // inv_beta_dt2 is only read when assembling the tangent.
        let (f_int0, _) = self.assemble(&u, false, 0.0)?;
        let residual0 = &self.external - &f_int0;
        let (a0_free, _) = self
            .solver
            .solve(&self.mass_free, &residual0, &self.solver_options)?;
        let mut a = DVector::zeros(self.total_dofs);
        for (k, &dof) in self.free_dofs.iter().enumerate() {
            a[dof] = a0_free[k];
        }
        Ok(DynamicState {
            displacement: DVector::zeros(self.total_dofs),
            velocity: DVector::zeros(self.total_dofs),
            acceleration: a,
        })
    }

    /// One Newmark step of the analysis's `dt` from `state` under the
    /// current forces. Returns the end-of-step state and the Newton
    /// iteration count; commits nothing (beyond rescue bookkeeping) —
    /// calling again with the same state and forces returns the
    /// identical result.
    ///
    /// The plain full-step Newton runs first, untouched — a step it
    /// converges is bit-identical to the pre-rescue stepper. Only when
    /// it FAILS does the rescue engage (measured need: both FSI3 study
    /// deaths were `ConvergenceFailed { iterations: 60 }` at a mid-swing
    /// load reversal — a full Newton step from the Newmark predictor
    /// under a reversed load leaves SVK's convergence region; a static
    /// load from rest was measured NOT to fail even at 1e6 N, because
    /// from a quiescent state the predictor is the current configuration
    /// and `M/(β Δt²)` regularizes the walk):
    ///
    /// 1. Newton again with a backtracking line search on `‖R‖`
    ///    (Armijo, α halved down to 2⁻⁸; a non-descending Newton
    ///    direction fails fast to level 2).
    /// 2. Step subdivision: 2, 4, 8, then 16 Newmark substeps of
    ///    `dt/n` from the same start state under the same (end-of-step)
    ///    load, each substep line-searched. The composed end state is
    ///    the step's result — same total interval, finer integration.
    ///
    /// Rescued steps are counted in [`Self::rescue_counts`]; if every
    /// level fails, the plain loop's original error is returned.
    pub fn step(&mut self, state: &DynamicState) -> FeaResult<(DynamicState, usize)> {
        self.step_with_dt(state, self.analysis.dt)
    }

    /// [`Self::step`] over an explicit interval `dt` instead of the
    /// analysis's own — the same plain Newton, line search and
    /// subdivision ladder (every level already takes `dt` as a
    /// parameter; the Newmark mass term follows it). The analysis's
    /// `dt` path is exactly `step`, float for float. A coupled march's
    /// coupling-level rescue uses this to repeat an interval as `n`
    /// substeps of `dt/n`.
    ///
    /// # Errors
    /// As [`Self::step`].
    pub fn step_with_dt(
        &mut self,
        state: &DynamicState,
        dt: f64,
    ) -> FeaResult<(DynamicState, usize)> {
        match self.newmark_newton(state, dt, false) {
            Ok(result) => Ok(result),
            Err(first_failure) => {
                if let Ok(result) = self.newmark_newton(state, dt, true) {
                    self.rescued_line_search += 1;
                    return Ok(result);
                }
                for n in [2usize, 4, 8, 16] {
                    if let Ok(result) = self.substep_march(state, dt, n) {
                        self.rescued_subdivision += 1;
                        return Ok(result);
                    }
                }
                Err(first_failure)
            }
        }
    }

    /// `n` Newmark substeps of `dt/n` from `state` under the current
    /// forces, each line-searched. The load is the step's own
    /// (end-of-step) load held constant across the substeps — the same
    /// closure the full step uses.
    fn substep_march(
        &mut self,
        state: &DynamicState,
        dt: f64,
        n: usize,
    ) -> FeaResult<(DynamicState, usize)> {
        let dt_sub = dt / n as f64;
        let mut current = state.clone();
        let mut total_iterations = 0usize;
        for _ in 0..n {
            let (next, iterations) = self.newmark_newton(&current, dt_sub, true)?;
            total_iterations += iterations;
            current = next;
        }
        Ok((current, total_iterations))
    }

    /// One Newmark step of `dt` from `state`: Newton on the end-of-step
    /// displacement from the predictor. With `line_search` false this is
    /// the original plain loop, float-op for float-op (`α = 1.0`
    /// multiplies exactly); with it true, each Newton direction is
    /// backtracked on the residual norm before acceptance.
    fn newmark_newton(
        &mut self,
        state: &DynamicState,
        dt: f64,
        line_search: bool,
    ) -> FeaResult<(DynamicState, usize)> {
        let gamma = self.analysis.gamma;
        let beta = self.analysis.beta;
        let criteria = &self.analysis.criteria;
        let force_scale = self.external.norm().max(1.0);

        let mut u_pred = DVector::zeros(self.total_dofs);
        for &dof in &self.free_dofs {
            u_pred[dof] = state.displacement[dof]
                + dt * state.velocity[dof]
                + dt * dt * (0.5 - beta) * state.acceleration[dof];
        }
        let inv_beta_dt2 = 1.0 / (beta * dt * dt);

        // Newton on the end-of-step displacement, starting from the
        // predictor (a_new = 0 there).
        let mut u_iter = u_pred.clone();
        let mut step_converged = false;
        let mut iterations = 0usize;
        for _ in 0..criteria.max_iterations {
            let mut a_new = DVector::zeros(self.total_dofs);
            for &dof in &self.free_dofs {
                a_new[dof] = inv_beta_dt2 * (u_iter[dof] - u_pred[dof]);
            }
            let (f_int, tangent) = self.assemble(&u_iter, true, inv_beta_dt2)?;
            let residual = &self.external - &f_int - self.mass_times(&a_new);
            let residual_norm = residual.norm();
            if line_search && !residual_norm.is_finite() {
                return Err(AnalysisError::ConvergenceFailed { iterations }.into());
            }
            if residual_norm < criteria.force_tolerance * force_scale {
                step_converged = true;
                break;
            }
            iterations += 1;
            let (delta, _) = self
                .solver
                .solve(&tangent, &residual, &self.solver_options)?;
            let alpha = if line_search {
                match self.backtrack(&u_iter, &delta, &u_pred, inv_beta_dt2, residual_norm) {
                    Some(alpha) => alpha,
                    None => return Err(AnalysisError::ConvergenceFailed { iterations }.into()),
                }
            } else {
                1.0
            };
            for (k, &dof) in self.free_dofs.iter().enumerate() {
                u_iter[dof] += alpha * delta[k];
            }
            if alpha * delta.norm() < criteria.displacement_tolerance * u_iter.norm().max(1.0) {
                step_converged = true;
                break;
            }
        }
        if !step_converged {
            return Err(AnalysisError::ConvergenceFailed { iterations }.into());
        }

        let mut a_new = DVector::zeros(self.total_dofs);
        let mut v_new = DVector::zeros(self.total_dofs);
        for &dof in &self.free_dofs {
            a_new[dof] = inv_beta_dt2 * (u_iter[dof] - u_pred[dof]);
            v_new[dof] = state.velocity[dof]
                + dt * ((1.0 - gamma) * state.acceleration[dof] + gamma * a_new[dof]);
        }
        Ok((
            DynamicState {
                displacement: u_iter,
                velocity: v_new,
                acceleration: a_new,
            },
            iterations,
        ))
    }

    /// Backtracking line search on the Newmark residual norm: the first
    /// `α ∈ {1, 1/2, …, 2⁻²⁹}` satisfying the Armijo decrease
    /// `‖R(u + α δ)‖ ≤ (1 − 10⁻⁴ α) ‖R(u)‖`; failing that, the best
    /// finite trial if it decreases the norm at all; `None` when the
    /// Newton direction yields no descent (the caller falls through to
    /// step subdivision rather than walking somewhere worse). The depth
    /// is deliberate: near a turning point the tangent
    /// `K_T + M/(β Δt²)` can be almost singular, the solved direction
    /// then enormous and inexact — with an exact Jacobian descent exists
    /// for small enough `α`, but 2⁻⁸ of a huge direction was measured
    /// still too large (the first rescue draft failed exactly here).
    fn backtrack(
        &mut self,
        u_iter: &DVector<f64>,
        delta: &DVector<f64>,
        u_pred: &DVector<f64>,
        inv_beta_dt2: f64,
        residual_norm: f64,
    ) -> Option<f64> {
        let mut best: Option<(f64, f64)> = None;
        let mut alpha = 1.0f64;
        for _ in 0..30 {
            let trial_norm = self.residual_norm_at(u_iter, delta, alpha, u_pred, inv_beta_dt2);
            if let Ok(trial_norm) = trial_norm {
                if trial_norm.is_finite() {
                    if trial_norm <= (1.0 - 1e-4 * alpha) * residual_norm {
                        return Some(alpha);
                    }
                    if best.is_none_or(|(_, b)| trial_norm < b) {
                        best = Some((alpha, trial_norm));
                    }
                }
            }
            alpha *= 0.5;
        }
        match best {
            Some((alpha, norm)) if norm < residual_norm => Some(alpha),
            _ => None,
        }
    }

    /// `‖F_ext − f_int(u + α δ) − M a(u + α δ)‖` — the line search's
    /// merit function (internal force only, no tangent).
    fn residual_norm_at(
        &self,
        u_iter: &DVector<f64>,
        delta: &DVector<f64>,
        alpha: f64,
        u_pred: &DVector<f64>,
        inv_beta_dt2: f64,
    ) -> FeaResult<f64> {
        let mut u_trial = u_iter.clone();
        let mut a_trial = DVector::zeros(self.total_dofs);
        for (k, &dof) in self.free_dofs.iter().enumerate() {
            u_trial[dof] += alpha * delta[k];
            a_trial[dof] = inv_beta_dt2 * (u_trial[dof] - u_pred[dof]);
        }
        let (f_int, _) = self.assemble(&u_trial, false, inv_beta_dt2)?;
        Ok((&self.external - &f_int - self.mass_times(&a_trial)).norm())
    }

    /// The DOF indices of a node, for reading [`DynamicState`] vectors.
    pub fn node_dofs(&self, node: NodeId) -> Vec<usize> {
        self.dof_numbering.get_node_dofs(node)
    }

    /// How many steps needed rescuing so far: `(line_search,
    /// subdivision)`. Zero on every healthy march — a nonzero count is a
    /// finding about the loads the stepper is being fed, worth reporting
    /// alongside a coupled march's bookkeeping.
    pub fn rescue_counts(&self) -> (usize, usize) {
        (self.rescued_line_search, self.rescued_subdivision)
    }

    /// Internal force and (optionally) tangent at a full displacement
    /// vector, reduced to the free DOFs. The tangent includes the Newmark
    /// mass term `M / (β Δt²)` at the caller's `inv_beta_dt2` (the
    /// rescue's substeps run a finer `dt` than the analysis's own).
    fn assemble(
        &self,
        solution: &DVector<f64>,
        with_tangent: bool,
        inv_beta_dt2: f64,
    ) -> FeaResult<(DVector<f64>, SparseMatrix)> {
        let dim = self.analysis.mesh.spatial_dimension;
        let num_free = self.free_dofs.len();
        let mut internal = DVector::zeros(num_free);
        let mut tangent = SparseMatrix::new(num_free, num_free);
        for cache in &self.caches {
            let material = self
                .analysis
                .materials
                .get_material(cache.material_id)
                .unwrap();
            let fe = StandardFiniteElement::new(cache.element_type, cache.coords.clone());
            let mut element_displacement = DVector::zeros(cache.dofs.len());
            for (local, &dof) in cache.dofs.iter().enumerate() {
                element_displacement[local] = solution[dof];
            }
            let (f_int, k_t) = if self.analysis.total_lagrangian {
                let (lambda, mu) = material.properties().lame_parameters();
                let constitutive = saint_venant_kirchhoff(lambda, mu, dim);
                total_lagrangian::internal_force_and_tangent(
                    &fe,
                    &cache.coords,
                    &element_displacement,
                    constitutive.as_ref(),
                    None,
                )?
            } else {
                let constitutive = reduced_constitutive(material, dim)?;
                ElementMatrixComputer::compute_internal_force_and_tangent(
                    &fe,
                    &cache.coords,
                    &element_displacement,
                    constitutive.as_ref(),
                    None,
                )?
            };
            for (local_row, &dof_row) in cache.dofs.iter().enumerate() {
                let Some(free_row) = self.free_index[dof_row] else {
                    continue;
                };
                internal[free_row] += f_int[local_row];
                if with_tangent {
                    for (local_col, &dof_col) in cache.dofs.iter().enumerate() {
                        if let Some(free_col) = self.free_index[dof_col] {
                            let value = k_t[(local_row, local_col)]
                                + inv_beta_dt2 * cache.mass[(local_row, local_col)];
                            if value != 0.0 {
                                tangent.add_entry(free_row, free_col, value)?;
                            }
                        }
                    }
                }
            }
        }
        if with_tangent {
            for (dof, &m) in self.added_mass.iter().enumerate() {
                if m != 0.0 {
                    if let Some(free) = self.free_index[dof] {
                        tangent.add_entry(free, free, inv_beta_dt2 * m)?;
                    }
                }
            }
            tangent.finalize()?;
        }
        Ok((internal, tangent))
    }

    /// M times a full-length vector, reduced to the free DOFs.
    fn mass_times(&self, a_full: &DVector<f64>) -> DVector<f64> {
        let num_free = self.free_dofs.len();
        let mut out = DVector::zeros(num_free);
        for cache in &self.caches {
            let mut a_e = DVector::zeros(cache.dofs.len());
            for (local, &dof) in cache.dofs.iter().enumerate() {
                a_e[local] = a_full[dof];
            }
            let m_a = &cache.mass * a_e;
            for (local, &dof) in cache.dofs.iter().enumerate() {
                if let Some(free) = self.free_index[dof] {
                    out[free] += m_a[local];
                }
            }
        }
        for (dof, &m) in self.added_mass.iter().enumerate() {
            if m != 0.0 {
                if let Some(free) = self.free_index[dof] {
                    out[free] += m * a_full[dof];
                }
            }
        }
        out
    }
}