//! The coupled time march shared by the FSI2 and FSI3 tests: rigid-flag
//! phase, release, and the per-step subiterated (fluid ↔ flag Newmark)
//! coupling with its measured robustness machinery. Code motion from
//! the FSI2 test; the case parameters (inflow, solid density, modulus)
//! come from the [`BenchmarkCase`].

use std::cell::RefCell;
use std::io::Write as _;

use nalgebra::Vector3;
use rtx_cfd::solvers::incompressible::FlowField;
use rtx_fea::analysis::{
    AnalysisConfig, ConvergenceCriteria, DynamicState, NonlinearDynamicAnalysis,
};
use rtx_fea::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::{MaterialId, NodeId};
use rtx_fsi::{IqnIls, Subiterated};

use super::{BenchmarkCase, Fsi2Harness, clamp_left, crossing_frequency, env_or, mid_amp};

/// Everything a march run is parameterised by. `from_env` reads the
/// `RTX_<PREFIX>_*` knobs over a set of defaults.
#[derive(Debug, Clone)]
pub struct MarchConfig {
    pub ny: usize,
    pub flag_nx: usize,
    pub t_release: f64,
    pub t_end: f64,
    /// Fluid substeps per coupled step.
    pub subcycle: usize,
    /// Per-step interface tolerance: max(`tol_floor`, `rtol` x that
    /// step's own interface increment). The floor is MEASURED per
    /// configuration (`fsi2_interface_noise.rs`), not wished: the
    /// step-to-step scatter of the accepted interface feeds the no-slip
    /// closure a wall-velocity noise of tolerance / dt_c, so a tighter
    /// coupling needs a proportionally tighter (and, as measured,
    /// reachable) floor.
    pub tol_floor: f64,
    pub rtol: f64,
    pub max_subiterations: usize,
    /// `"aitken"` (per-step scalar Aitken) or `"iqn"` (a persistent
    /// IQN-ILS whose secant history carries across steps).
    pub coupler: String,
    /// IQN secant history retained across steps.
    pub reuse: usize,
    /// Traction smoothing radius in multiples of the cell size (0 = off;
    /// measured to change nothing that matters — see the probe).
    pub smooth_in_h: f64,
    pub csv_path: Option<String>,
    /// IQN's relaxation on the very first pass, before any secant
    /// information exists. Must CONTRACT a repulsive added-mass map: for
    /// a per-pass gain `-g` the first update multiplies the residual by
    /// `|1 - omega (1 + g)|`, so 0.5 diverges past g = 3 while 0.2 holds
    /// to g = 9 (FSI3, density ratio 1, needed this — its first release
    /// pass at 0.5 drove the flag's Newton to failure).
    pub initial_relaxation: f64,
    /// Print every pass's interface residual for the first `trace_steps`
    /// coupled steps (diagnostics; 0 = off).
    pub trace_steps: usize,
    /// C^1 interface motion (constant acceleration across the step from
    /// the previous end velocity) instead of a constant velocity with a
    /// jump at the step boundary. See `Fsi2Harness::advance_subcycled`.
    pub c1_interface: bool,
    /// The per-step predictor: `"structure"` steps the flag alone under
    /// the committed load (FSI2's, exact for a heavy flag), `"kinematic"`
    /// extrapolates the CONVERGED interface velocity, d + dt v (velocity
    /// only — see the predictor code for why not acceleration). At unit
    /// density ratio the structure-alone predictor ignores an added mass
    /// comparable to the flag's own and overshoots the motion 2–5x
    /// (measured: 6.4e-4 predicted vs 1.4e-4 converged at FSI3's
    /// release), and the C^1 ramp toward that excess draws a 5–6x load
    /// (8,300 N vs ~1,300 N) — a violent first pass every step, one of
    /// which pushed the flag's Newton onto a wrong branch.
    pub predictor: String,
    /// Release with the acceleration zeroed instead of the
    /// structure-alone consistent initial acceleration M⁻¹F. That
    /// acceleration ignores the added mass — at unit density ratio it
    /// is wildly wrong (light tip nodes under −529 N of lift) and Newmark
    /// average acceleration then carries it as a sign-alternating mode.
    pub quiescent_release: bool,
}

impl MarchConfig {
    /// Read `RTX_{prefix}_{NY,T_RELEASE,T_END,SUBCYCLE,TOL,RTOL,MAXSUB,
    /// FLAG_NX,SMOOTH,COUPLER,REUSE,CSV}` over `defaults`.
    pub fn from_env(prefix: &str, defaults: MarchConfig) -> MarchConfig {
        let key = |name: &str| format!("RTX_{prefix}_{name}");
        let num = |name: &str, default: f64| env_or(&key(name), default);
        MarchConfig {
            ny: num("NY", defaults.ny as f64) as usize,
            flag_nx: num("FLAG_NX", defaults.flag_nx as f64) as usize,
            t_release: num("T_RELEASE", defaults.t_release),
            t_end: num("T_END", defaults.t_end),
            subcycle: num("SUBCYCLE", defaults.subcycle as f64) as usize,
            tol_floor: num("TOL", defaults.tol_floor),
            rtol: num("RTOL", defaults.rtol),
            max_subiterations: num("MAXSUB", defaults.max_subiterations as f64) as usize,
            coupler: std::env::var(key("COUPLER")).unwrap_or(defaults.coupler),
            reuse: num("REUSE", defaults.reuse as f64) as usize,
            smooth_in_h: num("SMOOTH", defaults.smooth_in_h),
            csv_path: std::env::var(key("CSV")).ok().or(defaults.csv_path),
            initial_relaxation: num("OMEGA0", defaults.initial_relaxation),
            trace_steps: num("TRACE", defaults.trace_steps as f64) as usize,
            c1_interface: num("C1", f64::from(u8::from(defaults.c1_interface))) != 0.0,
            predictor: std::env::var(key("PREDICTOR")).unwrap_or(defaults.predictor),
            quiescent_release: num("QUIESCENT", f64::from(u8::from(defaults.quiescent_release)))
                != 0.0,
        }
    }
}

/// Statistics over a trailing window of a march.
#[derive(Debug, Clone, Copy)]
pub struct WindowStats {
    pub t_start: f64,
    pub uy_mid: f64,
    pub uy_amp: f64,
    pub ux_mid: f64,
    pub ux_amp: f64,
    pub frequency: Option<f64>,
    pub drag_mid: f64,
    pub drag_amp: f64,
    pub lift_mid: f64,
    pub lift_amp: f64,
    /// uy amplitude over the first / last quarter of the coupled march.
    pub amp_early: f64,
    pub amp_late: f64,
}

/// What a march produced: the trajectories and the coupling bookkeeping.
#[derive(Debug, Clone)]
pub struct MarchResult {
    pub dt: f64,
    pub coupled_steps: usize,
    pub times: Vec<f64>,
    pub ux: Vec<f64>,
    pub uy: Vec<f64>,
    pub force_times: Vec<f64>,
    pub drag: Vec<f64>,
    pub lift: Vec<f64>,
    pub rigid_drag: f64,
    pub rigid_lift: f64,
    pub mean_subiterations: f64,
    pub max_subiterations: usize,
    pub stalled_steps: usize,
    pub retried_steps: usize,
    pub worst_stall: f64,
    pub worst_conservation: f64,
    pub skipped: usize,
    pub spiked: usize,
    pub final_state_finite: bool,
    pub elapsed: f64,
}

impl MarchResult {
    /// Measure over the last `seconds` of the march (or the last half,
    /// if the march is shorter).
    pub fn window(&self, seconds: f64) -> WindowStats {
        let t_end = *self.times.last().unwrap_or(&0.0);
        let start = self
            .times
            .iter()
            .position(|&t| t >= t_end - seconds)
            .unwrap_or(self.times.len() / 2);
        let (uy_mid, uy_amp) = mid_amp(&self.uy[start..]);
        let (ux_mid, ux_amp) = mid_amp(&self.ux[start..]);
        let frequency = crossing_frequency(&self.times[start..], &self.uy[start..]);
        let force_start = self
            .force_times
            .iter()
            .position(|&t| t >= t_end - seconds)
            .unwrap_or(self.force_times.len() / 2);
        let (drag_mid, drag_amp) = mid_amp(&self.drag[force_start..]);
        let (lift_mid, lift_amp) = mid_amp(&self.lift[force_start..]);
        let quarter = self.uy.len() / 4;
        let (_, amp_early) = mid_amp(&self.uy[..quarter.max(1)]);
        let (_, amp_late) = mid_amp(&self.uy[self.uy.len() - quarter.max(1)..]);
        WindowStats {
            t_start: self.times.get(start).copied().unwrap_or(0.0),
            uy_mid,
            uy_amp,
            ux_mid,
            ux_amp,
            frequency,
            drag_mid,
            drag_amp,
            lift_mid,
            lift_amp,
            amp_early,
            amp_late,
        }
    }
}

/// Run the coupled march for a benchmark case.
///
/// Phase 1 marches the rigid flag to `t_release` (the fluid harness is
/// checked against the case's rigid-flag drag). Phase 2 releases the
/// flag at rest under the sampled load and marches to `t_end`: per step,
/// a structure-alone predictor, then the coupler drives the pass
/// (subcycled fluid on the candidate interface → sampled load → one flag
/// Newmark step from the committed state) to a fixed point. Robustness
/// machinery, each piece measured before it was written (see the FSI2
/// test's module docs): stall acceptance at the noise floor, IQN
/// history reset + one retry from the predictor, increment-scaled
/// acceptance for the rare violent step; genuine runaway still panics.
#[allow(clippy::too_many_lines)]
pub fn run_march(case: BenchmarkCase, config: &MarchConfig) -> MarchResult {
    let MarchConfig {
        ny,
        flag_nx,
        t_release,
        t_end,
        subcycle,
        tol_floor,
        rtol,
        max_subiterations: max_subiterations_budget,
        ref coupler,
        reuse,
        smooth_in_h,
        ref csv_path,
        initial_relaxation,
        trace_steps,
        c1_interface,
        ref predictor,
        quiescent_release,
    } = *config;

    let (harness, mut solver, mut field) = Fsi2Harness::build_case(case, ny, flag_nx, smooth_in_h);
    let dt_fluid = harness.dt_fluid;
    let dt = dt_fluid * subcycle as f64;
    let interface = &harness.interface;
    let zero_d = vec![0.0; 2 * interface.wetted.len()];

    // Phase 1: rigid flag to t_release.
    let start = std::time::Instant::now();
    let rigid_steps = (t_release / dt_fluid).round() as usize;
    for _ in 0..rigid_steps {
        futures::executor::block_on(solver.advance(&mut field, dt_fluid)).unwrap();
    }
    // The fluid harness check: surface drag on cylinder + flag against
    // the rigid-flag CFD value on this geometry.
    let (rigid_drag, rigid_lift) = harness.measure_force(&solver, &field);
    println!(
        "  {} rigid phase: {rigid_steps} steps to t = {t_release:.1} s in {:.0} s wall; \
         surface drag {rigid_drag:.1} (rigid-flag reference {:.1}), lift {rigid_lift:.1}",
        case.name,
        start.elapsed().as_secs_f64(),
        case.rigid_drag_reference
    );

    // The flag: nonlinear Newmark stepper at the coupled dt.
    let mut db = MaterialDatabase::new();
    db.add_material(
        MaterialId(0),
        LinearElastic::new(case.e_s, case.nu_s).with_density(case.rho_s),
        None,
    );
    // A deep Newton budget: a mid-swing subiteration can hand the flag a
    // large sudden load change (the coupled lift swings hundreds of N
    // within a period); typical steps converge in 1-2 iterations, and a
    // t = 25.8 s failure at the default budget of 25 is what set this.
    let analysis = NonlinearDynamicAnalysis::new(
        harness.mesh.clone(),
        db,
        clamp_left(&harness.mesh),
        dt,
        1,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian()
    .with_convergence_criteria(ConvergenceCriteria {
        max_iterations: 60,
        ..ConvergenceCriteria::default()
    });
    let flag = RefCell::new(analysis.stepper().unwrap());
    let wetted_dofs: Vec<[usize; 2]> = interface
        .wetted
        .iter()
        .map(|&id| {
            let dofs = flag.borrow().node_dofs(id);
            [dofs[0], dofs[1]]
        })
        .collect();
    let a_dofs = flag.borrow().node_dofs(harness.a_node);
    let extract = |state: &DynamicState| -> Vec<f64> {
        let mut d = vec![0.0; 2 * wetted_dofs.len()];
        for (k, dofs) in wetted_dofs.iter().enumerate() {
            d[2 * k] = state.displacement[dofs[0]];
            d[2 * k + 1] = state.displacement[dofs[1]];
        }
        d
    };
    let extract_velocity = |state: &DynamicState| -> Vec<f64> {
        let mut v = vec![0.0; 2 * wetted_dofs.len()];
        for (k, dofs) in wetted_dofs.iter().enumerate() {
            v[2 * k] = state.velocity[dofs[0]];
            v[2 * k + 1] = state.velocity[dofs[1]];
        }
        v
    };

    // Phase 2: release. The flag starts at rest under the current fluid
    // load (consistent initial acceleration — the step response about the
    // steady deflection is the seed perturbation for the instability).
    let (nodal0, conservation0, _) = harness.sample_load(&solver, &field, &zero_d);
    flag.borrow_mut().set_nodal_forces(&nodal0);
    let mut flag_state = flag.borrow_mut().rest_state().unwrap();
    if quiescent_release {
        flag_state.acceleration.fill(0.0);
    }
    let mut committed_nodal = nodal0;
    let mut worst_conservation = conservation0;

    let solver = RefCell::new(solver);
    let field = RefCell::new(field);

    // The interface driver: per-step Aitken, or a persistent IQN-ILS
    // whose secant history carries across steps.
    let mut iqn = (coupler == "iqn").then(|| {
        IqnIls::new(max_subiterations_budget, 1.0)
            .unwrap()
            .with_reuse(reuse)
            .with_initial_relaxation(initial_relaxation)
            .unwrap()
    });

    let coupled_steps = ((t_end - t_release) / dt).round() as usize;
    let mut times = Vec::with_capacity(coupled_steps);
    let mut ux_series = Vec::with_capacity(coupled_steps);
    let mut uy_series = Vec::with_capacity(coupled_steps);
    let mut total_subiterations = 0usize;
    let mut max_subiterations = 0usize;
    let mut total_skipped = 0usize;
    let mut stalled_steps = 0usize;
    let mut retried_steps = 0usize;
    let mut worst_stall = 0.0f64;
    let mut force_times: Vec<f64> = Vec::new();
    let mut drag_series: Vec<f64> = Vec::new();
    let mut lift_series: Vec<f64> = Vec::new();
    let mut csv = csv_path
        .as_ref()
        .map(|p| std::fs::File::create(p).expect("csv path"));

    let phase_start = std::time::Instant::now();
    for step in 0..coupled_steps {
        let d_n = extract(&flag_state);
        let v_n: Option<Vec<f64>> = c1_interface.then(|| extract_velocity(&flag_state));

        // Predictor (see `MarchConfig::predictor`).
        let d_predicted = if predictor == "kinematic" {
            // Velocity only. The converged velocity is clean (Newmark's
            // trapezoidal update sums consecutive accelerations), but the
            // ACCELERATION is not: average acceleration carries an
            // inconsistent initial acceleration as a sign-alternating
            // mode step after step, and extrapolating it predicted 22 mm
            // at FSI3's release (converged: 0.14 mm).
            let v = extract_velocity(&flag_state);
            d_n.iter().zip(&v).map(|(d, v)| d + dt * v).collect()
        } else {
            // The structure alone under the committed load.
            flag.borrow_mut().set_nodal_forces(&committed_nodal);
            let (predicted, _) = flag.borrow_mut().step(&flag_state).unwrap();
            extract(&predicted)
        };

        let fluid_saved = solver.borrow().snapshot();
        let field_saved = field.borrow().clone();
        type PassResult = (
            FlowField,
            DynamicState,
            Vec<(NodeId, Vector3<f64>)>,
            f64,
            usize,
        );
        let latest: RefCell<Option<PassResult>> = RefCell::new(None);

        let pass = |d_candidate: &[f64]| -> Vec<f64> {
            // Subcycled fluid steps from the SAME start-of-step state,
            // geometry interpolated to each substep's end time, interface
            // velocity of THIS candidate constant over the step.
            let mut solver_ref = solver.borrow_mut();
            solver_ref.restore(&fluid_saved);
            let mut trial_field = field_saved.clone();
            harness.advance_subcycled(
                &mut solver_ref,
                &mut trial_field,
                &d_n,
                d_candidate,
                subcycle,
                v_n.as_deref(),
            );
            // Load on the candidate geometry, flag answers from the
            // committed state.
            let (nodal, conservation, skipped) =
                harness.sample_load(&solver_ref, &trial_field, d_candidate);
            if step < trace_steps {
                let load: f64 = nodal.iter().map(|(_, f)| f.norm()).sum();
                let peak = nodal.iter().map(|(_, f)| f.norm()).fold(0.0, f64::max);
                println!(
                    "    step {step} pass: candidate |d| = {:.3e}, sampled load: total nodal |F| \
                     = {load:.2}, peak nodal |F| = {peak:.2}",
                    d_candidate.iter().map(|v| v * v).sum::<f64>().sqrt()
                );
            }
            let mut flag_ref = flag.borrow_mut();
            flag_ref.set_nodal_forces(&nodal);
            let (candidate_state, _) = flag_ref.step(&flag_state).unwrap();
            let d_new = extract(&candidate_state);
            if step < trace_steps {
                let residual: f64 = d_new
                    .iter()
                    .zip(d_candidate)
                    .map(|(a, b)| (a - b) * (a - b))
                    .sum::<f64>()
                    .sqrt();
                println!(
                    "    step {step} pass: |d_new - d_candidate| = {residual:.3e}, |d_new| = {:.3e}",
                    d_new.iter().map(|v| v * v).sum::<f64>().sqrt()
                );
            }
            *latest.borrow_mut() =
                Some((trial_field, candidate_state, nodal, conservation, skipped));
            d_new
        };

        let increment: f64 = d_predicted
            .iter()
            .zip(&d_n)
            .map(|(a, b)| (a - b) * (a - b))
            .sum::<f64>()
            .sqrt();
        let tol_step = tol_floor.max(rtol * increment);
        // Acceptance beyond the tolerance: 5x the tolerance (noise
        // bouncing over a well-predicted step) or an order below the
        // step's own increment (the rare violent step near peak motion —
        // the s = 1 FSI2 run died at residual = 9% of its increment).
        // Counted as stalls and bounded by the caller.
        let acceptable = (5.0 * tol_step).max(0.1 * increment);
        let mut outcome = if let Some(iqn) = iqn.as_mut() {
            iqn.set_tolerance(tol_step).unwrap();
            iqn.solve(&d_predicted, pass)
        } else {
            Subiterated::aitken(max_subiterations_budget, tol_step)
                .unwrap()
                .solve(&d_predicted, pass)
        };
        // Stale-history recovery: cross-step secant columns assume the
        // interface Jacobian drifts slowly; during a rapid resonant
        // growth they can steer the quasi-Newton update into an
        // overshoot the divergence guard reads as added mass (measured:
        // a first residual of 1e-4 driven to 1e-3 by the first update at
        // 2.7x the previously seen amplitude). The map itself converges
        // deeply from a clean start, so: reset the history and retry the
        // step ONCE from the predictor. Aitken carries no history — a
        // retry would repeat the identical iteration — so IQN-only.
        if let (Err(e), Some(iqn_ref)) = (&outcome, iqn.as_mut()) {
            let recoverable = matches!(
                e,
                rtx_fsi::FsiError::CouplingNotConverged { residual, .. }
                | rtx_fsi::FsiError::CouplingDiverged { residual, .. }
                if *residual >= acceptable
            );
            if recoverable {
                iqn_ref.reset_history();
                retried_steps += 1;
                outcome = iqn_ref.solve(&d_predicted, pass);
            }
        }
        match outcome {
            Ok(converged) => {
                total_subiterations += converged.iterations;
                max_subiterations = max_subiterations.max(converged.iterations);
            }
            Err(
                rtx_fsi::FsiError::CouplingNotConverged {
                    iterations,
                    residual,
                    ..
                }
                | rtx_fsi::FsiError::CouplingDiverged {
                    iterations,
                    residual,
                },
            ) if residual < acceptable => {
                stalled_steps += 1;
                worst_stall = worst_stall.max(residual);
                total_subiterations += iterations;
                max_subiterations = max_subiterations.max(iterations);
            }
            Err(e) => panic!("{} coupling failed at step {step}: {e:?}", case.name),
        }
        // `latest` holds the response to the accepted interface (the last
        // pass) — commit it directly; the fluid, mask and flag are
        // consistent with that interface without an extra pass.
        let (new_field, new_flag_state, nodal, conservation, skipped) =
            latest.borrow_mut().take().expect("pass ran");
        *field.borrow_mut() = new_field;
        flag_state = new_flag_state;
        committed_nodal = nodal;
        worst_conservation = worst_conservation.max(conservation);
        total_skipped += skipped;

        let t = t_release + (step + 1) as f64 * dt;
        let ux = flag_state.displacement[a_dofs[0]];
        let uy = flag_state.displacement[a_dofs[1]];
        times.push(t);
        ux_series.push(ux);
        uy_series.push(uy);
        if (step + 1) % 10 == 0 {
            let (drag, lift) = harness.measure_force(&solver.borrow(), &field.borrow());
            force_times.push(t);
            drag_series.push(drag);
            lift_series.push(lift);
            if let Some(file) = csv.as_mut() {
                writeln!(file, "{t:.6},{ux:.6e},{uy:.6e},{drag:.6e},{lift:.6e}").unwrap();
            }
        } else if let Some(file) = csv.as_mut() {
            writeln!(file, "{t:.6},{ux:.6e},{uy:.6e},,").unwrap();
        }
        if (step + 1) % 1000 == 0 {
            let window = &uy_series[uy_series.len().saturating_sub(1000)..];
            let (w_mid, w_amp) = mid_amp(window);
            println!(
                "  t = {t:.3} s ({step} steps): uy(A) = {uy:.3e} (window mid {w_mid:.3e} \
                 amp {w_amp:.3e}), {:.1} subit/step, {:.0} s wall",
                total_subiterations as f64 / (step + 1) as f64,
                phase_start.elapsed().as_secs_f64()
            );
        }
    }

    MarchResult {
        dt,
        coupled_steps,
        times,
        ux: ux_series,
        uy: uy_series,
        force_times,
        drag: drag_series,
        lift: lift_series,
        rigid_drag,
        rigid_lift,
        mean_subiterations: total_subiterations as f64 / coupled_steps.max(1) as f64,
        max_subiterations,
        stalled_steps,
        retried_steps,
        worst_stall,
        worst_conservation,
        skipped: total_skipped,
        spiked: harness.spiked_total.get(),
        final_state_finite: flag_state.displacement.iter().all(|v| v.is_finite()),
        elapsed: start.elapsed().as_secs_f64(),
    }
}