//! R8-g: the sparse Newton-tangent path of the nonlinear Newmark stepper
//! ([`TangentSolver::SparseLdlt`]) against the default banded LU.
//!
//! Suite (fast, run by default):
//!
//! 1. `sparse_matches_banded_csm3_2d` — the 2-D 35×2 Quad8 CSM3 march
//!    (the FSI2 harness's flag), 60 steps: every step's displacement
//!    within 1e-10 of the banded path's (relative to the peak).
//! 2. `sparse_matches_banded_csm1_3d_plane_strain` — CSM1 static on the
//!    3-D 35×2×1 plane-strain flag.
//! 3. `sparse_matches_banded_csm3_3d_free` — a free-edge 3-D flag
//!    (12×2×2), 40 CSM3 steps.
//! 4. `modified_newton_reuses_the_factor` — `reuse = 4`: fewer
//!    factorisations than solves, same march to the Newton tolerance.
//! 5. `tangent_knob_parses` — the `RTX_FEA_TANGENT` values.
//!
//! Instruments (`#[ignore]`, env-driven, write under `R8G_OUT`):
//!
//! * `r8g_g1_equivalence` — CSM1 static and a CSM3 march per
//!   configuration with both solvers side by side; per-step CSV.
//! * `r8g_g2_cost` — wall time per Newton iteration and per step, per
//!   solver mode and configuration, with the sparse path's breakdown.

use std::io::Write as _;
use std::time::Instant;

use nalgebra::{DVector, Vector3};
use rtx_fea::analysis::flag3d::{Flag3d, Flag3dSpec, LateralFaces};
use rtx_fea::analysis::{
    ConvergenceCriteria, DynamicState, NonlinearDynamicAnalysis, TangentSolver,
};
use rtx_fea::assembly::dof_mapping::DofComponent;
use rtx_fea::boundary::dirichlet::{DirichletBC, DirichletType};
use rtx_fea::boundary::{BoundaryCondition, BoundaryConditionSet, SpatialFunction};
use rtx_fea::elements::{ElementMatrixComputer, StandardFiniteElement};
use rtx_fea::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};

const E_MOD: f64 = 1.4e6;
const NU: f64 = 0.4;
const RHO_CSM: f64 = 1000.0;
const G: f64 = 2.0;

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)
}

/// The 2-D flag, as the FSI2 harness builds it (copied from
/// `flag3d_structure.rs`).
fn quad8_flag(nx: usize, ny: usize) -> Mesh {
    let (x0, x1, y0, y1) = (0.25, 0.6, 0.19, 0.21);
    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 clamp_2d(mesh: &Mesh) -> BoundaryConditionSet {
    let clamped: Vec<NodeId> = mesh
        .nodes
        .iter()
        .filter(|(_, node)| (node.position().x - 0.25).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
}

/// Consistent gravity nodal forces `∫ N_a ρ g dV`.
fn gravity_forces(mesh: &Mesh, rho: f64, g: f64) -> Vec<(NodeId, Vector3<f64>)> {
    let dim = mesh.spatial_dimension;
    let mut acc: std::collections::BTreeMap<NodeId, Vector3<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 f = ElementMatrixComputer::compute_body_force_vector(
            &fe,
            &coords,
            &|_| Vector3::new(0.0, -rho * g, 0.0),
            None,
        )
        .unwrap();
        for (a, id) in element.nodes.iter().enumerate() {
            let e = acc.entry(*id).or_insert_with(Vector3::zeros);
            for c in 0..dim {
                e[c] += f[a * dim + c];
            }
        }
    }
    acc.into_iter().collect()
}

fn static_criteria() -> ConvergenceCriteria {
    ConvergenceCriteria {
        force_tolerance: 1e-12,
        displacement_tolerance: 1e-14,
        max_iterations: 60,
        ..ConvergenceCriteria::default()
    }
}

/// The CSM3 analysis (gravity from rest) on the 2-D 35×2 flag.
fn csm3_2d(dt: f64, solver: TangentSolver) -> NonlinearDynamicAnalysis {
    let mesh = quad8_flag(35, 2);
    let mut analysis = NonlinearDynamicAnalysis::new(
        mesh.clone(),
        Flag3d::materials(E_MOD, NU, RHO_CSM),
        clamp_2d(&mesh),
        dt,
        1,
        Default::default(),
    )
    .with_total_lagrangian()
    .with_tangent_solver(solver);
    analysis.set_body_force(|_| Vector3::new(0.0, -RHO_CSM * G, 0.0));
    analysis
}

/// The CSM3 analysis on a 3-D flag (R8-b's instrument settings).
fn csm3_3d(
    flag: &Flag3d,
    lateral: LateralFaces,
    dt: f64,
    solver: TangentSolver,
) -> NonlinearDynamicAnalysis {
    let mut analysis = flag
        .dynamic_analysis(E_MOD, NU, RHO_CSM, lateral, dt, 1, 0.5)
        .with_tangent_solver(solver);
    analysis.set_body_force(|_| Vector3::new(0.0, -RHO_CSM * G, 0.0));
    analysis
}

/// `max_i |a_i − b_i| / max_i |b_i|`.
fn rel_max(a: &DVector<f64>, b: &DVector<f64>) -> f64 {
    let scale = b.amax().max(1e-300);
    (a - b).amax() / scale
}

/// March both analyses `steps` steps side by side; per step the
/// relative displacement difference. Returns (max rel, per-step rows).
fn march_pair(
    banded: &NonlinearDynamicAnalysis,
    sparse: &NonlinearDynamicAnalysis,
    steps: usize,
) -> (f64, Vec<(f64, f64, usize, usize)>) {
    let mut sb = banded.stepper().unwrap();
    let mut ss = sparse.stepper().unwrap();
    let mut stb = sb.rest_state().unwrap();
    let mut sts = ss.rest_state().unwrap();
    let mut worst: f64 = rel_max(&sts.acceleration, &stb.acceleration);
    let mut rows = Vec::new();
    for _ in 0..steps {
        let t0 = Instant::now();
        let (nb, ib) = sb.step(&stb).unwrap();
        let tb = t0.elapsed().as_secs_f64();
        let t1 = Instant::now();
        let (ns, is) = ss.step(&sts).unwrap();
        let ts = t1.elapsed().as_secs_f64();
        stb = nb;
        sts = ns;
        let r = rel_max(&sts.displacement, &stb.displacement);
        worst = worst.max(r);
        rows.push((r, tb / ts.max(1e-12), ib, is));
    }
    (worst, rows)
}

#[test]
fn sparse_matches_banded_csm3_2d() {
    let dt = 0.005;
    let (worst, rows) = march_pair(
        &csm3_2d(dt, TangentSolver::BandedLu),
        &csm3_2d(dt, TangentSolver::SPARSE),
        60,
    );
    let same_newton = rows.iter().all(|r| r.2 == r.3);
    println!("CSM3 2-D 35x2, 60 steps: max rel |Δu| {worst:.3e}, same Newton counts {same_newton}");
    assert!(worst < 1e-10, "sparse departs from banded: {worst:.3e}");
    assert!(same_newton);
}

#[test]
fn sparse_matches_banded_csm1_3d_plane_strain() {
    let flag = Flag3d::build(Flag3dSpec::turek_hron(0.05, 0.0, 35, 2, 1)).unwrap();
    let forces = gravity_forces(&flag.mesh, RHO_CSM, G);
    let solve = |solver: TangentSolver| {
        let analysis = flag
            .dynamic_analysis(E_MOD, NU, RHO_CSM, LateralFaces::PlaneStrain, 1e4, 1, 0.5)
            .with_convergence_criteria(static_criteria())
            .with_tangent_solver(solver);
        let mut stepper = analysis.stepper().unwrap();
        let n = stepper.rest_state().unwrap().displacement.len();
        let mut u = DVector::zeros(n);
        for s in 1..=5 {
            let scale = s as f64 / 5.0;
            let scaled: Vec<_> = forces.iter().map(|(id, f)| (*id, f * scale)).collect();
            stepper.set_nodal_forces(&scaled);
            let state = DynamicState {
                displacement: u.clone(),
                velocity: DVector::zeros(n),
                acceleration: DVector::zeros(n),
            };
            u = stepper.step(&state).unwrap().0.displacement;
        }
        let d = stepper.node_dofs(flag.point_a());
        (u.clone(), u[d[1]])
    };
    let (ub, ay_b) = solve(TangentSolver::BandedLu);
    let (us, ay_s) = solve(TangentSolver::SPARSE);
    let r = rel_max(&us, &ub);
    println!(
        "CSM1 3-D 35x2x1 ps: uy(A) banded {ay_b:.12e} sparse {ay_s:.12e}, max rel |Δu| {r:.3e}"
    );
    assert!(r < 1e-10, "CSM1 departs: {r:.3e}");
    assert!((ay_b + 65.1406e-3).abs() < 1e-6, "CSM1 uy(A) moved: {ay_b}");
}

#[test]
fn sparse_matches_banded_csm3_3d_free() {
    let flag = Flag3d::build(Flag3dSpec::turek_hron(0.1, -0.05, 12, 2, 2)).unwrap();
    let dt = 0.005;
    let (worst, _) = march_pair(
        &csm3_3d(&flag, LateralFaces::Free, dt, TangentSolver::BandedLu),
        &csm3_3d(&flag, LateralFaces::Free, dt, TangentSolver::SPARSE),
        40,
    );
    println!("CSM3 3-D 12x2x2 free, 40 steps: max rel |Δu| {worst:.3e}");
    assert!(worst < 1e-10, "sparse departs from banded: {worst:.3e}");
}

#[test]
fn modified_newton_reuses_the_factor() {
    let flag = Flag3d::build(Flag3dSpec::turek_hron(0.1, -0.05, 12, 2, 2)).unwrap();
    let dt = 0.005;
    let full = csm3_3d(&flag, LateralFaces::Free, dt, TangentSolver::SPARSE);
    let modified = csm3_3d(
        &flag,
        LateralFaces::Free,
        dt,
        TangentSolver::SparseLdlt { reuse: 4 },
    );
    let mut sf = full.stepper().unwrap();
    let mut sm = modified.stepper().unwrap();
    let mut stf = sf.rest_state().unwrap();
    let mut stm = sm.rest_state().unwrap();
    for _ in 0..40 {
        stf = sf.step(&stf).unwrap().0;
        stm = sm.step(&stm).unwrap().0;
    }
    let r = rel_max(&stm.displacement, &stf.displacement);
    let stats_f = sf.tangent_stats().unwrap();
    let stats_m = sm.tangent_stats().unwrap();
    println!(
        "modified Newton (reuse 4) vs full after 40 steps: rel {r:.3e}; full {stats_f:?}; modified {stats_m:?}"
    );
    assert_eq!(stats_f.factorizations, stats_f.solves);
    assert!(stats_m.factorizations < stats_m.solves);
    assert_eq!(stats_m.analyses, 1);
    // Newton's own displacement tolerance is 1e-6 (relative).
    assert!(r < 1e-4, "modified Newton drifts: {r:.3e}");
}

#[test]
fn tangent_knob_parses() {
    // Only the parser (the knob is read when a stepper is built).
    assert_eq!(TangentSolver::default(), TangentSolver::BandedLu);
    assert_eq!(
        TangentSolver::SPARSE,
        TangentSolver::SparseLdlt { reuse: 1 }
    );
}

// ---------------------------------------------------------------------------
// Instruments
// ---------------------------------------------------------------------------

/// `NXxNYxNZ:span:free|ps` entries (as in `flag3d_structure.rs`); `2d`
/// = the 2-D 35×2 Quad8 flag.
fn parse_configs(spec: &str) -> Vec<Option<(usize, usize, usize, f64, LateralFaces)>> {
    spec.split(',')
        .map(|entry| {
            if entry.trim() == "2d" {
                return None;
            }
            let mut parts = entry.trim().split(':');
            let mesh = parts.next().unwrap();
            let span: f64 = parts.next().unwrap().parse().unwrap();
            let lateral = match parts.next().unwrap() {
                "free" => LateralFaces::Free,
                "ps" => LateralFaces::PlaneStrain,
                other => panic!("lateral {other}"),
            };
            let n: Vec<usize> = mesh.split('x').map(|t| t.parse().unwrap()).collect();
            Some((n[0], n[1], n[2], span, lateral))
        })
        .collect()
}

fn tag(cfg: &Option<(usize, usize, usize, f64, LateralFaces)>) -> String {
    match cfg {
        None => "2d_35x2".to_string(),
        Some((nx, ny, nz, span, lateral)) => {
            let l = if *lateral == LateralFaces::Free {
                "free"
            } else {
                "ps"
            };
            format!("{nx}x{ny}x{nz}_s{span}_{l}")
        }
    }
}

fn build_flag(cfg: &(usize, usize, usize, f64, LateralFaces)) -> Flag3d {
    let (nx, ny, nz, span, _) = *cfg;
    Flag3d::build(Flag3dSpec::turek_hron(span, -0.5 * span, nx, ny, nz)).unwrap()
}

fn analysis_for(
    cfg: &Option<(usize, usize, usize, f64, LateralFaces)>,
    flag: Option<&Flag3d>,
    dt: f64,
    solver: TangentSolver,
) -> NonlinearDynamicAnalysis {
    match cfg {
        None => csm3_2d(dt, solver),
        Some(c) => csm3_3d(flag.unwrap(), c.4, dt, solver),
    }
}

#[test]
#[ignore = "instrument: G1 — sparse vs banded, CSM1 static and a CSM3 march"]
fn r8g_g1_equivalence() {
    let out = env_str("R8G_OUT", ".");
    let configs = parse_configs(&env_str("R8G_CONFIGS", "35x2x1:0.05:ps,35x2x8:0.41:free"));
    let steps = env_num("R8G_STEPS", 200.0) as usize;
    let dt = env_num("R8G_DT", 0.005);
    let do_csm1 = env_str("R8G_CSM1", "1") == "1";
    let mut table = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(format!("{out}/g1_table.txt"))
        .unwrap();
    for cfg in &configs {
        let name = tag(cfg);
        let flag = cfg.as_ref().map(build_flag);
        if do_csm1 {
            if let (Some(c), Some(flag)) = (cfg, flag.as_ref()) {
                let forces = gravity_forces(&flag.mesh, RHO_CSM, G);
                let mut results = Vec::new();
                for solver in [TangentSolver::BandedLu, TangentSolver::SPARSE] {
                    let start = Instant::now();
                    let analysis = flag
                        .dynamic_analysis(E_MOD, NU, RHO_CSM, c.4, 1e4, 1, 0.5)
                        .with_convergence_criteria(static_criteria())
                        .with_tangent_solver(solver);
                    let mut stepper = analysis.stepper().unwrap();
                    let n = stepper.rest_state().unwrap().displacement.len();
                    let mut u = DVector::zeros(n);
                    let mut newton = 0;
                    for s in 1..=5 {
                        let scale = s as f64 / 5.0;
                        let scaled: Vec<_> =
                            forces.iter().map(|(id, f)| (*id, f * scale)).collect();
                        stepper.set_nodal_forces(&scaled);
                        let state = DynamicState {
                            displacement: u.clone(),
                            velocity: DVector::zeros(n),
                            acceleration: DVector::zeros(n),
                        };
                        let (next, it) = stepper.step(&state).unwrap();
                        newton += it;
                        u = next.displacement;
                    }
                    let d = stepper.node_dofs(flag.point_a());
                    results.push((
                        u.clone(),
                        u[d[0]],
                        u[d[1]],
                        newton,
                        start.elapsed().as_secs_f64(),
                    ));
                }
                let r = rel_max(&results[1].0, &results[0].0);
                let line = format!(
                    "G1 CSM1 {name}: A banded ux {:.9e} uy {:.9e} [{} Newton, {:.1} s] | sparse ux \
                     {:.9e} uy {:.9e} [{} Newton, {:.1} s] | max rel |Δu| {r:.3e}",
                    results[0].1,
                    results[0].2,
                    results[0].3,
                    results[0].4,
                    results[1].1,
                    results[1].2,
                    results[1].3,
                    results[1].4
                );
                println!("{line}");
                writeln!(table, "{line}").unwrap();
            }
        }
        let banded = analysis_for(cfg, flag.as_ref(), dt, TangentSolver::BandedLu);
        let sparse = analysis_for(cfg, flag.as_ref(), dt, TangentSolver::SPARSE);
        let start = Instant::now();
        let (worst, rows) = march_pair(&banded, &sparse, steps);
        let path = format!("{out}/g1_csm3_{name}_dt{dt}.csv");
        let mut csv = std::fs::File::create(&path).unwrap();
        writeln!(csv, "step,rel_u,speedup,newton_banded,newton_sparse").unwrap();
        for (k, (r, speed, ib, is)) in rows.iter().enumerate() {
            writeln!(csv, "{},{r:.6e},{speed:.3},{ib},{is}", k + 1).unwrap();
        }
        let differ = rows.iter().filter(|r| r.2 != r.3).count();
        let line = format!(
            "G1 CSM3 {name} dt {dt}: {steps} steps, max rel |Δu| {worst:.3e}, Newton-count \
             differences {differ}, {:.0} s → {path}",
            start.elapsed().as_secs_f64()
        );
        println!("{line}");
        writeln!(table, "{line}").unwrap();
    }
}

#[test]
#[ignore = "instrument: G2 — cost per Newton iteration and per step"]
fn r8g_g2_cost() {
    let out = env_str("R8G_OUT", ".");
    let configs = parse_configs(&env_str(
        "R8G_CONFIGS",
        "2d,35x2x8:0.41:free,35x2x16:0.41:free",
    ));
    let modes: Vec<String> = env_str("R8G_MODES", "banded,sparse,sparse:3")
        .split(',')
        .map(str::to_string)
        .collect();
    let steps = env_num("R8G_STEPS", 40.0) as usize;
    let warm = env_num("R8G_WARM", 20.0) as usize;
    let dt = env_num("R8G_DT", 0.005);
    let mut table = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(format!("{out}/g2_table.txt"))
        .unwrap();
    for cfg in &configs {
        let name = tag(cfg);
        let flag = cfg.as_ref().map(build_flag);
        let mut reference: Option<DVector<f64>> = None;
        for mode in &modes {
            let solver = match mode.as_str() {
                "banded" => TangentSolver::BandedLu,
                "sparse" => TangentSolver::SPARSE,
                other => TangentSolver::SparseLdlt {
                    reuse: other.strip_prefix("sparse:").unwrap().parse().unwrap(),
                },
            };
            let analysis = analysis_for(cfg, flag.as_ref(), dt, solver);
            let t_build = Instant::now();
            let mut stepper = analysis.stepper().unwrap();
            let mut state = stepper.rest_state().unwrap();
            let build = t_build.elapsed().as_secs_f64();
            // Warm-up steps (off the clock) move the flag off rest.
            for _ in 0..warm {
                state = stepper.step(&state).unwrap().0;
            }
            let before = stepper.tangent_stats();
            let mut times = Vec::with_capacity(steps);
            let mut newton = 0usize;
            for _ in 0..steps {
                let t = Instant::now();
                let (next, it) = stepper.step(&state).unwrap();
                times.push(t.elapsed().as_secs_f64());
                newton += it;
                state = next;
            }
            let total: f64 = times.iter().sum();
            let mut sorted = times.clone();
            sorted.sort_by(f64::total_cmp);
            let median = sorted[sorted.len() / 2];
            let drift = reference
                .as_ref()
                .map_or(f64::NAN, |r| rel_max(&state.displacement, r));
            if reference.is_none() {
                reference = Some(state.displacement.clone());
            }
            let breakdown = match (before, stepper.tangent_stats()) {
                (Some(b), Some(a)) => format!(
                    " | assemblies {} ({:.4} s each), factorisations {} ({:.4} s each incl. \
                     scatter), solves {} ({:.4} s each), fallbacks {}, nnz(L) {}",
                    a.assemblies - b.assemblies,
                    (a.assembly_seconds - b.assembly_seconds)
                        / (a.assemblies - b.assemblies).max(1) as f64,
                    a.factorizations - b.factorizations,
                    (a.factor_seconds - b.factor_seconds)
                        / (a.factorizations - b.factorizations).max(1) as f64,
                    a.solves - b.solves,
                    (a.solve_seconds - b.solve_seconds) / (a.solves - b.solves).max(1) as f64,
                    a.fallbacks,
                    a.nnz_l
                ),
                _ => String::new(),
            };
            let line = format!(
                "G2 {name} {mode} dt {dt}: dofs {} | build+rest {build:.3} s | {steps} steps after \
                 {warm} warm: {:.4} s/step mean, {median:.4} median, {:.2} Newton/step, {:.4} \
                 s/Newton | final rel to first mode {drift:.2e} | threads {}{breakdown}",
                state.displacement.len(),
                total / steps as f64,
                newton as f64 / steps as f64,
                total / newton.max(1) as f64,
                rayon::current_num_threads()
            );
            println!("{line}");
            writeln!(table, "{line}").unwrap();
        }
    }
}