//! Modal analysis against closed-form natural frequencies.
//!
//! This is the end-to-end check that mesh, DOF numbering, constraints,
//! element matrices, global assembly and the eigensolver are all correct
//! *together*. Each has its own unit tests; none of those would catch a
//! mismatch between them, such as a mass matrix assembled in a different DOF
//! order than its stiffness.
//!
//! # Why axial modes and not a cantilever
//!
//! The obvious benchmark is the bending frequency of a cantilever,
//! `β₁L = 1.8751`. It is the wrong first test here. Bilinear `Quad4` elements
//! suffer **shear locking** in bending: their assumed displacement field
//! cannot represent pure bending without spurious shear strain, so a coarse
//! mesh is far too stiff and reports frequencies well above the true value.
//! A cantilever test would fail for a reason that has nothing to do with
//! whether the code under test is correct, and tuning the tolerance until it
//! passed would destroy its value as evidence.
//!
//! Longitudinal (axial) vibration has no such problem. The exact solution of
//! the 1-D wave equation for a fixed-free bar is
//!
//! ```text
//!     f_n = (2n - 1) / (4L) * sqrt(E / rho),    n = 1, 2, 3, ...
//! ```
//!
//! and a plane-stress mesh with transverse motion suppressed reduces to
//! exactly that problem. Linear elements with a consistent mass matrix
//! converge to it from above at `O(h²)`, so a modest mesh lands within a
//! fraction of a percent — tight enough that a real error cannot hide.
//!
//! Bending is still checked below, but as a *convergence* statement rather
//! than a single tolerance, which is the honest way to assert on an element
//! that is known to lock.

use rtx_fea::analysis::{Analysis, AnalysisConfig, AnalysisData, ModalAnalysis};
use rtx_fea::assembly::DofComponent;
use rtx_fea::boundary::{BoundaryCondition, BoundaryConditionSet, DirichletBC};
use rtx_fea::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};

const E: f64 = 200e9;
const RHO: f64 = 8000.0;
const LENGTH: f64 = 1.0;
const HEIGHT: f64 = 0.05;

/// A rectangular `nx` by `ny` grid of `Quad4` elements spanning
/// `[0, LENGTH] x [0, HEIGHT]`, returned with its node grid so tests can pick
/// out edges to constrain.
fn bar_mesh(nx: usize, ny: usize) -> (Mesh, Vec<Vec<NodeId>>) {
    let mut mesh = Mesh::new(2).unwrap();

    let mut grid = vec![vec![NodeId(0); ny + 1]; nx + 1];
    for (i, column) in grid.iter_mut().enumerate() {
        for (j, slot) in column.iter_mut().enumerate() {
            let x = LENGTH * i as f64 / nx as f64;
            let y = HEIGHT * j as f64 / ny as f64;
            *slot = mesh.add_node(Node::new_2d(x, y));
        }
    }

    for i in 0..nx {
        for j in 0..ny {
            // Counter-clockwise, so the Jacobian determinant is positive.
            let nodes = vec![
                grid[i][j],
                grid[i + 1][j],
                grid[i + 1][j + 1],
                grid[i][j + 1],
            ];
            let element = Element::new(ElementType::Quad4, nodes, MaterialId(0)).unwrap();
            mesh.add_element(element).unwrap();
        }
    }

    (mesh, grid)
}

/// Poisson's ratio is zero throughout.
///
/// This is a deliberate modelling choice, not a convenience: with `nu = 0`
/// the axial and transverse responses decouple exactly, so the plane-stress
/// model reduces to the 1-D bar the closed form describes. A non-zero
/// Poisson's ratio would introduce a real physical difference between the two
/// and the comparison would no longer be exact.
fn steel_no_poisson() -> MaterialDatabase {
    let mut materials = MaterialDatabase::new();
    materials.add_material(
        MaterialId(0),
        LinearElastic::new(E, 0.0).with_density(RHO),
        Some("steel".to_string()),
    );
    materials
}

fn frequencies_of(results: &rtx_fea::analysis::AnalysisResults) -> Vec<f64> {
    match results
        .additional_data
        .get("frequencies")
        .expect("frequencies missing")
    {
        AnalysisData::Vector(v) => v.iter().copied().collect(),
        other => panic!("frequencies had unexpected type {other:?}"),
    }
}

/// Longitudinal modes of a fixed-free bar against `f_n = (2n-1)/(4L)·√(E/ρ)`.
#[test]
fn axial_modes_match_the_closed_form_bar() {
    let nx = 24;
    let ny = 2;
    let (mesh, grid) = bar_mesh(nx, ny);

    let mut bcs = BoundaryConditionSet::new();

    // Clamp the left edge axially.
    let left_edge: Vec<NodeId> = grid[0].clone();
    bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
        left_edge,
        vec![DofComponent::DisplacementX],
        0.0,
    )));

    // Suppress transverse motion everywhere, reducing the plane-stress model
    // to the 1-D bar the closed form describes. Without this the spectrum is
    // interleaved with bending modes and the comparison is meaningless.
    let all_nodes: Vec<NodeId> = grid.iter().flatten().copied().collect();
    bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
        all_nodes,
        vec![DofComponent::DisplacementY],
        0.0,
    )));

    let num_modes = 3;
    let mut analysis = ModalAnalysis::new(
        mesh,
        steel_no_poisson(),
        num_modes,
        AnalysisConfig::default(),
    )
    .with_boundary_conditions(bcs);

    let results = analysis.run().expect("modal analysis failed");
    let computed = frequencies_of(&results);

    let wave_speed = (E / RHO).sqrt();
    for n in 1..=num_modes {
        let exact = (2 * n - 1) as f64 / (4.0 * LENGTH) * wave_speed;
        let got = computed[n - 1];
        let relative_error = (got - exact).abs() / exact;

        assert!(
            relative_error < 0.01,
            "mode {n}: computed {got:.4} Hz against exact {exact:.4} Hz \
             ({:.3}% error)",
            relative_error * 100.0
        );

        // Linear elements with a consistent mass matrix are stiffer than the
        // continuum, so the discrete frequency must come in high. Landing
        // below the exact value means something is wrong even if the
        // magnitude looks plausible.
        assert!(
            got >= exact * (1.0 - 1e-9),
            "mode {n}: computed {got:.4} Hz is below the exact {exact:.4} Hz; \
             a consistent-mass discretisation cannot be softer than the continuum"
        );
    }
}

/// Refining the mesh must drive the axial error down, and at the expected
/// second-order rate.
///
/// A single tolerance check can be satisfied by a wrong formula with a
/// compensating error. A convergence *rate* cannot: it pins the
/// discretisation itself.
#[test]
fn axial_frequency_converges_at_second_order() {
    let wave_speed = (E / RHO).sqrt();
    let exact = wave_speed / (4.0 * LENGTH);

    let mut errors = Vec::new();
    for nx in [4usize, 8, 16] {
        let (mesh, grid) = bar_mesh(nx, 1);

        let mut bcs = BoundaryConditionSet::new();
        bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
            grid[0].clone(),
            vec![DofComponent::DisplacementX],
            0.0,
        )));
        bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
            grid.iter().flatten().copied().collect(),
            vec![DofComponent::DisplacementY],
            0.0,
        )));

        let mut analysis =
            ModalAnalysis::new(mesh, steel_no_poisson(), 1, AnalysisConfig::default())
                .with_boundary_conditions(bcs);
        let results = analysis.run().expect("modal analysis failed");
        errors.push((frequencies_of(&results)[0] - exact).abs() / exact);
    }

    for window in errors.windows(2) {
        let rate = (window[0] / window[1]).log2();
        assert!(
            rate > 1.7,
            "halving the element size reduced the error by only 2^{rate:.2}; \
             expected close to second order. errors: {errors:?}"
        );
    }
}

/// An unconstrained structure has rigid-body modes, so `K` is singular.
///
/// The failure must be a clear error rather than a set of near-zero
/// eigenvalues that look like real low-frequency modes.
#[test]
fn unconstrained_structure_is_rejected_rather_than_silently_wrong() {
    let (mesh, _) = bar_mesh(4, 1);
    let mut analysis = ModalAnalysis::new(mesh, steel_no_poisson(), 2, AnalysisConfig::default());

    let error = analysis
        .run()
        .expect_err("an unconstrained structure must not yield frequencies");
    let message = error.to_string().to_lowercase();
    assert!(
        message.contains("singular") || message.contains("shift"),
        "error should name the singular stiffness or the shift remedy, got: {error}"
    );
}

/// Every reported natural frequency must be real and positive.
///
/// A constrained, positive-definite structure has no zero-frequency mode. A
/// zero or NaN here means the constraints did not reach the assembled system
/// or the eigenvalues came back negative.
#[test]
fn frequencies_are_real_and_positive() {
    let (mesh, grid) = bar_mesh(6, 2);

    let mut bcs = BoundaryConditionSet::new();
    bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
        grid[0].clone(),
        vec![DofComponent::DisplacementX, DofComponent::DisplacementY],
        0.0,
    )));

    let mut analysis = ModalAnalysis::new(mesh, steel_no_poisson(), 4, AnalysisConfig::default())
        .with_boundary_conditions(bcs);
    let results = analysis.run().expect("modal analysis failed");

    let frequencies = frequencies_of(&results);
    assert_eq!(frequencies.len(), 4);
    for (i, f) in frequencies.iter().enumerate() {
        assert!(
            f.is_finite() && *f > 0.0,
            "mode {} frequency is {f}, which is not a physical frequency",
            i + 1
        );
    }

    // Ascending, since modes are named by index.
    for pair in frequencies.windows(2) {
        assert!(
            pair[1] >= pair[0],
            "frequencies are not ascending: {frequencies:?}"
        );
    }
}

/// Cantilever bending, asserted as convergence rather than as a tolerance.
///
/// `Quad4` locks in bending, so the coarse-mesh frequency is far too high.
/// What must still hold is that refinement moves it monotonically *towards*
/// the Euler-Bernoulli value `f₁ = (β₁L)²/(2πL²)·√(EI/ρA)` with
/// `β₁L = 1.8751` — and that it approaches from above, which is the signature
/// of locking rather than of a bug.
#[test]
fn cantilever_bending_converges_towards_euler_bernoulli_from_above() {
    let beta_l: f64 = 1.8751;
    // Plane stress with unit thickness: A = h, I = h³/12.
    let area = HEIGHT;
    let second_moment = HEIGHT.powi(3) / 12.0;
    let exact = beta_l.powi(2) / (2.0 * std::f64::consts::PI * LENGTH.powi(2))
        * (E * second_moment / (RHO * area)).sqrt();

    let mut computed = Vec::new();
    for (nx, ny) in [(8usize, 2usize), (16, 4), (32, 8)] {
        let (mesh, grid) = bar_mesh(nx, ny);

        let mut bcs = BoundaryConditionSet::new();
        bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
            grid[0].clone(),
            vec![DofComponent::DisplacementX, DofComponent::DisplacementY],
            0.0,
        )));

        let mut analysis =
            ModalAnalysis::new(mesh, steel_no_poisson(), 1, AnalysisConfig::default())
                .with_boundary_conditions(bcs);
        let results = analysis.run().expect("modal analysis failed");
        computed.push(frequencies_of(&results)[0]);
    }

    for (i, f) in computed.iter().enumerate() {
        assert!(
            *f > exact * 0.95,
            "mesh {i}: {f:.3} Hz is below the Euler-Bernoulli value {exact:.3} Hz; \
             a locking element cannot be softer than the beam theory it approximates"
        );
    }

    for pair in computed.windows(2) {
        assert!(
            pair[1] <= pair[0] * 1.001,
            "refining the mesh increased the bending frequency ({:?}); \
             locking must relax under refinement, not worsen",
            computed
        );
    }

    let coarse_error = (computed[0] - exact).abs() / exact;
    let fine_error = (computed[computed.len() - 1] - exact).abs() / exact;
    assert!(
        fine_error < coarse_error,
        "refinement did not reduce the bending error: {coarse_error:.4} -> {fine_error:.4} \
         against exact {exact:.3} Hz, computed {computed:?}"
    );
}