ripopt 0.6.0

use crate::options::SolverOptions;

/// Check if constraint `i` is an equality constraint (g_l ≈ g_u).
#[inline]
pub fn is_equality_constraint(g_l: f64, g_u: f64) -> bool {
    g_l.is_finite() && g_u.is_finite() && (g_l - g_u).abs() < 1e-15
}

/// Result of a convergence check.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConvergenceStatus {
    /// Not converged.
    NotConverged,
    /// Converged within desired tolerance.
    Converged,
    /// Converged within acceptable tolerance.
    Acceptable,
    /// Diverging (objective growing unboundedly).
    Diverging,
}

/// Information needed to check convergence.
pub struct ConvergenceInfo {
    /// Primal infeasibility: max |c_i(x)| for violated constraints.
    pub primal_inf: f64,
    /// Dual infeasibility: ||grad_f + J^T y - z||_inf (using z_opt for scaled check).
    pub dual_inf: f64,
    /// Dual infeasibility using iterative z (for unscaled gate).
    pub dual_inf_unscaled: f64,
    /// Complementarity: max |x_i * z_i - mu|.
    pub compl_inf: f64,
    /// Current barrier parameter.
    pub mu: f64,
    /// Current objective value.
    pub objective: f64,
    /// Sum of absolute values of all multipliers (y, z_l, z_u).
    /// Used for Ipopt-style dual scaling.
    pub multiplier_sum: f64,
    /// Total number of multiplier components (n + m for scaling denominator).
    pub multiplier_count: usize,
}

/// Check convergence of the IPM algorithm.
///
/// Returns the convergence status based on current optimality measures.
pub fn check_convergence(
    info: &ConvergenceInfo,
    options: &SolverOptions,
    consecutive_acceptable: usize,
) -> ConvergenceStatus {
    // Ipopt-style dual scaling: s_d = max(s_max, sum|mults| / count) / s_max
    // This scales the tolerance to account for large multiplier magnitudes.
    // Cap s_d to prevent false convergence when multipliers explode.
    let s_max: f64 = 100.0;
    let s_d_max: f64 = 1e4; // Don't let tolerance scale more than 10000x
    let s_d = if info.multiplier_count > 0 {
        ((s_max.max(info.multiplier_sum / info.multiplier_count as f64)) / s_max).min(s_d_max)
    } else {
        1.0
    };

    let primal_tol = options.tol;
    let dual_tol = options.tol * s_d;
    let compl_tol = options.tol * s_d;

    // Strict convergence: BOTH scaled AND unscaled must pass.
    // Scaled check uses z_opt for dual_inf (fast convergence for LPs/QPs).
    // Unscaled check uses iterative z for dual_inf (catches false convergence
    // where z_opt absorbs gradient residuals into bound multipliers).
    let scaled_ok = info.primal_inf <= primal_tol
        && info.dual_inf <= dual_tol
        && info.compl_inf <= compl_tol;
    let unscaled_ok = info.primal_inf <= options.constr_viol_tol
        && info.dual_inf_unscaled <= options.dual_inf_tol
        && info.compl_inf <= options.compl_inf_tol;
    if scaled_ok && unscaled_ok {
        return ConvergenceStatus::Converged;
    }

    // Internal near-tolerance check: BOTH scaled AND unscaled must pass.
    // Used to detect "close but not optimal" states that trigger promotion strategies.
    const NEAR_TOL_ITERS: usize = 10;
    let near_primal_tol = 100.0 * options.tol;
    let near_dual_tol = 100.0 * options.tol * s_d;
    let near_compl_tol = 100.0 * options.tol * s_d;

    let near_scaled_ok = info.primal_inf <= near_primal_tol
        && info.dual_inf <= near_dual_tol
        && info.compl_inf <= near_compl_tol;
    let near_unscaled_ok = info.primal_inf <= 10.0 * options.constr_viol_tol
        && info.dual_inf_unscaled <= 10.0 * options.dual_inf_tol
        && info.compl_inf <= 10.0 * options.compl_inf_tol;

    if near_scaled_ok && near_unscaled_ok
        && consecutive_acceptable >= NEAR_TOL_ITERS
    {
        return ConvergenceStatus::Acceptable;
    }

    // Check divergence (use 1e50 — constrained problems can have large feasible objectives,
    // and transient excursions to large |obj| can occur during interior point iterations)
    if info.objective.abs() > 1e50 {
        return ConvergenceStatus::Diverging;
    }

    ConvergenceStatus::NotConverged
}

/// Compute primal infeasibility (constraint violation).
/// Takes constraint values g(x), and constraint bounds g_l, g_u.
pub fn primal_infeasibility(g: &[f64], g_l: &[f64], g_u: &[f64]) -> f64 {
    let mut max_viol = 0.0f64;
    for i in 0..g.len() {
        if g[i] < g_l[i] {
            max_viol = max_viol.max(g_l[i] - g[i]);
        }
        if g[i] > g_u[i] {
            max_viol = max_viol.max(g[i] - g_u[i]);
        }
    }
    max_viol
}

/// Compute dual infeasibility: ||grad_f - J^T * lambda - z_l + z_u||_inf.
///
/// `grad_f`: gradient of objective
/// `jac_rows`, `jac_cols`, `jac_vals`: Jacobian in COO format
/// `lambda`: constraint multipliers
/// `z_l`, `z_u`: bound multipliers
/// `n`: number of variables
#[allow(clippy::too_many_arguments)]
pub fn dual_infeasibility(
    grad_f: &[f64],
    jac_rows: &[usize],
    jac_cols: &[usize],
    jac_vals: &[f64],
    lambda: &[f64],
    z_l: &[f64],
    z_u: &[f64],
    n: usize,
) -> f64 {
    let mut residual = vec![0.0; n];

    // Start with gradient of objective
    residual[..n].copy_from_slice(&grad_f[..n]);

    // Add J^T * lambda (Ipopt convention: L = f + y^T g)
    for (idx, (&row, &col)) in jac_rows.iter().zip(jac_cols.iter()).enumerate() {
        residual[col] += jac_vals[idx] * lambda[row];
    }

    // Subtract z_l and add z_u (bound multipliers)
    for i in 0..n {
        residual[i] -= z_l[i];
        residual[i] += z_u[i];
    }

    residual.iter().map(|r| r.abs()).fold(0.0f64, f64::max)
}

/// Compute component-wise scaled dual infeasibility.
///
/// Uses `|r_i| / (1 + |grad_f_i|)` per component, which makes the metric
/// insensitive to gradient magnitude across variables. This prevents
/// poorly-scaled problems from having artificially large unscaled dual
/// infeasibility even when the scaled version is small.
#[allow(clippy::too_many_arguments)]
pub fn dual_infeasibility_scaled(
    grad_f: &[f64],
    jac_rows: &[usize],
    jac_cols: &[usize],
    jac_vals: &[f64],
    lambda: &[f64],
    z_l: &[f64],
    z_u: &[f64],
    n: usize,
) -> f64 {
    let mut residual = vec![0.0; n];

    // Start with gradient of objective
    residual[..n].copy_from_slice(&grad_f[..n]);

    // Add J^T * lambda (Ipopt convention: L = f + y^T g)
    for (idx, (&row, &col)) in jac_rows.iter().zip(jac_cols.iter()).enumerate() {
        residual[col] += jac_vals[idx] * lambda[row];
    }

    // Subtract z_l and add z_u (bound multipliers)
    for i in 0..n {
        residual[i] -= z_l[i];
        residual[i] += z_u[i];
    }

    // Component-wise scaling: divide by (1 + |grad_f_i|)
    residual
        .iter()
        .enumerate()
        .map(|(i, r)| r.abs() / (1.0 + grad_f[i].abs()))
        .fold(0.0f64, f64::max)
}

/// Compute complementarity error for bound constraints.
/// compl = max_i |x_i * z_l_i| where x_i is near lower bound,
///         max_i |s_u_i * z_u_i| where x_i is near upper bound.
///
/// For the barrier method: complementarity = max(|(x-x_l)*z_l - mu|, |(x_u-x)*z_u - mu|).
pub fn complementarity_error(
    x: &[f64],
    x_l: &[f64],
    x_u: &[f64],
    z_l: &[f64],
    z_u: &[f64],
    mu: f64,
) -> f64 {
    let mut max_err = 0.0f64;
    let n = x.len();
    for i in 0..n {
        if x_l[i].is_finite() {
            let slack = x[i] - x_l[i];
            max_err = max_err.max((slack * z_l[i] - mu).abs());
        }
        if x_u[i].is_finite() {
            let slack = x_u[i] - x[i];
            max_err = max_err.max((slack * z_u[i] - mu).abs());
        }
    }
    max_err
}

/// Compute full complementarity error including constraint slack complementarity.
///
/// In addition to variable bound complementarity (x-x_l)*z_l and (x_u-x)*z_u,
/// this also checks constraint slack complementarity for inequality constraints:
/// - Lower-bounded: (g(x) - g_l) * max(y\[i\], 0)
/// - Upper-bounded: (g_u - g(x)) * max(-y\[i\], 0)
/// - Equality constraints are skipped.
#[allow(clippy::too_many_arguments)]
pub fn complementarity_error_full(
    x: &[f64],
    x_l: &[f64],
    x_u: &[f64],
    z_l: &[f64],
    z_u: &[f64],
    g: &[f64],
    g_l: &[f64],
    g_u: &[f64],
    y: &[f64],
    mu: f64,
) -> f64 {
    // Start with variable bound complementarity
    let mut max_err = complementarity_error(x, x_l, x_u, z_l, z_u, mu);

    // Add constraint slack complementarity for inequality constraints
    let m = g.len();
    for i in 0..m {
        if is_equality_constraint(g_l[i], g_u[i]) {
            continue;
        }
        if g_l[i].is_finite() {
            let slack = g[i] - g_l[i];
            let mult = y[i].max(0.0);
            max_err = max_err.max((slack * mult - mu).abs());
        }
        if g_u[i].is_finite() {
            let slack = g_u[i] - g[i];
            let mult = (-y[i]).max(0.0);
            max_err = max_err.max((slack * mult - mu).abs());
        }
    }
    max_err
}

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

    #[test]
    fn test_primal_infeasibility_feasible() {
        let g = vec![1.5, 3.0];
        let g_l = vec![1.0, 2.0];
        let g_u = vec![2.0, 4.0];
        assert_eq!(primal_infeasibility(&g, &g_l, &g_u), 0.0);
    }

    #[test]
    fn test_primal_infeasibility_violated() {
        let g = vec![0.5, 5.0];
        let g_l = vec![1.0, 2.0];
        let g_u = vec![2.0, 4.0];
        assert_eq!(primal_infeasibility(&g, &g_l, &g_u), 1.0);
    }

    #[test]
    fn test_convergence_optimal() {
        let info = ConvergenceInfo {
            primal_inf: 1e-10,
            dual_inf: 1e-10,
            dual_inf_unscaled: 1e-10,
            compl_inf: 1e-10,
            mu: 1e-11,
            objective: 17.0,
            multiplier_sum: 0.0,
            multiplier_count: 0,
        };
        let opts = SolverOptions::default();
        assert_eq!(
            check_convergence(&info, &opts, 0),
            ConvergenceStatus::Converged
        );
    }

    #[test]
    fn test_convergence_not_converged() {
        let info = ConvergenceInfo {
            primal_inf: 1e-3,
            dual_inf: 1e-3,
            dual_inf_unscaled: 1e-3,
            compl_inf: 1e-3,
            mu: 0.01,
            objective: 17.0,
            multiplier_sum: 0.0,
            multiplier_count: 0,
        };
        let opts = SolverOptions::default();
        assert_eq!(
            check_convergence(&info, &opts, 0),
            ConvergenceStatus::NotConverged
        );
    }

    #[test]
    fn test_convergence_diverging() {
        let info = ConvergenceInfo {
            primal_inf: 1e-3,
            dual_inf: 1e-3,
            dual_inf_unscaled: 1e-3,
            compl_inf: 1e-3,
            mu: 1e-11,
            objective: 1e51,
            multiplier_sum: 0.0,
            multiplier_count: 0,
        };
        let opts = SolverOptions::default();
        assert_eq!(
            check_convergence(&info, &opts, 0),
            ConvergenceStatus::Diverging
        );
    }

    #[test]
    fn test_convergence_acceptable() {
        let info = ConvergenceInfo {
            primal_inf: 1e-7,
            dual_inf: 1e-7,
            dual_inf_unscaled: 1e-7,
            compl_inf: 1e-7,
            mu: 1e-8,
            objective: 5.0,
            multiplier_sum: 0.0,
            multiplier_count: 0,
        };
        let opts = SolverOptions::default();
        // Need enough consecutive near-tolerance iterations (hardcoded NEAR_TOL_ITERS=10)
        assert_eq!(
            check_convergence(&info, &opts, 10),
            ConvergenceStatus::Acceptable
        );
    }

    #[test]
    fn test_convergence_acceptable_insufficient_count() {
        let info = ConvergenceInfo {
            primal_inf: 1e-7,
            dual_inf: 1e-7,
            dual_inf_unscaled: 1e-7,
            compl_inf: 1e-7,
            mu: 1e-8,
            objective: 5.0,
            multiplier_sum: 0.0,
            multiplier_count: 0,
        };
        let opts = SolverOptions::default();
        // Not enough consecutive iterations (9 < 10)
        assert_eq!(
            check_convergence(&info, &opts, 9),
            ConvergenceStatus::NotConverged
        );
    }

    #[test]
    fn test_convergence_dual_scaling() {
        // Large multipliers should scale the dual tolerance
        let info = ConvergenceInfo {
            primal_inf: 1e-10,
            dual_inf: 5e-5, // Would fail without scaling
            dual_inf_unscaled: 5e-5,
            compl_inf: 1e-10,
            mu: 1e-11,
            objective: 1.0,
            multiplier_sum: 1e6, // Large multipliers
            multiplier_count: 10,
        };
        let opts = SolverOptions::default();
        // s_d = max(100, 1e6/10)/100 = 1e5/100 = 1000
        // dual_tol = 1e-8 * 1000 = 1e-5
        // 5e-5 > 1e-5, so not converged
        assert_eq!(
            check_convergence(&info, &opts, 0),
            ConvergenceStatus::NotConverged
        );

        // With slightly smaller dual_inf it should pass
        let info2 = ConvergenceInfo {
            dual_inf: 5e-6,
            dual_inf_unscaled: 5e-6,
            ..info
        };
        assert_eq!(
            check_convergence(&info2, &opts, 0),
            ConvergenceStatus::Converged
        );
    }

    #[test]
    fn test_convergence_unscaled_gate_blocks_false_convergence() {
        // z_opt says converged (dual_inf small), but iterative z says not (dual_inf_unscaled large)
        let info = ConvergenceInfo {
            primal_inf: 1e-10,
            dual_inf: 1e-10, // z_opt-based: looks converged
            dual_inf_unscaled: 1.5, // iterative z: clearly not converged (> dual_inf_tol=1.0)
            compl_inf: 1e-10,
            mu: 1e-11,
            objective: 1.0,
            multiplier_sum: 0.0,
            multiplier_count: 0,
        };
        let opts = SolverOptions::default();
        // Should NOT converge: unscaled gate blocks it
        assert_eq!(
            check_convergence(&info, &opts, 0),
            ConvergenceStatus::NotConverged
        );

        // Now with unscaled also passing
        let info2 = ConvergenceInfo {
            dual_inf_unscaled: 0.5, // below dual_inf_tol=1.0
            ..info
        };
        assert_eq!(
            check_convergence(&info2, &opts, 0),
            ConvergenceStatus::Converged
        );
    }

    #[test]
    fn test_complementarity_error_no_bounds() {
        let x = vec![1.0, 2.0];
        let x_l = vec![f64::NEG_INFINITY; 2];
        let x_u = vec![f64::INFINITY; 2];
        let z_l = vec![0.0; 2];
        let z_u = vec![0.0; 2];
        let err = complementarity_error(&x, &x_l, &x_u, &z_l, &z_u, 0.0);
        assert!((err).abs() < 1e-15);
    }

    #[test]
    fn test_complementarity_error_at_optimality() {
        // At optimality: (x - x_l) * z_l = mu
        let mu = 0.01;
        let x = vec![1.1]; // slack = 0.1
        let x_l = vec![1.0];
        let x_u = vec![f64::INFINITY];
        let z_l = vec![mu / 0.1]; // z_l = mu/slack = 0.1
        let z_u = vec![0.0];
        let err = complementarity_error(&x, &x_l, &x_u, &z_l, &z_u, mu);
        assert!(err < 1e-12, "At optimality, complementarity error should be ~0, got {}", err);
    }

    #[test]
    fn test_complementarity_error_away_from_optimality() {
        let mu = 0.01;
        let x = vec![1.5]; // slack = 0.5
        let x_l = vec![1.0];
        let x_u = vec![f64::INFINITY];
        let z_l = vec![1.0]; // z_l * slack = 0.5 >> mu
        let z_u = vec![0.0];
        let err = complementarity_error(&x, &x_l, &x_u, &z_l, &z_u, mu);
        // err = |0.5 * 1.0 - 0.01| = 0.49
        assert!((err - 0.49).abs() < 1e-12);
    }

    #[test]
    fn test_dual_infeasibility_stationarity() {
        // Exact stationarity: grad_f + J^T * lambda - z_l + z_u = 0
        let n = 2;
        let grad_f = vec![1.0, 2.0];
        let jac_rows = vec![0, 0];
        let jac_cols = vec![0, 1];
        let jac_vals = vec![1.0, 1.0]; // J = [1, 1]
        let lambda = vec![-0.5]; // J^T * lambda = [-0.5, -0.5]
        // residual = [1.0 + (-0.5), 2.0 + (-0.5)] - z_l + z_u
        // Need z_l, z_u such that residual = 0
        let z_l = vec![0.5, 1.5];
        let z_u = vec![0.0, 0.0];
        let di = dual_infeasibility(&grad_f, &jac_rows, &jac_cols, &jac_vals, &lambda, &z_l, &z_u, n);
        assert!(di < 1e-12, "Exact stationarity should give 0, got {}", di);

        // Nonzero case
        let z_l2 = vec![0.0, 0.0];
        let di2 = dual_infeasibility(&grad_f, &jac_rows, &jac_cols, &jac_vals, &lambda, &z_l2, &z_u, n);
        assert!(di2 > 0.1, "Non-stationary should give positive dual_inf");
    }

    #[test]
    fn test_dual_infeasibility_scaled_insensitive_to_gradient_magnitude() {
        let n = 2;
        // Large gradient magnitudes
        let grad_f = vec![1e6, 1e6];
        let jac_rows = vec![];
        let jac_cols = vec![];
        let jac_vals: Vec<f64> = vec![];
        let lambda: Vec<f64> = vec![];
        // Residual = grad_f - z_l + z_u = [1e6 - 0, 1e6 - 0] = [1e6, 1e6]
        let z_l = vec![0.0, 0.0];
        let z_u = vec![0.0, 0.0];

        // Unscaled: max |r_i| = 1e6
        let di = dual_infeasibility(&grad_f, &jac_rows, &jac_cols, &jac_vals, &lambda, &z_l, &z_u, n);
        assert!(di > 1e5, "Unscaled should be large: {}", di);

        // Scaled: max |r_i| / (1 + |grad_f_i|) ≈ 1e6 / (1 + 1e6) ≈ 1.0
        let di_s = dual_infeasibility_scaled(&grad_f, &jac_rows, &jac_cols, &jac_vals, &lambda, &z_l, &z_u, n);
        assert!(di_s < 1.1, "Scaled should be ~1.0: {}", di_s);
    }

    #[test]
    fn test_dual_infeasibility_scaled_stationarity() {
        let n = 2;
        let grad_f = vec![1.0, 2.0];
        let jac_rows = vec![0, 0];
        let jac_cols = vec![0, 1];
        let jac_vals = vec![1.0, 1.0];
        let lambda = vec![-0.5];
        let z_l = vec![0.5, 1.5];
        let z_u = vec![0.0, 0.0];
        // At stationarity, both scaled and unscaled should be ~0
        let di_s = dual_infeasibility_scaled(&grad_f, &jac_rows, &jac_cols, &jac_vals, &lambda, &z_l, &z_u, n);
        assert!(di_s < 1e-12, "Scaled stationarity should give 0, got {}", di_s);
    }
}