//! Verification of the POD/ECSW model-order reduction pipeline.
//!
//! A clamped nonlinear block under a load sweep provides the full-order
//! snapshots; the reduced model is then asked for a load level it has never
//! seen and compared against the full-order solve at that level. Three
//! separate claims, each with its own instrument:
//!
//! 1. **POD-Galerkin accuracy**: the reduced solution at the unseen load
//!    must sit close to the full solve — the subspace captures the load
//!    sweep's solution manifold.
//! 2. **ECSW fidelity**: sampling must not cost accuracy — the hyper-reduced
//!    solution must stay within a small multiple of the POD-Galerkin error,
//!    while assembling a strict subset of the elements. The training
//!    residual itself is asserted against the requested tolerance (that is
//!    the energy-conservation property: the weighted sample reproduces the
//!    virtual work of the full element set over the training states).
//! 3. **Nonnegativity**: every ECSW weight is positive by construction —
//!    a negative weight would let a sampled element produce energy.

use nalgebra::{DMatrix, DVector, Vector3, Vector6};
use rtx_fea::analysis::{Analysis, AnalysisConfig, NonlinearConfig, NonlinearStaticAnalysis};
use rtx_fea::assembly::dof_mapping::{AdvancedDofNumbering, DofComponent, DofMappingStrategy};
use rtx_fea::boundary::{BoundaryCondition, BoundaryConditionSet, DirichletBC};
use rtx_fea::elements::{ElementMatrixComputer, StandardFiniteElement};
use rtx_fea::materials::{
    Material, MaterialDatabase, MaterialProperties, MaterialResponse, MaterialState,
};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};
use rtx_fea::mor::{ReducedNonlinearModel, pod_basis, train_ecsw};

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

/// Same energy-derived nonlinear material as `tests/nonlinear_static.rs`:
/// `W = 1/2 e'De + alpha/3 I1^3`, stress and tangent exact derivatives.
struct CubicEnergy {
    properties: MaterialProperties,
    d: DMatrix<f64>,
}

impl CubicEnergy {
    fn new() -> Self {
        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;
        }
        Self {
            properties: MaterialProperties::isotropic_elastic(E, NU, 1.0),
            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"
    }
}

/// Cantilever-ish block: 6 x 2 x 2 Hex8 elements spanning [0,3] x [0,1] x
/// [0,1], clamped on the x = 0 face.
fn block_mesh() -> Mesh {
    let (nx, ny, nz) = (6usize, 2usize, 2usize);
    let mut mesh = Mesh::new(3).unwrap();
    let mut grid = vec![vec![vec![NodeId(0); nz + 1]; ny + 1]; nx + 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(
                    3.0 * i as f64 / nx as f64,
                    j as f64 / ny as f64,
                    k as f64 / nz as f64,
                ));
            }
        }
    }
    for i in 0..nx {
        for j in 0..ny {
            for k in 0..nz {
                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 clamped_face_bcs(mesh: &Mesh) -> BoundaryConditionSet {
    let clamped: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| node.position().x.abs() < 1e-12)
        .map(|(&id, _)| id)
        .collect();
    let mut bcs = BoundaryConditionSet::new();
    bcs.add_condition(BoundaryCondition::Dirichlet(DirichletBC::fixed(
        clamped,
        vec![
            DofComponent::DisplacementX,
            DofComponent::DisplacementY,
            DofComponent::DisplacementZ,
        ],
        0.0,
    )));
    bcs
}

/// Body force at unit load factor: transverse plus a little axial so the
/// solution manifold is not one-dimensional. Sized for a tip deflection of
/// a few percent of the length: the cubic energy's tangent goes indefinite
/// at large compressive I1, so an over-ambitious load has no equilibrium to
/// converge to.
fn base_force(p: Vector3<f64>) -> Vector3<f64> {
    Vector3::new(1e-4 * p.x, -4e-4, 1e-4 * (p.z - 0.5))
}

fn database() -> MaterialDatabase {
    let mut materials = MaterialDatabase::new();
    materials.add_material(MaterialId(0), CubicEnergy::new(), None);
    materials
}

/// Full-order solve at load factor `lambda`; returns the full DOF vector.
fn full_solve(mesh: &Mesh, lambda: f64) -> DVector<f64> {
    let mut analysis = NonlinearStaticAnalysis::new(
        mesh.clone(),
        database(),
        clamped_face_bcs(mesh),
        NonlinearConfig::default(),
        AnalysisConfig::default(),
    );
    analysis.set_body_force(move |p| base_force(p) * lambda);
    let results = analysis.run().unwrap();
    assert!(results.convergence.converged);
    results.displacements
}

/// The numbering the analysis uses internally, rebuilt identically: same
/// strategy, same constraint source, so global DOF indices agree.
fn numbering(mesh: &Mesh) -> AdvancedDofNumbering {
    let mut dof_numbering =
        AdvancedDofNumbering::displacement_only(mesh, DofMappingStrategy::Sequential).unwrap();
    for (&node_id, node) in &mesh.nodes {
        if node.position().x.abs() < 1e-12 {
            for component in [
                DofComponent::DisplacementX,
                DofComponent::DisplacementY,
                DofComponent::DisplacementZ,
            ] {
                let dof = dof_numbering.get_dof(node_id, component).unwrap();
                dof_numbering.constrain_dof(dof).unwrap();
            }
        }
    }
    dof_numbering
}

fn to_free(full: &DVector<f64>, dof_numbering: &AdvancedDofNumbering) -> DVector<f64> {
    DVector::from_iterator(
        dof_numbering.free_dofs.len(),
        dof_numbering.free_dofs.iter().map(|&dof| full[dof]),
    )
}

/// Consistent body force over the free DOFs at load factor `lambda`.
fn external_force_free(
    mesh: &Mesh,
    dof_numbering: &AdvancedDofNumbering,
    lambda: f64,
) -> DVector<f64> {
    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 mut external = DVector::zeros(dof_numbering.free_dofs.len());
    let force = move |p: Vector3<f64>| base_force(p) * lambda;
    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, &force, None)
                .unwrap();
        let dofs: Vec<usize> = element
            .nodes
            .iter()
            .flat_map(|node| dof_numbering.get_node_dofs(*node))
            .collect();
        for (row, &dof) in dofs.iter().enumerate() {
            if let Some(free) = free_index[dof] {
                external[free] += local[row];
            }
        }
    }
    external
}

#[test]
fn ecsw_reduced_model_reproduces_an_unseen_load_case() {
    let mesh = block_mesh();
    let dof_numbering = numbering(&mesh);
    let materials = database();

    // Snapshots from a load sweep; the test load 0.625 is not in it.
    let training_loads = [0.25, 0.5, 0.75, 1.0];
    let snapshots: Vec<DVector<f64>> = training_loads
        .iter()
        .map(|&lambda| to_free(&full_solve(&mesh, lambda), &dof_numbering))
        .collect();

    let basis = pod_basis(&snapshots, 1e-10).unwrap();
    println!("  POD modes: {}", basis.ncols());

    let ecsw = train_ecsw(&mesh, &materials, &dof_numbering, &basis, &snapshots, 1e-4).unwrap();
    println!(
        "  ECSW: {} of {} elements, training residual {:.3e}",
        ecsw.weights.len(),
        mesh.num_elements(),
        ecsw.training_residual
    );

    // The energy-conservation property: the weighted sample reproduces the
    // reduced internal force over the training set to the requested
    // tolerance, and every weight is nonnegative.
    assert!(ecsw.training_residual <= 1e-4 + 1e-12);
    assert!(ecsw.weights.iter().all(|&(_, w)| w > 0.0));
    assert!(
        ecsw.weights.len() < mesh.num_elements(),
        "ECSW selected every element — no reduction"
    );

    // Unseen load case.
    let lambda = 0.625;
    let full = to_free(&full_solve(&mesh, lambda), &dof_numbering);
    let external = external_force_free(&mesh, &dof_numbering, lambda);

    let galerkin = ReducedNonlinearModel::new(&mesh, &materials, &dof_numbering, basis.clone())
        .unwrap()
        .solve(&external, 1e-9, 30)
        .unwrap();
    let sampled = ReducedNonlinearModel::new(&mesh, &materials, &dof_numbering, basis.clone())
        .unwrap()
        .with_ecsw(&ecsw)
        .solve(&external, 1e-9, 30)
        .unwrap();

    let scale = full.norm();
    let galerkin_error = (&galerkin - &full).norm() / scale;
    let sampled_error = (&sampled - &full).norm() / scale;
    let sampling_gap = (&sampled - &galerkin).norm() / scale;
    println!(
        "  relative errors: POD-Galerkin {galerkin_error:.3e}, ECSW {sampled_error:.3e}, \
         sampling gap {sampling_gap:.3e}"
    );

    // Negative control: the same 5 elements with their weights forced to 1
    // assemble only a fraction of the internal force, so the model must be
    // badly wrong — this is what pins the accuracy on the WEIGHTS rather
    // than on the subset happening to be representative.
    let unweighted = rtx_fea::mor::EcswModel {
        weights: ecsw.weights.iter().map(|&(id, _)| (id, 1.0)).collect(),
        training_residual: f64::NAN,
    };
    let control = ReducedNonlinearModel::new(&mesh, &materials, &dof_numbering, basis.clone())
        .unwrap()
        .with_ecsw(&unweighted)
        .solve(&external, 1e-9, 30)
        .unwrap();
    let control_error = (&control - &full).norm() / scale;
    println!("  unweighted-subset control error: {control_error:.3e}");

    // Measured: POD-Galerkin 3.09e-7, ECSW 3.08e-7, sampling gap 1.6e-8 —
    // the hyper-reduced model is indistinguishable from the full ROM, and
    // both reproduce the unseen full-order solve to the subspace's own
    // accuracy (the load sweep's manifold is analytic in the load factor,
    // which is why two modes carry it so far). The unweighted control reads
    // 2 orders of magnitude worse or more.
    assert!(
        galerkin_error < 1e-5,
        "POD-Galerkin error {galerkin_error:.3e}"
    );
    assert!(sampled_error < 1e-5, "ECSW error {sampled_error:.3e}");
    assert!(
        sampling_gap < 1e-6,
        "hyper-reduction cost {sampling_gap:.3e} against the full ROM"
    );
    assert!(
        control_error > 20.0 * sampled_error,
        "unweighted subset error {control_error:.3e} is not clearly worse than \
         ECSW's {sampled_error:.3e} — the weights are not doing anything"
    );
}

#[test]
fn ecsw_residual_on_the_training_set_reproduces_the_training_residual() {
    use rtx_fea::mor::{Formulation, ecsw_residual};

    let mesh = block_mesh();
    let dof_numbering = numbering(&mesh);
    let materials = database();
    let snapshots: Vec<DVector<f64>> = [0.25, 0.5, 0.75, 1.0]
        .iter()
        .map(|&lambda| to_free(&full_solve(&mesh, lambda), &dof_numbering))
        .collect();
    let basis = pod_basis(&snapshots, 1e-10).unwrap();
    let ecsw = train_ecsw(&mesh, &materials, &dof_numbering, &basis, &snapshots, 1e-4).unwrap();

    // The held-out evaluator, applied to the training set itself, must
    // reproduce the residual the training reported — same C, same b, same
    // weights. This is what licenses using it on a held-out split.
    let replayed = ecsw_residual(
        &mesh,
        &materials,
        &dof_numbering,
        &basis,
        &snapshots,
        &ecsw,
        Formulation::SmallStrain,
    )
    .unwrap();
    assert!(
        (replayed - ecsw.training_residual).abs() < 1e-12,
        "residual evaluator {replayed:.6e} disagrees with training residual {:.6e}",
        ecsw.training_residual
    );
}

/// Full-order total-Lagrangian solve at load factor `lambda`. Convergence
/// is tightened far past the defaults because this solve is the REFERENCE
/// for an identity-basis equivalence assertion — at the default 1e-6
/// criteria the full solve stops ~1.7e-4 (relative) short of the
/// tight-tolerance reduced solve, which is convergence looseness, not a
/// formulation disagreement (measured while writing the test).
fn full_solve_total_lagrangian(mesh: &Mesh, lambda: f64) -> DVector<f64> {
    let mut analysis = NonlinearStaticAnalysis::new(
        mesh.clone(),
        database(),
        clamped_face_bcs(mesh),
        NonlinearConfig {
            convergence_criteria: rtx_fea::analysis::ConvergenceCriteria {
                force_tolerance: 1e-12,
                displacement_tolerance: 1e-12,
                energy_tolerance: 1e-15,
                max_iterations: 60,
            },
            ..NonlinearConfig::default()
        },
        AnalysisConfig::default(),
    )
    .with_total_lagrangian();
    analysis.set_body_force(move |p| base_force(p) * lambda);
    let results = analysis.run().unwrap();
    assert!(results.convergence.converged);
    results.displacements
}

#[test]
fn total_lagrangian_operators_match_the_full_total_lagrangian_solve() {
    use rtx_fea::mor::Formulation;

    let mesh = block_mesh();
    let dof_numbering = numbering(&mesh);
    let materials = database();
    let lambda = 1.0;
    let full_tl = to_free(&full_solve_total_lagrangian(&mesh, lambda), &dof_numbering);
    let external = external_force_free(&mesh, &dof_numbering, lambda);

    // Identity basis: the reduced Newton spans the whole free-DOF space,
    // so the ONLY thing under test is the operators' internal force — if
    // they assembled the small-strain force instead of the
    // total-Lagrangian one, the equilibrium they find would be a
    // different solution.
    let n_free = dof_numbering.free_dofs.len();
    let identity = DMatrix::<f64>::identity(n_free, n_free);

    let reduced_tl = ReducedNonlinearModel::new_formulated(
        &mesh,
        &materials,
        &dof_numbering,
        identity.clone(),
        Formulation::TotalLagrangian,
    )
    .unwrap()
    .solve(&external, 1e-11, 60)
    .unwrap();
    let reduced_ss = ReducedNonlinearModel::new_formulated(
        &mesh,
        &materials,
        &dof_numbering,
        identity,
        Formulation::SmallStrain,
    )
    .unwrap()
    .solve(&external, 1e-11, 60)
    .unwrap();

    let scale = full_tl.norm();
    let tl_error = (&reduced_tl - &full_tl).norm() / scale;
    let ss_gap = (&reduced_ss - &full_tl).norm() / scale;
    println!("  identity-basis TL error {tl_error:.3e}; small-strain gap {ss_gap:.3e}");

    // Note the full TL solve uses the SVK constitutive from the material's
    // Lamé parameters (the analysis's TL branch bypasses compute_response),
    // so the identity-basis reduced TL solve targets the same equations.
    assert!(
        tl_error < 1e-6,
        "identity-basis total-Lagrangian solve should reproduce the full \
         TL solve: relative error {tl_error:.3e}"
    );
    assert!(
        ss_gap > 20.0 * tl_error.max(1e-9),
        "small-strain operators land indistinguishably close to the TL \
         solution ({ss_gap:.3e}) — the formulation switch is not being \
         exercised at this load"
    );
}