//! The reduced Newmark driver (`mor::dynamic`) against the full-order
//! stepper — phase 4's machinery, verified before it touches the flag.
//!
//! 1. Identity-basis equivalence: with `V = I` over the free DOFs the
//!    reduced Newmark IS the full Newmark (same predictor, same
//!    residual, same tangent, different linear-solver rounding), so a
//!    march must reproduce the full stepper's trajectory to near
//!    machine precision when both Newton loops are run tight.
//! 2. The reduced mass carries the right physics: a rigid-translation
//!    vector must see exactly the total mass ρ·A.
//! 3. A truncated POD basis built from full-order snapshots must track
//!    the full trajectory it was trained on to within a band far above
//!    its projection error but far below any wrong-dynamics answer.

use nalgebra::{DVector, Vector3};
use rtx_fea::analysis::{AnalysisConfig, ConvergenceCriteria, NonlinearDynamicAnalysis};
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::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};
use rtx_fea::mor::{Formulation, ReducedNewmark, ReducedNonlinearModel, pod_basis};

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

/// `nx` by `ny` Quad8 serendipity mesh of `[x0, x1] x [y0, y1]` (as in
/// `newton_rescue.rs` and the FSI harness).
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).with_density(RHO),
        None,
    );
    db
}

fn clamp_left(mesh: &Mesh, x_left: f64) -> BoundaryConditionSet {
    let clamped: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| (node.position().x - x_left).abs() < 1e-12)
        .map(|(&id, _)| id)
        .collect();
    let mut set = BoundaryConditionSet::new();
    for component in [DofComponent::DisplacementX, DofComponent::DisplacementY] {
        set.add_condition(BoundaryCondition::Dirichlet(DirichletBC {
            nodes: clamped.clone(),
            components: vec![component],
            condition_type: DirichletType::Spatial(SpatialFunction(Box::new(|_| 0.0))),
            time_range: None,
            ramping_factor: 1.0,
            gradual_enforcement: false,
        }));
    }
    set
}

/// The numbering the analysis uses internally, rebuilt identically
/// (Sequential strategy, same clamp criterion) — the
/// `fsi3_ecsw_offline` pattern, self-checked there against the dump.
fn clamped_numbering(mesh: &Mesh, x_left: f64) -> AdvancedDofNumbering {
    let mut numbering =
        AdvancedDofNumbering::displacement_only(mesh, DofMappingStrategy::Sequential).unwrap();
    for (&node_id, node) in &mesh.nodes {
        if (node.position().x - x_left).abs() < 1e-12 {
            for component in [DofComponent::DisplacementX, DofComponent::DisplacementY] {
                let dof = numbering.get_dof(node_id, component).unwrap();
                numbering.constrain_dof(dof).unwrap();
            }
        }
    }
    numbering
}

fn tip_node(mesh: &Mesh, x: f64, y: f64) -> NodeId {
    mesh.nodes
        .iter()
        .find(|(_, node)| {
            (node.position().x - x).abs() < 1e-12 && (node.position().y - y).abs() < 1e-12
        })
        .map(|(&id, _)| id)
        .expect("tip node")
}

/// A nodal force as a free-DOF vector under `numbering`.
fn free_force(numbering: &AdvancedDofNumbering, loads: &[(NodeId, Vector3<f64>)]) -> DVector<f64> {
    let mut free_index = vec![None; numbering.total_dofs];
    for (i, &dof) in numbering.free_dofs.iter().enumerate() {
        free_index[dof] = Some(i);
    }
    let mut force = DVector::zeros(numbering.free_dofs.len());
    for (node, f) in loads {
        for (component, &dof) in numbering.get_node_dofs(*node).iter().enumerate() {
            if let Some(free) = free_index[dof] {
                force[free] += f[component];
            }
        }
    }
    force
}

/// Tight Newton on both sides so the converged states differ only by
/// linear-solver rounding, not by the stopping tolerance.
fn tight_criteria() -> ConvergenceCriteria {
    ConvergenceCriteria {
        force_tolerance: 1e-10,
        displacement_tolerance: 1e-12,
        energy_tolerance: 1e-14,
        max_iterations: 50,
    }
}

#[test]
fn identity_basis_matches_full_stepper() {
    let (x0, x1) = (0.0, 0.35);
    let mesh = quad8_rect_mesh(x0, x1, 0.0, 0.02, 4, 1);
    let tip = tip_node(&mesh, x1, 0.01);
    let dt = 1e-3;

    // Full order.
    let analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, x0),
        dt,
        1,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian()
    .with_convergence_criteria(tight_criteria());
    let mut stepper = analysis.stepper().unwrap();
    let load = [(tip, Vector3::new(0.0, -40.0, 0.0))];
    stepper.set_nodal_forces(&load);
    let mut full_state = stepper.rest_state().unwrap();

    // Reduced with V = I over the free DOFs.
    let db = materials();
    let numbering = clamped_numbering(&mesh, x0);
    let n_free = numbering.free_dofs.len();
    let basis = nalgebra::DMatrix::identity(n_free, n_free);
    let model = ReducedNonlinearModel::new_formulated(
        &mesh,
        &db,
        &numbering,
        basis,
        Formulation::TotalLagrangian,
    )
    .unwrap();
    let newmark = ReducedNewmark::new(&model, dt)
        .unwrap()
        .with_convergence_criteria(tight_criteria());
    let external = free_force(&numbering, &load);
    let mut reduced_state = newmark.rest_state(&external).unwrap();

    // The consistent initial accelerations must already agree.
    let a_full: DVector<f64> = DVector::from_iterator(
        n_free,
        numbering
            .free_dofs
            .iter()
            .map(|&d| full_state.acceleration[d]),
    );
    let a0_err = (&a_full - &reduced_state.q_ddot).norm() / a_full.norm();
    assert!(
        a0_err < 1e-9,
        "rest-state accelerations differ: {a0_err:.3e}"
    );

    let mut worst = 0.0f64;
    for step in 0..15 {
        let (next_full, _) = stepper.step(&full_state).expect("full step");
        let (next_reduced, _) = newmark
            .step(&reduced_state, &external)
            .expect("reduced step");
        full_state = next_full;
        reduced_state = next_reduced;

        let d_full: DVector<f64> = DVector::from_iterator(
            n_free,
            numbering
                .free_dofs
                .iter()
                .map(|&d| full_state.displacement[d]),
        );
        let d_reduced = model.expand(&reduced_state.q);
        let rel = (&d_full - &d_reduced).norm() / d_full.norm().max(1e-30);
        worst = worst.max(rel);
        assert!(
            rel < 1e-9,
            "step {step}: identity-basis reduced Newmark left the full trajectory \
             (rel {rel:.3e})"
        );
    }
    println!("  identity-basis march: worst per-step rel deviation {worst:.3e} over 15 steps");
}

#[test]
fn reduced_mass_is_total_mass_on_rigid_translation() {
    // Unconstrained 2x1 mesh: every DOF free, identity basis — the
    // reduced mass IS the assembled mass. A rigid x-translation stores
    // the total mass ρ·A; any quadrature or expansion slip breaks the
    // total while keeping the matrix symmetric positive definite.
    let (w, h) = (2.0, 0.5);
    let mesh = quad8_rect_mesh(0.0, w, 0.0, h, 2, 1);
    let db = materials();
    let numbering =
        AdvancedDofNumbering::displacement_only(&mesh, DofMappingStrategy::Sequential).unwrap();
    let n_free = numbering.free_dofs.len();
    let basis = nalgebra::DMatrix::identity(n_free, n_free);
    let model = ReducedNonlinearModel::new_formulated(
        &mesh,
        &db,
        &numbering,
        basis,
        Formulation::TotalLagrangian,
    )
    .unwrap();
    let mass = model.reduced_mass().unwrap();

    // x-translation: 1 on every x DOF (free DOF order follows the
    // numbering; even/odd split by component index within a node).
    let mut e_x = DVector::zeros(n_free);
    let mut free_index = vec![None; numbering.total_dofs];
    for (i, &dof) in numbering.free_dofs.iter().enumerate() {
        free_index[dof] = Some(i);
    }
    for (&node_id, _) in &mesh.nodes {
        let dofs = numbering.get_node_dofs(node_id);
        if let Some(free) = free_index[dofs[0]] {
            e_x[free] = 1.0;
        }
    }
    let total = (e_x.transpose() * &mass * &e_x)[(0, 0)];
    let expected = RHO * w * h;
    let rel = (total - expected).abs() / expected;
    assert!(
        rel < 1e-9,
        "rigid x-translation sees {total:.6} against total mass {expected:.6} (rel {rel:.3e})"
    );
}

#[test]
fn truncated_basis_tracks_its_training_trajectory() {
    let (x0, x1) = (0.0, 0.35);
    let mesh = quad8_rect_mesh(x0, x1, 0.0, 0.02, 4, 1);
    let tip = tip_node(&mesh, x1, 0.01);
    let dt = 1e-3;
    let steps = 60;

    // Full-order march under a smoothly ramped tip load; snapshot every
    // step.
    let analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, x0),
        dt,
        1,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian()
    .with_convergence_criteria(tight_criteria());
    let mut stepper = analysis.stepper().unwrap();
    let db = materials();
    let numbering = clamped_numbering(&mesh, x0);
    let n_free = numbering.free_dofs.len();
    let load_at = |k: usize| {
        let t = (k as f64) * dt;
        [(tip, Vector3::new(0.0, -30.0 * (35.0 * t).sin(), 0.0))]
    };

    stepper.set_nodal_forces(&load_at(0));
    let mut full_state = stepper.rest_state().unwrap();
    let mut snapshots: Vec<DVector<f64>> = Vec::with_capacity(steps);
    let mut full_trajectory: Vec<DVector<f64>> = Vec::with_capacity(steps);
    for k in 0..steps {
        stepper.set_nodal_forces(&load_at(k + 1));
        let (next, _) = stepper.step(&full_state).expect("full step");
        full_state = next;
        let d: DVector<f64> = DVector::from_iterator(
            n_free,
            numbering
                .free_dofs
                .iter()
                .map(|&d| full_state.displacement[d]),
        );
        snapshots.push(d.clone());
        full_trajectory.push(d);
    }

    // POD basis from the trajectory, truncated hard.
    let basis_full = pod_basis(&snapshots, 1e-14).unwrap();
    let modes = basis_full.ncols().min(6);
    let basis = basis_full.columns(0, modes).into_owned();
    // Same normalization as the tracking metric below (absolute error
    // rms over the max displacement scale) so the two are comparable.
    let scale = snapshots
        .iter()
        .map(nalgebra::DVector::norm)
        .fold(0.0f64, f64::max);
    let projection_rms = {
        let mut sum = 0.0;
        for d in &snapshots {
            let err = (d - &basis * (basis.transpose() * d)).norm();
            sum += err * err;
        }
        (sum / snapshots.len() as f64).sqrt() / scale
    };

    // Reduced march under the same loads from the same rest state.
    let model = ReducedNonlinearModel::new_formulated(
        &mesh,
        &db,
        &numbering,
        basis,
        Formulation::TotalLagrangian,
    )
    .unwrap();
    let newmark = ReducedNewmark::new(&model, dt)
        .unwrap()
        .with_convergence_criteria(tight_criteria());
    let mut reduced_state = newmark
        .rest_state(&free_force(&numbering, &load_at(0)))
        .unwrap();
    let mut sum_sq = 0.0;
    let mut scale_sq = 0.0f64;
    for (k, d_full) in full_trajectory.iter().enumerate() {
        let external = free_force(&numbering, &load_at(k + 1));
        let (next, _) = newmark
            .step(&reduced_state, &external)
            .expect("reduced step");
        reduced_state = next;
        let err = (d_full - model.expand(&reduced_state.q)).norm();
        sum_sq += err * err;
        scale_sq = scale_sq.max(d_full.norm_squared());
    }
    let tracking_rms = (sum_sq / full_trajectory.len() as f64).sqrt() / scale_sq.sqrt();
    println!(
        "  {modes}-mode reduced march: tracking rms {tracking_rms:.3e} \
         (projection rms {projection_rms:.3e})"
    );
    // The reduced march may legitimately exceed pure projection error
    // (closure: the dynamics leave the subspace and come back), but a
    // wrong mass, wrong formulation, or wrong Newmark constant lands
    // orders of magnitude higher.
    assert!(
        tracking_rms < 1e-2,
        "reduced march does not track its own training trajectory: rms {tracking_rms:.3e} \
         (projection floor {projection_rms:.3e})"
    );
}