//! P6-b (`docs/overset_metal_campaign.md` §5.17): the fictitious added mass
//! for the partitioned coupling loop. The stepper carries a lumped mass
//! `M_f` on chosen nodes; the coupling loop adds the load `M_f ü_k` from
//! the previous subiterate. Pin: solving `(M + M_f) ü = F + M_f ü_A` with
//! `ü_A` the plain step's own acceleration reproduces the plain step (the
//! fixed point is unchanged — the added terms cancel at convergence), while
//! the added mass WITHOUT its compensating load changes the step (the hook
//! is live). The Newmark effective matrix carries `M_f` (a step from rest
//! under a nodal load has the acceleration of a heavier body).

use nalgebra::Vector3;
use rtx_fea::analysis::{AnalysisConfig, NonlinearDynamicAnalysis};
use rtx_fea::assembly::dof_mapping::DofComponent;
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};

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

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
}

fn point_a(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("tracking point must be a mesh node")
}

#[test]
fn fictitious_mass_cancels_with_its_load_and_is_live_without_it() {
    let mesh = quad8_rect_mesh(0.25, 0.6, 0.19, 0.21, 10, 2);
    let dt = 0.005;
    let a_node = point_a(&mesh, 0.6, 0.2);
    let tip_force = Vector3::new(0.0, -0.5, 0.0);
    let analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        materials(),
        clamp_left(&mesh, 0.25),
        dt,
        1,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian();

    // A: the plain step from a moving state (three plain steps in, so the
    // acceleration is not the rest one).
    let mut plain = analysis.stepper().unwrap();
    plain.set_nodal_forces(&[(a_node, tip_force)]);
    let mut state = plain.rest_state().unwrap();
    for _ in 0..3 {
        state = plain.step(&state).unwrap().0;
    }
    let (step_a, _) = plain.step(&state).unwrap();
    let a_dofs = plain.node_dofs(a_node);

    // The wetted nodes of this flag: everything not clamped. M_f on all of
    // them, the size of the flag's own mass per node (a hard case for the
    // cancellation: the added term is O(1) of the inertia).
    let wetted: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, n)| (n.position().x - 0.25).abs() > 1e-12)
        .map(|(&id, _)| id)
        .collect();
    let m_f = 0.5; // kg per node (the flag's 0.02 × 0.35 × 1000 = 7 kg over ~60 nodes)
    let entries: Vec<(NodeId, f64)> = wetted.iter().map(|&n| (n, m_f)).collect();

    // B: the added mass with the compensating load M_f ü_A (ü_A = A's own
    // acceleration at each wetted node): the same step to solver tolerance.
    let mut fict = analysis.stepper().unwrap();
    fict.set_added_lumped_mass(&entries);
    let mut loads = vec![(a_node, tip_force)];
    for &n in &wetted {
        let d = fict.node_dofs(n);
        loads.push((
            n,
            Vector3::new(
                m_f * step_a.acceleration[d[0]],
                m_f * step_a.acceleration[d[1]],
                0.0,
            ),
        ));
    }
    fict.set_nodal_forces(&loads);
    let (step_b, _) = fict.step(&state).unwrap();
    let uy_a = step_a.displacement[a_dofs[1]];
    let uy_b = step_b.displacement[a_dofs[1]];
    let du =
        (step_a.displacement.clone() - &step_b.displacement).norm() / step_a.displacement.norm();
    println!(
        "  cancellation: tip uy A {uy_a:.6e} vs B {uy_b:.6e}; relative displacement difference {du:.2e}"
    );
    assert!(
        du < 1e-6,
        "the compensated added mass changed the step: {du:.2e}"
    );

    // C: the added mass WITHOUT the compensating load: a heavier body, a
    // visibly different step (the hook is live in the effective matrix).
    let mut heavy = analysis.stepper().unwrap();
    heavy.set_added_lumped_mass(&entries);
    heavy.set_nodal_forces(&[(a_node, tip_force)]);
    let (step_c, _) = heavy.step(&state).unwrap();
    let dc =
        (step_a.displacement.clone() - &step_c.displacement).norm() / step_a.displacement.norm();
    println!("  uncompensated: relative displacement difference {dc:.2e}");
    assert!(
        dc > 1e-3,
        "the added mass alone did not change the step: {dc:.2e}"
    );

    // D: zero added mass is the plain stepper bit for bit.
    let mut zero = analysis.stepper().unwrap();
    zero.set_added_lumped_mass(&wetted.iter().map(|&n| (n, 0.0)).collect::<Vec<_>>());
    zero.set_nodal_forces(&[(a_node, tip_force)]);
    let (step_d, _) = zero.step(&state).unwrap();
    assert_eq!(
        step_a.displacement, step_d.displacement,
        "zero added mass is not bit-identical"
    );
}