//! Coupling-level rescue — rung A of omni-cortex
//! `docs/coupling_rescue_campaign.md`.
//!
//! A coupled step the march would otherwise abandon (a fatal stall:
//! the coupling ended above its acceptance window) or reject (a
//! committed tip jump beyond the march's own running statistic — the
//! runaway signature every recorded death shows for 5–37 steps before
//! its panic) is REPEATED over the same interval as n = 2, 4, 8, 16,
//! 32 coupled substeps of dt/n: Mayr, Wall and Gee's reduced-step
//! repetition (κ = 0.5, at most five repetitions) applied to the
//! partitioned loop, one layer above the structural Newton rescue.
//!
//! Each substep is a complete coupled step at dt/n: its own predictor
//! from the substep's start state, the fluid subcycled at
//! dt_fluid/n from the fluid state the previous substep committed, the
//! flag stepped at dt/n (`NonlinearDynamicStepper::step_with_dt`, whose
//! Newton and rescue ladder already take dt as a parameter), the C¹
//! interface velocity chained through the substep's own start
//! velocity, and a FRESH coupler (the secant history of the other dt
//! is not this map's). Tolerances follow the march's own rule at the
//! substep's own increment. A substep that ends above its acceptance
//! window fails the rung; the next rung restores the start-of-step
//! fluid snapshot and tries twice as many substeps.
//!
//! **Default OFF.** With the knob off nothing here runs — the march's
//! own step logic is not routed through this module, on purpose: the
//! path the digit-identity protocol certifies stays float-for-float
//! untouched, and this duplicate of it only exists past a trigger.

use std::cell::{Cell, RefCell};

use nalgebra::Vector3;
use rtx_cfd::solvers::incompressible::{EmbeddedPisoSolver, EmbeddedSolverState, FlowField};
use rtx_fea::analysis::{DynamicState, NonlinearDynamicStepper};
use rtx_fea::mesh::NodeId;
use rtx_fsi::{FsiError, IqnIls, Subiterated};

use super::Fsi2Harness;
use super::march::MarchConfig;

/// The substep ladder: five repetitions at κ = 0.5.
pub const LADDER: [usize; 5] = [2, 4, 8, 16, 32];

/// Kinematic trigger: a committed tip jump above this multiple of the
/// march's running p95 of |Δuy| rejects the step. Calibrated on all 74
/// recorded s = 1 marches (`scripts/fsi_burst_scan.py`): every march
/// that reached its horizon stays ≤ 2.4×; 27 of 28 deaths exceed 5×.
pub const TIP_JUMP_FACTOR: f64 = 3.0;

/// The running statistic's window (committed, non-rescued steps); the
/// kinematic trigger is inert until the window is full.
pub const TIP_JUMP_WINDOW: usize = 2000;

/// More rescues than this within one second of march is a runaway the
/// substeps only delay — the march dies loudly instead of hiding it.
pub const RATE_CAP_PER_SECOND: usize = 20;

/// One rescue, as recorded in the march result and printed as it fires.
#[derive(Debug, Clone)]
pub struct RescueRecord {
    pub step: usize,
    pub t: f64,
    /// `"fatal stall"` or `"tip jump"`.
    pub trigger: &'static str,
    /// The rejected step's residual (fatal stall) or |Δuy| (tip jump).
    pub before: f64,
    /// Substeps that carried the interval; 0 = every rung failed (the
    /// step was then kept as the coupling accepted it, or the march
    /// died — see the trigger).
    pub n: usize,
    /// Coupling passes spent across all rungs tried.
    pub passes: usize,
    /// Tip uy the rejected step had committed (NaN for a fatal stall,
    /// which never commits) and the tip the rescue committed.
    pub tip_rejected: f64,
    pub tip_rescued: f64,
}

/// What a successful ladder rung hands back to the march.
pub struct RescueOutcome {
    pub state: DynamicState,
    pub nodal: Vec<(NodeId, Vector3<f64>)>,
    pub n: usize,
    pub passes: usize,
    pub stalls: usize,
    pub worst_residual: f64,
    pub worst_conservation: f64,
    pub skipped: usize,
}

/// Everything one interval repetition needs, borrowed from the march.
pub struct Interval<'a, 'b> {
    pub harness: &'a Fsi2Harness,
    pub solver: &'a RefCell<EmbeddedPisoSolver>,
    pub field: &'a RefCell<FlowField>,
    pub flag: &'a RefCell<NonlinearDynamicStepper<'b>>,
    pub wetted_dofs: &'a [[usize; 2]],
    /// The start-of-step fluid snapshot and field the march already
    /// saves for its coupling passes.
    pub fluid_saved: &'a EmbeddedSolverState,
    pub field_saved: &'a FlowField,
    pub start_state: &'a DynamicState,
    pub start_nodal: &'a [(NodeId, Vector3<f64>)],
    pub config: &'a MarchConfig,
    /// The full coupled step the interval spans.
    pub dt: f64,
}

fn extract(state: &DynamicState, wetted: &[[usize; 2]]) -> Vec<f64> {
    let mut d = vec![0.0; 2 * wetted.len()];
    for (k, dofs) in wetted.iter().enumerate() {
        d[2 * k] = state.displacement[dofs[0]];
        d[2 * k + 1] = state.displacement[dofs[1]];
    }
    d
}

fn extract_velocity(state: &DynamicState, wetted: &[[usize; 2]]) -> Vec<f64> {
    let mut v = vec![0.0; 2 * wetted.len()];
    for (k, dofs) in wetted.iter().enumerate() {
        v[2 * k] = state.velocity[dofs[0]];
        v[2 * k + 1] = state.velocity[dofs[1]];
    }
    v
}

/// Repeat the interval up the ladder; the first rung whose every
/// substep converges or stalls inside its acceptance window wins.
///
/// # Errors
/// The last rung's coupling error when all five fail. The fluid and
/// field are then left as the last failed rung left them — the caller
/// restores or dies.
pub fn substep_interval(iv: &Interval<'_, '_>) -> Result<RescueOutcome, FsiError> {
    let mut last_err = None;
    let mut passes_so_far = 0usize;
    for &n in &LADDER {
        match attempt(iv, n) {
            Ok(mut outcome) => {
                outcome.passes += passes_so_far;
                return Ok(outcome);
            }
            Err((e, passes)) => {
                passes_so_far += passes;
                last_err = Some(e);
            }
        }
    }
    Err(last_err.expect("ladder is non-empty"))
}

/// More coarse episodes than this within one second of march is a
/// runaway the coarsening only delays — the march dies loudly.
pub const COARSE_EPISODE_CAP_PER_SECOND: usize = 5;

/// Rung C: ONE coupled step over `factor × dt` with the fluid
/// subcycled at `factor ×` the configured subcycle — the fluid dt is
/// unchanged, so this is exactly the s = 2 interpolated closure
/// (factor 2) that marches through the crossing where s = 1 dies
/// (campaign doc §11, P0.3). From the saved start, like a rung.
///
/// # Errors
/// The coupling error when the coarse step ends above its acceptance
/// window; the fluid and field are then mid-failure — the caller
/// restores or dies.
pub fn coarse_step(iv: &Interval<'_, '_>, factor: usize) -> Result<RescueOutcome, FsiError> {
    attempt_with(iv, 1, iv.dt * factor as f64, iv.config.subcycle * factor).map_err(|(e, _)| e)
}

/// One rung: `n` coupled substeps of `dt/n` from the saved start.
fn attempt(iv: &Interval<'_, '_>, n: usize) -> Result<RescueOutcome, (FsiError, usize)> {
    attempt_with(iv, n, iv.dt / n as f64, iv.config.subcycle)
}

/// `n` coupled steps of `dt_sub`, each with `subcycle_sub` fluid
/// substeps, from the saved start.
fn attempt_with(
    iv: &Interval<'_, '_>,
    n: usize,
    dt_sub: f64,
    subcycle_sub: usize,
) -> Result<RescueOutcome, (FsiError, usize)> {
    let cfg = iv.config;
    let dt_fluid_sub = dt_sub / subcycle_sub as f64;

    iv.solver.borrow_mut().restore(iv.fluid_saved);
    *iv.field.borrow_mut() = iv.field_saved.clone();
    let mut state = iv.start_state.clone();
    let mut nodal: Vec<(NodeId, Vector3<f64>)> = iv.start_nodal.to_vec();

    let mut passes = 0usize;
    let mut stalls = 0usize;
    let mut worst_residual = 0.0f64;
    let mut worst_conservation = 0.0f64;
    let mut skipped = 0usize;

    for _ in 0..n {
        let d_n = extract(&state, iv.wetted_dofs);
        let v_n: Option<Vec<f64>> = cfg
            .c1_interface
            .then(|| extract_velocity(&state, iv.wetted_dofs));
        let d_predicted: Vec<f64> = if cfg.predictor == "kinematic" {
            let v = extract_velocity(&state, iv.wetted_dofs);
            d_n.iter().zip(&v).map(|(d, v)| d + dt_sub * v).collect()
        } else {
            let mut flag = iv.flag.borrow_mut();
            flag.set_nodal_forces(&nodal);
            let (predicted, _) = flag
                .step_with_dt(&state, dt_sub)
                .expect("structure-alone predictor at the substep dt");
            extract(&predicted, iv.wetted_dofs)
        };

        let sub_saved = iv.solver.borrow().snapshot();
        let sub_field = iv.field.borrow().clone();
        type PassResult = (
            FlowField,
            DynamicState,
            Vec<(NodeId, Vector3<f64>)>,
            f64,
            usize,
        );
        let latest: RefCell<Option<PassResult>> = RefCell::new(None);
        let pass_count = Cell::new(0usize);
        let state_ref = &state;
        let pass = |d_candidate: &[f64]| -> Vec<f64> {
            pass_count.set(pass_count.get() + 1);
            let mut solver_ref = iv.solver.borrow_mut();
            solver_ref.restore(&sub_saved);
            let mut trial = sub_field.clone();
            iv.harness.advance_subcycled_with(
                &mut solver_ref,
                &mut trial,
                &d_n,
                d_candidate,
                subcycle_sub,
                v_n.as_deref(),
                dt_fluid_sub,
            );
            let (nodal_c, conservation, skipped_c) =
                iv.harness.sample_load(&solver_ref, &trial, d_candidate);
            let mut flag = iv.flag.borrow_mut();
            flag.set_nodal_forces(&nodal_c);
            let (candidate, _) = flag
                .step_with_dt(state_ref, dt_sub)
                .expect("flag step at the substep dt");
            let d_new = extract(&candidate, iv.wetted_dofs);
            *latest.borrow_mut() = Some((trial, candidate, nodal_c, conservation, skipped_c));
            d_new
        };

        let increment: f64 = d_predicted
            .iter()
            .zip(&d_n)
            .map(|(a, b)| (a - b) * (a - b))
            .sum::<f64>()
            .sqrt();
        let tol_step = cfg.tol_floor.max(cfg.rtol * increment);
        let acceptable = (cfg.stall_accept * tol_step).max(0.1 * increment);
        let outcome = if cfg.coupler == "iqn" {
            let mut iqn = IqnIls::new(cfg.max_subiterations, tol_step)
                .map_err(|e| (e, passes))?
                .with_reuse(cfg.reuse)
                .with_initial_relaxation(cfg.initial_relaxation)
                .map_err(|e| (e, passes))?;
            iqn.solve(&d_predicted, pass)
        } else {
            Subiterated::aitken(cfg.max_subiterations, tol_step)
                .map_err(|e| (e, passes))?
                .solve(&d_predicted, pass)
        };
        passes += pass_count.get();
        match outcome {
            Ok(c) => worst_residual = worst_residual.max(c.residual),
            Err(
                FsiError::CouplingNotConverged { residual, .. }
                | FsiError::CouplingDiverged { residual, .. },
            ) if residual < acceptable => {
                stalls += 1;
                worst_residual = worst_residual.max(residual);
            }
            Err(e) => return Err((e, passes)),
        }
        let (new_field, new_state, new_nodal, conservation, skipped_c) =
            latest.borrow_mut().take().expect("pass ran");
        // As in the march: the solver's state after the last pass IS
        // the response to the accepted interface; commit it directly.
        *iv.field.borrow_mut() = new_field;
        state = new_state;
        nodal = new_nodal;
        worst_conservation = worst_conservation.max(conservation);
        skipped += skipped_c;
    }

    Ok(RescueOutcome {
        state,
        nodal,
        n,
        passes,
        stalls,
        worst_residual,
        worst_conservation,
        skipped,
    })
}