//! ECSW×FSI campaign, phase 2: the offline POD/ECSW measurements on the
//! harvested FSI3 flag trajectory (omni-cortex `next_session.md`, the
//! campaign section; snapshots from the phase-1 `FSNP` dump).
//!
//! Three measurements, pre-registered in the campaign doc — deliberately
//! NOT a static-solve-vs-dynamic-state comparison (the recorded
//! trajectory is dynamic; the dynamic reduced Newmark is phase 4):
//!
//! 1. **POD projection error vs mode count** — does the flapping flag's
//!    solution manifold live in a small subspace? (The campaign's
//!    load-bearing question.)
//! 2. **ECSW training + HELD-OUT residual** (total-Lagrangian operators,
//!    matching the march's `with_total_lagrangian`) — does a small
//!    element sample reproduce the reduced internal force on states it
//!    never trained on?
//! 3. **Assembly wall-clock, full-active vs sampled** — the
//!    per-Newton-iteration cost ECSW actually reduces.
//!
//! Env-gated study: set `RTX_ECSW_SNAP=<path to .fsnp>`; without it the
//! test prints a skip note and passes (committed default is a no-op).

mod fsi2_harness;

use fsi2_harness::{FLAG_X0, FSI3, flag_mesh};
use nalgebra::{DMatrix, DVector};
use rtx_fea::assembly::dof_mapping::{AdvancedDofNumbering, DofComponent, DofMappingStrategy};
use rtx_fea::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::MaterialId;
use rtx_fea::mor::{
    Formulation, ReducedNonlinearModel, ecsw_residual, pod_basis, train_ecsw_formulated,
};

/// One FSNP record (the phase-1 dump format; see `march::write_snapshot`).
struct Snapshot {
    #[allow(dead_code)]
    t: f64,
    displacement: Vec<f64>,
}

fn read_fsnp(path: &str) -> (usize, Vec<Snapshot>) {
    let data = std::fs::read(path).expect("snapshot file");
    assert_eq!(&data[0..4], b"FSNP", "bad magic");
    assert_eq!(u32::from_le_bytes(data[4..8].try_into().unwrap()), 1);
    let n_dofs = u64::from_le_bytes(data[8..16].try_into().unwrap()) as usize;
    let mut off = 16usize;
    let mut records = Vec::new();
    let f64_at = |data: &[u8], off: usize| -> f64 {
        f64::from_le_bytes(data[off..off + 8].try_into().unwrap())
    };
    while off < data.len() {
        let t = f64_at(&data, off);
        off += 8;
        let displacement: Vec<f64> = (0..n_dofs).map(|i| f64_at(&data, off + 8 * i)).collect();
        off += 8 * n_dofs * 3; // skip velocity + acceleration for this phase
        let n_forces = u64::from_le_bytes(data[off..off + 8].try_into().unwrap()) as usize;
        off += 8 + n_forces * 32;
        records.push(Snapshot { t, displacement });
    }
    assert_eq!(off, data.len(), "trailing bytes");
    (n_dofs, records)
}

/// The numbering the flag's analysis uses internally, rebuilt identically
/// (same strategy, same clamp criterion as `clamp_left`), so full-DOF
/// indices in the dump line up.
fn numbering(mesh: &rtx_fea::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 - FLAG_X0).abs() < 1e-9 {
            for component in [DofComponent::DisplacementX, DofComponent::DisplacementY] {
                let dof = dof_numbering.get_dof(node_id, component).unwrap();
                dof_numbering.constrain_dof(dof).unwrap();
            }
        }
    }
    dof_numbering
}

#[test]
fn fsi3_ecsw_offline_study() {
    let Ok(snap_path) = std::env::var("RTX_ECSW_SNAP") else {
        println!("  RTX_ECSW_SNAP not set — offline ECSW study skipped");
        return;
    };
    let (n_dofs, records) = read_fsnp(&snap_path);
    println!(
        "  {} snapshots, {} full DOFs, t in [{:.4}, {:.4}]",
        records.len(),
        n_dofs,
        records.first().unwrap().t,
        records.last().unwrap().t
    );

    // The flag exactly as the march builds it (FSI3 defaults).
    let mesh = flag_mesh(35, 2);
    let mut materials = MaterialDatabase::new();
    materials.add_material(
        MaterialId(0),
        LinearElastic::new(FSI3.e_s, FSI3.nu_s).with_density(FSI3.rho_s),
        None,
    );
    let dof_numbering = numbering(&mesh);
    let n_free = dof_numbering.free_dofs.len();
    assert_eq!(dof_numbering.total_dofs, n_dofs, "DOF count mismatch");

    // Numbering self-check: the dump's clamped DOFs must be exactly zero
    // at every snapshot — if the rebuilt numbering disagreed with the
    // analysis's internal one, real displacements would land on
    // "constrained" indices and this fails loudly.
    let mut is_free = vec![false; dof_numbering.total_dofs];
    for &dof in &dof_numbering.free_dofs {
        is_free[dof] = true;
    }
    let worst_clamped = records
        .iter()
        .flat_map(|r| {
            r.displacement
                .iter()
                .enumerate()
                .filter(|(dof, _)| !is_free[*dof])
                .map(|(_, v)| v.abs())
        })
        .fold(0.0f64, f64::max);
    println!("  numbering self-check: worst clamped-DOF displacement {worst_clamped:.2e}");
    assert!(
        worst_clamped < 1e-14,
        "clamped DOFs carry displacement {worst_clamped:.2e} — numbering mismatch"
    );

    let to_free = |full: &[f64]| -> DVector<f64> {
        DVector::from_iterator(n_free, dof_numbering.free_dofs.iter().map(|&dof| full[dof]))
    };

    // Interleaved split so both halves span release AND settled cycle.
    let train: Vec<DVector<f64>> = records
        .iter()
        .step_by(2)
        .map(|r| to_free(&r.displacement))
        .collect();
    let held_out: Vec<DVector<f64>> = records
        .iter()
        .skip(1)
        .step_by(2)
        .map(|r| to_free(&r.displacement))
        .collect();

    // ---- Measurement 1: POD projection error vs mode count. ----
    let basis_full = pod_basis(&train, 1e-14).unwrap();
    println!(
        "  POD basis: {} modes at energy tolerance 1e-14 (of {n_free} free DOFs)",
        basis_full.ncols()
    );
    let proj_errors = |basis: &DMatrix<f64>, set: &[DVector<f64>]| -> (f64, f64) {
        let mut max = 0.0f64;
        let mut sum_sq = 0.0f64;
        for d in set {
            let q = basis.transpose() * d;
            let err = (d - basis * q).norm() / d.norm().max(1e-30);
            max = max.max(err);
            sum_sq += err * err;
        }
        (max, (sum_sq / set.len() as f64).sqrt())
    };
    println!("  modes | train max / rms      | held-out max / rms");
    for m in [2usize, 4, 6, 8, 12, 16, 20, 30] {
        if m > basis_full.ncols() {
            break;
        }
        let v = basis_full.columns(0, m).into_owned();
        let (tr_max, tr_rms) = proj_errors(&v, &train);
        let (ho_max, ho_rms) = proj_errors(&v, &held_out);
        println!("  {m:5} | {tr_max:.3e} / {tr_rms:.3e} | {ho_max:.3e} / {ho_rms:.3e}");
    }

    // ---- Measurement 2: ECSW training + held-out residual (TL). ----
    // Thinned sets: the NNLS system is (snapshots x modes) rows.
    let ecsw_held: Vec<DVector<f64>> = held_out.iter().step_by(4).cloned().collect();
    // Cycle-only variant: restrict to the settled cycle (t > 6), the
    // narrower manifold the coupled phase-4 model would actually live on.
    let cycle_train: Vec<DVector<f64>> = records
        .iter()
        .filter(|r| r.t > 6.0)
        .step_by(8)
        .map(|r| to_free(&r.displacement))
        .collect();
    for (m, tol, thin) in [
        (8usize, 1e-3f64, 4usize),
        (8, 1e-2, 4),
        (8, 3e-2, 4),
        (8, 1e-1, 4),
        (4, 1e-1, 4),
        (8, 1e-2, 16),
        (12, 1e-2, 4),
        (8, 1e-2, 0), // thin = 0 marks the cycle-only training set
    ] {
        let ecsw_train: Vec<DVector<f64>> = if thin == 0 {
            cycle_train.clone()
        } else {
            train.iter().step_by(thin).cloned().collect()
        };
        let v = basis_full.columns(0, m).into_owned();
        let model = train_ecsw_formulated(
            &mesh,
            &materials,
            &dof_numbering,
            &v,
            &ecsw_train,
            tol,
            Formulation::TotalLagrangian,
        )
        .unwrap();
        let held = ecsw_residual(
            &mesh,
            &materials,
            &dof_numbering,
            &v,
            &ecsw_held,
            &model,
            Formulation::TotalLagrangian,
        )
        .unwrap();
        println!(
            "  ECSW m={m} tol={tol:.0e} thin={thin} ({} train states): {} of {} elements, \
             training residual {:.3e}, held-out residual {held:.3e}",
            ecsw_train.len(),
            model.weights.len(),
            mesh.num_elements(),
            model.training_residual
        );

        // ---- Measurement 3: assembly wall-clock, full vs sampled. ----
        let full_model = ReducedNonlinearModel::new_formulated(
            &mesh,
            &materials,
            &dof_numbering,
            v.clone(),
            Formulation::TotalLagrangian,
        )
        .unwrap();
        let sampled_model = ReducedNonlinearModel::new_formulated(
            &mesh,
            &materials,
            &dof_numbering,
            v.clone(),
            Formulation::TotalLagrangian,
        )
        .unwrap()
        .with_ecsw(&model);
        let time_per_eval = |model: &ReducedNonlinearModel, label: &str| -> f64 {
            let states = &ecsw_held;
            // Warm-up pass, then best of 3 timed sweeps.
            for d in states.iter().take(8) {
                let _ = model.assemble_reduced_force(d).unwrap();
            }
            let mut best = f64::INFINITY;
            for _ in 0..3 {
                let start = std::time::Instant::now();
                for d in states {
                    let _ = model.assemble_reduced_force(d).unwrap();
                }
                best = best.min(start.elapsed().as_secs_f64() / states.len() as f64);
            }
            println!(
                "    {label}: {:.1} us/assembly over {} states ({} elements)",
                best * 1e6,
                states.len(),
                model.active_elements()
            );
            best
        };
        let t_full = time_per_eval(&full_model, "full-active");
        let t_sampled = time_per_eval(&sampled_model, "ECSW-sampled");
        println!("    assembly speedup {:.1}x", t_full / t_sampled);
    }
}