//! Verification of the total-Lagrangian St. Venant–Kirchhoff path
//! (`elements::total_lagrangian`, `NonlinearStaticAnalysis::with_total_lagrangian`).
//!
//! In the order the claims must be established:
//!
//! 1. At zero displacement the TL tangent is the small-strain plane-strain
//!    stiffness, to rounding.
//! 2. The tangent is the derivative of the internal force — finite
//!    differences at a finite random displacement, for Quad4, Quad8, Hex8.
//!    A plausible-but-wrong geometric stiffness fails this and nothing else.
//! 3. A finite rigid rotation produces no internal force (Green–Lagrange
//!    strain is objective); the small-strain routine does produce one — the
//!    negative control that shows the test has teeth.
//! 4. Manufactured solutions at finite strain recover the element orders
//!    (Quad4 ~2, Quad8 ~3 in L2), with the body force obtained by finite
//!    differences of the exact first Piola–Kirchhoff stress.
//! 5. Turek–Hron CSM1 and CSM2 (flag under gravity, clamped at the cylinder):
//!    reference `u_x(A) = −7.18777e-3, u_y(A) = −66.1029e-3` (CSM1) and
//!    `−0.469006e-3, −16.9740e-3` (CSM2), FEATFLOW tables.

use nalgebra::{DMatrix, DVector, Vector3};
use rtx_fea::analysis::{Analysis, AnalysisConfig, NonlinearConfig, 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::total_lagrangian::{internal_force_and_tangent, saint_venant_kirchhoff};
use rtx_fea::elements::{ElementMatrixComputer, FiniteElement, StandardFiniteElement};
use rtx_fea::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};
use std::f64::consts::PI;

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

fn lame() -> (f64, f64) {
    let mu = E_MOD / (2.0 * (1.0 + NU));
    let lambda = E_MOD * NU / ((1.0 + NU) * (1.0 - 2.0 * NU));
    (lambda, mu)
}

fn plane_strain_d() -> DMatrix<f64> {
    let (lambda, mu) = lame();
    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;
    d
}

fn quad4_coords() -> Vec<Vector3<f64>> {
    // A deliberately non-rectangular quadrilateral.
    vec![
        Vector3::new(0.0, 0.0, 0.0),
        Vector3::new(1.1, 0.1, 0.0),
        Vector3::new(0.9, 1.0, 0.0),
        Vector3::new(-0.1, 0.8, 0.0),
    ]
}

fn quad8_coords() -> Vec<Vector3<f64>> {
    let c = quad4_coords();
    let mid = |a: usize, b: usize| 0.5 * (c[a] + c[b]);
    vec![
        c[0],
        c[1],
        c[2],
        c[3],
        mid(0, 1),
        mid(1, 2),
        mid(2, 3),
        mid(3, 0),
    ]
}

fn hex8_coords() -> Vec<Vector3<f64>> {
    vec![
        Vector3::new(0.0, 0.0, 0.0),
        Vector3::new(1.0, 0.0, 0.1),
        Vector3::new(1.1, 1.0, 0.0),
        Vector3::new(0.0, 0.9, 0.0),
        Vector3::new(0.1, 0.0, 1.0),
        Vector3::new(1.0, 0.1, 1.0),
        Vector3::new(1.0, 1.0, 1.1),
        Vector3::new(0.0, 1.0, 0.9),
    ]
}

/// Deterministic pseudo-random displacement of amplitude `amp`.
fn pseudo_random(n: usize, amp: f64, seed: u64) -> DVector<f64> {
    let mut x = seed;
    DVector::from_iterator(
        n,
        (0..n).map(|_| {
            x = x
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            amp * (((x >> 33) as f64) / (1u64 << 31) as f64 - 1.0)
        }),
    )
}

/// 1. Zero displacement: TL tangent == small-strain plane-strain stiffness.
#[test]
fn at_zero_displacement_the_tangent_is_the_plane_strain_stiffness() {
    let coords = quad4_coords();
    let fe = StandardFiniteElement::new(ElementType::Quad4, coords.clone());
    let (lambda, mu) = lame();
    let svk = saint_venant_kirchhoff(lambda, mu, 2);
    let u0 = DVector::zeros(8);
    let (f_int, k_tl) = internal_force_and_tangent(&fe, &coords, &u0, svk.as_ref(), None).unwrap();
    let d = plane_strain_d();
    let linear = move |strain: &DVector<f64>| Ok((&d * strain, d.clone()));
    let (_, k_lin) =
        ElementMatrixComputer::compute_internal_force_and_tangent(&fe, &coords, &u0, &linear, None)
            .unwrap();
    assert!(
        f_int.norm() == 0.0,
        "internal force at zero displacement: {}",
        f_int.norm()
    );
    let diff = (&k_tl - &k_lin).norm() / k_lin.norm();
    assert!(
        diff < 1e-13,
        "TL tangent at u = 0 differs from the plane-strain stiffness by {diff:.3e}"
    );
}

/// 2. The tangent is the derivative of the internal force (central
/// differences at a finite displacement; step chosen so the FD error is far
/// below the tolerance for a smooth polynomial residual).
#[test]
fn tangent_is_the_derivative_of_the_internal_force() {
    let (lambda, mu) = lame();
    for (name, element_type, coords, dim) in [
        ("Quad4", ElementType::Quad4, quad4_coords(), 2usize),
        ("Quad8", ElementType::Quad8, quad8_coords(), 2),
        ("Hex8", ElementType::Hex8, hex8_coords(), 3),
    ] {
        let fe = StandardFiniteElement::new(element_type, coords.clone());
        let svk = saint_venant_kirchhoff(lambda, mu, dim);
        let n = coords.len() * dim;
        // 20% strain-level displacement: genuinely finite.
        let u = pseudo_random(n, 0.2, 7);
        let (_, k) = internal_force_and_tangent(&fe, &coords, &u, svk.as_ref(), None).unwrap();
        let eps = 1e-6;
        let mut k_fd = DMatrix::zeros(n, n);
        for j in 0..n {
            let mut up = u.clone();
            let mut um = u.clone();
            up[j] += eps;
            um[j] -= eps;
            let (fp, _) =
                internal_force_and_tangent(&fe, &coords, &up, svk.as_ref(), None).unwrap();
            let (fm, _) =
                internal_force_and_tangent(&fe, &coords, &um, svk.as_ref(), None).unwrap();
            let col = (fp - fm) / (2.0 * eps);
            k_fd.set_column(j, &col);
        }
        let rel = (&k - &k_fd).norm() / k.norm();
        assert!(
            rel < 1e-7,
            "{name}: tangent vs finite-difference derivative of f_int: relative {rel:.3e}"
        );
        let asym = (&k - k.transpose()).norm() / k.norm();
        assert!(asym < 1e-12, "{name}: tangent not symmetric: {asym:.3e}");
    }
}

/// 3. Finite rigid rotation: no internal force in the TL routine; the
/// small-strain routine sees strain and pushes back (negative control).
#[test]
fn rigid_rotation_produces_no_internal_force_and_the_small_strain_routine_fails_this() {
    let coords = quad8_coords();
    let fe = StandardFiniteElement::new(ElementType::Quad8, coords.clone());
    let theta: f64 = 0.6; // 34 degrees
    let (s, c) = theta.sin_cos();
    let mut u = DVector::zeros(16);
    for (a, x) in coords.iter().enumerate() {
        u[2 * a] = c * x.x - s * x.y - x.x;
        u[2 * a + 1] = s * x.x + c * x.y - x.y;
    }
    let (lambda, mu) = lame();
    let svk = saint_venant_kirchhoff(lambda, mu, 2);
    let (f_tl, _) = internal_force_and_tangent(&fe, &coords, &u, svk.as_ref(), None).unwrap();
    let d = plane_strain_d();
    let linear = move |strain: &DVector<f64>| Ok((&d * strain, d.clone()));
    let (f_small, _) =
        ElementMatrixComputer::compute_internal_force_and_tangent(&fe, &coords, &u, &linear, None)
            .unwrap();
    let scale = E_MOD * u.norm();
    assert!(
        f_tl.norm() < 1e-10 * scale,
        "TL internal force under a rigid rotation: {:.3e} (scale {scale:.3e})",
        f_tl.norm()
    );
    assert!(
        f_small.norm() > 1e-2 * scale,
        "the small-strain routine should see a rigid rotation as strain; got {:.3e}",
        f_small.norm()
    );
}

// ---------------------------------------------------------------------------
// Manufactured solution at finite strain
// ---------------------------------------------------------------------------

const AMP: f64 = 0.03;

fn u_exact(p: Vector3<f64>) -> Vector3<f64> {
    let (x, y) = (p.x, p.y);
    Vector3::new(
        AMP * (PI * x).sin() * (PI * y).sin(),
        AMP * (PI * x).sin() * (PI * y).sin() * 0.5 + AMP * 0.3 * x * y,
        0.0,
    )
}

fn grad_u(p: Vector3<f64>) -> DMatrix<f64> {
    let (x, y) = (p.x, p.y);
    let sx = (PI * x).sin();
    let cx = (PI * x).cos();
    let sy = (PI * y).sin();
    let cy = (PI * y).cos();
    let mut h = DMatrix::zeros(2, 2);
    h[(0, 0)] = AMP * PI * cx * sy;
    h[(0, 1)] = AMP * PI * sx * cy;
    h[(1, 0)] = 0.5 * AMP * PI * cx * sy + AMP * 0.3 * y;
    h[(1, 1)] = 0.5 * AMP * PI * sx * cy + AMP * 0.3 * x;
    h
}

/// First Piola–Kirchhoff stress `P = F S` of the manufactured field.
fn piola(p: Vector3<f64>) -> DMatrix<f64> {
    let (lambda, mu) = lame();
    let f = DMatrix::identity(2, 2) + grad_u(p);
    let e = 0.5 * (f.transpose() * &f - DMatrix::identity(2, 2));
    let s = lambda * e.trace() * DMatrix::identity(2, 2) + 2.0 * mu * &e;
    f * s
}

/// Body force per unit reference volume `b = −Div P`, by central
/// differences of the analytic `P` (step 1e-6: FD error ~1e-12 relative).
fn body_force(p: Vector3<f64>) -> Vector3<f64> {
    let eps = 1e-6;
    let mut div = Vector3::zeros();
    for j in 0..2 {
        let mut dp = Vector3::zeros();
        dp[j] = eps;
        let plus = piola(p + dp);
        let minus = piola(p - dp);
        for i in 0..2 {
            div[i] += (plus[(i, j)] - minus[(i, j)]) / (2.0 * eps);
        }
    }
    -div
}

fn on_unit_square_boundary(p: Vector3<f64>) -> bool {
    (0..2).any(|d| p[d].abs() < 1e-12 || (p[d] - 1.0).abs() < 1e-12)
}

fn quad4_mesh(n: usize) -> Mesh {
    let mut mesh = Mesh::new(2).unwrap();
    let mut grid = vec![vec![NodeId(0); n + 1]; n + 1];
    for (i, column) in grid.iter_mut().enumerate() {
        for (j, slot) in column.iter_mut().enumerate() {
            *slot = mesh.add_node(Node::new_2d(i as f64 / n as f64, j as f64 / n as f64));
        }
    }
    for i in 0..n {
        for j in 0..n {
            let nodes = vec![
                grid[i][j],
                grid[i + 1][j],
                grid[i + 1][j + 1],
                grid[i][j + 1],
            ];
            mesh.add_element(Element::new(ElementType::Quad4, nodes, MaterialId(0)).unwrap())
                .unwrap();
        }
    }
    mesh
}

/// `nx` by `ny` Quad8 mesh of `[x0, x1] x [y0, y1]` (serendipity lattice with
/// cell centres left out), corners counter-clockwise then mid-edges.
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), None);
    db
}

fn dirichlet(
    nodes: Vec<NodeId>,
    component: DofComponent,
    f: impl Fn(Vector3<f64>) -> f64 + Send + Sync + 'static,
) -> BoundaryCondition {
    BoundaryCondition::Dirichlet(DirichletBC {
        nodes,
        components: vec![component],
        condition_type: DirichletType::Spatial(SpatialFunction(Box::new(move |p| f(*p)))),
        time_range: None,
        ramping_factor: 1.0,
        gradual_enforcement: false,
    })
}

fn exact_boundary_conditions(mesh: &Mesh) -> BoundaryConditionSet {
    let boundary: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| on_unit_square_boundary(node.position()))
        .map(|(&id, _)| id)
        .collect();
    let mut set = BoundaryConditionSet::new();
    set.add_condition(dirichlet(
        boundary.clone(),
        DofComponent::DisplacementX,
        |p| u_exact(p).x,
    ));
    set.add_condition(dirichlet(boundary, DofComponent::DisplacementY, |p| {
        u_exact(p).y
    }));
    set
}

/// Quadrature-integrated L2 error against the manufactured field.
fn l2_error(mesh: &Mesh, dof_numbering: &AdvancedDofNumbering, solution: &DVector<f64>) -> f64 {
    let mut squared = 0.0;
    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 rule = fe.quadrature_rule(Some(4)).unwrap();
        let dofs: Vec<usize> = element
            .nodes
            .iter()
            .flat_map(|node| dof_numbering.get_node_dofs(*node))
            .collect();
        for point in &rule.points {
            let shape = fe.shape_functions(&point.coords).unwrap();
            let jac = fe.jacobian(&point.coords, &coords).unwrap();
            let physical = fe.map_to_physical(&point.coords, &coords).unwrap();
            let mut uh = Vector3::zeros();
            for a in 0..element.nodes.len() {
                let n = shape.value(a).unwrap();
                uh.x += n * solution[dofs[2 * a]];
                uh.y += n * solution[dofs[2 * a + 1]];
            }
            let exact = u_exact(physical.coords);
            squared += (uh - exact).norm_squared() * point.weight * jac.determinant().abs();
        }
    }
    squared.sqrt()
}

fn solve_mms(mesh: Mesh) -> f64 {
    let dof_numbering =
        AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();
    let bcs = exact_boundary_conditions(&mesh);
    let config = NonlinearConfig {
        max_load_steps: 5,
        ..NonlinearConfig::default()
    };
    let mut analysis = NonlinearStaticAnalysis::new(
        mesh.clone(),
        materials(),
        bcs,
        config,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian();
    analysis.set_body_force(body_force);
    let results = analysis.run().unwrap();
    assert!(
        results.convergence.converged,
        "Newton did not converge on the manufactured problem"
    );
    l2_error(&mesh, &dof_numbering, &results.displacements)
}

fn observed_order(errors: &[f64]) -> Vec<f64> {
    errors.windows(2).map(|w| (w[0] / w[1]).log2()).collect()
}

/// 4. Manufactured finite-strain solution: Quad4 at order 2, Quad8 at 3.
#[test]
fn manufactured_finite_strain_solution_converges_at_the_element_orders() {
    // Sanity on the manufactured data: max |grad u| ~ AMP*pi ~ 0.25 — finite.
    let quad4_errors: Vec<f64> = [4usize, 8, 16]
        .iter()
        .map(|&n| solve_mms(quad4_mesh(n)))
        .collect();
    let quad8_errors: Vec<f64> = [2usize, 4, 8, 16]
        .iter()
        .map(|&n| solve_mms(quad8_rect_mesh(0.0, 1.0, 0.0, 1.0, n, n)))
        .collect();
    let o4 = observed_order(&quad4_errors);
    let o8 = observed_order(&quad8_errors);
    println!("  Quad4 L2 errors {quad4_errors:?} orders {o4:?}");
    println!("  Quad8 L2 errors {quad8_errors:?} orders {o8:?}");
    assert!(o4.last().unwrap() > &1.8, "Quad4 order {o4:?}");
    assert!(o8.last().unwrap() > &2.7, "Quad8 order {o8:?}");
}

// ---------------------------------------------------------------------------
// Turek–Hron CSM1 / CSM2
// ---------------------------------------------------------------------------

struct Csm {
    ux_a: f64,
    uy_a: f64,
    iterations: usize,
}

/// The flag `[0.25, 0.6] x [0.19, 0.21]`, clamped at `x = 0.25`, under
/// gravity `g = 2` downward, density 1000, plane-strain SVK with the given
/// shear modulus and `nu = 0.4`. Returns the displacement of
/// `A = (0.6, 0.2)`.
fn run_csm(mu_s: f64, nx: usize, ny: usize, load_steps: usize) -> Csm {
    let e_mod = 2.0 * mu_s * (1.0 + NU);
    let mesh = quad8_rect_mesh(0.25, 0.6, 0.19, 0.21, nx, ny);
    let clamped: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| (node.position().x - 0.25).abs() < 1e-12)
        .map(|(&id, _)| id)
        .collect();
    let point_a = mesh
        .nodes
        .iter()
        .find(|(_, node)| {
            (node.position().x - 0.6).abs() < 1e-12 && (node.position().y - 0.2).abs() < 1e-12
        })
        .map(|(&id, _)| id)
        .expect("point A (0.6, 0.2) must be a mesh node");
    let mut bcs = BoundaryConditionSet::new();
    bcs.add_condition(dirichlet(
        clamped.clone(),
        DofComponent::DisplacementX,
        |_| 0.0,
    ));
    bcs.add_condition(dirichlet(clamped, DofComponent::DisplacementY, |_| 0.0));

    let mut db = MaterialDatabase::new();
    db.add_material(MaterialId(0), LinearElastic::new(e_mod, NU), None);
    let dof_numbering =
        AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();
    let config = NonlinearConfig {
        max_load_steps: load_steps,
        ..NonlinearConfig::default()
    };
    let mut analysis =
        NonlinearStaticAnalysis::new(mesh, db, bcs, config, AnalysisConfig::default())
            .with_total_lagrangian();
    analysis.set_body_force(|_| Vector3::new(0.0, -1000.0 * 2.0, 0.0));
    let results = analysis.run().unwrap();
    assert!(results.convergence.converged, "CSM Newton did not converge");
    let dofs = dof_numbering.get_node_dofs(point_a);
    Csm {
        ux_a: results.displacements[dofs[0]],
        uy_a: results.displacements[dofs[1]],
        iterations: results.convergence.iterations,
    }
}

#[test]
fn turek_hron_csm1_and_csm2_deflections() {
    // (mu_s, reference ux, reference uy) in metres.
    let cases = [
        ("CSM1", 0.5e6, -7.18777e-3, -66.1029e-3),
        ("CSM2", 2.0e6, -0.469006e-3, -16.9740e-3),
    ];
    // Measured (Quad8, plane-strain SVK, 5 load steps): CSM1 35x2
    // (-7.006e-3, -65.14e-3), 70x4 (-7.060e-3, -65.43e-3) — converging
    // from below onto the reference (-7.188e-3, -66.10e-3), 1.0% short in
    // u_y and 1.8% in u_x at 5 mm elements. `TL_FINE=1` runs 140x8 as well
    // (too slow for the suite) for the convergence record.
    let fine_mesh = if std::env::var("TL_FINE").is_ok() {
        (140, 8)
    } else {
        (70, 4)
    };
    for (name, mu_s, ref_ux, ref_uy) in cases {
        let coarse = run_csm(mu_s, 35, 2, 5);
        let fine = run_csm(mu_s, fine_mesh.0, fine_mesh.1, 5);
        println!(
            "  {name}: 35x2 Quad8 u(A) = ({:.5e}, {:.5e}); {}x{} Quad8 u(A) = ({:.5e}, {:.5e}) \
             [{} Newton iterations]; reference ({ref_ux:.5e}, {ref_uy:.5e})",
            coarse.ux_a,
            coarse.uy_a,
            fine_mesh.0,
            fine_mesh.1,
            fine.ux_a,
            fine.uy_a,
            fine.iterations
        );
        let rel = |a: f64, b: f64| ((a - b) / b).abs();
        assert!(
            rel(fine.uy_a, ref_uy) < 0.015,
            "{name}: u_y(A) = {:.5e} vs reference {ref_uy:.5e}",
            fine.uy_a
        );
        assert!(
            rel(fine.ux_a, ref_ux) < 0.03,
            "{name}: u_x(A) = {:.5e} vs reference {ref_ux:.5e}",
            fine.ux_a
        );
        assert!(
            rel(fine.uy_a, ref_uy) <= rel(coarse.uy_a, ref_uy) + 1e-4,
            "{name}: refinement moved away from the reference"
        );
    }
}