//! Verification of `NonlinearStaticAnalysis`.
//!
//! Two independent instruments, per the crate's verification protocol:
//!
//! 1. **Equivalence with the linear path.** With a linear-elastic material
//!    the Newton loop must land on exactly the solution the linear assembly
//!    produces — same `B` matrices, same quadrature, same solver — and it
//!    must get there in one Newton step, because the residual of a linear
//!    problem after one exact tangent solve is zero. Any disagreement is a
//!    defect in the nonlinear assembly, since everything else is shared.
//!
//! 2. **Manufactured solution with a genuinely nonlinear material.** The
//!    material `CubicEnergy` below derives from the stored energy
//!    `W = 1/2 eps' D eps + (alpha/3) I1^3`, so its stress
//!    `sigma = D eps + alpha I1^2 m` and consistent tangent
//!    `D_T = D + 2 alpha I1 m m'` (with `m = [1,1,1,0,0,0]'`) are exact by
//!    construction and the tangent is symmetric. The body force
//!    `f = -div sigma(eps(u_exact))` is computed by central differences of
//!    the closed-form stress field, the same trick
//!    `tests/mms_elastostatics.rs` uses to cross-check its hand-derived
//!    force. The observed L2 order must be 2 — and it can only get there if
//!    the nonlinear term is actually solved, because the forcing contains it.

use nalgebra::{DMatrix, DVector, Matrix3, Vector3, Vector6};
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::{ElementMatrixComputer, FiniteElement, StandardFiniteElement};
use rtx_fea::materials::{
    LinearElastic, Material, MaterialDatabase, MaterialProperties, MaterialResponse, MaterialState,
};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};
use rtx_fea::solvers::{LinearSolver, LuDirect, SolverOptions};

const E: f64 = 1.0;
const NU: f64 = 0.3;
const ALPHA: f64 = 3.0;
const AMP: f64 = 0.15;

// ---------------------------------------------------------------------------
// The manufactured field and its gradient, hand-differentiated
// ---------------------------------------------------------------------------

fn u_exact(p: Vector3<f64>) -> Vector3<f64> {
    use std::f64::consts::PI;
    let (x, y, z) = (p.x, p.y, p.z);
    Vector3::new(
        AMP * (PI * x).sin() * (PI * y).cos() * (PI * z).cos(),
        AMP * (PI * x).cos() * (PI * y).sin() * (PI * z).cos(),
        -2.0 * AMP * (PI * x).cos() * (PI * y).cos() * (PI * z).sin(),
    )
}

fn grad_u(p: Vector3<f64>) -> Matrix3<f64> {
    use std::f64::consts::PI;
    let (x, y, z) = (p.x, p.y, p.z);
    let (sx, cx) = ((PI * x).sin(), (PI * x).cos());
    let (sy, cy) = ((PI * y).sin(), (PI * y).cos());
    let (sz, cz) = ((PI * z).sin(), (PI * z).cos());
    let a = AMP * PI;
    Matrix3::new(
        a * cx * cy * cz,
        -a * sx * sy * cz,
        -a * sx * cy * sz,
        -a * sx * sy * cz,
        a * cx * cy * cz,
        -a * cx * sy * sz,
        2.0 * a * sx * cy * sz,
        2.0 * a * cx * sy * sz,
        -2.0 * a * cx * cy * cz,
    )
}

/// Small-strain tensor in Voigt-6 (engineering shear), matching the element
/// `B` matrix convention.
fn strain_voigt(p: Vector3<f64>) -> Vector6<f64> {
    let g = grad_u(p);
    Vector6::new(
        g[(0, 0)],
        g[(1, 1)],
        g[(2, 2)],
        g[(1, 2)] + g[(2, 1)],
        g[(0, 2)] + g[(2, 0)],
        g[(0, 1)] + g[(1, 0)],
    )
}

fn elastic_d() -> DMatrix<f64> {
    let lambda = E * NU / ((1.0 + NU) * (1.0 - 2.0 * NU));
    let mu = E / (2.0 * (1.0 + NU));
    let mut d = DMatrix::zeros(6, 6);
    for i in 0..3 {
        for j in 0..3 {
            d[(i, j)] = lambda;
        }
        d[(i, i)] += 2.0 * mu;
        d[(i + 3, i + 3)] = mu;
    }
    d
}

/// Stress of the `CubicEnergy` material at a point of the exact field.
fn sigma_exact(p: Vector3<f64>) -> Vector6<f64> {
    let eps = strain_voigt(p);
    let d = elastic_d();
    let i1 = eps[0] + eps[1] + eps[2];
    let mut sigma = Vector6::zeros();
    for i in 0..6 {
        for j in 0..6 {
            sigma[i] += d[(i, j)] * eps[j];
        }
    }
    for i in 0..3 {
        sigma[i] += ALPHA * i1 * i1;
    }
    sigma
}

/// Body force `f_i = -d sigma_ij / d x_j`, by central differences of the
/// closed-form stress. Voigt row (i, j) lookup: the full tensor from Voigt-6.
fn body_force(p: Vector3<f64>) -> Vector3<f64> {
    let h = 1e-6;
    let tensor = |q: Vector3<f64>| -> Matrix3<f64> {
        let s = sigma_exact(q);
        Matrix3::new(s[0], s[5], s[4], s[5], s[1], s[3], s[4], s[3], s[2])
    };
    let mut f = Vector3::zeros();
    for j in 0..3 {
        let mut dq = Vector3::zeros();
        dq[j] = h;
        let ds = (tensor(p + dq) - tensor(p - dq)) / (2.0 * h);
        for i in 0..3 {
            f[i] -= ds[(i, j)];
        }
    }
    f
}

// ---------------------------------------------------------------------------
// The nonlinear test material
// ---------------------------------------------------------------------------

/// `W = 1/2 eps' D eps + (alpha/3) I1^3`; stress and tangent are exact
/// derivatives of the energy, so the tangent is consistent by construction.
struct CubicEnergy {
    properties: MaterialProperties,
    d: DMatrix<f64>,
}

impl CubicEnergy {
    fn new() -> Self {
        Self {
            properties: MaterialProperties::isotropic_elastic(E, NU, 1.0),
            d: elastic_d(),
        }
    }
}

impl Material for CubicEnergy {
    fn properties(&self) -> &MaterialProperties {
        &self.properties
    }

    fn compute_response(
        &self,
        strain: &Vector6<f64>,
        state: &MaterialState,
        _dt: f64,
    ) -> rtx_fea::error::FeaResult<MaterialResponse> {
        let i1 = strain[0] + strain[1] + strain[2];
        let mut stress = Vector6::zeros();
        for i in 0..6 {
            for j in 0..6 {
                stress[i] += self.d[(i, j)] * strain[j];
            }
        }
        for i in 0..3 {
            stress[i] += ALPHA * i1 * i1;
        }

        let mut tangent = self.d.clone();
        for i in 0..3 {
            for j in 0..3 {
                tangent[(i, j)] += 2.0 * ALPHA * i1;
            }
        }

        Ok(MaterialResponse::new(stress, tangent, state.clone()))
    }

    fn elastic_tangent(&self) -> rtx_fea::error::FeaResult<DMatrix<f64>> {
        Ok(self.d.clone())
    }

    fn material_type(&self) -> &'static str {
        "CubicEnergy"
    }
}

// ---------------------------------------------------------------------------
// Meshing and boundary conditions
// ---------------------------------------------------------------------------

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

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

/// The exact field prescribed on the whole boundary: one spatial Dirichlet
/// condition per displacement component, over the boundary nodes.
fn exact_boundary_conditions(mesh: &Mesh) -> BoundaryConditionSet {
    let boundary_nodes: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| on_boundary(node.position()))
        .map(|(&id, _)| id)
        .collect();

    let mut set = BoundaryConditionSet::new();
    let components = [
        (DofComponent::DisplacementX, 0usize),
        (DofComponent::DisplacementY, 1),
        (DofComponent::DisplacementZ, 2),
    ];
    for (component, axis) in components {
        set.add_condition(BoundaryCondition::Dirichlet(DirichletBC {
            nodes: boundary_nodes.clone(),
            components: vec![component],
            condition_type: DirichletType::Spatial(SpatialFunction(Box::new(move |p| {
                u_exact(*p)[axis]
            }))),
            time_range: None,
            ramping_factor: 1.0,
            gradual_enforcement: false,
        }));
    }
    set
}

/// Quadrature-integrated L2 error of a solved displacement field against the
/// exact one, mirroring `tests/mms_elastostatics.rs`.
fn l2_error(mesh: &Mesh, dof_numbering: &AdvancedDofNumbering, solution: &DVector<f64>) -> f64 {
    let mut squared = 0.0;
    for element in mesh.elements.values() {
        let node_coords: Vec<Vector3<f64>> = element
            .nodes
            .iter()
            .map(|id| mesh.get_node(*id).unwrap().position())
            .collect();
        let fe = StandardFiniteElement::new(element.element_type, node_coords.clone());
        let rule = fe.quadrature_rule(None).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 jacobian = fe.jacobian(&point.coords, &node_coords).unwrap();
            let physical = fe.map_to_physical(&point.coords, &node_coords).unwrap();

            let mut uh = Vector3::zeros();
            for node_index in 0..element.nodes.len() {
                let n = shape.value(node_index).unwrap();
                for d in 0..3 {
                    uh[d] += n * solution[dofs[node_index * 3 + d]];
                }
            }
            let exact = u_exact(physical.coords);
            squared += (uh - exact).norm_squared() * point.weight * jacobian.determinant().abs();
        }
    }
    squared.sqrt()
}

fn nonlinear_config() -> NonlinearConfig {
    NonlinearConfig::default()
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

/// With a linear material, one Newton step must land on exactly the solution
/// of the directly assembled linear system.
#[test]
fn linear_material_reproduces_the_linear_solution_in_one_step() {
    let mesh = hex8_mesh(3);
    let bcs = exact_boundary_conditions(&mesh);

    // Nonlinear path.
    let mut materials = MaterialDatabase::new();
    materials.add_material(MaterialId(0), LinearElastic::new(E, NU), None);
    let mut analysis = NonlinearStaticAnalysis::new(
        mesh.clone(),
        materials,
        bcs,
        nonlinear_config(),
        AnalysisConfig::default(),
    );
    analysis.set_body_force(body_force);
    let results = analysis.run().unwrap();
    assert!(results.convergence.converged);

    // A linear problem after one exact tangent solve has zero residual; the
    // config's single load step should therefore take exactly one iteration.
    assert!(
        results.convergence.iterations <= 1,
        "linear problem took {} Newton iterations",
        results.convergence.iterations
    );

    // Direct linear solve with the same numbering, constraints and loads.
    let mut dof_numbering =
        AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();
    let mut prescribed = Vec::new();
    for (&node_id, node) in &mesh.nodes {
        let p = node.position();
        if !on_boundary(p) {
            continue;
        }
        let value = u_exact(p);
        for (i, component) in [
            DofComponent::DisplacementX,
            DofComponent::DisplacementY,
            DofComponent::DisplacementZ,
        ]
        .into_iter()
        .enumerate()
        {
            let dof = dof_numbering.get_dof(node_id, component).unwrap();
            dof_numbering.constrain_dof(dof).unwrap();
            prescribed.push((dof, value[i]));
        }
    }
    let total = dof_numbering.total_dofs;
    let mut free_index = vec![None; total];
    for (i, &dof) in dof_numbering.free_dofs.iter().enumerate() {
        free_index[dof] = Some(i);
    }
    let num_free = dof_numbering.free_dofs.len();

    let mut linear = DVector::zeros(total);
    for &(dof, value) in &prescribed {
        linear[dof] = value;
    }

    // K_ff x = f_f - K_fc u_c, assembled per element.
    let mut stiffness = rtx_fea::assembly::SparseMatrix::new(num_free, num_free);
    let mut rhs = DVector::zeros(num_free);
    for element in mesh.elements.values() {
        let node_coords: Vec<Vector3<f64>> = element
            .nodes
            .iter()
            .map(|id| mesh.get_node(*id).unwrap().position())
            .collect();
        let fe = StandardFiniteElement::new(element.element_type, node_coords.clone());
        let k = ElementMatrixComputer::compute_stiffness_matrix(&fe, &node_coords, E, NU, None)
            .unwrap()
            .matrix;
        let f =
            ElementMatrixComputer::compute_body_force_vector(&fe, &node_coords, &body_force, None)
                .unwrap();

        let dofs: Vec<usize> = element
            .nodes
            .iter()
            .flat_map(|node| dof_numbering.get_node_dofs(*node))
            .collect();
        for (row, &dof_row) in dofs.iter().enumerate() {
            let Some(free_row) = free_index[dof_row] else {
                continue;
            };
            rhs[free_row] += f[row];
            for (col, &dof_col) in dofs.iter().enumerate() {
                match free_index[dof_col] {
                    Some(free_col) => {
                        if k[(row, col)] != 0.0 {
                            stiffness
                                .add_entry(free_row, free_col, k[(row, col)])
                                .unwrap();
                        }
                    }
                    None => rhs[free_row] -= k[(row, col)] * linear[dof_col],
                }
            }
        }
    }
    stiffness.finalize().unwrap();
    let mut solver = LuDirect::new();
    let (x, _) = solver
        .solve(&stiffness, &rhs, &SolverOptions::default())
        .unwrap();
    for (i, &dof) in dof_numbering.free_dofs.iter().enumerate() {
        linear[dof] = x[i];
    }

    let max_diff = results
        .displacements
        .iter()
        .zip(linear.iter())
        .map(|(a, b)| (a - b).abs())
        .fold(0.0_f64, f64::max);
    assert!(
        max_diff < 1e-10,
        "nonlinear path differs from the linear solution by {max_diff:.3e}"
    );
}

/// Manufactured solution with the genuinely nonlinear material: the observed
/// L2 order must be 2, which only happens if the nonlinear term is solved —
/// the forcing contains it.
///
/// Measured (2 -> 4 -> 8 Hex8): L2 error 6.032e-2, 1.780e-2, 4.595e-3 —
/// observed orders 1.76 and 1.95, climbing to the theoretical 2, with Newton
/// converging in a handful of iterations at every resolution.
#[test]
fn nonlinear_mms_converges_at_second_order() {
    let resolutions = [2usize, 4, 8];
    let mut errors = Vec::new();

    for &n in &resolutions {
        let mesh = hex8_mesh(n);
        let bcs = exact_boundary_conditions(&mesh);
        let mut materials = MaterialDatabase::new();
        materials.add_material(MaterialId(0), CubicEnergy::new(), None);

        let mut analysis = NonlinearStaticAnalysis::new(
            mesh.clone(),
            materials,
            bcs,
            nonlinear_config(),
            AnalysisConfig::default(),
        );
        analysis.set_body_force(body_force);
        let results = analysis.run().unwrap();
        assert!(results.convergence.converged);

        // Full Newton with a consistent tangent converges quadratically: a
        // handful of iterations per load step, not dozens. A wrong tangent
        // still creeps to the answer — this is what catches it.
        let steps = nonlinear_config().max_load_steps.max(1);
        assert!(
            results.convergence.iterations <= 8 * steps,
            "Newton took {} iterations over {steps} load steps — the tangent \
             is not consistent with the stress",
            results.convergence.iterations
        );

        let dof_numbering =
            AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();
        errors.push(l2_error(&mesh, &dof_numbering, &results.displacements));
    }

    let rates: Vec<f64> = errors
        .windows(2)
        .map(|pair| (pair[0] / pair[1]).log2())
        .collect();
    for (i, &n) in resolutions.iter().enumerate() {
        let rate = if i == 0 {
            String::from("   -")
        } else {
            format!("{:4.2}", rates[i - 1])
        };
        println!("  n = {n}   L2 error = {:.6e}   order = {rate}", errors[i]);
    }

    assert!(errors.windows(2).all(|pair| pair[1] < pair[0]));
    for &rate in &rates {
        assert!(
            (1.7..2.4).contains(&rate),
            "observed order {rate:.2}, expected 2 for Hex8; errors {errors:?}"
        );
    }
}

/// No load and homogeneous boundary data must produce the zero solution.
#[test]
fn zero_problem_stays_zero() {
    let mesh = hex8_mesh(2);
    let boundary_nodes: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| on_boundary(node.position()))
        .map(|(&id, _)| id)
        .collect();
    let mut bcs = BoundaryConditionSet::new();
    bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
        boundary_nodes,
        vec![
            DofComponent::DisplacementX,
            DofComponent::DisplacementY,
            DofComponent::DisplacementZ,
        ],
        0.0,
    )));

    let mut materials = MaterialDatabase::new();
    materials.add_material(MaterialId(0), CubicEnergy::new(), None);
    let mut analysis = NonlinearStaticAnalysis::new(
        mesh,
        materials,
        bcs,
        nonlinear_config(),
        AnalysisConfig::default(),
    );
    let results = analysis.run().unwrap();
    assert!(results.convergence.converged);
    assert!(results.displacements.norm() < 1e-12);
}