slatec 0.1.0

//! Safe two-stage wrappers for `CHKDER` and `DCKDER`.

use alloc::{vec, vec::Vec};
use core::fmt;

use slatec_core::to_fortran_integer;

use super::JacobianMut;

/// Failure while preparing or evaluating a SLATEC Jacobian consistency check.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum JacobianCheckError {
    /// The square system has no variables or equations.
    EmptySystem,
    /// A component of the evaluation point is NaN or infinite.
    NonFinitePoint {
        /// Zero-based invalid component.
        index: usize,
    },
    /// A Rust function evaluation left a component non-finite or unwritten.
    FunctionReturnedNonFinite {
        /// Whether the failure occurred at the original or perturbed point.
        evaluation: &'static str,
        /// Zero-based residual component.
        index: usize,
    },
    /// The user Jacobian left a logical entry non-finite or unwritten.
    JacobianReturnedNonFinite {
        /// Zero-based row.
        row: usize,
        /// Zero-based column.
        column: usize,
    },
    /// The system dimension cannot be represented by Fortran `INTEGER`.
    IntegerOverflow,
    /// Checked square matrix allocation arithmetic overflowed.
    WorkspaceOverflow,
    /// Native checker output violated its documented score contract.
    NativeContractViolation,
}

impl fmt::Display for JacobianCheckError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::EmptySystem => write!(formatter, "Jacobian checks require a nonempty system"),
            Self::NonFinitePoint { index } => {
                write!(
                    formatter,
                    "Jacobian-check point at index {index} must be finite"
                )
            }
            Self::FunctionReturnedNonFinite { evaluation, index } => write!(
                formatter,
                "function evaluation at {evaluation} left component {index} non-finite"
            ),
            Self::JacobianReturnedNonFinite { row, column } => write!(
                formatter,
                "Jacobian entry ({row}, {column}) is non-finite or unwritten"
            ),
            Self::IntegerOverflow => write!(
                formatter,
                "Jacobian-check dimension exceeds Fortran INTEGER"
            ),
            Self::WorkspaceOverflow => write!(formatter, "Jacobian-check matrix size overflowed"),
            Self::NativeContractViolation => write!(
                formatter,
                "native Jacobian checker returned an invalid score"
            ),
        }
    }
}

#[cfg(feature = "std")]
impl std::error::Error for JacobianCheckError {}

/// Componentwise consistency evidence produced by `CHKDER` or `DCKDER`.
#[derive(Clone, Debug, PartialEq)]
pub struct JacobianCheckResult<T = f64> {
    /// One score per equation, from zero (inconsistent) to one (consistent).
    pub scores: Vec<T>,
    /// Zero-based rows with scores below SLATEC's documented `0.5` heuristic.
    pub suspicious_rows: Vec<usize>,
}

fn check_input_f64(point: &[f64]) -> Result<usize, JacobianCheckError> {
    if point.is_empty() {
        return Err(JacobianCheckError::EmptySystem);
    }
    if let Some((index, _)) = point
        .iter()
        .enumerate()
        .find(|(_, value)| !value.is_finite())
    {
        return Err(JacobianCheckError::NonFinitePoint { index });
    }
    point
        .len()
        .checked_mul(point.len())
        .ok_or(JacobianCheckError::WorkspaceOverflow)
}

fn check_input_f32(point: &[f32]) -> Result<usize, JacobianCheckError> {
    if point.is_empty() {
        return Err(JacobianCheckError::EmptySystem);
    }
    if let Some((index, _)) = point
        .iter()
        .enumerate()
        .find(|(_, value)| !value.is_finite())
    {
        return Err(JacobianCheckError::NonFinitePoint { index });
    }
    point
        .len()
        .checked_mul(point.len())
        .ok_or(JacobianCheckError::WorkspaceOverflow)
}

/// Checks a square double-precision analytic Jacobian with SLATEC `DCKDER`.
///
/// `point` is Fortran `X`. The function callback fills `FVEC` at `X` and at
/// the perturbed `XP` generated by mode 1. The Jacobian callback fills the
/// dense column-major `FJAC` at `X`; its residual argument is the valid
/// function value at that point. Mode 2 returns one score per equation:
/// `1` indicates consistency, `0` inconsistency, and `0.5` is the original
/// routine's heuristic boundary. Severe cancellation can make scores
/// inconclusive, especially when point components are unusually small.
///
/// The wrapper allocates five vectors plus an `N × N` matrix but requires no
/// callback FFI, thread-local state, or process lock: both Rust callbacks run
/// before direct `DCKDER` calls. It therefore requires `alloc` and
/// `nonlinear-jacobian-check`, but not `std`. A Rust callback panic remains a
/// normal Rust panic and never crosses Fortran.
///
/// # Errors
///
/// Returns [`JacobianCheckError`] for an empty/non-finite point, checked size
/// overflow, non-finite callback output, or invalid native scores.
///
/// # Example
///
/// ```no_run
/// # fn main() -> Result<(), slatec::nonlinear::JacobianCheckError> {
/// use slatec::nonlinear::check_jacobian;
/// let result = check_jacobian(
///     &[1.5],
///     |x, f| f[0] = x[0] * x[0] - 2.0,
///     |x, _, mut j| j.set(0, 0, 2.0 * x[0]).unwrap(),
/// )?;
/// assert!(result.scores[0] > 0.5);
/// # Ok(()) }
/// ```
pub fn check_jacobian<F, J>(
    point: &[f64],
    mut function: F,
    mut jacobian: J,
) -> Result<JacobianCheckResult<f64>, JacobianCheckError>
where
    F: FnMut(&[f64], &mut [f64]),
    J: FnMut(&[f64], &[f64], JacobianMut<'_, f64>),
{
    let matrix_len = check_input_f64(point)?;
    let dimension = point.len();
    let mut x = point.to_vec();
    let mut fvec = vec![f64::NAN; dimension];
    function(point, &mut fvec);
    if let Some((index, _)) = fvec
        .iter()
        .enumerate()
        .find(|(_, value)| !value.is_finite())
    {
        return Err(JacobianCheckError::FunctionReturnedNonFinite {
            evaluation: "X",
            index,
        });
    }
    let mut fjac = vec![f64::NAN; matrix_len];
    let view = JacobianMut::new(&mut fjac, dimension, dimension, dimension)
        .ok_or(JacobianCheckError::WorkspaceOverflow)?;
    jacobian(point, &fvec, view);
    for column in 0..dimension {
        for row in 0..dimension {
            if !fjac[row + column * dimension].is_finite() {
                return Err(JacobianCheckError::JacobianReturnedNonFinite { row, column });
            }
        }
    }
    let mut xp = vec![0.0; dimension];
    let mut fvecp = vec![0.0; dimension];
    let mut scores = vec![0.0; dimension];
    let mut n = to_fortran_integer(dimension).map_err(|_| JacobianCheckError::IntegerOverflow)?;
    let mut m = n;
    let mut ldfjac = n;
    let mut mode = 1;
    // SAFETY: all arrays have their exact reviewed lengths, N=M=LDFJAC is a
    // positive checked Fortran integer, and DCKDER mode 1 writes only XP.
    unsafe {
        slatec_sys::nonlinear::dckder(
            &mut m,
            &mut n,
            x.as_mut_ptr(),
            fvec.as_mut_ptr(),
            fjac.as_mut_ptr(),
            &mut ldfjac,
            xp.as_mut_ptr(),
            fvecp.as_mut_ptr(),
            &mut mode,
            scores.as_mut_ptr(),
        );
    }
    fvecp.fill(f64::NAN);
    function(&xp, &mut fvecp);
    if let Some((index, _)) = fvecp
        .iter()
        .enumerate()
        .find(|(_, value)| !value.is_finite())
    {
        return Err(JacobianCheckError::FunctionReturnedNonFinite {
            evaluation: "XP",
            index,
        });
    }
    mode = 2;
    // SAFETY: mode 2 receives the same validated storage plus finite FVEC,
    // FVECP, and every logical FJAC entry.
    unsafe {
        slatec_sys::nonlinear::dckder(
            &mut m,
            &mut n,
            x.as_mut_ptr(),
            fvec.as_mut_ptr(),
            fjac.as_mut_ptr(),
            &mut ldfjac,
            xp.as_mut_ptr(),
            fvecp.as_mut_ptr(),
            &mut mode,
            scores.as_mut_ptr(),
        );
    }
    if scores
        .iter()
        .any(|score| !score.is_finite() || !(0.0..=1.0).contains(score))
    {
        return Err(JacobianCheckError::NativeContractViolation);
    }
    let suspicious_rows = scores
        .iter()
        .enumerate()
        .filter_map(|(index, score)| (*score < 0.5).then_some(index))
        .collect();
    Ok(JacobianCheckResult {
        scores,
        suspicious_rows,
    })
}

/// Single-precision `CHKDER` counterpart of [`check_jacobian`].
///
/// Original SLATEC routine: `CHKDER`. It uses the same two-stage perturbation,
/// column-major Jacobian, score interpretation, allocation policy, and error
/// checks with `f32` values.
///
/// # Example
///
/// ```no_run
/// # fn main() -> Result<(), slatec::nonlinear::JacobianCheckError> {
/// use slatec::nonlinear::check_jacobian_f32;
/// let result = check_jacobian_f32(
///     &[1.5_f32], |x, f| f[0] = x[0] * x[0] - 2.0,
///     |x, _, mut j| j.set(0, 0, 2.0 * x[0]).unwrap(),
/// )?;
/// assert!(result.scores[0] > 0.5);
/// # Ok(()) }
/// ```
pub fn check_jacobian_f32<F, J>(
    point: &[f32],
    mut function: F,
    mut jacobian: J,
) -> Result<JacobianCheckResult<f32>, JacobianCheckError>
where
    F: FnMut(&[f32], &mut [f32]),
    J: FnMut(&[f32], &[f32], JacobianMut<'_, f32>),
{
    let matrix_len = check_input_f32(point)?;
    let dimension = point.len();
    let mut x = point.to_vec();
    let mut fvec = vec![f32::NAN; dimension];
    function(point, &mut fvec);
    if let Some((index, _)) = fvec
        .iter()
        .enumerate()
        .find(|(_, value)| !value.is_finite())
    {
        return Err(JacobianCheckError::FunctionReturnedNonFinite {
            evaluation: "X",
            index,
        });
    }
    let mut fjac = vec![f32::NAN; matrix_len];
    let view = JacobianMut::new(&mut fjac, dimension, dimension, dimension)
        .ok_or(JacobianCheckError::WorkspaceOverflow)?;
    jacobian(point, &fvec, view);
    for column in 0..dimension {
        for row in 0..dimension {
            if !fjac[row + column * dimension].is_finite() {
                return Err(JacobianCheckError::JacobianReturnedNonFinite { row, column });
            }
        }
    }
    let mut xp = vec![0.0; dimension];
    let mut fvecp = vec![0.0; dimension];
    let mut scores = vec![0.0; dimension];
    let mut n = to_fortran_integer(dimension).map_err(|_| JacobianCheckError::IntegerOverflow)?;
    let mut m = n;
    let mut ldfjac = n;
    let mut mode = 1;
    // SAFETY: equivalent to the reviewed DCKDER mode-1 call above for f32.
    unsafe {
        slatec_sys::nonlinear::chkder(
            &mut m,
            &mut n,
            x.as_mut_ptr(),
            fvec.as_mut_ptr(),
            fjac.as_mut_ptr(),
            &mut ldfjac,
            xp.as_mut_ptr(),
            fvecp.as_mut_ptr(),
            &mut mode,
            scores.as_mut_ptr(),
        );
    }
    fvecp.fill(f32::NAN);
    function(&xp, &mut fvecp);
    if let Some((index, _)) = fvecp
        .iter()
        .enumerate()
        .find(|(_, value)| !value.is_finite())
    {
        return Err(JacobianCheckError::FunctionReturnedNonFinite {
            evaluation: "XP",
            index,
        });
    }
    mode = 2;
    // SAFETY: equivalent to the reviewed DCKDER mode-2 call above for f32.
    unsafe {
        slatec_sys::nonlinear::chkder(
            &mut m,
            &mut n,
            x.as_mut_ptr(),
            fvec.as_mut_ptr(),
            fjac.as_mut_ptr(),
            &mut ldfjac,
            xp.as_mut_ptr(),
            fvecp.as_mut_ptr(),
            &mut mode,
            scores.as_mut_ptr(),
        );
    }
    if scores
        .iter()
        .any(|score| !score.is_finite() || !(0.0..=1.0).contains(score))
    {
        return Err(JacobianCheckError::NativeContractViolation);
    }
    let suspicious_rows = scores
        .iter()
        .enumerate()
        .filter_map(|(index, score)| (*score < 0.5).then_some(index))
        .collect();
    Ok(JacobianCheckResult {
        scores,
        suspicious_rows,
    })
}

#[cfg(test)]
mod tests {
    use super::{JacobianCheckError, check_input_f64};

    #[test]
    fn checker_input_validation_is_alloc_only() {
        assert_eq!(check_input_f64(&[]), Err(JacobianCheckError::EmptySystem));
        assert_eq!(
            check_input_f64(&[f64::NAN]),
            Err(JacobianCheckError::NonFinitePoint { index: 0 })
        );
        assert_eq!(check_input_f64(&[1.0, 2.0]), Ok(4));
    }
}