//! Rung S2: the nonlinear Newmark analysis
//! (`analysis::nonlinear_dynamic`), verified in two steps:
//!
//! 1. **Linear limit**: under a load small enough that finite-strain terms
//!    vanish (strains ~1e-9), the total-Lagrangian nonlinear stepper must
//!    reproduce the verified linear `NewmarkStepper` marching the same
//!    consistent mass and (plane-strain) stiffness, step for step.
//!
//! 2. **Turek–Hron CSM3**: the flag under gravity switched on at rest,
//!    plane-strain St. Venant–Kirchhoff, Newmark average acceleration,
//!    Δt = 0.005 — the benchmark's own time step. Reference (FEATFLOW):
//!    `ux(A) = −14.305 ± 14.305 [1.0995 Hz]`,
//!    `uy(A) = −63.607 ± 65.160 [1.0995 Hz]`.
//!    Measured values and bands are in the test body; the mesh is the
//!    35×2 Quad8 the static CSM tests bounded at ~1.5%.

use nalgebra::{DMatrix, DVector, Vector3};
use rtx_fea::analysis::{
    Analysis, AnalysisConfig, NewmarkStepper, NonlinearConfig, NonlinearDynamicAnalysis,
    NonlinearStaticAnalysis,
};
use rtx_fea::assembly::dof_mapping::{AdvancedDofNumbering, DofComponent, DofMappingStrategy};
use rtx_fea::boundary::dirichlet::{DirichletBC, DirichletType};
use rtx_fea::boundary::{BoundaryCondition, BoundaryConditionSet, SpatialFunction};
use rtx_fea::elements::{ElementMatrixComputer, FiniteElement, StandardFiniteElement};
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;
const G: f64 = 2.0;

/// `nx` by `ny` Quad8 mesh of `[x0, x1] x [y0, y1]` (serendipity lattice),
/// as in `tests/total_lagrangian_svk.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")
}

fn mid_amp(series: &[f64]) -> (f64, f64) {
    let max = series.iter().copied().fold(f64::MIN, f64::max);
    let min = series.iter().copied().fold(f64::MAX, f64::min);
    (0.5 * (max + min), 0.5 * (max - min))
}

fn crossing_frequency(times: &[f64], series: &[f64]) -> f64 {
    let (mean, _) = mid_amp(series);
    let mut crossings: Vec<f64> = Vec::new();
    for k in 1..series.len() {
        let (a, b) = (series[k - 1] - mean, series[k] - mean);
        if a < 0.0 && b >= 0.0 {
            crossings.push(times[k - 1] + (a / (a - b)) * (times[k] - times[k - 1]));
        }
    }
    assert!(crossings.len() >= 3, "too few oscillation periods");
    (crossings.len() - 1) as f64 / (crossings.last().unwrap() - crossings.first().unwrap())
}

/// 1. Linear limit: tiny gravity, TL nonlinear stepper vs the linear
/// `NewmarkStepper` on the same dense M, K, F.
#[test]
fn linear_limit_matches_the_linear_newmark_stepper() {
    let mesh = quad8_rect_mesh(0.25, 0.6, 0.19, 0.21, 10, 2);
    let scale = 1e-6;
    let dt = 0.005;
    let steps = 120;
    let a_node = point_a(&mesh, 0.6, 0.2);

    // Nonlinear TL run.
    let mut analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, 0.25),
        dt,
        steps,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian();
    analysis.set_body_force(move |_| Vector3::new(0.0, -RHO * G * scale, 0.0));
    analysis.track_node(a_node);
    let results = analysis.run().unwrap();
    let uy_nonlinear: Vec<f64> = results.tracked[0].iter().map(|u| u[1]).collect();

    // Linear reference: dense free-free M, K (plane strain, the TL tangent
    // at u = 0), consistent F; the verified linear stepper.
    let mut numbering =
        AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();
    for (&id, node) in mesh.nodes.iter() {
        if (node.position().x - 0.25).abs() < 1e-12 {
            for component in [DofComponent::DisplacementX, DofComponent::DisplacementY] {
                let dof = numbering.get_dof(id, component).unwrap();
                numbering.constrain_dof(dof).unwrap();
            }
        }
    }
    let free = numbering.free_dofs.clone();
    let mut free_index = vec![None; numbering.total_dofs];
    for (k, &dof) in free.iter().enumerate() {
        free_index[dof] = Some(k);
    }
    let n = free.len();
    let mut mass = DMatrix::zeros(n, n);
    let mut stiffness = DMatrix::zeros(n, n);
    let mut force = DVector::zeros(n);
    let mu = E_MOD / (2.0 * (1.0 + NU));
    let lambda = E_MOD * NU / ((1.0 + NU) * (1.0 - 2.0 * NU));
    let mut d = DMatrix::zeros(3, 3);
    d[(0, 0)] = lambda + 2.0 * mu;
    d[(1, 1)] = lambda + 2.0 * mu;
    d[(0, 1)] = lambda;
    d[(1, 0)] = lambda;
    d[(2, 2)] = mu;
    let linear = move |strain: &DVector<f64>| Ok((&d * strain, d.clone()));
    for element in mesh.elements.values() {
        let coords: Vec<Vector3<f64>> = element
            .nodes
            .iter()
            .map(|id| mesh.get_node(*id).unwrap().position())
            .collect();
        let fe = StandardFiniteElement::new(element.element_type, coords.clone());
        let dofs: Vec<usize> = element
            .nodes
            .iter()
            .flat_map(|node| numbering.get_node_dofs(*node))
            .collect();
        let zero = DVector::zeros(dofs.len());
        let (_, k_e) = ElementMatrixComputer::compute_internal_force_and_tangent(
            &fe, &coords, &zero, &linear, None,
        )
        .unwrap();
        let m_scalar =
            ElementMatrixComputer::compute_consistent_mass_matrix(&fe, &coords, RHO, None).unwrap();
        let f_e = ElementMatrixComputer::compute_body_force_vector(
            &fe,
            &coords,
            &|_| Vector3::new(0.0, -RHO * G * scale, 0.0),
            None,
        )
        .unwrap();
        for (lr, &dr) in dofs.iter().enumerate() {
            let Some(fr) = free_index[dr] else { continue };
            force[fr] += f_e[lr];
            for (lc, &dc) in dofs.iter().enumerate() {
                if let Some(fc) = free_index[dc] {
                    stiffness[(fr, fc)] += k_e[(lr, lc)];
                    let (a, b) = (lr / 2, lc / 2);
                    if lr % 2 == lc % 2 {
                        mass[(fr, fc)] += m_scalar.matrix[(a, b)];
                    }
                }
            }
        }
    }
    let damping = DMatrix::zeros(n, n);
    let stepper = NewmarkStepper::average_acceleration(mass, damping, stiffness, dt).unwrap();
    let mut state = stepper
        .initial_state(DVector::zeros(n), DVector::zeros(n), &force)
        .unwrap();
    let a_dofs = numbering.get_node_dofs(a_node);
    let a_free = free_index[a_dofs[1]].expect("point A is free");
    let mut uy_linear = Vec::with_capacity(steps);
    for _ in 0..steps {
        state = stepper.step(&state, &force).unwrap();
        uy_linear.push(state.displacement[a_free]);
    }

    let amplitude = uy_linear.iter().fold(0.0f64, |acc, &x| acc.max(x.abs()));
    let max_diff = uy_nonlinear
        .iter()
        .zip(&uy_linear)
        .fold(0.0f64, |acc, (a, b)| acc.max((a - b).abs()));
    println!(
        "  linear limit: amplitude {amplitude:.3e}, max |nonlinear - linear| {max_diff:.3e} \
         ({:.2e} relative)",
        max_diff / amplitude
    );
    assert!(
        max_diff < 1e-4 * amplitude,
        "nonlinear stepper deviates from the linear one in the linear limit: \
         {max_diff:.3e} vs amplitude {amplitude:.3e}"
    );
}

/// 2. Turek–Hron CSM3. Bands are the measured ones for the 35×2 Quad8 mesh
/// (recorded in the printed line; the static CSM tests bound this mesh's
/// spatial error at ~1.5%).
#[test]
fn turek_hron_csm3_oscillation() {
    let mesh = quad8_rect_mesh(0.25, 0.6, 0.19, 0.21, 35, 2);
    let dt = 0.005;
    let steps = 1200; // 6 s, ~6.6 oscillation periods (2000 steps measured the same bands in 372 s)
    let a_node = point_a(&mesh, 0.6, 0.2);

    let mut analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, 0.25),
        dt,
        steps,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian();
    analysis.set_body_force(|_| Vector3::new(0.0, -RHO * G, 0.0));
    analysis.track_node(a_node);
    let start = std::time::Instant::now();
    let results = analysis.run().unwrap();
    let seconds = start.elapsed().as_secs_f64();

    let ux: Vec<f64> = results.tracked[0].iter().map(|u| u[0]).collect();
    let uy: Vec<f64> = results.tracked[0].iter().map(|u| u[1]).collect();
    let (ux_mean, ux_amp) = mid_amp(&ux);
    let (uy_mean, uy_amp) = mid_amp(&uy);
    let frequency = crossing_frequency(&results.times, &uy);
    let half = uy.len() / 2;
    let (_, amp_first) = mid_amp(&uy[..half]);
    let (_, amp_second) = mid_amp(&uy[half..]);

    println!(
        "  CSM3 35x2 Quad8, dt = {dt}: ux(A) {:.3} ± {:.3} mm, uy(A) {:.3} ± {:.3} mm, \
         f = {frequency:.4} Hz; half-window uy amps {:.3}/{:.3} mm; \
         {} Newton iterations total (max {}/step); {seconds:.0} s.  \
         Reference: ux −14.305 ± 14.305, uy −63.607 ± 65.160 [1.0995 Hz]",
        ux_mean * 1e3,
        ux_amp * 1e3,
        uy_mean * 1e3,
        uy_amp * 1e3,
        amp_first * 1e3,
        amp_second * 1e3,
        results.total_iterations,
        results.max_iterations_per_step,
    );

    let rel = |a: f64, b: f64| ((a - b) / b).abs();
    // Undamped average acceleration: the amplitude must persist.
    assert!(
        (amp_first - amp_second).abs() < 0.02 * amp_second,
        "the undamped oscillation is losing amplitude: {amp_first:.4} vs {amp_second:.4}"
    );
    assert!(
        rel(frequency, 1.0995) < 0.02,
        "frequency {frequency:.4} vs reference 1.0995"
    );
    assert!(
        rel(uy_mean, -63.607e-3) < 0.03,
        "uy mean {:.4e} vs reference -63.607e-3",
        uy_mean
    );
    assert!(
        rel(uy_amp, 65.160e-3) < 0.03,
        "uy amplitude {:.4e} vs reference 65.160e-3",
        uy_amp
    );
    assert!(
        rel(ux_mean, -14.305e-3) < 0.05 && rel(ux_amp, 14.305e-3) < 0.05,
        "ux {:.4e} ± {:.4e} vs reference -14.305e-3 ± 14.305e-3",
        ux_mean,
        ux_amp
    );
    assert!(
        results.max_iterations_per_step <= 5,
        "Newton needed {} iterations in one step",
        results.max_iterations_per_step
    );
}

/// 3. The stepper under a nodal step load — the FSI2 seam. Three claims:
///
/// a. Driving the stepper by hand (set the nodal force, step, commit) is
///    bit-identical to `run()` with the same force set on the analysis —
///    one code path, verified from both ends.
/// b. `step` commits nothing: repeating a step from the same state under
///    the same force is bit-identical; changing the force between the
///    repeats changes the answer (the subiteration a coupling loop needs).
/// c. The undamped step response oscillates about the static deflection:
///    `mid_amp` of the tip trajectory must match the *static* nonlinear
///    analysis under the identical nodal force — a different code path —
///    in both mean and amplitude (`u(t) ≈ u_s (1 − cos ωt)` while the
///    first mode dominates a tip-loaded cantilever).
#[test]
fn stepper_nodal_step_load_oscillates_about_the_static_deflection() {
    let mesh = quad8_rect_mesh(0.25, 0.6, 0.19, 0.21, 10, 2);
    let dt = 0.005;
    let steps = 400; // 2 s: two periods of the ~1 Hz first mode
    let a_node = point_a(&mesh, 0.6, 0.2);
    let tip_force = Vector3::new(0.0, -0.1, 0.0);

    // run() with the force set on the analysis.
    let mut analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, 0.25),
        dt,
        steps,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian();
    analysis.set_nodal_forces(vec![(a_node, tip_force)]);
    analysis.track_node(a_node);
    let results = analysis.run().unwrap();
    let uy_run: Vec<f64> = results.tracked[0].iter().map(|u| u[1]).collect();

    // The same march, driven by hand through the stepper.
    let mut stepper = analysis.stepper().unwrap();
    stepper.set_nodal_forces(&[(a_node, tip_force)]);
    let mut state = stepper.rest_state().unwrap();
    let a_dofs = stepper.node_dofs(a_node);
    let mut uy_manual = Vec::with_capacity(steps);
    for step in 0..steps {
        if step == 7 {
            // b. Re-running the same step is bit-identical; a different
            // force from the same state gives a different answer and
            // leaves no trace once the force is restored.
            let (first, _) = stepper.step(&state).unwrap();
            let (again, _) = stepper.step(&state).unwrap();
            assert_eq!(
                first.displacement, again.displacement,
                "re-running a step from the same state changed the answer"
            );
            stepper.set_nodal_forces(&[(a_node, 2.0 * tip_force)]);
            let (other, _) = stepper.step(&state).unwrap();
            // One step's response to an extra force is ~ ΔF β Δt² / m_modal
            // (≈ 1% of the accumulated displacement here), downward.
            let moved = other.displacement[a_dofs[1]] - first.displacement[a_dofs[1]];
            assert!(
                moved < -1e-3 * first.displacement[a_dofs[1]].abs(),
                "doubling the interface force did not move the step down: \
                 delta {moved:.3e} vs u {:.3e}",
                first.displacement[a_dofs[1]]
            );
            stepper.set_nodal_forces(&[(a_node, tip_force)]);
        }
        let (new_state, _) = stepper.step(&state).unwrap();
        state = new_state;
        uy_manual.push(state.displacement[a_dofs[1]]);
    }
    // a. One code path, verified from both ends.
    assert_eq!(
        uy_run, uy_manual,
        "manual stepper drive deviates from run()"
    );

    // c. Static deflection under the identical nodal force, from the
    // nonlinear *static* analysis.
    let mut static_analysis = NonlinearStaticAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, 0.25),
        NonlinearConfig::default(),
        AnalysisConfig::default(),
    )
    .with_total_lagrangian();
    static_analysis.set_nodal_forces(vec![(a_node, tip_force)]);
    let static_results = static_analysis.run().unwrap();
    assert!(static_results.convergence.converged);
    let numbering =
        AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();
    let uy_static = static_results.displacements[numbering.get_node_dofs(a_node)[1]];

    let (uy_mean, uy_amp) = mid_amp(&uy_manual);
    println!(
        "  step load: static uy = {uy_static:.4e}, dynamic mid ± amp = \
         {uy_mean:.4e} ± {uy_amp:.4e}"
    );
    assert!(
        uy_static < -1e-4,
        "static deflection suspiciously small: {uy_static:.3e}"
    );
    let rel = |a: f64, b: f64| ((a - b) / b).abs();
    assert!(
        rel(uy_mean, uy_static) < 0.03,
        "oscillation midpoint {uy_mean:.4e} vs static deflection {uy_static:.4e}"
    );
    assert!(
        rel(uy_amp, -uy_static) < 0.06,
        "oscillation amplitude {uy_amp:.4e} vs |static| {:.4e}",
        -uy_static
    );
}