//! Jacobian row/column equilibration (Ruiz-style scaling).
//!
//! The Newton state vector mixes quantities of vastly different magnitudes:
//! mass flow (ṁ ≈ 1 kg/s), pressure (P ≈ 1e5–4e6 Pa), enthalpy (h ≈ 1e5–5e5
//! J/kg) and — once moist-air edges arrive — humidity ratio (W ≈ 0.005–0.025).
//! Without rescaling, the Jacobian condition number can reach 1e10+, so the LU
//! solve loses significant digits and Newton stalls or diverges.
//!
//! This module applies a **solution-preserving** diagonal scaling. Given the
//! row/column factors `(d_r, d_c)` from [`equilibrate`], the scaled matrix
//! `Ĵ = diag(d_r) · J · diag(d_c)` has row and column ∞-norms ≈ 1. The Newton
//! system is then solved as
//!
//! ```text
//!   Ĵ · y = diag(d_r) · (−r)
//!   Δx    = diag(d_c) · y          (see `unscale_dx`)
//! ```
//!
//! which is mathematically identical to the unscaled solve but numerically far
//! more robust on stiff, mixed-unit, >50-variable systems.
//!
//! Row scaling uses the inverse row ∞-norm; column scaling the inverse column
//! ∞-norm. Both are applied iteratively (Ruiz equilibration) until the norms are
//! within tolerance of 1. This data-driven scheme subsumes the nominal
//! per-variable column scaling described in the design doc (§3.5.2) without
//! needing hand-tuned reference magnitudes.
//!
//! **Ref:** IDAES scaling theory,
//! `idaes-pse.readthedocs.io/en/stable/explanations/scaling_toolbox/scaling_theory.html`
//! and the Ruiz (2001) iterative equilibration algorithm. Design context:
//! `_bmad-output/planning-artifacts/complex-machine-design.md#3.5`.

use nalgebra::DMatrix;

/// Maximum Ruiz sweeps before giving up on the unit-norm target.
const MAX_ITERS: usize = 20;
/// Convergence tolerance on the row/column ∞-norm deviation from 1.
const TOL: f64 = 1e-3;

/// Computes Ruiz row/column equilibration factors for a matrix.
///
/// Iteratively rescales rows and columns by the square root of their current
/// ∞-norm until the norms are within tolerance of 1 (or [`MAX_ITERS`] is
/// reached). Returns `(d_r, d_c)` such that `diag(d_r) · m · diag(d_c)` is
/// equilibrated (row/column ∞-norms ≈ 1).
///
/// Zero rows/columns are left untouched (factor `1.0`), so the rank — and
/// therefore singularity detection in a subsequent LU — is preserved exactly.
///
/// # Example
///
/// ```rust
/// use entropyk_solver::scaling::equilibrate;
/// use nalgebra::DMatrix;
///
/// // Badly scaled diagonal spanning 12 orders of magnitude.
/// let m = DMatrix::from_row_slice(2, 2, &[1e6, 0.0, 0.0, 1e-6]);
/// let (d_r, d_c) = equilibrate(&m);
/// assert_eq!(d_r.len(), 2);
/// assert_eq!(d_c.len(), 2);
/// ```
pub fn equilibrate(m: &DMatrix<f64>) -> (Vec<f64>, Vec<f64>) {
    let n_rows = m.nrows();
    let n_cols = m.ncols();
    let mut d_r = vec![1.0_f64; n_rows];
    let mut d_c = vec![1.0_f64; n_cols];

    if n_rows == 0 || n_cols == 0 {
        return (d_r, d_c);
    }

    for _ in 0..MAX_ITERS {
        let mut max_deviation = 0.0_f64;

        // Row scaling: divide each row factor by sqrt of its current ∞-norm.
        for i in 0..n_rows {
            let mut row_max = 0.0_f64;
            for j in 0..n_cols {
                let v = (d_r[i] * m[(i, j)] * d_c[j]).abs();
                if v > row_max {
                    row_max = v;
                }
            }
            if row_max > 0.0 {
                d_r[i] /= row_max.sqrt();
                max_deviation = max_deviation.max((1.0 - row_max).abs());
            }
        }

        // Column scaling: divide each column factor by sqrt of its current ∞-norm.
        for j in 0..n_cols {
            let mut col_max = 0.0_f64;
            for i in 0..n_rows {
                let v = (d_r[i] * m[(i, j)] * d_c[j]).abs();
                if v > col_max {
                    col_max = v;
                }
            }
            if col_max > 0.0 {
                d_c[j] /= col_max.sqrt();
                max_deviation = max_deviation.max((1.0 - col_max).abs());
            }
        }

        if max_deviation < TOL {
            break;
        }
    }

    (d_r, d_c)
}

/// Un-scales a solved step back to physical units.
///
/// After solving the equilibrated system `diag(d_r) · J · diag(d_c) · y =
/// diag(d_r) · (−r)`, the true Newton step is recovered by reapplying the column
/// scaling: `Δx[j] = d_c[j] · y[j]`.
///
/// If `y` and `d_c` differ in length, the shorter length governs the result
/// (defensive — callers always pass matching lengths).
///
/// # Example
///
/// ```rust
/// use entropyk_solver::scaling::unscale_dx;
///
/// let y = [2.0, -3.0];
/// let d_c = [0.5, 10.0];
/// let dx = unscale_dx(&y, &d_c);
/// assert_eq!(dx, vec![1.0, -30.0]);
/// ```
pub fn unscale_dx(y: &[f64], d_c: &[f64]) -> Vec<f64> {
    y.iter().zip(d_c.iter()).map(|(&yj, &cj)| cj * yj).collect()
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    /// After equilibration, scaled row/column ∞-norms should be ≈ 1.
    #[test]
    fn test_equilibrate_unit_norms() {
        let m = DMatrix::from_row_slice(2, 2, &[1e6, 1e3, 1e-3, 1e-6]);
        let (d_r, d_c) = equilibrate(&m);

        for i in 0..2 {
            let row_max = (0..2)
                .map(|j| (d_r[i] * m[(i, j)] * d_c[j]).abs())
                .fold(0.0_f64, f64::max);
            assert!(
                (row_max - 1.0).abs() < 0.05,
                "row {} ∞-norm not unit: {}",
                i,
                row_max
            );
        }
        for j in 0..2 {
            let col_max = (0..2)
                .map(|i| (d_r[i] * m[(i, j)] * d_c[j]).abs())
                .fold(0.0_f64, f64::max);
            assert!(
                (col_max - 1.0).abs() < 0.05,
                "col {} ∞-norm not unit: {}",
                j,
                col_max
            );
        }
    }

    /// `unscale_dx` must be the exact inverse of column scaling: given a step
    /// expressed in scaled coordinates `y = D_c^{-1} · x`, reapplying `d_c`
    /// recovers `x` exactly.
    #[test]
    fn test_unscale_dx_inverts_column_scaling() {
        let x = [3.0_f64, -7.0, 0.25];
        let d_c = [0.5_f64, 10.0, 2.0];
        // y = D_c^{-1} · x
        let y: Vec<f64> = x.iter().zip(d_c.iter()).map(|(&xj, &cj)| xj / cj).collect();
        let recovered = unscale_dx(&y, &d_c);
        for (got, want) in recovered.iter().zip(x.iter()) {
            assert_relative_eq!(got, want, epsilon = 1e-12);
        }
    }

    /// Zero rows and columns must keep factor 1.0 so the rank is preserved and
    /// singularity detection in the downstream LU is unchanged.
    #[test]
    fn test_zero_row_and_column_keep_unit_factor() {
        // Row 1 is all zero; column 1 is all zero.
        let m = DMatrix::from_row_slice(2, 2, &[5.0, 0.0, 0.0, 0.0]);
        let (d_r, d_c) = equilibrate(&m);
        assert_relative_eq!(d_r[1], 1.0, epsilon = 1e-12);
        assert_relative_eq!(d_c[1], 1.0, epsilon = 1e-12);
    }

    /// Equilibration must slash the condition number of a badly-scaled matrix.
    #[test]
    fn test_equilibrate_reduces_condition_number() {
        // Row 0 lives at ~1e6, row 1 at ~1; dynamic range keeps the small
        // singular value well above underflow so κ is measured reliably.
        let m = DMatrix::from_row_slice(2, 2, &[1.0e6, 2.0e6, 3.0, 1.0]);

        let cond_before = condition_number(&m).unwrap();

        let (d_r, d_c) = equilibrate(&m);
        let mut scaled = m.clone();
        for i in 0..2 {
            for j in 0..2 {
                scaled[(i, j)] *= d_r[i] * d_c[j];
            }
        }
        let cond_after = condition_number(&scaled).unwrap();

        assert!(
            cond_after < cond_before / 100.0 && cond_after < 10.0,
            "Ruiz should slash κ: before={:.3e}, after={:.3e}",
            cond_before,
            cond_after
        );
    }

    /// Empty matrices return empty/unit factors without panicking.
    #[test]
    fn test_empty_matrix() {
        let m: DMatrix<f64> = DMatrix::zeros(0, 0);
        let (d_r, d_c) = equilibrate(&m);
        assert!(d_r.is_empty());
        assert!(d_c.is_empty());
    }

    /// Local κ helper (SVD σ_max/σ_min) for the test above only.
    fn condition_number(m: &DMatrix<f64>) -> Option<f64> {
        if m.nrows() == 0 || m.ncols() == 0 {
            return None;
        }
        let svd = m.clone().svd(false, false);
        let sv = svd.singular_values;
        if sv.len() == 0 {
            return None;
        }
        let sigma_max = sv.max();
        let sigma_min = sv
            .iter()
            .filter(|&&s| s > 0.0)
            .min_by(|a, b| a.partial_cmp(b).unwrap())
            .copied();
        sigma_min.map(|min| sigma_max / min)
    }
}