//! The FSI2 flag's in-vacuo structure model, as the coupled march builds
//! it (track 1 round 4, `smodes`): the question is whether the lift
//! excess's 5f = 9.67 Hz content is the flag's 4th bending mode, and
//! whether the 35×2 Quad8 mesh or the Newmark γ = 0.7 dissipation puts
//! that mode where it is.
//!
//! Both tests are `#[ignore]`d instruments (test-only code; nothing in the
//! march changes):
//!
//! * `fsi2_flag_modes_dump` writes the linearised operators of the
//!   march's structure — the total-Lagrangian St. Venant–Kirchhoff tangent
//!   at u = 0 (plane strain, the element's default quadrature) and the
//!   consistent mass, on the free DoFs of the root-clamped flag — for
//!   each `SMODES_MESHES` entry (`35x2,70x4,140x8` by default) as COO
//!   files, for a generalised eigen-solve outside (scipy `eigsh`).
//! * `fsi2_flag_free_vibration` releases the march's own
//!   `NonlinearDynamicStepper` (35×2 by default) from a small-amplitude
//!   mode shape (`SMODES_MODE`: lines `node_id ux uy`; `SMODES_OMEGA2` =
//!   its ω² for a consistent start acceleration) at `SMODES_GAMMA` /
//!   β = (γ + ½)²/4 and `SMODES_DT`, and writes the probe node's uy(t) —
//!   the measured per-period decay of that mode in the real stepper.

mod fsi2_harness;

use std::io::Write as _;

use fsi2_harness::{FLAG_X0, FSI2, clamp_left, flag_mesh};
use nalgebra::{DVector, Vector3};
use rtx_fea::analysis::{
    AnalysisConfig, ConvergenceCriteria, DynamicState, NonlinearDynamicAnalysis,
};
use rtx_fea::elements::total_lagrangian::{internal_force_and_tangent, saint_venant_kirchhoff};
use rtx_fea::elements::{ElementMatrixComputer, StandardFiniteElement};
use rtx_fea::materials::{LinearElastic, Material as _, MaterialDatabase};
use rtx_fea::mesh::{MaterialId, NodeId};

fn env_str(name: &str, default: &str) -> String {
    std::env::var(name).unwrap_or_else(|_| default.to_string())
}

fn env_num(name: &str, default: f64) -> f64 {
    std::env::var(name)
        .map(|v| v.parse().expect(name))
        .unwrap_or(default)
}

fn parse_mesh(spec: &str) -> (usize, usize) {
    let (a, b) = spec.split_once('x').expect("mesh spec NXxNY");
    (a.trim().parse().unwrap(), b.trim().parse().unwrap())
}

#[test]
#[ignore = "instrument: writes the flag's K and M for an outside eigen-solve"]
fn fsi2_flag_modes_dump() {
    let out = env_str("SMODES_OUT", ".");
    let meshes = env_str("SMODES_MESHES", "35x2,70x4,140x8");
    let (lambda, mu) = LinearElastic::new(FSI2.e_s, FSI2.nu_s)
        .with_density(FSI2.rho_s)
        .properties()
        .lame_parameters();
    println!(
        "FSI2 flag: E {:.4e} ν {} ρ_s {} → λ {lambda:.6e} μ {mu:.6e} (plane strain SVK)",
        FSI2.e_s, FSI2.nu_s, FSI2.rho_s
    );
    let constitutive = saint_venant_kirchhoff(lambda, mu, 2);
    for spec in meshes.split(',') {
        let (nx, ny) = parse_mesh(spec);
        let mesh = flag_mesh(nx, ny);
        let mut ids: Vec<NodeId> = mesh.nodes.keys().copied().collect();
        ids.sort();
        // Free DoFs: every node off the clamped root edge, (x, y) each.
        let mut free = std::collections::HashMap::new();
        let mut dof_lines = Vec::new();
        for id in &ids {
            let p = mesh.get_node(*id).unwrap().position();
            if (p.x - FLAG_X0).abs() < 1e-9 {
                continue;
            }
            for c in 0..2 {
                free.insert((*id, c), dof_lines.len());
                dof_lines.push(format!(
                    "{} {} {:.12e} {:.12e} {c}",
                    dof_lines.len(),
                    id.0,
                    p.x,
                    p.y
                ));
            }
        }
        let mut k_trip: std::collections::BTreeMap<(usize, usize), f64> = Default::default();
        let mut m_trip: std::collections::BTreeMap<(usize, usize), f64> = Default::default();
        for element in mesh.elements.values() {
            let coords: Vec<Vector3<f64>> = element
                .nodes
                .iter()
                .map(|id| mesh.get_node(*id).unwrap().position())
                .collect();
            let fe = StandardFiniteElement::new(element.element_type, coords.clone());
            let zero = DVector::zeros(2 * element.nodes.len());
            let (_, k_e) =
                internal_force_and_tangent(&fe, &coords, &zero, constitutive.as_ref(), None)
                    .unwrap();
            let m_s = ElementMatrixComputer::compute_consistent_mass_matrix(
                &fe, &coords, FSI2.rho_s, None,
            )
            .unwrap();
            let local: Vec<Option<usize>> = element
                .nodes
                .iter()
                .flat_map(|n| [free.get(&(*n, 0)).copied(), free.get(&(*n, 1)).copied()])
                .collect();
            for (a, ga) in local.iter().enumerate() {
                let Some(ga) = ga else { continue };
                for (b, gb) in local.iter().enumerate() {
                    let Some(gb) = gb else { continue };
                    *k_trip.entry((*ga, *gb)).or_default() += k_e[(a, b)];
                    if a % 2 == b % 2 {
                        *m_trip.entry((*ga, *gb)).or_default() += m_s.matrix[(a / 2, b / 2)];
                    }
                }
            }
        }
        let tag = format!("{nx}x{ny}");
        let write = |name: &str, trip: &std::collections::BTreeMap<(usize, usize), f64>| {
            let mut f = std::fs::File::create(format!("{out}/{name}_{tag}.coo")).unwrap();
            for ((i, j), v) in trip {
                if *v != 0.0 {
                    writeln!(f, "{i} {j} {v:.17e}").unwrap();
                }
            }
        };
        write("k", &k_trip);
        write("m", &m_trip);
        std::fs::write(format!("{out}/dofs_{tag}.txt"), dof_lines.join("\n") + "\n").unwrap();
        let mass: f64 = m_trip
            .iter()
            .filter(|((i, _), _)| i % 2 == 0)
            .map(|(_, v)| v)
            .sum();
        println!(
            "  {tag}: {} free DoFs, K nnz {}, M nnz {}, free-node x-mass {mass:.4} kg/m (flag 0.35·0.02·1e4 = 70 minus the root column's share)",
            dof_lines.len(),
            k_trip.len(),
            m_trip.len()
        );
    }
}

#[test]
#[ignore = "instrument: the march's stepper released from a mode shape"]
fn fsi2_flag_free_vibration() {
    let (nx, ny) = parse_mesh(&env_str("SMODES_MESH", "35x2"));
    let gamma = env_num("SMODES_GAMMA", 0.5);
    let beta = (gamma + 0.5).powi(2) / 4.0;
    let dt = env_num("SMODES_DT", 1.601e-4);
    let steps = env_num("SMODES_STEPS", 2000.0) as usize;
    let amp = env_num("SMODES_AMP", 1e-6);
    let omega2 = env_num("SMODES_OMEGA2", 0.0);
    let probe = NodeId(env_num("SMODES_PROBE", 0.0) as usize);
    let mode_path = std::env::var("SMODES_MODE").expect("SMODES_MODE");
    let out = std::env::var("SMODES_CSV").expect("SMODES_CSV");

    let mesh = flag_mesh(nx, ny);
    let mut db = MaterialDatabase::new();
    db.add_material(
        MaterialId(0),
        LinearElastic::new(FSI2.e_s, FSI2.nu_s).with_density(FSI2.rho_s),
        None,
    );
    let analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        db,
        clamp_left(&mesh),
        dt,
        1,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian()
    .with_convergence_criteria(ConvergenceCriteria {
        max_iterations: 60,
        ..ConvergenceCriteria::default()
    })
    .with_newmark_parameters(gamma, beta);
    let mut stepper = analysis.stepper().unwrap();
    let rest = stepper.rest_state().unwrap();
    let n = rest.displacement.len();
    let mut u = DVector::zeros(n);
    for line in std::fs::read_to_string(&mode_path).unwrap().lines() {
        let f: Vec<f64> = line
            .split_whitespace()
            .map(|t| t.parse().unwrap())
            .collect();
        if f.len() < 3 {
            continue;
        }
        let dofs = stepper.node_dofs(NodeId(f[0] as usize));
        u[dofs[0]] = amp * f[1];
        u[dofs[1]] = amp * f[2];
    }
    let mut state = DynamicState {
        acceleration: &u * (-omega2),
        velocity: DVector::zeros(n),
        displacement: u,
    };
    let probe_dofs = stepper.node_dofs(probe);
    let mut csv = std::fs::File::create(&out).unwrap();
    writeln!(csv, "t,uy,ux").unwrap();
    writeln!(
        csv,
        "0,{:.17e},{:.17e}",
        state.displacement[probe_dofs[1]], state.displacement[probe_dofs[0]]
    )
    .unwrap();
    for k in 1..=steps {
        let (next, _) = stepper.step(&state).unwrap();
        state = next;
        writeln!(
            csv,
            "{:.9e},{:.17e},{:.17e}",
            k as f64 * dt,
            state.displacement[probe_dofs[1]],
            state.displacement[probe_dofs[0]]
        )
        .unwrap();
    }
    println!("free vibration {nx}x{ny} γ {gamma} β {beta:.4} dt {dt:.4e}: {steps} steps → {out}");
}