Entropyk/crates/solver/src/jacobian.rs

//! Jacobian matrix assembly and solving for Newton-Raphson.
//!
//! This module provides the `JacobianMatrix` type, which wraps `nalgebra::DMatrix<f64>`
//! and provides methods for:
//!
//! - Building from sparse entries (from `JacobianBuilder`)
//! - Solving linear systems J·Δx = -r via LU decomposition
//! - Computing numerical Jacobians via finite differences
//!
//! # Example
//!
//! ```rust
//! use entropyk_solver::jacobian::JacobianMatrix;
//!
//! // Build from sparse entries
//! let entries = vec![(0, 0, 2.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 3.0)];
//! let jacobian = JacobianMatrix::from_builder(&entries, 2, 2);
//!
//! // Solve J·Δx = -r
//! let residuals = vec![1.0, 2.0];
//! let delta = jacobian.solve(&residuals).expect("non-singular");
//! ```

use nalgebra::{DMatrix, DVector};

/// Wrapper around `nalgebra::DMatrix<f64>` for Jacobian operations.
///
/// The Jacobian matrix J represents the partial derivatives of the residual vector
/// with respect to the state vector:
///
/// $$J_{ij} = \frac{\partial r_i}{\partial x_j}$$
///
/// For Newton-Raphson, we solve the linear system:
///
/// $$J \cdot \Delta x = -r$$
#[derive(Debug, Clone)]
pub struct JacobianMatrix(DMatrix<f64>);

impl JacobianMatrix {
    /// Builds a Jacobian matrix from sparse entries.
    ///
    /// Each entry is a tuple `(row, col, value)`. The matrix is zero-initialized
    /// and then filled with the provided entries.
    ///
    /// # Arguments
    ///
    /// * `entries` - Slice of `(row, col, value)` tuples
    /// * `n_rows` - Number of rows (equations)
    /// * `n_cols` - Number of columns (state variables)
    ///
    /// # Example
    ///
    /// ```rust
    /// use entropyk_solver::jacobian::JacobianMatrix;
    ///
    /// let entries = vec![(0, 0, 1.0), (1, 1, 2.0)];
    /// let j = JacobianMatrix::from_builder(&entries, 2, 2);
    /// ```
    pub fn from_builder(entries: &[(usize, usize, f64)], n_rows: usize, n_cols: usize) -> Self {
        let mut matrix = DMatrix::zeros(n_rows, n_cols);
        for &(row, col, value) in entries {
            if row < n_rows && col < n_cols {
                matrix[(row, col)] += value;
            }
        }
        JacobianMatrix(matrix)
    }

    /// Updates an existing Jacobian matrix from sparse entries in-place.
    ///
    /// The matrix is first zeroed out, then filled with the new entries.
    /// This avoids re-allocating memory during iterations, satisfying the
    /// zero-allocation architecture constraint.
    ///
    /// # Arguments
    ///
    /// * `entries` - Slice of `(row, col, value)` tuples
    pub fn update_from_builder(&mut self, entries: &[(usize, usize, f64)]) {
        self.0.fill(0.0);
        let n_rows = self.0.nrows();
        let n_cols = self.0.ncols();
        for &(row, col, value) in entries {
            if row < n_rows && col < n_cols {
                self.0[(row, col)] += value;
            }
        }
    }

    /// Creates a zero Jacobian matrix with the given dimensions.
    pub fn zeros(n_rows: usize, n_cols: usize) -> Self {
        JacobianMatrix(DMatrix::zeros(n_rows, n_cols))
    }

    /// Returns the number of rows (equations).
    pub fn nrows(&self) -> usize {
        self.0.nrows()
    }

    /// Returns the number of columns (state variables).
    pub fn ncols(&self) -> usize {
        self.0.ncols()
    }

    /// Solves the linear system J·Δx = -r and returns Δx.
    ///
    /// Uses **Ruiz equilibration** (iterative row/column scaling) followed by LU
    /// decomposition with partial pivoting. Equilibration rescales the rows and
    /// columns so their ∞-norms approach 1, which dramatically lowers the
    /// condition number of badly-scaled Jacobians — exactly the situation that
    /// arises when a thermodynamic system mixes pressures (~1e6), enthalpies
    /// (~1e5) and dimensionless controls (~1) in the same matrix. The scaling is
    /// solution-preserving (it is undone on the returned step), so the result is
    /// mathematically identical to an unscaled solve but numerically far more
    /// robust on stiff, >50-variable systems.
    ///
    /// Returns `None` if the matrix is singular (no unique solution exists).
    ///
    /// # Arguments
    ///
    /// * `residuals` - The residual vector r (length must equal `nrows()`)
    ///
    /// # Returns
    ///
    /// * `Some(Δx)` - The Newton step (length = `ncols()`)
    /// * `None` - If the Jacobian is singular
    ///
    /// # Example
    ///
    /// ```rust
    /// use entropyk_solver::jacobian::JacobianMatrix;
    ///
    /// let entries = vec![(0, 0, 2.0), (1, 1, 1.0)];
    /// let j = JacobianMatrix::from_builder(&entries, 2, 2);
    ///
    /// let r = vec![4.0, 3.0];
    /// let delta = j.solve(&r).expect("non-singular");
    /// assert!((delta[0] - (-2.0)).abs() < 1e-10);
    /// assert!((delta[1] - (-3.0)).abs() < 1e-10);
    /// ```
    pub fn solve(&self, residuals: &[f64]) -> Option<Vec<f64>> {
        if residuals.len() != self.0.nrows() {
            tracing::warn!(
                "residual length {} != Jacobian rows {}",
                residuals.len(),
                self.0.nrows()
            );
            return None;
        }

        // For square systems, use Ruiz-equilibrated LU decomposition.
        if self.0.nrows() == self.0.ncols() {
            let n = self.0.nrows();

            // Diagonal scalings D_r, D_c such that the scaled matrix
            // Ĵ = diag(d_r) · J · diag(d_c) has near-unit row/column ∞-norms.
            let (d_r, d_c) = crate::scaling::equilibrate(&self.0);

            // Build the scaled matrix in place on a single clone (same number of
            // allocations as the previous unscaled path).
            let mut scaled = self.0.clone();
            for i in 0..n {
                for j in 0..n {
                    scaled[(i, j)] *= d_r[i] * d_c[j];
                }
            }

            let lu = scaled.lu();

            // Scaled right-hand side: D_r · (-r).
            let neg_r_scaled = DVector::from_iterator(n, (0..n).map(|i| -residuals[i] * d_r[i]));

            match lu.solve(&neg_r_scaled) {
                // Undo the column scaling to recover the true step: Δx = D_c · y.
                Some(y) => {
                    let delta = crate::scaling::unscale_dx(y.as_slice(), &d_c);
                    // A NaN/Inf entry anywhere in the Jacobian (or RHS) silently
                    // slips through LU as a non-finite step. Reject it here so the
                    // caller treats the iteration as a clean failure instead of
                    // propagating NaN into the state for another iteration.
                    if delta.iter().all(|v| v.is_finite()) {
                        Some(delta)
                    } else {
                        tracing::warn!(
                            "LU solve produced a non-finite step - Jacobian may contain NaN/Inf"
                        );
                        None
                    }
                }
                None => {
                    tracing::warn!("LU solve failed - Jacobian may be singular");
                    None
                }
            }
        } else {
            // For non-square systems, use least-squares (SVD)
            // This is a fallback for overdetermined/underdetermined systems
            tracing::debug!(
                "Non-square Jacobian ({}x{}) - using least-squares",
                self.0.nrows(),
                self.0.ncols()
            );

            let r_vec = DVector::from_row_slice(residuals);
            let neg_r = -r_vec;

            // Use SVD for robust least-squares solution
            let svd = self.0.clone().svd(true, true);
            match svd.solve(&neg_r, 1e-10) {
                Ok(delta) => {
                    let v: Vec<f64> = delta.iter().copied().collect();
                    if v.iter().all(|x| x.is_finite()) {
                        Some(v)
                    } else {
                        tracing::warn!(
                            "SVD solve produced a non-finite step - Jacobian may contain NaN/Inf"
                        );
                        None
                    }
                }
                Err(e) => {
                    tracing::warn!("SVD solve failed - Jacobian may be rank-deficient: {}", e);
                    None
                }
            }
        }
    }

    /// Returns the Ruiz row/column equilibration factors `(d_r, d_c)`.
    ///
    /// The scaled matrix `diag(d_r) · J · diag(d_c)` has row and column
    /// ∞-norms close to 1. This is the same scaling applied internally by
    /// [`solve`](Self::solve); it is exposed for diagnostics (e.g. inspecting how
    /// ill-scaled a Jacobian is, or estimating the conditioning improvement).
    ///
    /// # Example
    ///
    /// ```rust
    /// use entropyk_solver::jacobian::JacobianMatrix;
    ///
    /// // Badly scaled diagonal: entries span 12 orders of magnitude.
    /// let entries = vec![(0, 0, 1e6), (1, 1, 1e-6)];
    /// let j = JacobianMatrix::from_builder(&entries, 2, 2);
    /// let (d_r, d_c) = j.ruiz_scaling_factors();
    /// assert_eq!(d_r.len(), 2);
    /// assert_eq!(d_c.len(), 2);
    /// ```
    pub fn ruiz_scaling_factors(&self) -> (Vec<f64>, Vec<f64>) {
        crate::scaling::equilibrate(&self.0)
    }

    /// Estimates the condition number of the Jacobian matrix.
    ///
    /// The condition number κ = σ_max / σ_min indicates how ill-conditioned
    /// the matrix is. Values > 1e10 indicate an ill-conditioned system that
    /// may cause numerical instability in the solver.
    ///
    /// Uses SVD decomposition to compute singular values. This is an O(n³)
    /// operation and should only be used for diagnostics.
    ///
    /// # Returns
    ///
    /// * `Some(κ)` - The condition number (ratio of largest to smallest singular value)
    /// * `None` - If the matrix is rank-deficient (σ_min = 0)
    ///
    /// # Example
    ///
    /// ```rust
    /// use entropyk_solver::jacobian::JacobianMatrix;
    ///
    /// // Well-conditioned matrix
    /// let entries = vec![(0, 0, 2.0), (1, 1, 1.0)];
    /// let j = JacobianMatrix::from_builder(&entries, 2, 2);
    /// let cond = j.estimate_condition_number().unwrap();
    /// assert!(cond < 10.0, "Expected low condition number, got {}", cond);
    ///
    /// // Ill-conditioned matrix (nearly singular)
    /// let bad_entries = vec![(0, 0, 1.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 1.0000001)];
    /// let bad_j = JacobianMatrix::from_builder(&bad_entries, 2, 2);
    /// let bad_cond = bad_j.estimate_condition_number().unwrap();
    /// assert!(bad_cond > 1e7, "Expected high condition number, got {}", bad_cond);
    /// ```
    pub fn estimate_condition_number(&self) -> Option<f64> {
        // Handle empty matrices
        if self.0.nrows() == 0 || self.0.ncols() == 0 {
            return None;
        }

        // Use SVD to get singular values
        let svd = self.0.clone().svd(true, true);

        // Get singular values
        let singular_values = svd.singular_values;

        if singular_values.len() == 0 {
            return None;
        }

        let sigma_max = singular_values.max();
        let sigma_min = singular_values
            .iter()
            .filter(|&&s| s > 0.0)
            .min_by(|a, b| a.partial_cmp(b).unwrap())
            .copied();

        match sigma_min {
            Some(min) => Some(sigma_max / min),
            None => None, // Matrix is rank-deficient
        }
    }

    /// Computes a numerical Jacobian via finite differences.
    ///
    /// For each state variable x_j, perturbs by epsilon and computes:
    ///
    /// $$J_{ij} \approx \frac{r_i(x + \epsilon e_j) - r_i(x)}{\epsilon}$$
    ///
    /// # Arguments
    ///
    /// * `compute_residuals` - Function that computes residuals from state
    /// * `state` - Current state vector
    /// * `residuals` - Current residual vector (avoid recomputing)
    /// * `epsilon` - Perturbation size (typically 1e-8)
    ///
    /// # Returns
    ///
    /// A `JacobianMatrix` with the numerical derivatives.
    ///
    /// # Example
    ///
    /// ```rust
    /// use entropyk_solver::jacobian::JacobianMatrix;
    ///
    /// let state: Vec<f64> = vec![1.0, 2.0];
    /// let residuals: Vec<f64> = vec![state[0] * state[0], state[1] * 2.0];
    /// let compute_residuals = |s: &[f64], r: &mut [f64]| {
    ///     r[0] = s[0] * s[0];
    ///     r[1] = s[1] * 2.0;
    ///     Ok(())
    /// };
    ///
    /// let j = JacobianMatrix::numerical(
    ///     compute_residuals,
    ///     &state,
    ///     &residuals,
    ///     1e-8
    /// ).unwrap();
    /// ```
    pub fn numerical<F>(
        compute_residuals: F,
        state: &[f64],
        residuals: &[f64],
        relative_epsilon: f64,
    ) -> Result<Self, String>
    where
        F: Fn(&[f64], &mut [f64]) -> Result<(), String>,
    {
        let n = state.len();
        let m = residuals.len();
        let mut matrix = DMatrix::zeros(m, n);

        for j in 0..n {
            // Use relative epsilon scaled to the state variable magnitude.
            // For variables like P~350,000 Pa or h~400,000 J/kg, an absolute
            // epsilon of 1e-8 is below library numerical precision and gives zero
            // finite differences. A relative epsilon of ~1e-5 gives physically
            // meaningful perturbations across all thermodynamic property scales.
            let eps_j = state[j].abs().max(1.0) * relative_epsilon;

            let mut state_perturbed = state.to_vec();
            state_perturbed[j] += eps_j;

            let mut residuals_perturbed = vec![0.0; m];
            compute_residuals(&state_perturbed, &mut residuals_perturbed)?;

            for i in 0..m {
                matrix[(i, j)] = (residuals_perturbed[i] - residuals[i]) / eps_j;
            }
        }

        Ok(JacobianMatrix(matrix))
    }

    /// Returns a reference to the underlying matrix.
    pub fn as_matrix(&self) -> &DMatrix<f64> {
        &self.0
    }

    /// Returns a mutable reference to the underlying matrix.
    pub fn as_matrix_mut(&mut self) -> &mut DMatrix<f64> {
        &mut self.0
    }

    /// Gets an element at (row, col).
    pub fn get(&self, row: usize, col: usize) -> Option<f64> {
        if row < self.0.nrows() && col < self.0.ncols() {
            Some(self.0[(row, col)])
        } else {
            None
        }
    }

    /// Sets an element at (row, col).
    pub fn set(&mut self, row: usize, col: usize, value: f64) {
        if row < self.0.nrows() && col < self.0.ncols() {
            self.0[(row, col)] = value;
        }
    }

    /// Returns the Frobenius norm of the matrix.
    pub fn norm(&self) -> f64 {
        self.0.norm()
    }

    /// Returns the condition number (ratio of largest to smallest singular value).
    ///
    /// Returns `None` if the matrix is rank-deficient.
    pub fn condition_number(&self) -> Option<f64> {
        let svd = self.0.clone().svd(false, false);
        let singular_values = svd.singular_values;

        let max_sv = singular_values.max();
        let min_sv = singular_values.min();

        if min_sv > 1e-14 {
            Some(max_sv / min_sv)
        } else {
            None
        }
    }

    /// Returns the block structure of the Jacobian matrix for a multi-circuit system.
    ///
    /// For a system with N circuits, each circuit's equations and state variables
    /// form a contiguous block in the Jacobian (assuming the state vector layout
    /// `[P_edge0, h_edge0, P_edge1, h_edge1, ...]` is ordered by circuit).
    ///
    /// Returns one tuple per circuit: `(row_start, row_end, col_start, col_end)`,
    /// where rows correspond to equations and columns to state variables.
    ///
    /// # Notes
    ///
    /// - For uncoupled circuits, the blocks do not overlap and off-block entries
    ///   are zero (verified by [`is_block_diagonal`](Self::is_block_diagonal)).
    /// - Row/col ranges are inclusive-start, exclusive-end: `row_start..row_end`.
    ///
    /// # AC: #6
    pub fn block_structure(
        &self,
        system: &crate::system::System,
    ) -> Vec<(usize, usize, usize, usize)> {
        let n_circuits = system.circuit_count();
        let mut blocks = Vec::with_capacity(n_circuits);

        for circuit_idx in 0..n_circuits {
            let circuit_id = circuit_idx as u8;

            // Collect state-variable indices for this circuit.
            // State layout: [P_edge0, h_edge0, P_edge1, h_edge1, ...], so for edge i:
            //   col p_idx = 2*i, col h_idx = 2*i+1.
            // The equation rows mirror the same layout, so row = col for square systems.
            let indices: Vec<usize> = system
                .circuit_edges(crate::CircuitId(circuit_id.into()))
                .flat_map(|edge| {
                    let (p_idx, h_idx) = system.edge_state_indices(edge);
                    [p_idx, h_idx]
                })
                .collect();

            if indices.is_empty() {
                continue;
            }

            let col_start = *indices.iter().min().unwrap();
            let col_end = *indices.iter().max().unwrap() + 1; // exclusive
                                                              // Equations mirror state layout for square systems
            let row_start = col_start;
            let row_end = col_end;

            blocks.push((row_start, row_end, col_start, col_end));
        }

        blocks
    }

    /// Returns `true` if the Jacobian has block-diagonal structure for a multi-circuit system.
    ///
    /// Checks that all entries **outside** the circuit blocks (as returned by
    /// [`block_structure`](Self::block_structure)) have absolute value ≤ `tolerance`.
    ///
    /// For uncoupled multi-circuit systems, the Jacobian is block-diagonal because
    /// equations in one circuit do not depend on state variables in another circuit.
    ///
    /// # Arguments
    ///
    /// * `system` — The system whose circuit decomposition defines the expected blocks.
    /// * `tolerance` — Maximum allowed absolute value for off-block entries.
    ///
    /// # AC: #6
    pub fn is_block_diagonal(&self, system: &crate::system::System, tolerance: f64) -> bool {
        let blocks = self.block_structure(system);
        let nrows = self.0.nrows();
        let ncols = self.0.ncols();

        // Map each row to its corresponding block column range (if any)
        // This optimizes the check from O(N^2 * C) to O(N^2)
        let mut row_block_cols = vec![None; nrows];
        for &(rs, re, cs, ce) in &blocks {
            for block in &mut row_block_cols[rs..re] {
                *block = Some((cs, ce));
            }
        }

        for (row, block) in row_block_cols.iter().enumerate().take(nrows) {
            for col in 0..ncols {
                let in_block = match *block {
                    Some((cs, ce)) => col >= cs && col < ce,
                    None => false,
                };

                if !in_block {
                    let val = self.0[(row, col)].abs();
                    if val > tolerance {
                        tracing::debug!(
                            row = row,
                            col = col,
                            value = val,
                            tolerance = tolerance,
                            "Off-block nonzero entry found — not block-diagonal"
                        );
                        return false;
                    }
                }
            }
        }

        true
    }
}

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

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

    /// A NaN entry in the Jacobian must yield `None` (clean failure) rather than
    /// a `Some(..)` step containing NaN that would poison the next iteration.
    #[test]
    fn test_solve_with_nan_entry_returns_none() {
        let entries = vec![(0, 0, f64::NAN), (1, 1, 1.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);
        assert!(j.solve(&[1.0, 1.0]).is_none());
    }

    #[test]
    fn test_from_builder_simple() {
        let entries = vec![(0, 0, 1.0), (0, 1, 2.0), (1, 0, 3.0), (1, 1, 4.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        assert_eq!(j.nrows(), 2);
        assert_eq!(j.ncols(), 2);
        assert_relative_eq!(j.get(0, 0).unwrap(), 1.0);
        assert_relative_eq!(j.get(0, 1).unwrap(), 2.0);
        assert_relative_eq!(j.get(1, 0).unwrap(), 3.0);
        assert_relative_eq!(j.get(1, 1).unwrap(), 4.0);
    }

    #[test]
    fn test_from_builder_accumulates() {
        // Multiple entries for the same position should accumulate
        let entries = vec![(0, 0, 1.0), (0, 0, 2.0), (0, 0, 3.0)];
        let j = JacobianMatrix::from_builder(&entries, 1, 1);

        assert_relative_eq!(j.get(0, 0).unwrap(), 6.0);
    }

    #[test]
    fn test_from_builder_out_of_bounds_ignored() {
        let entries = vec![(0, 0, 1.0), (5, 5, 100.0)]; // (5, 5) is out of bounds
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        assert_relative_eq!(j.get(0, 0).unwrap(), 1.0);
        assert_eq!(j.get(5, 5), None); // Out of bounds
    }

    #[test]
    fn test_solve_identity() {
        let entries = vec![(0, 0, 1.0), (1, 1, 1.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let r = vec![3.0, 4.0];
        let delta = j.solve(&r).expect("identity is non-singular");

        assert_relative_eq!(delta[0], -3.0, epsilon = 1e-10);
        assert_relative_eq!(delta[1], -4.0, epsilon = 1e-10);
    }

    #[test]
    fn test_solve_diagonal() {
        let entries = vec![(0, 0, 2.0), (1, 1, 4.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let r = vec![6.0, 8.0];
        let delta = j.solve(&r).expect("diagonal is non-singular");

        assert_relative_eq!(delta[0], -3.0, epsilon = 1e-10);
        assert_relative_eq!(delta[1], -2.0, epsilon = 1e-10);
    }

    #[test]
    fn test_solve_full_matrix() {
        // J = [[2, 1], [1, 3]]
        // J·Δx = -r where r = [1, 2]
        // Solution: Δx = [-0.2, -0.6]
        let entries = vec![(0, 0, 2.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 3.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let r = vec![1.0, 2.0];
        let delta = j.solve(&r).expect("non-singular");

        // Verify: J·Δx = -r
        assert_relative_eq!(2.0 * delta[0] + 1.0 * delta[1], -1.0, epsilon = 1e-10);
        assert_relative_eq!(1.0 * delta[0] + 3.0 * delta[1], -2.0, epsilon = 1e-10);
    }

    #[test]
    fn test_solve_singular_returns_none() {
        // Singular matrix: [[1, 1], [1, 1]]
        let entries = vec![(0, 0, 1.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 1.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let r = vec![1.0, 2.0];
        let result = j.solve(&r);

        assert!(result.is_none(), "Singular matrix should return None");
    }

    #[test]
    fn test_solve_zero_matrix_returns_none() {
        let entries: Vec<(usize, usize, f64)> = vec![];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let r = vec![1.0, 2.0];
        let result = j.solve(&r);

        assert!(result.is_none(), "Zero matrix should return None");
    }

    #[test]
    fn test_numerical_jacobian_linear() {
        // r[0] = 2*x0 + 3*x1
        // r[1] = x0 - x1
        // J = [[2, 3], [1, -1]]
        let state = vec![1.0, 2.0];
        let residuals = vec![2.0 * state[0] + 3.0 * state[1], state[0] - state[1]];

        let compute_residuals = |s: &[f64], r: &mut [f64]| {
            r[0] = 2.0 * s[0] + 3.0 * s[1];
            r[1] = s[0] - s[1];
            Ok(())
        };

        let j = JacobianMatrix::numerical(compute_residuals, &state, &residuals, 1e-8).unwrap();

        assert_relative_eq!(j.get(0, 0).unwrap(), 2.0, epsilon = 1e-6);
        assert_relative_eq!(j.get(0, 1).unwrap(), 3.0, epsilon = 1e-6);
        assert_relative_eq!(j.get(1, 0).unwrap(), 1.0, epsilon = 1e-6);
        assert_relative_eq!(j.get(1, 1).unwrap(), -1.0, epsilon = 1e-6);
    }

    #[test]
    fn test_numerical_jacobian_quadratic() {
        // r[0] = x0^2
        // r[1] = x1^3
        // J = [[2*x0, 0], [0, 3*x1^2]]
        let state: Vec<f64> = vec![2.0, 3.0];
        let residuals: Vec<f64> = vec![state[0].powi(2), state[1].powi(3)];

        let compute_residuals = |s: &[f64], r: &mut [f64]| {
            r[0] = s[0].powi(2);
            r[1] = s[1].powi(3);
            Ok(())
        };

        let j = JacobianMatrix::numerical(compute_residuals, &state, &residuals, 1e-8).unwrap();

        assert_relative_eq!(j.get(0, 0).unwrap(), 4.0, epsilon = 1e-5); // 2*2
        assert_relative_eq!(j.get(0, 1).unwrap(), 0.0, epsilon = 1e-5);
        assert_relative_eq!(j.get(1, 0).unwrap(), 0.0, epsilon = 1e-5);
        assert_relative_eq!(j.get(1, 1).unwrap(), 27.0, epsilon = 1e-4); // 3*3^2
    }

    #[test]
    fn test_condition_number() {
        // Well-conditioned identity
        let entries = vec![(0, 0, 1.0), (1, 1, 1.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);
        let cond = j.condition_number().unwrap();
        assert_relative_eq!(cond, 1.0, epsilon = 1e-10);

        // Ill-conditioned (nearly singular)
        let entries = vec![(0, 0, 1.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 1.0 + 1e-10)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);
        let cond = j.condition_number();
        assert!(cond.unwrap() > 1e9, "Should be ill-conditioned");
    }

    #[test]
    fn test_norm() {
        let entries = vec![(0, 0, 3.0), (0, 1, 4.0)];
        let j = JacobianMatrix::from_builder(&entries, 1, 2);
        // Frobenius norm = sqrt(3^2 + 4^2) = 5
        assert_relative_eq!(j.norm(), 5.0, epsilon = 1e-10);
    }

    #[test]
    fn test_zeros() {
        let j = JacobianMatrix::zeros(3, 4);
        assert_eq!(j.nrows(), 3);
        assert_eq!(j.ncols(), 4);
        assert_relative_eq!(j.norm(), 0.0);
    }

    #[test]
    fn test_set_and_get() {
        let mut j = JacobianMatrix::zeros(2, 2);
        j.set(0, 0, 5.0);
        j.set(1, 1, 7.0);

        assert_relative_eq!(j.get(0, 0).unwrap(), 5.0);
        assert_relative_eq!(j.get(1, 1).unwrap(), 7.0);
        assert_relative_eq!(j.get(0, 1).unwrap(), 0.0);
    }

    #[test]
    fn test_solve_wrong_residual_length() {
        let entries = vec![(0, 0, 1.0), (1, 1, 1.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let r = vec![1.0]; // Wrong length
        let result = j.solve(&r);

        assert!(result.is_none(), "Wrong residual length should return None");
    }

    #[test]
    fn test_numerical_vs_analytical_agree() {
        // For a simple function, numerical and analytical Jacobians should match
        // r[0] = x0^2 + x0*x1
        // r[1] = sin(x0) + cos(x1)
        // J = [[2*x0 + x1, x0], [cos(x0), -sin(x1)]]

        let state: Vec<f64> = vec![0.5, 1.0];
        let residuals: Vec<f64> = vec![
            state[0].powi(2) + state[0] * state[1],
            state[0].sin() + state[1].cos(),
        ];

        let compute_residuals = |s: &[f64], r: &mut [f64]| {
            r[0] = s[0].powi(2) + s[0] * s[1];
            r[1] = s[0].sin() + s[1].cos();
            Ok(())
        };

        let j_num = JacobianMatrix::numerical(compute_residuals, &state, &residuals, 1e-8).unwrap();

        // Analytical values
        let j00 = 2.0 * state[0] + state[1]; // 2*0.5 + 1.0 = 2.0
        let j01 = state[0]; // 0.5
        let j10 = state[0].cos(); // cos(0.5)
        let j11 = -state[1].sin(); // -sin(1.0)

        assert_relative_eq!(j_num.get(0, 0).unwrap(), j00, epsilon = 1e-5);
        assert_relative_eq!(j_num.get(0, 1).unwrap(), j01, epsilon = 1e-5);
        assert_relative_eq!(j_num.get(1, 0).unwrap(), j10, epsilon = 1e-5);
        assert_relative_eq!(j_num.get(1, 1).unwrap(), j11, epsilon = 1e-5);
    }

    #[test]
    fn test_ruiz_equilibration_unit_norms() {
        // After equilibration, scaled row/column ∞-norms should be ≈ 1.
        let entries = vec![(0, 0, 1e6), (0, 1, 1e3), (1, 0, 1e-3), (1, 1, 1e-6)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);
        let (d_r, d_c) = j.ruiz_scaling_factors();

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

    #[test]
    fn test_ruiz_reduces_condition_number() {
        // Badly-scaled but SVD-resolvable matrix (dynamic range ~1e6 keeps the
        // small singular value above underflow, so the condition number is
        // measured reliably). Row 0 lives at ~1e6, row 1 at ~1.
        let entries = vec![(0, 0, 1.0e6), (0, 1, 2.0e6), (1, 0, 3.0), (1, 1, 1.0)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let cond_before = j.estimate_condition_number().unwrap();

        // Apply the Ruiz scaling and recompute the condition number.
        let (d_r, d_c) = j.ruiz_scaling_factors();
        let m = j.as_matrix();
        let mut scaled_entries = Vec::new();
        for i in 0..2 {
            for jj in 0..2 {
                scaled_entries.push((i, jj, d_r[i] * m[(i, jj)] * d_c[jj]));
            }
        }
        let scaled = JacobianMatrix::from_builder(&scaled_entries, 2, 2);
        let cond_after = scaled.estimate_condition_number().unwrap();

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

    #[test]
    fn test_solve_illconditioned_matches_known_solution() {
        // Badly-scaled system with a known solution. J·x = b where
        // J = [[1e6, 2e6], [3e-6, 4e-6]], x = [1, -1]  =>  b = [-1e6, -1e-6].
        // solve() returns Δx for J·Δx = -r, so set r = -b to get Δx = x.
        let entries = vec![(0, 0, 1e6), (0, 1, 2e6), (1, 0, 3e-6), (1, 1, 4e-6)];
        let j = JacobianMatrix::from_builder(&entries, 2, 2);

        let b = [-1e6, -1e-6];
        let r = [-b[0], -b[1]];
        let delta = j.solve(&r).expect("non-singular");

        assert_relative_eq!(delta[0], 1.0, epsilon = 1e-9);
        assert_relative_eq!(delta[1], -1.0, epsilon = 1e-9);
    }
}