//! Code verification by the Method of Manufactured Solutions.
//!
//! Every other test in this crate asserts that some quantity is *close enough*
//! to a value someone believed. MMS asserts something stronger and more
//! objective: that the discretisation converges to the exact solution at the
//! rate the theory predicts.
//!
//! The method needs no benchmark data. Choose any sufficiently smooth field
//! `u`, substitute it into the governing equations, and whatever they fail to
//! balance is the body force `f = -∇·σ(u)` that makes `u` the exact solution of
//! the problem with that load. Solve with `f` applied and `u` imposed on the
//! boundary, measure the error, refine, and read off the observed order
//! `log2(e_h / e_{h/2})`.
//!
//! This matters here specifically. Every defect found in this crate returned a
//! plausible wrong answer and passed the tests written for it — element
//! matrices of zeros, a quadrature rule with no points, an assembly transposed
//! per node. None of them survives an order-of-accuracy measurement: a wrong
//! discretisation either fails to converge or converges at the wrong rate, and
//! neither can be tuned away.
//!
//! # The manufactured solution
//!
//! ```text
//!     u(x, y) = sin(pi x) sin(pi y)
//!     v(x, y) = x^2 (1 - x) y (1 - y)
//! ```
//!
//! Chosen so that the field is smooth but **not** in the element space — a
//! bilinear `Quad4` cannot represent either component exactly, so the error is
//! genuinely discretisation error rather than round-off. The two components are
//! deliberately different in form, and the shear strain
//! `du/dy + dv/dx` is non-zero, so the shear block of the constitutive matrix
//! is exercised rather than silently skipped.

use nalgebra::{DVector, Vector3};
use rtx_fea::assembly::{
    AdvancedDofNumbering, AssemblyOptions, DofComponent, DofMappingStrategy, GlobalAssembler,
};
use rtx_fea::elements::{
    ElementMatrixComputer, FiniteElement, NaturalCoords, StandardFiniteElement,
};
use rtx_fea::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};
use rtx_fea::solvers::{SolverFactory, SolverOptions};
use std::f64::consts::PI;

const E: f64 = 1.0;
const NU: f64 = 0.3;

/// Exact displacement field.
fn exact(x: f64, y: f64) -> (f64, f64) {
    (
        (PI * x).sin() * (PI * y).sin(),
        x * x * (1.0 - x) * y * (1.0 - y),
    )
}

/// Body force `f = -div(sigma(u))` for plane stress.
///
/// Derived by hand from the manufactured field above with
/// `sigma_xx = c(e_xx + nu e_yy)`, `sigma_yy = c(nu e_xx + e_yy)`,
/// `sigma_xy = mu gamma_xy`, where `c = E / (1 - nu^2)` and
/// `mu = E / (2(1 + nu))` — the same plane-stress constitutive matrix
/// `ElementMatrixComputer::elasticity_matrix` builds.
///
/// Strains:
/// ```text
///     e_xx = pi cos(pi x) sin(pi y)
///     e_yy = x^2 (1 - x) (1 - 2y)
///     gamma_xy = pi sin(pi x) cos(pi y) + (2x - 3x^2) y (1 - y)
/// ```
fn body_force(x: f64, y: f64) -> (f64, f64) {
    let c = E / (1.0 - NU * NU);
    let mu = E / (2.0 * (1.0 + NU));

    let (sx, cx) = ((PI * x).sin(), (PI * x).cos());
    let (sy, cy) = ((PI * y).sin(), (PI * y).cos());

    // d(e_xx)/dx, d(e_yy)/dx, d(e_xx)/dy, d(e_yy)/dy
    let dexx_dx = -PI * PI * sx * sy;
    let deyy_dx = (2.0 * x - 3.0 * x * x) * (1.0 - 2.0 * y);
    let dexx_dy = PI * PI * cx * cy;
    let deyy_dy = -2.0 * x * x * (1.0 - x);

    // d(gamma_xy)/dy and d(gamma_xy)/dx
    let dgxy_dy = -PI * PI * sx * sy + (2.0 * x - 3.0 * x * x) * (1.0 - 2.0 * y);
    let dgxy_dx = PI * PI * cx * cy + (2.0 - 6.0 * x) * y * (1.0 - y);

    let dsxx_dx = c * (dexx_dx + NU * deyy_dx);
    let dsyy_dy = c * (NU * dexx_dy + deyy_dy);
    let dsxy_dy = mu * dgxy_dy;
    let dsxy_dx = mu * dgxy_dx;

    (-(dsxx_dx + dsxy_dy), -(dsxy_dx + dsyy_dy))
}

/// Unit square meshed with `n` by `n` `Quad4` elements.
fn unit_square_mesh(n: usize) -> (Mesh, Vec<Vec<NodeId>>) {
    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() {
            let x = i as f64 / n as f64;
            let y = j as f64 / n as f64;
            *slot = mesh.add_node(Node::new_2d(x, y));
        }
    }
    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, grid)
}

/// Solve the manufactured problem on an `n` by `n` mesh and return the
/// discrete L2 displacement error.
fn l2_error(n: usize) -> f64 {
    let (mesh, grid) = unit_square_mesh(n);

    let mut materials = MaterialDatabase::new();
    materials.add_material(MaterialId(0), LinearElastic::new(E, NU), None);

    let mut dof_numbering =
        AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();

    // Impose the exact solution on the whole boundary.
    let mut prescribed = Vec::new();
    for i in 0..=n {
        for j in 0..=n {
            if i != 0 && i != n && j != 0 && j != n {
                continue;
            }
            let node = grid[i][j];
            let (ux, uy) = exact(i as f64 / n as f64, j as f64 / n as f64);
            for (component, value) in [
                (DofComponent::DisplacementX, ux),
                (DofComponent::DisplacementY, uy),
            ] {
                let dof = dof_numbering.get_dof(node, component).unwrap();
                dof_numbering.constrain_dof(dof).unwrap();
                prescribed.push((dof, value));
            }
        }
    }

    let simple = dof_numbering.to_dof_numbering();
    let assembler = GlobalAssembler::new(AssemblyOptions::default(), materials);
    let mut system = assembler.assemble_system(&mesh, &simple, 0.0).unwrap();

    // Consistent body force: integral of N_i f over each element.
    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 local = ElementMatrixComputer::compute_body_force_vector(
            &fe,
            &node_coords,
            &|p: Vector3<f64>| {
                let (fx, fy) = body_force(p.x, p.y);
                Vector3::new(fx, fy, 0.0)
            },
            None,
        )
        .unwrap();

        let mut global_dofs = Vec::new();
        for node in &element.nodes {
            global_dofs.extend_from_slice(simple.get_node_dofs(*node).unwrap());
        }
        for (local_index, &global_dof) in global_dofs.iter().enumerate() {
            system.force_vector[global_dof] += local[local_index];
        }
    }

    for (dof, value) in prescribed {
        system.set_prescribed_value(dof, value).unwrap();
    }

    let (free_stiffness, free_force) = system.extract_free_system().unwrap();
    let properties = SolverFactory::analyze_matrix(&free_stiffness);
    let options = SolverOptions::default();
    let mut solver = SolverFactory::create_linear_solver(&properties, &options);
    let (free_solution, _) = solver
        .solve(&free_stiffness, &free_force, &options)
        .unwrap();

    // Expand to the full DOF vector, with the prescribed values in place.
    let mut solution = DVector::zeros(system.dof_numbering.total_dofs);
    for (i, &dof) in system.dof_numbering.free_dofs.iter().enumerate() {
        solution[dof] = free_solution[i];
    }
    for (&dof, &value) in system.prescribed_values() {
        solution[dof] = value;
    }

    // L2 error by the elements' own quadrature.
    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(Some(3)).unwrap();

        let dofs: Vec<usize> = element
            .nodes
            .iter()
            .flat_map(|node| simple.get_node_dofs(*node).unwrap().to_vec())
            .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 = (0.0, 0.0);
            for (node_index, _) in element.nodes.iter().enumerate() {
                let n = shape.value(node_index).unwrap();
                uh.0 += n * solution[dofs[node_index * 2]];
                uh.1 += n * solution[dofs[node_index * 2 + 1]];
            }

            let (ux, uy) = exact(physical.x(), physical.y());
            let weight = point.weight * jacobian.determinant().abs();
            squared += ((uh.0 - ux).powi(2) + (uh.1 - uy).powi(2)) * weight;
        }
    }

    squared.sqrt()
}

/// The displacement error must fall at second order under refinement.
///
/// `Quad4` interpolates bilinearly, so the theoretical L2 displacement rate is
/// 2. A rate near 1 means something in the chain is only first-order accurate;
/// a rate near 0 means the solution is not converging to the exact one at all,
/// which is what every stub found in this crate would produce.
#[test]
fn quad4_displacement_converges_at_second_order() {
    // 8 -> 16 -> 32 gives two independent rate estimates, which is enough to
    // establish the order. A 64 x 64 rung adds a third at roughly ten times
    // the cost, and it reads 2.00 as well.
    let resolutions = [8usize, 16, 32];
    let errors: Vec<f64> = resolutions.iter().map(|&n| l2_error(n)).collect();

    let rates: Vec<f64> = errors
        .windows(2)
        .map(|pair| (pair[0] / pair[1]).log2())
        .collect();

    // Printed so the measurement is auditable rather than merely asserted.
    for (i, &n) in resolutions.iter().enumerate() {
        let rate = if i == 0 {
            String::from("    -")
        } else {
            format!("{:5.2}", rates[i - 1])
        };
        println!(
            "  n = {n:3}   L2 error = {:.6e}   observed order = {rate}",
            errors[i]
        );
    }

    for (i, &rate) in rates.iter().enumerate() {
        assert!(
            rate > 1.8 && rate < 2.2,
            "refinement {} -> {}: observed order {rate:.3}, expected 2 for Quad4. \
             Errors: {errors:?}",
            resolutions[i],
            resolutions[i + 1]
        );
    }
}

/// The error must actually be small, not merely converging at the right rate.
///
/// A scheme converging at second order toward the wrong answer would satisfy
/// the rate test alone.
#[test]
fn quad4_error_is_small_on_a_fine_mesh() {
    let error = l2_error(32);
    assert!(
        error < 1e-3,
        "L2 error {error:.3e} on a 32x32 mesh is too large for a second-order \
         method on a smooth solution"
    );
}