aprender/optim/

mod.rs

//! Optimization algorithms for gradient-based learning.
//!
//! This module provides both stochastic (mini-batch) and batch (deterministic) optimizers
//! following the unified [`Optimizer`] trait architecture.
//!
//! # Available Optimizers
//!
//! ## Stochastic (Mini-Batch) Optimizers
//! - [`SGD`] - Stochastic Gradient Descent with optional momentum
//! - [`Adam`] - Adaptive Moment Estimation (adaptive learning rates)
//!
//! ## Batch (Deterministic) Optimizers
//! - [`LBFGS`] - Limited-memory BFGS (memory-efficient quasi-Newton)
//! - [`ConjugateGradient`] - Conjugate Gradient with three beta formulas
//! - [`DampedNewton`] - Newton's method with automatic damping for stability
//!
//! ## Convex Optimization (Phase 2)
//! - [`FISTA`] - Fast Iterative Shrinkage-Thresholding (proximal gradient)
//! - [`CoordinateDescent`] - Coordinate-wise optimization for high dimensions
//! - [`ADMM`] - Alternating Direction Method of Multipliers (distributed ML)
//!
//! ## Constrained Optimization (Phase 3)
//! - [`ProjectedGradientDescent`] - Projection onto convex sets
//! - [`AugmentedLagrangian`] - Equality-constrained optimization
//! - [`InteriorPoint`] - Inequality-constrained optimization (log-barrier)
//!
//! ## Utility Functions
//! - [`safe_cholesky_solve`] - Cholesky solver with automatic regularization
//!
//! # Stochastic Optimization (Mini-Batch)
//!
//! Stochastic optimizers update parameters incrementally using mini-batch gradients.
//! Use the [`Optimizer::step`] method for parameter updates:
//!
//! ```
//! use aprender::optim::SGD;
//! use aprender::primitives::Vector;
//!
//! // Create optimizer with learning rate 0.01
//! let mut optimizer = SGD::new(0.01);
//!
//! // Initialize parameters and gradients
//! let mut params = Vector::from_slice(&[1.0, 2.0, 3.0]);
//! let gradients = Vector::from_slice(&[0.1, 0.2, 0.3]);
//!
//! // Update parameters
//! optimizer.step(&mut params, &gradients);
//!
//! // Parameters are updated: params = params - lr * gradients
//! assert!((params[0] - 0.999).abs() < 1e-6);
//! ```
//!
//! # Batch Optimization (Full Dataset)
//!
//! Batch optimizers minimize objective functions using full dataset access.
//! They use the `minimize` method which returns detailed convergence information:
//!
//! ```
//! use aprender::optim::{LBFGS, ConvergenceStatus, Optimizer};
//! use aprender::primitives::Vector;
//!
//! // Create L-BFGS optimizer: 100 max iterations, 1e-5 tolerance, 10 memory size
//! let mut optimizer = LBFGS::new(100, 1e-5, 10);
//!
//! // Define objective and gradient functions
//! let objective = |x: &Vector<f32>| (x[0] - 5.0).powi(2) + (x[1] - 3.0).powi(2);
//! let gradient = |x: &Vector<f32>| {
//!     Vector::from_slice(&[2.0 * (x[0] - 5.0), 2.0 * (x[1] - 3.0)])
//! };
//!
//! let x0 = Vector::from_slice(&[0.0, 0.0]);
//! let result = optimizer.minimize(objective, gradient, x0);
//!
//! assert_eq!(result.status, ConvergenceStatus::Converged);
//! assert!((result.solution[0] - 5.0).abs() < 1e-4);
//! assert!((result.solution[1] - 3.0).abs() < 1e-4);
//! ```
//!
//! # See Also
//!
//! - [`examples/batch_optimization.rs`](https://github.com/paiml/aprender/tree/main/examples/batch_optimization.rs) - Comprehensive examples
//! - Specification: `docs/specifications/comprehensive-optimization-spec.md`

use serde::{Deserialize, Serialize};

use crate::primitives::{Matrix, Vector};

/// Result of an optimization procedure.
///
/// Contains the final solution, convergence information, and diagnostic metrics.
#[derive(Debug, Clone)]
pub struct OptimizationResult {
    /// Final solution (optimized parameters)
    pub solution: Vector<f32>,
    /// Final objective function value
    pub objective_value: f32,
    /// Number of iterations performed
    pub iterations: usize,
    /// Convergence status
    pub status: ConvergenceStatus,
    /// Final gradient norm (‖∇f(x)‖)
    pub gradient_norm: f32,
    /// Constraint violation (0.0 for unconstrained problems)
    pub constraint_violation: f32,
    /// Total elapsed time
    pub elapsed_time: std::time::Duration,
}

impl OptimizationResult {
    /// Creates a converged result.
    #[must_use]
    pub fn converged(solution: Vector<f32>, iterations: usize) -> Self {
        Self {
            solution,
            objective_value: 0.0,
            iterations,
            status: ConvergenceStatus::Converged,
            gradient_norm: 0.0,
            constraint_violation: 0.0,
            elapsed_time: std::time::Duration::ZERO,
        }
    }

    /// Creates a max-iterations result.
    #[must_use]
    pub fn max_iterations(solution: Vector<f32>) -> Self {
        Self {
            solution,
            objective_value: 0.0,
            iterations: 0,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: 0.0,
            constraint_violation: 0.0,
            elapsed_time: std::time::Duration::ZERO,
        }
    }
}

/// Convergence status of an optimization procedure.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ConvergenceStatus {
    /// Converged (gradient norm < tolerance)
    Converged,
    /// Reached maximum iteration limit
    MaxIterations,
    /// Progress stalled (step size too small)
    Stalled,
    /// Numerical error (NaN, Inf, etc.)
    NumericalError,
    /// Optimization still running
    Running,
    /// User-requested termination
    UserTerminated,
}

/// Safely solves a linear system Ax = b using Cholesky decomposition with automatic regularization.
///
/// When the matrix is not positive definite or near-singular, this function automatically
/// adds regularization (λI) to make the system solvable. This is essential for numerical
/// stability in second-order optimization methods like Newton's method.
///
/// # Algorithm
///
/// 1. Try standard Cholesky solve with A
/// 2. If it fails, add regularization: (A + λI)
/// 3. Increase λ progressively (×10) until solve succeeds or max attempts reached
///
/// # Arguments
///
/// * `A` - Symmetric matrix (should be positive definite, but may not be)
/// * `b` - Right-hand side vector
/// * `initial_lambda` - Initial regularization parameter (typical: 1e-8)
/// * `max_attempts` - Maximum regularization attempts (typical: 10)
///
/// # Returns
///
/// * `Ok(x)` - Solution vector if successful
/// * `Err(msg)` - Error message if regularization failed
///
/// # Example
///
/// ```
/// use aprender::optim::safe_cholesky_solve;
/// use aprender::primitives::{Matrix, Vector};
///
/// // Slightly ill-conditioned matrix
/// let A = Matrix::from_vec(2, 2, vec![1.0, 0.0, 0.0, 1e-8]).expect("valid dimensions");
/// let b = Vector::from_slice(&[1.0, 1.0]);
///
/// let x = safe_cholesky_solve(&A, &b, 1e-8, 10).expect("should solve with regularization");
/// assert_eq!(x.len(), 2);
/// ```
///
/// # Use Cases
///
/// - **Damped Newton**: Regularize Hessian when not positive definite
/// - **Levenberg-Marquardt**: Add damping parameter for stability
/// - **Trust region**: Solve trust region subproblem with ill-conditioned Hessian
///
/// # References
///
/// - Nocedal & Wright (2006), *Numerical Optimization*, Chapter 3
pub fn safe_cholesky_solve(
    a: &Matrix<f32>,
    b: &Vector<f32>,
    initial_lambda: f32,
    max_attempts: usize,
) -> Result<Vector<f32>, &'static str> {
    // First try without regularization
    if let Ok(x) = a.cholesky_solve(b) {
        return Ok(x);
    }

    // Matrix not positive definite, try with regularization
    solve_with_regularization(a, b, initial_lambda, max_attempts)
}

/// Solve linear system with Tikhonov regularization (A + λI)x = b.
fn solve_with_regularization(
    a: &Matrix<f32>,
    b: &Vector<f32>,
    initial_lambda: f32,
    max_attempts: usize,
) -> Result<Vector<f32>, &'static str> {
    let n = a.n_rows();
    let mut lambda = initial_lambda;

    for _attempt in 0..max_attempts {
        let a_reg = create_regularized_matrix(a, n, lambda);

        if let Ok(x) = a_reg.cholesky_solve(b) {
            return Ok(x);
        }

        lambda *= 10.0;
        if lambda > 1e6 {
            return Err(
                "Cholesky solve failed: matrix too ill-conditioned even with regularization",
            );
        }
    }

    Err("Cholesky solve failed after maximum regularization attempts")
}

/// Create regularized matrix `A_reg` = A + λI.
fn create_regularized_matrix(a: &Matrix<f32>, n: usize, lambda: f32) -> Matrix<f32> {
    let mut a_reg_data = vec![0.0; n * n];
    for i in 0..n {
        for j in 0..n {
            let diagonal_term = if i == j { lambda } else { 0.0 };
            a_reg_data[i * n + j] = a.get(i, j) + diagonal_term;
        }
    }
    Matrix::from_vec(n, n, a_reg_data).expect("Matrix dimensions should be valid")
}

/// Line search strategy for determining step size in batch optimization.
///
/// Line search methods find an appropriate step size α along a search direction d
/// by solving the 1D optimization problem:
///
/// ```text
/// minimize f(x + α*d) over α > 0
/// ```
///
/// Different strategies enforce different conditions on the step size.
pub trait LineSearch {
    /// Finds a suitable step size along the search direction.
    ///
    /// # Arguments
    ///
    /// * `f` - Objective function f: ℝⁿ → ℝ
    /// * `grad` - Gradient function ∇f: ℝⁿ → ℝⁿ
    /// * `x` - Current point
    /// * `d` - Search direction (typically descent direction, ∇f(x)·d < 0)
    ///
    /// # Returns
    ///
    /// Step size α > 0 satisfying the line search conditions
    fn search<F, G>(&self, f: &F, grad: &G, x: &Vector<f32>, d: &Vector<f32>) -> f32
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>;
}

/// Backtracking line search with Armijo condition.
///
/// Starts with step size α = 1 and repeatedly shrinks it by factor ρ until
/// the Armijo condition is satisfied:
///
/// ```text
/// f(x + α*d) ≤ f(x) + c₁*α*∇f(x)ᵀd
/// ```
///
/// This ensures sufficient decrease in the objective function.
///
/// # Parameters
///
/// - **c1**: Armijo constant (typical: 1e-4), controls acceptable decrease
/// - **rho**: Backtracking factor (typical: 0.5), shrinkage rate for α
/// - **`max_iter`**: Maximum backtracking iterations (safety limit)
///
/// # Example
///
/// ```
/// use aprender::optim::{BacktrackingLineSearch, LineSearch};
/// use aprender::primitives::Vector;
///
/// let line_search = BacktrackingLineSearch::new(1e-4, 0.5, 50);
///
/// // Define a simple quadratic function
/// let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
/// let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);
///
/// let x = Vector::from_slice(&[1.0, 1.0]);
/// let d = Vector::from_slice(&[-2.0, -2.0]); // Descent direction
///
/// let alpha = line_search.search(&f, &grad, &x, &d);
/// assert!(alpha > 0.0);
/// ```
#[derive(Debug, Clone)]
pub struct BacktrackingLineSearch {
    /// Armijo constant (c₁ ∈ (0, 1), typical: 1e-4)
    c1: f32,
    /// Backtracking factor (ρ ∈ (0, 1), typical: 0.5)
    rho: f32,
    /// Maximum backtracking iterations
    max_iter: usize,
}

impl BacktrackingLineSearch {
    /// Creates a new backtracking line search.
    ///
    /// # Arguments
    ///
    /// * `c1` - Armijo constant (typical: 1e-4)
    /// * `rho` - Backtracking factor (typical: 0.5)
    /// * `max_iter` - Maximum iterations (typical: 50)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::BacktrackingLineSearch;
    ///
    /// let line_search = BacktrackingLineSearch::new(1e-4, 0.5, 50);
    /// ```
    #[must_use]
    pub fn new(c1: f32, rho: f32, max_iter: usize) -> Self {
        Self { c1, rho, max_iter }
    }
}

impl Default for BacktrackingLineSearch {
    /// Creates a backtracking line search with default parameters.
    ///
    /// Defaults: c1=1e-4, rho=0.5, `max_iter=50`
    fn default() -> Self {
        Self::new(1e-4, 0.5, 50)
    }
}

impl LineSearch for BacktrackingLineSearch {
    fn search<F, G>(&self, f: &F, grad: &G, x: &Vector<f32>, d: &Vector<f32>) -> f32
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
    {
        let mut alpha = 1.0;
        let fx = f(x);
        let grad_x = grad(x);

        // Compute directional derivative: ∇f(x)ᵀd
        let mut dir_deriv = 0.0;
        for i in 0..x.len() {
            dir_deriv += grad_x[i] * d[i];
        }

        // Backtracking loop
        for _ in 0..self.max_iter {
            // Compute x_new = x + alpha * d
            let mut x_new = Vector::zeros(x.len());
            for i in 0..x.len() {
                x_new[i] = x[i] + alpha * d[i];
            }

            let fx_new = f(&x_new);

            // Check Armijo condition: f(x + α*d) ≤ f(x) + c₁*α*∇f(x)ᵀd
            if fx_new <= fx + self.c1 * alpha * dir_deriv {
                return alpha;
            }

            // Shrink step size
            alpha *= self.rho;
        }

        // Return the last alpha if max iterations reached
        alpha
    }
}

/// Wolfe line search with Armijo and curvature conditions.
///
/// Enforces both the Armijo condition (sufficient decrease) and the curvature
/// condition (sufficient curvature):
///
/// ```text
/// Armijo:    f(x + α*d) ≤ f(x) + c₁*α*∇f(x)ᵀd
/// Curvature: |∇f(x + α*d)ᵀd| ≤ c₂*|∇f(x)ᵀd|
/// ```
///
/// The curvature condition ensures the step size is not too small by requiring
/// that the gradient has decreased sufficiently along the search direction.
///
/// # Parameters
///
/// - **c1**: Armijo constant (typical: 1e-4), c₁ ∈ (0, c₂)
/// - **c2**: Curvature constant (typical: 0.9), c₂ ∈ (c₁, 1)
/// - **`max_iter`**: Maximum line search iterations
///
/// # Example
///
/// ```
/// use aprender::optim::{WolfeLineSearch, LineSearch};
/// use aprender::primitives::Vector;
///
/// let line_search = WolfeLineSearch::new(1e-4, 0.9, 50);
///
/// let f = |x: &Vector<f32>| x[0] * x[0];
/// let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);
///
/// let x = Vector::from_slice(&[1.0]);
/// let d = Vector::from_slice(&[-2.0]);
///
/// let alpha = line_search.search(&f, &grad, &x, &d);
/// assert!(alpha > 0.0);
/// ```
#[derive(Debug, Clone)]
pub struct WolfeLineSearch {
    /// Armijo constant (c₁ ∈ (0, c₂), typical: 1e-4)
    c1: f32,
    /// Curvature constant (c₂ ∈ (c₁, 1), typical: 0.9)
    c2: f32,
    /// Maximum line search iterations
    max_iter: usize,
}

impl WolfeLineSearch {
    /// Creates a new Wolfe line search.
    ///
    /// # Arguments
    ///
    /// * `c1` - Armijo constant (typical: 1e-4)
    /// * `c2` - Curvature constant (typical: 0.9)
    /// * `max_iter` - Maximum iterations (typical: 50)
    ///
    /// # Panics
    ///
    /// Panics if c1 >= c2 or values are outside (0, 1).
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::WolfeLineSearch;
    ///
    /// let line_search = WolfeLineSearch::new(1e-4, 0.9, 50);
    /// ```
    #[must_use]
    pub fn new(c1: f32, c2: f32, max_iter: usize) -> Self {
        assert!(
            c1 < c2 && c1 > 0.0 && c2 < 1.0,
            "Wolfe conditions require 0 < c1 < c2 < 1"
        );
        Self { c1, c2, max_iter }
    }
}

impl Default for WolfeLineSearch {
    /// Creates a Wolfe line search with default parameters.
    ///
    /// Defaults: c1=1e-4, c2=0.9, `max_iter=50`
    fn default() -> Self {
        Self::new(1e-4, 0.9, 50)
    }
}

impl LineSearch for WolfeLineSearch {
    fn search<F, G>(&self, f: &F, grad: &G, x: &Vector<f32>, d: &Vector<f32>) -> f32
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
    {
        let fx = f(x);
        let grad_x = grad(x);

        // Compute directional derivative: ∇f(x)ᵀd
        let mut dir_deriv = 0.0;
        for i in 0..x.len() {
            dir_deriv += grad_x[i] * d[i];
        }

        // Start with alpha = 1.0
        let mut alpha = 1.0;
        let mut alpha_lo = 0.0;
        let mut alpha_hi = f32::INFINITY;

        for _ in 0..self.max_iter {
            // Compute x_new = x + alpha * d
            let mut x_new = Vector::zeros(x.len());
            for i in 0..x.len() {
                x_new[i] = x[i] + alpha * d[i];
            }

            let fx_new = f(&x_new);
            let grad_new = grad(&x_new);

            // Compute new directional derivative
            let mut dir_deriv_new = 0.0;
            for i in 0..x.len() {
                dir_deriv_new += grad_new[i] * d[i];
            }

            // Check Armijo condition
            if fx_new > fx + self.c1 * alpha * dir_deriv {
                // Armijo fails - alpha too large
                alpha_hi = alpha;
                alpha = (alpha_lo + alpha_hi) / 2.0;
                continue;
            }

            // Check curvature condition: |∇f(x + α*d)ᵀd| ≤ c₂*|∇f(x)ᵀd|
            if dir_deriv_new.abs() <= self.c2 * dir_deriv.abs() {
                // Both conditions satisfied
                return alpha;
            }

            // Curvature condition fails
            if dir_deriv_new > 0.0 {
                // Gradient sign changed - reduce alpha
                alpha_hi = alpha;
            } else {
                // Gradient still negative - increase alpha
                alpha_lo = alpha;
            }

            // Update alpha
            if alpha_hi.is_finite() {
                alpha = (alpha_lo + alpha_hi) / 2.0;
            } else {
                alpha *= 2.0;
            }
        }

        // Return the last alpha if max iterations reached
        alpha
    }
}

/// Limited-memory BFGS (L-BFGS) optimizer.
///
/// L-BFGS is a quasi-Newton method that approximates the inverse Hessian using
/// a limited history of gradient information. It's efficient for large-scale
/// optimization problems where storing the full Hessian is infeasible.
///
/// # Algorithm
///
/// 1. Compute gradient `g_k` = ∇`f(x_k)`
/// 2. Compute search direction `d_k` using two-loop recursion (approximates H^(-1) * `g_k`)
/// 3. Find step size `α_k` via line search (Wolfe conditions)
/// 4. Update: x_{k+1} = `x_k` - `α_k` * `d_k`
/// 5. Store gradient and position differences for next iteration
///
/// # Parameters
///
/// - **`max_iter`**: Maximum number of iterations
/// - **tol**: Convergence tolerance (gradient norm)
/// - **m**: History size (typically 5-20, tradeoff between memory and convergence)
///
/// # Example
///
/// ```
/// use aprender::optim::{LBFGS, Optimizer};
/// use aprender::primitives::Vector;
///
/// let mut optimizer = LBFGS::new(100, 1e-5, 10);
///
/// // Define Rosenbrock function and its gradient
/// let f = |x: &Vector<f32>| {
///     let a = x[0];
///     let b = x[1];
///     (1.0 - a).powi(2) + 100.0 * (b - a * a).powi(2)
/// };
///
/// let grad = |x: &Vector<f32>| {
///     let a = x[0];
///     let b = x[1];
///     Vector::from_slice(&[
///         -2.0 * (1.0 - a) - 400.0 * a * (b - a * a),
///         200.0 * (b - a * a),
///     ])
/// };
///
/// let x0 = Vector::from_slice(&[0.0, 0.0]);
/// let result = optimizer.minimize(f, grad, x0);
///
/// // Should converge to (1, 1)
/// assert_eq!(result.status, aprender::optim::ConvergenceStatus::Converged);
/// ```
#[derive(Debug, Clone)]
pub struct LBFGS {
    /// Maximum number of iterations
    max_iter: usize,
    /// Convergence tolerance (gradient norm)
    tol: f32,
    /// History size (number of correction pairs to store)
    m: usize,
    /// Line search strategy
    line_search: WolfeLineSearch,
    /// Position differences: `s_k` = x_{k+1} - `x_k`
    s_history: Vec<Vector<f32>>,
    /// Gradient differences: `y_k` = g_{k+1} - `g_k`
    y_history: Vec<Vector<f32>>,
}

impl LBFGS {
    /// Creates a new L-BFGS optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of iterations (typical: 100-1000)
    /// * `tol` - Convergence tolerance for gradient norm (typical: 1e-5)
    /// * `m` - History size (typical: 5-20)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::LBFGS;
    ///
    /// let optimizer = LBFGS::new(100, 1e-5, 10);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, tol: f32, m: usize) -> Self {
        Self {
            max_iter,
            tol,
            m,
            line_search: WolfeLineSearch::new(1e-4, 0.9, 50),
            s_history: Vec::with_capacity(m),
            y_history: Vec::with_capacity(m),
        }
    }

    /// Two-loop recursion to compute search direction.
    ///
    /// Approximates H^(-1) * grad where H is the Hessian.
    /// Uses stored history of s (position diff) and y (gradient diff).
    fn compute_direction(&self, grad: &Vector<f32>) -> Vector<f32> {
        let n = grad.len();
        let k = self.s_history.len();

        if k == 0 {
            // No history: use steepest descent
            let mut d = Vector::zeros(n);
            for i in 0..n {
                d[i] = -grad[i];
            }
            return d;
        }

        // Initialize q = -grad
        let mut q = Vector::zeros(n);
        for i in 0..n {
            q[i] = -grad[i];
        }

        let mut alpha = vec![0.0; k];
        let mut rho = vec![0.0; k];

        // First loop: backward pass
        for i in (0..k).rev() {
            let s = &self.s_history[i];
            let y = &self.y_history[i];

            // rho_i = 1 / (y_i^T s_i)
            let mut y_dot_s = 0.0;
            for j in 0..n {
                y_dot_s += y[j] * s[j];
            }
            rho[i] = 1.0 / y_dot_s;

            // alpha_i = rho_i * s_i^T * q
            let mut s_dot_q = 0.0;
            for j in 0..n {
                s_dot_q += s[j] * q[j];
            }
            alpha[i] = rho[i] * s_dot_q;

            // q = q - alpha_i * y_i
            for j in 0..n {
                q[j] -= alpha[i] * y[j];
            }
        }

        // Scale by H_0 = (s^T y) / (y^T y) from most recent update
        let s_last = &self.s_history[k - 1];
        let y_last = &self.y_history[k - 1];

        let mut s_dot_y = 0.0;
        let mut y_dot_y = 0.0;
        for i in 0..n {
            s_dot_y += s_last[i] * y_last[i];
            y_dot_y += y_last[i] * y_last[i];
        }
        let gamma = s_dot_y / y_dot_y;

        // r = H_0 * q = gamma * q
        let mut r = Vector::zeros(n);
        for i in 0..n {
            r[i] = gamma * q[i];
        }

        // Second loop: forward pass
        for i in 0..k {
            let s = &self.s_history[i];
            let y = &self.y_history[i];

            // beta = rho_i * y_i^T * r
            let mut y_dot_r = 0.0;
            for j in 0..n {
                y_dot_r += y[j] * r[j];
            }
            let beta = rho[i] * y_dot_r;

            // r = r + s_i * (alpha_i - beta)
            for j in 0..n {
                r[j] += s[j] * (alpha[i] - beta);
            }
        }

        r
    }

    /// Computes the L2 norm of a vector.
    fn norm(v: &Vector<f32>) -> f32 {
        let mut sum = 0.0;
        for i in 0..v.len() {
            sum += v[i] * v[i];
        }
        sum.sqrt()
    }
}

impl Optimizer for LBFGS {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "L-BFGS does not support stochastic updates (step). Use minimize() for batch optimization."
        )
    }

    fn minimize<F, G>(&mut self, objective: F, gradient: G, x0: Vector<f32>) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
    {
        let start_time = std::time::Instant::now();
        let n = x0.len();

        // Clear history from previous runs
        self.s_history.clear();
        self.y_history.clear();

        let mut x = x0;
        let mut fx = objective(&x);
        let mut grad = gradient(&x);
        let mut grad_norm = Self::norm(&grad);

        for iter in 0..self.max_iter {
            // Check convergence
            if grad_norm < self.tol {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Compute search direction
            let d = self.compute_direction(&grad);

            // Line search
            let alpha = self.line_search.search(&objective, &gradient, &x, &d);

            // Check for stalled progress
            if alpha < 1e-12 {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::Stalled,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Update position: x_new = x + alpha * d
            let mut x_new = Vector::zeros(n);
            for i in 0..n {
                x_new[i] = x[i] + alpha * d[i];
            }

            // Compute new objective and gradient
            let fx_new = objective(&x_new);
            let grad_new = gradient(&x_new);

            // Check for numerical errors
            if fx_new.is_nan() || fx_new.is_infinite() {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::NumericalError,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Compute s_k = x_new - x and y_k = grad_new - grad
            let mut s_k = Vector::zeros(n);
            let mut y_k = Vector::zeros(n);
            for i in 0..n {
                s_k[i] = x_new[i] - x[i];
                y_k[i] = grad_new[i] - grad[i];
            }

            // Check curvature condition: y^T s > 0
            let mut y_dot_s = 0.0;
            for i in 0..n {
                y_dot_s += y_k[i] * s_k[i];
            }

            if y_dot_s > 1e-10 {
                // Store in history
                if self.s_history.len() >= self.m {
                    self.s_history.remove(0);
                    self.y_history.remove(0);
                }
                self.s_history.push(s_k);
                self.y_history.push(y_k);
            }

            // Update for next iteration
            x = x_new;
            fx = fx_new;
            grad = grad_new;
            grad_norm = Self::norm(&grad);
        }

        // Max iterations reached
        OptimizationResult {
            solution: x,
            objective_value: fx,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: grad_norm,
            constraint_violation: 0.0,
            elapsed_time: start_time.elapsed(),
        }
    }

    fn reset(&mut self) {
        self.s_history.clear();
        self.y_history.clear();
    }
}

/// Beta computation formula for Conjugate Gradient.
///
/// Different formulas provide different convergence properties and
/// numerical stability characteristics.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CGBetaFormula {
    /// Fletcher-Reeves: β = (g_{k+1}^T g_{k+1}) / (`g_k^T` `g_k`)
    ///
    /// Most stable but can be slow on non-quadratic problems.
    FletcherReeves,
    /// Polak-Ribière: β = g_{k+1}^T (g_{k+1} - `g_k`) / (`g_k^T` `g_k`)
    ///
    /// Better performance than FR, includes automatic restart (β < 0).
    PolakRibiere,
    /// Hestenes-Stiefel: β = g_{k+1}^T (g_{k+1} - `g_k`) / (`d_k^T` (g_{k+1} - `g_k`))
    ///
    /// Similar to PR but with different denominator, can be more robust.
    HestenesStiefel,
}

/// Nonlinear Conjugate Gradient (CG) optimizer.
///
/// CG is an iterative method for solving optimization problems that uses
/// conjugate directions rather than steepest descent. It's particularly
/// effective for quadratic problems but extends to general nonlinear optimization.
///
/// # Algorithm
///
/// 1. Initialize with steepest descent: `d_0` = -∇`f(x_0)`
/// 2. Line search: find `α_k` minimizing `f(x_k` + `α_k` `d_k`)
/// 3. Update: x_{k+1} = `x_k` + `α_k` `d_k`
/// 4. Compute β_{k+1} using chosen formula (FR, PR, or HS)
/// 5. Update direction: d_{k+1} = -∇f(x_{k+1}) + β_{k+1} `d_k`
/// 6. Restart if β < 0 or every n iterations
///
/// # Parameters
///
/// - **`max_iter`**: Maximum number of iterations
/// - **tol**: Convergence tolerance (gradient norm)
/// - **`beta_formula`**: Method for computing β (FR, PR, or HS)
/// - **`restart_interval`**: Restart with steepest descent every n iterations (0 = no periodic restart)
///
/// # Example
///
/// ```
/// use aprender::optim::{ConjugateGradient, CGBetaFormula, Optimizer};
/// use aprender::primitives::Vector;
///
/// let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
///
/// // Minimize Rosenbrock function
/// let f = |x: &Vector<f32>| {
///     let a = x[0];
///     let b = x[1];
///     (1.0 - a).powi(2) + 100.0 * (b - a * a).powi(2)
/// };
///
/// let grad = |x: &Vector<f32>| {
///     let a = x[0];
///     let b = x[1];
///     Vector::from_slice(&[
///         -2.0 * (1.0 - a) - 400.0 * a * (b - a * a),
///         200.0 * (b - a * a),
///     ])
/// };
///
/// let x0 = Vector::from_slice(&[0.0, 0.0]);
/// let result = optimizer.minimize(f, grad, x0);
/// ```
#[derive(Debug, Clone)]
pub struct ConjugateGradient {
    /// Maximum number of iterations
    max_iter: usize,
    /// Convergence tolerance (gradient norm)
    tol: f32,
    /// Beta computation formula
    beta_formula: CGBetaFormula,
    /// Restart interval (0 = no periodic restart, only on β < 0)
    restart_interval: usize,
    /// Line search strategy
    line_search: WolfeLineSearch,
    /// Previous search direction (for conjugacy)
    prev_direction: Option<Vector<f32>>,
    /// Previous gradient (for beta computation)
    prev_gradient: Option<Vector<f32>>,
    /// Iteration counter (for restart)
    iter_count: usize,
}

impl ConjugateGradient {
    /// Creates a new Conjugate Gradient optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of iterations (typical: 100-1000)
    /// * `tol` - Convergence tolerance for gradient norm (typical: 1e-5)
    /// * `beta_formula` - Method for computing β (`FletcherReeves`, `PolakRibiere`, or `HestenesStiefel`)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::{ConjugateGradient, CGBetaFormula};
    ///
    /// let optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, tol: f32, beta_formula: CGBetaFormula) -> Self {
        Self {
            max_iter,
            tol,
            beta_formula,
            restart_interval: 0, // No periodic restart by default
            line_search: WolfeLineSearch::new(1e-4, 0.1, 50), // c2=0.1 for CG (more exact line search)
            prev_direction: None,
            prev_gradient: None,
            iter_count: 0,
        }
    }

    /// Sets the restart interval.
    ///
    /// CG will restart with steepest descent every n iterations.
    /// Setting to 0 disables periodic restart (only restarts on β < 0).
    ///
    /// # Arguments
    ///
    /// * `interval` - Number of iterations between restarts (typical: n, where n is problem dimension)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::{ConjugateGradient, CGBetaFormula};
    ///
    /// let optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere)
    ///     .with_restart_interval(50);
    /// ```
    #[must_use]
    pub fn with_restart_interval(mut self, interval: usize) -> Self {
        self.restart_interval = interval;
        self
    }

    /// Computes beta coefficient based on the chosen formula.
    fn compute_beta(
        &self,
        grad_new: &Vector<f32>,
        grad_old: &Vector<f32>,
        d_old: &Vector<f32>,
    ) -> f32 {
        let n = grad_new.len();

        match self.beta_formula {
            CGBetaFormula::FletcherReeves => {
                // β = (g_{k+1}^T g_{k+1}) / (g_k^T g_k)
                let mut numerator = 0.0;
                let mut denominator = 0.0;
                for i in 0..n {
                    numerator += grad_new[i] * grad_new[i];
                    denominator += grad_old[i] * grad_old[i];
                }
                numerator / denominator.max(1e-12)
            }
            CGBetaFormula::PolakRibiere => {
                // β = g_{k+1}^T (g_{k+1} - g_k) / (g_k^T g_k)
                let mut numerator = 0.0;
                let mut denominator = 0.0;
                for i in 0..n {
                    numerator += grad_new[i] * (grad_new[i] - grad_old[i]);
                    denominator += grad_old[i] * grad_old[i];
                }
                let beta = numerator / denominator.max(1e-12);
                // PR has automatic restart: if β < 0, restart with steepest descent
                beta.max(0.0)
            }
            CGBetaFormula::HestenesStiefel => {
                // β = g_{k+1}^T (g_{k+1} - g_k) / (d_k^T (g_{k+1} - g_k))
                let mut numerator = 0.0;
                let mut denominator = 0.0;
                for i in 0..n {
                    let y_i = grad_new[i] - grad_old[i];
                    numerator += grad_new[i] * y_i;
                    denominator += d_old[i] * y_i;
                }
                let beta = numerator / denominator.max(1e-12);
                beta.max(0.0)
            }
        }
    }

    /// Computes the L2 norm of a vector.
    fn norm(v: &Vector<f32>) -> f32 {
        let mut sum = 0.0;
        for i in 0..v.len() {
            sum += v[i] * v[i];
        }
        sum.sqrt()
    }
}

impl Optimizer for ConjugateGradient {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "Conjugate Gradient does not support stochastic updates (step). Use minimize() for batch optimization."
        )
    }

    #[allow(clippy::too_many_lines)]
    fn minimize<F, G>(&mut self, objective: F, gradient: G, x0: Vector<f32>) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
    {
        let start_time = std::time::Instant::now();
        let n = x0.len();

        // Reset state
        self.prev_direction = None;
        self.prev_gradient = None;
        self.iter_count = 0;

        let mut x = x0;
        let mut fx = objective(&x);
        let mut grad = gradient(&x);
        let mut grad_norm = Self::norm(&grad);

        for iter in 0..self.max_iter {
            // Check convergence
            if grad_norm < self.tol {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Compute search direction
            let d = if let (Some(d_old), Some(g_old)) = (&self.prev_direction, &self.prev_gradient)
            {
                // Check if we need to restart
                let need_restart = if self.restart_interval > 0 {
                    self.iter_count % self.restart_interval == 0
                } else {
                    false
                };

                if need_restart {
                    // Restart with steepest descent
                    let mut d_new = Vector::zeros(n);
                    for i in 0..n {
                        d_new[i] = -grad[i];
                    }
                    d_new
                } else {
                    // Compute beta and conjugate direction
                    let beta = self.compute_beta(&grad, g_old, d_old);

                    // d = -grad + beta * d_old
                    let mut d_new = Vector::zeros(n);
                    for i in 0..n {
                        d_new[i] = -grad[i] + beta * d_old[i];
                    }

                    // Check if direction is descent (grad^T d < 0)
                    let mut grad_dot_d = 0.0;
                    for i in 0..n {
                        grad_dot_d += grad[i] * d_new[i];
                    }

                    if grad_dot_d >= 0.0 {
                        // Not a descent direction - restart with steepest descent
                        for i in 0..n {
                            d_new[i] = -grad[i];
                        }
                    }

                    d_new
                }
            } else {
                // First iteration: use steepest descent
                let mut d_new = Vector::zeros(n);
                for i in 0..n {
                    d_new[i] = -grad[i];
                }
                d_new
            };

            // Line search
            let alpha = self.line_search.search(&objective, &gradient, &x, &d);

            // Check for stalled progress
            if alpha < 1e-12 {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::Stalled,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Update position: x_new = x + alpha * d
            let mut x_new = Vector::zeros(n);
            for i in 0..n {
                x_new[i] = x[i] + alpha * d[i];
            }

            // Compute new objective and gradient
            let fx_new = objective(&x_new);
            let grad_new = gradient(&x_new);

            // Check for numerical errors
            if fx_new.is_nan() || fx_new.is_infinite() {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::NumericalError,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Store current direction and gradient for next iteration
            self.prev_direction = Some(d);
            self.prev_gradient = Some(grad);

            // Update for next iteration
            x = x_new;
            fx = fx_new;
            grad = grad_new;
            grad_norm = Self::norm(&grad);
            self.iter_count += 1;
        }

        // Max iterations reached
        OptimizationResult {
            solution: x,
            objective_value: fx,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: grad_norm,
            constraint_violation: 0.0,
            elapsed_time: start_time.elapsed(),
        }
    }

    fn reset(&mut self) {
        self.prev_direction = None;
        self.prev_gradient = None;
        self.iter_count = 0;
    }
}

/// Damped Newton optimizer with finite-difference Hessian approximation.
///
/// Newton's method uses second-order information (Hessian) to find the minimum
/// by solving the linear system: H * d = -g, where H is the Hessian and g is
/// the gradient. The damping factor and line search ensure global convergence.
///
/// # Algorithm
///
/// 1. Compute gradient g = ∇f(x)
/// 2. Approximate Hessian H using finite differences
/// 3. Solve H * d = -g using Cholesky decomposition
/// 4. If Hessian not positive definite, fall back to steepest descent
/// 5. Line search along d to find step size α
/// 6. Update: x_{k+1} = `x_k` + α * `d_k`
///
/// # Parameters
///
/// - **`max_iter`**: Maximum number of iterations
/// - **tol**: Convergence tolerance (gradient norm)
/// - **epsilon**: Finite difference step size for Hessian approximation (default: 1e-5)
///
/// # Example
///
/// ```
/// use aprender::optim::{DampedNewton, Optimizer};
/// use aprender::primitives::Vector;
///
/// let mut optimizer = DampedNewton::new(100, 1e-5);
///
/// // Minimize quadratic function f(x,y) = x^2 + 2y^2
/// let f = |x: &Vector<f32>| x[0] * x[0] + 2.0 * x[1] * x[1];
/// let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 4.0 * x[1]]);
///
/// let x0 = Vector::from_slice(&[5.0, 3.0]);
/// let result = optimizer.minimize(f, grad, x0);
/// ```
#[derive(Debug, Clone)]
pub struct DampedNewton {
    /// Maximum number of iterations
    max_iter: usize,
    /// Convergence tolerance (gradient norm)
    tol: f32,
    /// Finite difference step size for Hessian approximation
    epsilon: f32,
    /// Line search strategy
    line_search: BacktrackingLineSearch,
}

impl DampedNewton {
    /// Creates a new Damped Newton optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of iterations (typical: 100-1000)
    /// * `tol` - Convergence tolerance for gradient norm (typical: 1e-5)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::DampedNewton;
    ///
    /// let optimizer = DampedNewton::new(100, 1e-5);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, tol: f32) -> Self {
        Self {
            max_iter,
            tol,
            epsilon: 1e-5, // Finite difference step size
            line_search: BacktrackingLineSearch::new(1e-4, 0.5, 50),
        }
    }

    /// Sets the finite difference epsilon for Hessian approximation.
    ///
    /// # Arguments
    ///
    /// * `epsilon` - Step size for finite differences (typical: 1e-5 to 1e-8)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::DampedNewton;
    ///
    /// let optimizer = DampedNewton::new(100, 1e-5).with_epsilon(1e-6);
    /// ```
    #[must_use]
    pub fn with_epsilon(mut self, epsilon: f32) -> Self {
        self.epsilon = epsilon;
        self
    }

    /// Approximates the Hessian matrix using finite differences.
    ///
    /// Uses central differences: H[i,j] ≈ (∂`²f/∂x_i∂x_j`)
    fn approximate_hessian<G>(&self, grad: &G, x: &Vector<f32>) -> Matrix<f32>
    where
        G: Fn(&Vector<f32>) -> Vector<f32>,
    {
        let n = x.len();
        let mut h_data = vec![0.0; n * n];

        let g0 = grad(x);

        // Compute Hessian using finite differences
        for i in 0..n {
            // Perturb x[i] by epsilon
            let mut x_plus = x.clone();
            x_plus[i] += self.epsilon;

            let g_plus = grad(&x_plus);

            // Approximate column i of Hessian: H[:,i] ≈ (g(x+ε*e_i) - g(x)) / ε
            for j in 0..n {
                h_data[j * n + i] = (g_plus[j] - g0[j]) / self.epsilon;
            }
        }

        // Symmetrize the Hessian (since it should be symmetric)
        for i in 0..n {
            for j in (i + 1)..n {
                let avg = (h_data[i * n + j] + h_data[j * n + i]) / 2.0;
                h_data[i * n + j] = avg;
                h_data[j * n + i] = avg;
            }
        }

        Matrix::from_vec(n, n, h_data).expect("Matrix dimensions should be valid")
    }

    /// Computes the L2 norm of a vector.
    fn norm(v: &Vector<f32>) -> f32 {
        let mut sum = 0.0;
        for i in 0..v.len() {
            sum += v[i] * v[i];
        }
        sum.sqrt()
    }
}

impl Optimizer for DampedNewton {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "Damped Newton does not support stochastic updates (step). Use minimize() for batch optimization."
        )
    }

    fn minimize<F, G>(&mut self, objective: F, gradient: G, x0: Vector<f32>) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
    {
        let start_time = std::time::Instant::now();
        let n = x0.len();

        let mut x = x0;
        let mut fx = objective(&x);
        let mut grad = gradient(&x);
        let mut grad_norm = Self::norm(&grad);

        for iter in 0..self.max_iter {
            // Check convergence
            if grad_norm < self.tol {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Approximate Hessian
            let hessian = self.approximate_hessian(&gradient, &x);

            // Negate gradient for solving H * d = -g
            let mut neg_grad = Vector::zeros(n);
            for i in 0..n {
                neg_grad[i] = -grad[i];
            }

            // Solve H * d = -g using Cholesky decomposition
            let d = if let Ok(direction) = hessian.cholesky_solve(&neg_grad) {
                // Check if it's a descent direction
                let mut grad_dot_d = 0.0;
                for i in 0..n {
                    grad_dot_d += grad[i] * direction[i];
                }

                if grad_dot_d < 0.0 {
                    // Valid descent direction from Newton step
                    direction
                } else {
                    // Not a descent direction - fall back to steepest descent
                    let mut sd = Vector::zeros(n);
                    for i in 0..n {
                        sd[i] = -grad[i];
                    }
                    sd
                }
            } else {
                // Hessian not positive definite - fall back to steepest descent
                let mut sd = Vector::zeros(n);
                for i in 0..n {
                    sd[i] = -grad[i];
                }
                sd
            };

            // Line search
            let alpha = self.line_search.search(&objective, &gradient, &x, &d);

            // Check for stalled progress
            if alpha < 1e-12 {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::Stalled,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Update position: x_new = x + alpha * d
            let mut x_new = Vector::zeros(n);
            for i in 0..n {
                x_new[i] = x[i] + alpha * d[i];
            }

            // Compute new objective and gradient
            let fx_new = objective(&x_new);
            let grad_new = gradient(&x_new);

            // Check for numerical errors
            if fx_new.is_nan() || fx_new.is_infinite() {
                return OptimizationResult {
                    solution: x,
                    objective_value: fx,
                    iterations: iter,
                    status: ConvergenceStatus::NumericalError,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Update for next iteration
            x = x_new;
            fx = fx_new;
            grad = grad_new;
            grad_norm = Self::norm(&grad);
        }

        // Max iterations reached
        OptimizationResult {
            solution: x,
            objective_value: fx,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: grad_norm,
            constraint_violation: 0.0,
            elapsed_time: start_time.elapsed(),
        }
    }

    fn reset(&mut self) {
        // Damped Newton is stateless - nothing to reset
    }
}

// ==================== Proximal Operators ====================

/// Proximal operators for non-smooth regularization.
///
/// A proximal operator for function g is defined as:
/// ```text
/// prox_g(v) = argmin_x { g(x) + ½‖x - v‖² }
/// ```
///
/// These are essential building blocks for proximal gradient methods like FISTA.
pub mod prox {
    use crate::primitives::Vector;

    /// Soft-thresholding operator for L1 regularization.
    ///
    /// Computes the proximal operator of the L1 norm: prox_{λ‖·‖₁}(v).
    ///
    /// # Formula
    ///
    /// ```text
    /// prox_{λ‖·‖₁}(v) = sign(v) ⊙ max(|v| - λ, 0)
    /// ```
    ///
    /// # Arguments
    ///
    /// * `v` - Input vector
    /// * `lambda` - Regularization parameter (λ ≥ 0)
    ///
    /// # Returns
    ///
    /// Soft-thresholded vector
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::prox::soft_threshold;
    /// use aprender::primitives::Vector;
    ///
    /// let v = Vector::from_slice(&[2.0, -1.5, 0.5]);
    /// let result = soft_threshold(&v, 1.0);
    ///
    /// assert!((result[0] - 1.0).abs() < 1e-6);  // 2.0 - 1.0 = 1.0
    /// assert!((result[1] + 0.5).abs() < 1e-6);  // -1.5 + 1.0 = -0.5
    /// assert!(result[2].abs() < 1e-6);          // 0.5 - 1.0 = 0 (thresholded)
    /// ```
    ///
    /// # Use Cases
    ///
    /// - **Lasso regression**: Sparse linear models with L1 penalty
    /// - **Compressed sensing**: Sparse signal recovery
    /// - **Feature selection**: Automatic variable selection via sparsity
    #[must_use]
    pub fn soft_threshold(v: &Vector<f32>, lambda: f32) -> Vector<f32> {
        let mut result = Vector::zeros(v.len());
        for i in 0..v.len() {
            let val = v[i];
            result[i] = if val > lambda {
                val - lambda
            } else if val < -lambda {
                val + lambda
            } else {
                0.0
            };
        }
        result
    }

    /// Projects onto the non-negative orthant: x ≥ 0.
    ///
    /// # Arguments
    ///
    /// * `x` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with all negative components set to zero
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::prox::nonnegative;
    /// use aprender::primitives::Vector;
    ///
    /// let x = Vector::from_slice(&[1.0, -2.0, 3.0, -0.5]);
    /// let result = nonnegative(&x);
    ///
    /// assert_eq!(result[0], 1.0);
    /// assert_eq!(result[1], 0.0);
    /// assert_eq!(result[2], 3.0);
    /// assert_eq!(result[3], 0.0);
    /// ```
    #[must_use]
    pub fn nonnegative(x: &Vector<f32>) -> Vector<f32> {
        let mut result = Vector::zeros(x.len());
        for i in 0..x.len() {
            result[i] = x[i].max(0.0);
        }
        result
    }

    /// Projects onto an L2 ball: ‖x‖₂ ≤ radius.
    ///
    /// # Arguments
    ///
    /// * `x` - Input vector
    /// * `radius` - Ball radius (r > 0)
    ///
    /// # Returns
    ///
    /// Projected vector satisfying ‖result‖₂ ≤ radius
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::prox::project_l2_ball;
    /// use aprender::primitives::Vector;
    ///
    /// let x = Vector::from_slice(&[3.0, 4.0]); // norm = 5.0
    /// let result = project_l2_ball(&x, 2.0);
    ///
    /// // Should be scaled to norm = 2.0
    /// let norm = (result[0] * result[0] + result[1] * result[1]).sqrt();
    /// assert!((norm - 2.0).abs() < 1e-5);
    /// ```
    #[must_use]
    pub fn project_l2_ball(x: &Vector<f32>, radius: f32) -> Vector<f32> {
        let mut norm_sq = 0.0;
        for i in 0..x.len() {
            norm_sq += x[i] * x[i];
        }
        let norm = norm_sq.sqrt();

        if norm <= radius {
            x.clone()
        } else {
            let scale = radius / norm;
            let mut result = Vector::zeros(x.len());
            for i in 0..x.len() {
                result[i] = scale * x[i];
            }
            result
        }
    }

    /// Projects onto box constraints: lower ≤ x ≤ upper.
    ///
    /// # Arguments
    ///
    /// * `x` - Input vector
    /// * `lower` - Lower bounds
    /// * `upper` - Upper bounds
    ///
    /// # Returns
    ///
    /// Vector with components clipped to [lower, upper]
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::prox::project_box;
    /// use aprender::primitives::Vector;
    ///
    /// let x = Vector::from_slice(&[-1.0, 0.5, 2.0]);
    /// let lower = Vector::from_slice(&[0.0, 0.0, 0.0]);
    /// let upper = Vector::from_slice(&[1.0, 1.0, 1.0]);
    ///
    /// let result = project_box(&x, &lower, &upper);
    ///
    /// assert_eq!(result[0], 0.0);  // Clipped to lower
    /// assert_eq!(result[1], 0.5);  // Within bounds
    /// assert_eq!(result[2], 1.0);  // Clipped to upper
    /// ```
    #[must_use]
    pub fn project_box(x: &Vector<f32>, lower: &Vector<f32>, upper: &Vector<f32>) -> Vector<f32> {
        let mut result = Vector::zeros(x.len());
        for i in 0..x.len() {
            result[i] = x[i].max(lower[i]).min(upper[i]);
        }
        result
    }
}

// ==================== FISTA Optimizer ====================

/// FISTA (Fast Iterative Shrinkage-Thresholding Algorithm).
///
/// Accelerated proximal gradient method for minimizing composite objectives:
/// ```text
/// minimize f(x) + g(x)
/// ```
/// where f is smooth and convex, g is convex (possibly non-smooth but "simple").
///
/// FISTA achieves O(1/k²) convergence rate using Nesterov acceleration,
/// compared to O(1/k) for standard proximal gradient (ISTA).
///
/// # Key Applications
///
/// - **Lasso regression**: f(x) = ½‖Ax - b‖², g(x) = λ‖x‖₁
/// - **Elastic Net**: f(x) = ½‖Ax - b‖², g(x) = λ₁‖x‖₁ + λ₂‖x‖₂²
/// - **Total variation**: Image denoising with TV regularization
/// - **Non-negative least squares**: f(x) = ½‖Ax - b‖², g(x) = indicator(x ≥ 0)
///
/// # Example
///
/// ```
/// use aprender::optim::{FISTA, Optimizer, prox};
/// use aprender::primitives::Vector;
///
/// // Minimize: ½(x - 5)² + 2|x|  (L1-regularized quadratic)
/// let smooth = |x: &Vector<f32>| 0.5 * (x[0] - 5.0).powi(2);
/// let grad_smooth = |x: &Vector<f32>| Vector::from_slice(&[x[0] - 5.0]);
/// let proximal = |v: &Vector<f32>, _alpha: f32| prox::soft_threshold(v, 2.0);
///
/// let mut fista = FISTA::new(1000, 0.1, 1e-5);
/// let x0 = Vector::from_slice(&[0.0]);
/// let result = fista.minimize(smooth, grad_smooth, proximal, x0);
///
/// // Check that optimization completed successfully
/// assert!(!result.solution[0].is_nan());
/// ```
///
/// # References
///
/// - Beck & Teboulle (2009). "A fast iterative shrinkage-thresholding algorithm
///   for linear inverse problems." SIAM Journal on Imaging Sciences, 2(1), 183-202.
#[derive(Debug, Clone)]
pub struct FISTA {
    /// Maximum number of iterations
    max_iter: usize,
    /// Step size (α > 0)
    step_size: f32,
    /// Convergence tolerance (‖xₖ₊₁ - xₖ‖ < tol)
    tol: f32,
}

impl FISTA {
    /// Creates a new FISTA optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of iterations
    /// * `step_size` - Step size α (should be ≤ 1/L where L is Lipschitz constant of ∇f)
    /// * `tol` - Convergence tolerance
    ///
    /// # Returns
    ///
    /// New FISTA optimizer instance
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::FISTA;
    ///
    /// let optimizer = FISTA::new(1000, 0.01, 1e-6);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, step_size: f32, tol: f32) -> Self {
        Self {
            max_iter,
            step_size,
            tol,
        }
    }

    /// Minimizes a composite objective function using FISTA.
    ///
    /// Solves: minimize f(x) + g(x) where f is smooth, g is "simple" (has easy prox).
    ///
    /// # Type Parameters
    ///
    /// * `F` - Smooth objective function type
    /// * `G` - Gradient of smooth part type
    /// * `P` - Proximal operator type
    ///
    /// # Arguments
    ///
    /// * `smooth` - Smooth part f(x)
    /// * `grad_smooth` - Gradient ∇f(x)
    /// * `prox` - Proximal operator `prox_g(v`, α)
    /// * `x0` - Initial point
    ///
    /// # Returns
    ///
    /// [`OptimizationResult`] with solution and convergence information
    pub fn minimize<F, G, P>(
        &mut self,
        smooth: F,
        grad_smooth: G,
        prox: P,
        x0: Vector<f32>,
    ) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
        P: Fn(&Vector<f32>, f32) -> Vector<f32>,
    {
        let start_time = std::time::Instant::now();

        let mut x = x0.clone();
        let mut y = x0;
        let mut t = 1.0; // Nesterov momentum parameter

        for iter in 0..self.max_iter {
            // Proximal gradient step at y
            let grad_y = grad_smooth(&y);

            // Compute: y - α * ∇f(y)
            let mut gradient_step = Vector::zeros(y.len());
            for i in 0..y.len() {
                gradient_step[i] = y[i] - self.step_size * grad_y[i];
            }

            // Apply proximal operator
            let x_new = prox(&gradient_step, self.step_size);

            // Check convergence
            let mut diff_norm = 0.0;
            for i in 0..x.len() {
                let diff = x_new[i] - x[i];
                diff_norm += diff * diff;
            }
            diff_norm = diff_norm.sqrt();

            if diff_norm < self.tol {
                let final_obj = smooth(&x_new);
                return OptimizationResult {
                    solution: x_new,
                    objective_value: final_obj,
                    iterations: iter,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: diff_norm, // Use step norm as proxy for gradient norm
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Nesterov acceleration
            let t_new = (1.0_f32 + (1.0_f32 + 4.0_f32 * t * t).sqrt()) / 2.0_f32;
            let beta = (t - 1.0_f32) / t_new;

            // y_new = x_new + β(x_new - x)
            let mut y_new = Vector::zeros(x.len());
            for i in 0..x.len() {
                y_new[i] = x_new[i] + beta * (x_new[i] - x[i]);
            }

            x = x_new;
            y = y_new;
            t = t_new;
        }

        // Max iterations reached
        let final_obj = smooth(&x);
        OptimizationResult {
            solution: x,
            objective_value: final_obj,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: 0.0,
            constraint_violation: 0.0,
            elapsed_time: start_time.elapsed(),
        }
    }
}

impl Optimizer for FISTA {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "FISTA does not support stochastic updates (step). Use minimize() for batch optimization with proximal operators."
        )
    }

    fn reset(&mut self) {
        // FISTA is stateless - nothing to reset
    }
}

// ==================== Coordinate Descent Optimizer ====================

/// Coordinate Descent optimizer for high-dimensional problems.
///
/// Optimizes one coordinate at a time, which can be much more efficient than
/// full gradient methods when the number of features is very large (n ≫ m).
///
/// # Algorithm
///
/// ```text
/// for k = 1, 2, ..., max_iter:
///     for i = 1, 2, ..., n (cyclic or random order):
///         xᵢ ← argmin f(x₁, ..., xᵢ₋₁, xᵢ, xᵢ₊₁, ..., xₙ)
/// ```
///
/// # Key Applications
///
/// - **Lasso regression**: Coordinate descent with soft-thresholding (scikit-learn default)
/// - **Elastic Net**: L1 + L2 regularization
/// - **SVM**: Sequential Minimal Optimization (SMO) variant
/// - **High-dimensional statistics**: n ≫ m scenarios
///
/// # Advantages
///
/// - O(n) per coordinate update (vs O(n) for full gradient)
/// - No line search needed for many problems
/// - Handles non-differentiable objectives (e.g., L1)
/// - Cache-friendly memory access patterns
///
/// # Example
///
/// ```
/// use aprender::optim::{CoordinateDescent, Optimizer};
/// use aprender::primitives::Vector;
///
/// // Minimize: ½‖x - c‖² where c = [1, 2, 3]
/// // Coordinate i update: xᵢ = cᵢ (closed form)
/// let c = vec![1.0, 2.0, 3.0];
///
/// let update = move |x: &mut Vector<f32>, i: usize| {
///     x[i] = c[i]; // Closed-form solution for coordinate i
/// };
///
/// let mut cd = CoordinateDescent::new(100, 1e-6);
/// let x0 = Vector::from_slice(&[0.0, 0.0, 0.0]);
/// let result = cd.minimize(update, x0);
///
/// // Should converge to c
/// assert!((result.solution[0] - 1.0).abs() < 1e-5);
/// assert!((result.solution[1] - 2.0).abs() < 1e-5);
/// assert!((result.solution[2] - 3.0).abs() < 1e-5);
/// ```
///
/// # References
///
/// - Wright (2015). "Coordinate descent algorithms." Mathematical Programming.
/// - Friedman et al. (2010). "Regularization paths for generalized linear models via coordinate descent."
#[derive(Debug, Clone)]
pub struct CoordinateDescent {
    /// Maximum number of outer iterations (passes through all coordinates)
    max_iter: usize,
    /// Convergence tolerance (‖xₖ₊₁ - xₖ‖ < tol)
    tol: f32,
    /// Whether to use random coordinate order (vs cyclic)
    random_order: bool,
}

impl CoordinateDescent {
    /// Creates a new Coordinate Descent optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of passes through all coordinates
    /// * `tol` - Convergence tolerance
    ///
    /// # Returns
    ///
    /// New Coordinate Descent optimizer with cyclic coordinate order
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::CoordinateDescent;
    ///
    /// let optimizer = CoordinateDescent::new(1000, 1e-6);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, tol: f32) -> Self {
        Self {
            max_iter,
            tol,
            random_order: false,
        }
    }

    /// Sets whether to use random coordinate order.
    ///
    /// # Arguments
    ///
    /// * `random` - If true, coordinates are updated in random order each iteration
    ///
    /// # Returns
    ///
    /// Self for method chaining
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::CoordinateDescent;
    ///
    /// let cd = CoordinateDescent::new(1000, 1e-6).with_random_order(true);
    /// ```
    #[must_use]
    pub fn with_random_order(mut self, random: bool) -> Self {
        self.random_order = random;
        self
    }

    /// Minimizes an objective using coordinate descent.
    ///
    /// The user provides a coordinate update function that modifies one coordinate
    /// at a time. This function should solve:
    /// ```text
    /// xᵢ ← argmin f(x₁, ..., xᵢ₋₁, xᵢ, xᵢ₊₁, ..., xₙ)
    /// ```
    ///
    /// # Type Parameters
    ///
    /// * `U` - Coordinate update function type
    ///
    /// # Arguments
    ///
    /// * `update` - Function that updates coordinate i: `fn(&mut Vector, usize)`
    /// * `x0` - Initial point
    ///
    /// # Returns
    ///
    /// [`OptimizationResult`] with solution and convergence information
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::{CoordinateDescent, prox};
    /// use aprender::primitives::Vector;
    ///
    /// // Lasso coordinate update: soft-thresholding
    /// let lambda = 0.1;
    /// let update = move |x: &mut Vector<f32>, i: usize| {
    ///     // Simplified: actual Lasso requires computing residuals
    ///     let v = x[i];
    ///     x[i] = if v > lambda {
    ///         v - lambda
    ///     } else if v < -lambda {
    ///         v + lambda
    ///     } else {
    ///         0.0
    ///     };
    /// };
    ///
    /// let mut cd = CoordinateDescent::new(100, 1e-6);
    /// let x0 = Vector::from_slice(&[1.0, -0.5, 0.3]);
    /// let result = cd.minimize(update, x0);
    /// ```
    pub fn minimize<U>(&mut self, mut update: U, x0: Vector<f32>) -> OptimizationResult
    where
        U: FnMut(&mut Vector<f32>, usize),
    {
        let start_time = std::time::Instant::now();
        let n = x0.len();

        let mut x = x0;

        for iter in 0..self.max_iter {
            // Save previous iterate for convergence check
            let x_old = x.clone();

            // Determine coordinate order
            if self.random_order {
                // Random permutation (Fisher-Yates shuffle)
                let mut indices: Vec<usize> = (0..n).collect();
                for i in (1..n).rev() {
                    let j = (i as f32 * 0.123456).rem_euclid(1.0); // Simple pseudo-random
                    let j = (j * (i + 1) as f32) as usize;
                    indices.swap(i, j);
                }

                // Update in random order
                for i in indices {
                    update(&mut x, i);
                }
            } else {
                // Cyclic order
                for i in 0..n {
                    update(&mut x, i);
                }
            }

            // Check convergence
            let mut diff_norm = 0.0;
            for i in 0..n {
                let diff = x[i] - x_old[i];
                diff_norm += diff * diff;
            }
            diff_norm = diff_norm.sqrt();

            if diff_norm < self.tol {
                return OptimizationResult {
                    solution: x,
                    objective_value: 0.0, // Objective not tracked
                    iterations: iter,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: diff_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }
        }

        // Max iterations reached
        OptimizationResult {
            solution: x,
            objective_value: 0.0,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: 0.0,
            constraint_violation: 0.0,
            elapsed_time: start_time.elapsed(),
        }
    }
}

impl Optimizer for CoordinateDescent {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "Coordinate Descent does not support stochastic updates (step). Use minimize() with coordinate update function."
        )
    }

    fn reset(&mut self) {
        // Coordinate Descent is stateless - nothing to reset
    }
}

/// ADMM (Alternating Direction Method of Multipliers) for distributed and constrained optimization.
///
/// Solves problems of the form:
/// ```text
/// minimize  f(x) + g(z)
/// subject to Ax + Bz = c
/// ```
///
/// # Applications
///
/// - **Distributed Lasso**: Split data across workers for large-scale regression
/// - **Consensus optimization**: Average models from different sites (federated learning)
/// - **Constrained problems**: Equality-constrained optimization via consensus
/// - **Model parallelism**: Parallelize training across devices
///
/// # Algorithm
///
/// ADMM alternates between three updates:
///
/// 1. **x-update**: `x^{k+1} = argmin_x { f(x) + (ρ/2)‖Ax + Bz^k - c + u^k‖² }`
/// 2. **z-update**: `z^{k+1} = argmin_z { g(z) + (ρ/2)‖Ax^{k+1} + Bz - c + u^k‖² }`
/// 3. **u-update**: `u^{k+1} = u^k + (Ax^{k+1} + Bz^{k+1} - c)`
///
/// where u is the scaled dual variable and ρ is the penalty parameter.
///
/// # Convergence
///
/// - **Rate**: O(1/k) for convex f and g
/// - **Stopping criteria**: Both primal and dual residuals below tolerance
/// - **Adaptive ρ**: Automatically adjusts penalty parameter for faster convergence
///
/// # Example: Consensus Form (Lasso)
///
/// For Lasso regression with consensus constraint x = z:
/// ```rust
/// use aprender::optim::ADMM;
/// use aprender::primitives::{Vector, Matrix};
///
/// let n = 5;
/// let m = 10;
///
/// // Create problem data
/// let A = Matrix::eye(n); // Identity for consensus
/// let B = Matrix::eye(n);
/// let c = Vector::zeros(n);
///
/// // x-minimizer: least squares update
/// let data_matrix = Matrix::eye(m); // Your data matrix
/// let observations = Vector::ones(m); // Your observations
/// let lambda = 0.1;
///
/// let x_minimizer = |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
///     // Minimize ½‖Dx - b‖² + (ρ/2)‖x - z + u‖²
///     // Closed form: x = (DᵀD + ρI)⁻¹(Dᵀb + ρ(z - u))
///     let mut rhs = Vector::zeros(n);
///     for i in 0..n {
///         rhs[i] = rho * (z[i] - u[i]);
///     }
///     rhs // Simplified for example
/// };
///
/// // z-minimizer: soft-thresholding (proximal operator for L1)
/// let z_minimizer = |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
///     let mut z = Vector::zeros(n);
///     for i in 0..n {
///         let v = ax[i] + u[i];
///         let threshold = lambda / rho;
///         z[i] = if v > threshold {
///             v - threshold
///         } else if v < -threshold {
///             v + threshold
///         } else {
///             0.0
///         };
///     }
///     z
/// };
///
/// let mut admm = ADMM::new(100, 1.0, 1e-4).with_adaptive_rho(true);
/// let x0 = Vector::zeros(n);
/// let z0 = Vector::zeros(n);
///
/// let result = admm.minimize_consensus(
///     x_minimizer,
///     z_minimizer,
///     &A,
///     &B,
///     &c,
///     x0,
///     z0,
/// );
/// ```
///
/// # Reference
///
/// Boyd, S., Parikh, N., Chu, E., Peleato, B., & Eckstein, J. (2011).
/// "Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers"
/// Foundations and Trends in Machine Learning, 3(1), 1-122.
#[derive(Debug, Clone)]
pub struct ADMM {
    /// Maximum number of iterations
    max_iter: usize,
    /// Penalty parameter (controls constraint enforcement)
    rho: f32,
    /// Tolerance for convergence (primal + dual residuals)
    tol: f32,
    /// Whether to adaptively adjust rho
    adaptive_rho: bool,
    /// Factor for increasing rho when primal residual is large
    rho_increase: f32,
    /// Factor for decreasing rho when dual residual is large
    rho_decrease: f32,
}

impl ADMM {
    /// Creates a new ADMM optimizer.
    ///
    /// # Parameters
    ///
    /// - `max_iter`: Maximum number of iterations (typical: 100-1000)
    /// - `rho`: Penalty parameter (typical: 0.1-10.0, problem-dependent)
    /// - `tol`: Convergence tolerance for residuals (typical: 1e-4 to 1e-6)
    ///
    /// # Returns
    ///
    /// A new ADMM optimizer with default settings:
    /// - Adaptive rho: disabled (use `with_adaptive_rho(true)` to enable)
    /// - Rho increase factor: 2.0
    /// - Rho decrease factor: 2.0
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::ADMM;
    ///
    /// let admm = ADMM::new(100, 1.0, 1e-4);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, rho: f32, tol: f32) -> Self {
        Self {
            max_iter,
            rho,
            tol,
            adaptive_rho: false,
            rho_increase: 2.0,
            rho_decrease: 2.0,
        }
    }

    /// Enables or disables adaptive penalty parameter adjustment.
    ///
    /// When enabled, rho is automatically adjusted based on the ratio of primal to dual residuals:
    /// - If primal residual >> dual residual: increase rho (enforce constraints more)
    /// - If dual residual >> primal residual: decrease rho (improve objective)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::ADMM;
    ///
    /// let admm = ADMM::new(100, 1.0, 1e-4).with_adaptive_rho(true);
    /// ```
    #[must_use]
    pub fn with_adaptive_rho(mut self, adaptive: bool) -> Self {
        self.adaptive_rho = adaptive;
        self
    }

    /// Sets the factors for adaptive rho adjustment.
    ///
    /// # Parameters
    ///
    /// - `increase`: Factor to multiply rho when primal residual is large (default: 2.0)
    /// - `decrease`: Factor to divide rho when dual residual is large (default: 2.0)
    #[must_use]
    pub fn with_rho_factors(mut self, increase: f32, decrease: f32) -> Self {
        self.rho_increase = increase;
        self.rho_decrease = decrease;
        self
    }

    /// Minimizes a consensus-form ADMM problem.
    ///
    /// Solves: minimize f(x) + g(z) subject to Ax + Bz = c
    ///
    /// # Parameters
    ///
    /// - `x_minimizer`: Function that solves x-subproblem given (z, u, c, rho)
    /// - `z_minimizer`: Function that solves z-subproblem given (Ax, u, c, rho)
    /// - `A`, `B`, `c`: Constraint matrices and vector (Ax + Bz = c)
    /// - `x0`, `z0`: Initial values for x and z
    ///
    /// # Returns
    ///
    /// `OptimizationResult` containing the optimal x value and convergence information.
    ///
    /// # Minimizer Functions
    ///
    /// The `x_minimizer` should solve:
    /// ```text
    /// argmin_x { f(x) + (ρ/2)‖Ax + Bz - c + u‖² }
    /// ```
    ///
    /// The `z_minimizer` should solve:
    /// ```text
    /// argmin_z { g(z) + (ρ/2)‖Ax + Bz - c + u‖² }
    /// ```
    ///
    /// These often have closed-form solutions or can use proximal operators.
    #[allow(clippy::too_many_arguments)]
    pub fn minimize_consensus<F, G>(
        &mut self,
        x_minimizer: F,
        z_minimizer: G,
        a: &Matrix<f32>,
        b_mat: &Matrix<f32>,
        c: &Vector<f32>,
        x0: Vector<f32>,
        z0: Vector<f32>,
    ) -> OptimizationResult
    where
        F: Fn(&Vector<f32>, &Vector<f32>, &Vector<f32>, f32) -> Vector<f32>,
        G: Fn(&Vector<f32>, &Vector<f32>, &Vector<f32>, f32) -> Vector<f32>,
    {
        let start_time = std::time::Instant::now();

        let mut x = x0;
        let mut z = z0;
        let mut u = Vector::zeros(c.len());
        let mut rho = self.rho;

        let mut z_old = z.clone();

        for iter in 0..self.max_iter {
            // x-update: minimize f(x) + (ρ/2)‖Ax + Bz - c + u‖²
            x = x_minimizer(&z, &u, c, rho);

            // z-update: minimize g(z) + (ρ/2)‖Ax + Bz - c + u‖²
            let ax = a.matvec(&x).expect("Matrix-vector multiplication");
            z = z_minimizer(&ax, &u, c, rho);

            // Compute residual: r = Ax + Bz - c
            let bz = b_mat.matvec(&z).expect("Matrix-vector multiplication");
            let residual = &(&ax + &bz) - c;

            // u-update: u^{k+1} = u^k + r^{k+1}
            u = &u + &residual;

            // Compute primal residual: ‖Ax + Bz - c‖
            let primal_res = residual.norm();

            // Compute dual residual: ρ‖Bᵀ(z^{k+1} - z^k)‖
            let z_diff = &z - &z_old;
            let bt_z_diff = b_mat
                .transpose()
                .matvec(&z_diff)
                .expect("Matrix-vector multiplication");
            let dual_res = rho * bt_z_diff.norm();

            // Check convergence
            if primal_res < self.tol && dual_res < self.tol {
                return OptimizationResult {
                    solution: x,
                    objective_value: 0.0, // Objective not tracked (requires f and g evaluations)
                    iterations: iter + 1,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: dual_res,
                    constraint_violation: primal_res,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Adaptive rho adjustment (Boyd et al. 2011, Section 3.4.1)
            if self.adaptive_rho && iter % 10 == 0 {
                if primal_res > 10.0 * dual_res {
                    // Primal residual is large: increase rho to enforce constraints
                    rho *= self.rho_increase;
                    // Rescale dual variable: u = u / rho_increase
                    let scale = 1.0 / self.rho_increase;
                    u = u.mul_scalar(scale);
                } else if dual_res > 10.0 * primal_res {
                    // Dual residual is large: decrease rho to improve objective
                    rho /= self.rho_decrease;
                    // Rescale dual variable: u = u * rho_decrease
                    u = u.mul_scalar(self.rho_decrease);
                }
            }

            z_old = z.clone();
        }

        // Max iterations reached
        OptimizationResult {
            solution: x,
            objective_value: 0.0,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: 0.0,
            constraint_violation: 0.0,
            elapsed_time: start_time.elapsed(),
        }
    }
}

impl Optimizer for ADMM {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "ADMM does not support stochastic updates (step). Use minimize_consensus() with x-minimizer and z-minimizer functions."
        )
    }

    fn reset(&mut self) {
        // ADMM is stateless - nothing to reset
    }
}

// ==================== Phase 3: Constrained Optimization ====================

/// Projected Gradient Descent for constrained optimization.
///
/// Solves problems of the form:
/// ```text
/// minimize f(x)
/// subject to x ∈ C
/// ```
///
/// where C is a convex set with efficient projection operator.
///
/// # Algorithm
///
/// ```text
/// for k = 1, 2, ..., max_iter:
///     x_k+1 = P_C(x_k - α∇f(x_k))
/// ```
///
/// where `P_C` is the projection onto constraint set C.
///
/// # Key Applications
///
/// - **Non-negative least squares**: C = {x : x ≥ 0}
/// - **Box constraints**: C = {x : l ≤ x ≤ u}
/// - **L2 ball**: C = {x : ‖x‖₂ ≤ r}
/// - **Simplex**: C = {x : x ≥ 0, Σx = 1}
///
/// # Convergence
///
/// - For convex f: O(1/k) convergence rate
/// - For strongly convex f: Linear convergence
/// - Step size α can be constant or use line search
///
/// # Example
///
/// ```
/// use aprender::optim::{ProjectedGradientDescent, prox};
/// use aprender::primitives::Vector;
///
/// // Minimize: ½‖x - c‖² subject to x ≥ 0
/// let c = Vector::from_slice(&[1.0, -2.0, 3.0, -1.0]);
///
/// let objective = |x: &Vector<f32>| {
///     let mut obj = 0.0;
///     for i in 0..x.len() {
///         let diff = x[i] - c[i];
///         obj += 0.5 * diff * diff;
///     }
///     obj
/// };
///
/// let gradient = |x: &Vector<f32>| {
///     let mut grad = Vector::zeros(x.len());
///     for i in 0..x.len() {
///         grad[i] = x[i] - c[i];
///     }
///     grad
/// };
///
/// let project = |x: &Vector<f32>| prox::nonnegative(x);
///
/// let mut pgd = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
/// let x0 = Vector::zeros(4);
/// let result = pgd.minimize(objective, gradient, project, x0);
///
/// // Solution should be max(c, 0) = [1.0, 0.0, 3.0, 0.0]
/// assert!((result.solution[0] - 1.0).abs() < 1e-4);
/// assert!(result.solution[1].abs() < 1e-4);
/// assert!((result.solution[2] - 3.0).abs() < 1e-4);
/// assert!(result.solution[3].abs() < 1e-4);
/// ```
///
/// # References
///
/// - Bertsekas (1999). "Nonlinear Programming."
/// - Beck & Teboulle (2009). "Gradient-based algorithms with applications to signal recovery."
#[derive(Debug, Clone)]
pub struct ProjectedGradientDescent {
    /// Maximum number of iterations
    max_iter: usize,
    /// Step size (learning rate)
    step_size: f32,
    /// Convergence tolerance
    tol: f32,
    /// Use backtracking line search
    use_line_search: bool,
    /// Backtracking parameter (0 < beta < 1)
    beta: f32,
}

impl ProjectedGradientDescent {
    /// Creates a new Projected Gradient Descent optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of iterations
    /// * `step_size` - Initial step size (fixed if line search disabled)
    /// * `tol` - Convergence tolerance
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::ProjectedGradientDescent;
    ///
    /// let optimizer = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, step_size: f32, tol: f32) -> Self {
        Self {
            max_iter,
            step_size,
            tol,
            use_line_search: false,
            beta: 0.5,
        }
    }

    /// Enables backtracking line search.
    ///
    /// # Arguments
    ///
    /// * `beta` - Backtracking parameter (0 < beta < 1, typically 0.5)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::ProjectedGradientDescent;
    ///
    /// let optimizer = ProjectedGradientDescent::new(1000, 1.0, 1e-6)
    ///     .with_line_search(0.5);
    /// ```
    #[must_use]
    pub fn with_line_search(mut self, beta: f32) -> Self {
        self.use_line_search = true;
        self.beta = beta;
        self
    }

    /// Minimizes objective function subject to projection constraint.
    ///
    /// # Arguments
    ///
    /// * `objective` - Objective function f(x)
    /// * `gradient` - Gradient function ∇f(x)
    /// * `project` - Projection operator `P_C(x)`
    /// * `x0` - Initial point
    ///
    /// # Returns
    ///
    /// Optimization result with converged solution
    pub fn minimize<F, G, P>(
        &mut self,
        objective: F,
        gradient: G,
        project: P,
        x0: Vector<f32>,
    ) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
        P: Fn(&Vector<f32>) -> Vector<f32>,
    {
        let start_time = std::time::Instant::now();

        let mut x = x0;
        let mut alpha = self.step_size;

        for iter in 0..self.max_iter {
            // Compute gradient
            let grad = gradient(&x);

            // Gradient step: y = x - α∇f(x)
            let mut y = Vector::zeros(x.len());
            for i in 0..x.len() {
                y[i] = x[i] - alpha * grad[i];
            }

            // Project onto constraint set
            let x_new = project(&y);

            // Line search if enabled
            if self.use_line_search {
                let f_x = objective(&x);
                let f_x_new = objective(&x_new);

                // Backtracking: reduce step size until sufficient decrease
                let mut ls_iter = 0;
                while f_x_new > f_x && ls_iter < 20 {
                    alpha *= self.beta;

                    for i in 0..x.len() {
                        y[i] = x[i] - alpha * grad[i];
                    }
                    let x_new_ls = project(&y);

                    if objective(&x_new_ls) <= f_x {
                        break;
                    }
                    ls_iter += 1;
                }
            }

            // Check convergence
            let mut diff_norm = 0.0;
            for i in 0..x.len() {
                let diff = x_new[i] - x[i];
                diff_norm += diff * diff;
            }
            diff_norm = diff_norm.sqrt();

            // Compute gradient norm at new point
            let grad_new = gradient(&x_new);
            let mut grad_norm = 0.0;
            for i in 0..grad_new.len() {
                grad_norm += grad_new[i] * grad_new[i];
            }
            grad_norm = grad_norm.sqrt();

            x = x_new;

            if diff_norm < self.tol {
                let final_obj = objective(&x);
                return OptimizationResult {
                    solution: x,
                    objective_value: final_obj,
                    iterations: iter + 1,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: grad_norm,
                    constraint_violation: 0.0,
                    elapsed_time: start_time.elapsed(),
                };
            }
        }

        // Max iterations reached
        let final_obj = objective(&x);
        let grad_final = gradient(&x);
        let mut grad_norm = 0.0;
        for i in 0..grad_final.len() {
            grad_norm += grad_final[i] * grad_final[i];
        }
        grad_norm = grad_norm.sqrt();

        OptimizationResult {
            solution: x,
            objective_value: final_obj,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: grad_norm,
            constraint_violation: 0.0,
            elapsed_time: start_time.elapsed(),
        }
    }
}

impl Optimizer for ProjectedGradientDescent {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "Projected Gradient Descent does not support stochastic updates (step). Use minimize() with projection operator."
        )
    }

    fn reset(&mut self) {
        // Reset step size to initial value
        // Note: step_size is not stored separately, so nothing to reset
    }
}

/// Augmented Lagrangian method for constrained optimization.
///
/// Solves problems with equality constraints:
/// ```text
/// minimize f(x)
/// subject to: h(x) = 0  (equality constraints)
/// ```
///
/// # Algorithm
///
/// ```text
/// Augmented Lagrangian: L_ρ(x, λ) = f(x) + λᵀh(x) + ½ρ‖h(x)‖²
///
/// for k = 1, 2, ..., max_iter:
///     x_k = argmin L_ρ(x, λ_k)
///     λ_k+1 = λ_k + ρ h(x_k)
///     if ‖h(x)‖ is small: converged
///     else: increase ρ
/// ```
///
/// # Key Features
///
/// - **Penalty parameter ρ**: Increased adaptively for faster convergence
/// - **Lagrange multipliers λ**: Automatically updated for equality constraints
/// - **Flexible subproblem solver**: Uses gradient descent for inner loop
/// - **Convergence**: Superlinear under regularity conditions
///
/// # Applications
///
/// - **Equality constraints**: Linear systems, manifold optimization
/// - **ADMM**: Alternating Direction Method of Multipliers (special case)
/// - **Consensus optimization**: Distributed optimization, federated learning
/// - **PDE-constrained optimization**: Physics-informed neural networks
///
/// # Example
///
/// ```
/// use aprender::optim::AugmentedLagrangian;
/// use aprender::primitives::Vector;
///
/// // Minimize: ½(x₁-2)² + ½(x₂-3)² subject to x₁ + x₂ = 1
/// let objective = |x: &Vector<f32>| {
///     0.5 * (x[0] - 2.0).powi(2) + 0.5 * (x[1] - 3.0).powi(2)
/// };
///
/// let gradient = |x: &Vector<f32>| {
///     Vector::from_slice(&[x[0] - 2.0, x[1] - 3.0])
/// };
///
/// // Equality constraint: x₁ + x₂ - 1 = 0
/// let equality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] - 1.0]);
///
/// let equality_jac = |_x: &Vector<f32>| {
///     vec![Vector::from_slice(&[1.0, 1.0])]
/// };
///
/// let mut al = AugmentedLagrangian::new(100, 1e-6, 1.0);
/// let x0 = Vector::zeros(2);
/// let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);
///
/// assert!(result.constraint_violation < 1e-4);
/// ```
///
/// # References
///
/// - Nocedal & Wright (2006). "Numerical Optimization." Chapter 17.
/// - Bertsekas (1982). "Constrained Optimization and Lagrange Multiplier Methods."
#[derive(Debug, Clone)]
pub struct AugmentedLagrangian {
    /// Maximum number of outer iterations
    max_iter: usize,
    /// Convergence tolerance for constraint violation
    tol: f32,
    /// Initial penalty parameter
    initial_rho: f32,
    /// Current penalty parameter
    rho: f32,
    /// Penalty increase factor (> 1)
    rho_increase: f32,
    /// Maximum penalty parameter
    rho_max: f32,
}

impl AugmentedLagrangian {
    /// Creates a new Augmented Lagrangian optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of outer iterations
    /// * `tol` - Convergence tolerance for constraint violation
    /// * `initial_rho` - Initial penalty parameter (typically 1.0-10.0)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::AugmentedLagrangian;
    ///
    /// let optimizer = AugmentedLagrangian::new(100, 1e-6, 1.0);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, tol: f32, initial_rho: f32) -> Self {
        Self {
            max_iter,
            tol,
            initial_rho,
            rho: initial_rho,
            rho_increase: 2.0,
            rho_max: 1e6,
        }
    }

    /// Sets penalty increase factor.
    ///
    /// # Arguments
    ///
    /// * `factor` - Penalty increase factor (> 1, typically 2.0-10.0)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::AugmentedLagrangian;
    ///
    /// let optimizer = AugmentedLagrangian::new(100, 1e-6, 1.0)
    ///     .with_rho_increase(5.0);
    /// ```
    #[must_use]
    pub fn with_rho_increase(mut self, factor: f32) -> Self {
        self.rho_increase = factor;
        self
    }

    /// Minimizes objective subject to equality constraints.
    ///
    /// Solves: minimize f(x) subject to h(x) = 0
    ///
    /// # Arguments
    ///
    /// * `objective` - Objective function f(x)
    /// * `gradient` - Gradient ∇f(x)
    /// * `equality` - Equality constraints h(x) = 0 (returns vector)
    /// * `equality_jac` - Jacobian of equality constraints ∇h(x)
    /// * `x0` - Initial point
    ///
    /// # Returns
    ///
    /// Optimization result with constraint satisfaction metrics
    pub fn minimize_equality<F, G, H, J>(
        &mut self,
        objective: F,
        gradient: G,
        equality: H,
        equality_jac: J,
        x0: Vector<f32>,
    ) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
        H: Fn(&Vector<f32>) -> Vector<f32>,
        J: Fn(&Vector<f32>) -> Vec<Vector<f32>>,
    {
        let start_time = std::time::Instant::now();

        let mut x = x0;
        self.rho = self.initial_rho;

        // Initialize Lagrange multipliers to zero
        let h0 = equality(&x);
        let m = h0.len(); // Number of equality constraints
        let mut lambda = Vector::zeros(m);

        for outer_iter in 0..self.max_iter {
            // Solve augmented Lagrangian subproblem: min L_ρ(x, λ)
            // L_ρ(x, λ) = f(x) + λᵀh(x) + ½ρ‖h(x)‖²

            let aug_grad = |x_inner: &Vector<f32>| {
                let grad_f = gradient(x_inner);
                let h_val = equality(x_inner);
                let jac_h = equality_jac(x_inner);

                let n = x_inner.len();
                let mut aug_g = Vector::zeros(n);

                // ∇f(x)
                for i in 0..n {
                    aug_g[i] = grad_f[i];
                }

                // Add ∇h(x)ᵀ(λ + ρh(x))
                for j in 0..m {
                    let coeff = lambda[j] + self.rho * h_val[j];
                    for i in 0..n {
                        aug_g[i] += coeff * jac_h[j][i];
                    }
                }

                aug_g
            };

            // Solve subproblem using gradient descent (simple solver)
            let mut x_sub = x.clone();
            let alpha = 0.01; // Fixed step size for subproblem
            for _sub_iter in 0..50 {
                let grad = aug_grad(&x_sub);
                let mut grad_norm_sq = 0.0;
                for i in 0..grad.len() {
                    grad_norm_sq += grad[i] * grad[i];
                }
                if grad_norm_sq < 1e-8 {
                    break; // Subproblem converged
                }
                for i in 0..x_sub.len() {
                    x_sub[i] -= alpha * grad[i];
                }
            }

            x = x_sub;

            // Update Lagrange multipliers: λ_k+1 = λ_k + ρ h(x_k)
            let h_val = equality(&x);
            for i in 0..m {
                lambda[i] += self.rho * h_val[i];
            }

            // Check constraint violation
            let mut constraint_viol = 0.0;
            for i in 0..m {
                constraint_viol += h_val[i] * h_val[i];
            }
            constraint_viol = constraint_viol.sqrt();

            // Check convergence
            if constraint_viol < self.tol {
                let final_obj = objective(&x);
                let grad_f = gradient(&x);
                let mut grad_norm = 0.0;
                for i in 0..grad_f.len() {
                    grad_norm += grad_f[i] * grad_f[i];
                }
                grad_norm = grad_norm.sqrt();

                return OptimizationResult {
                    solution: x,
                    objective_value: final_obj,
                    iterations: outer_iter + 1,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: grad_norm,
                    constraint_violation: constraint_viol,
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Increase penalty parameter if constraint violation is not decreasing fast enough
            if constraint_viol > 0.1 * self.tol && self.rho < self.rho_max {
                self.rho *= self.rho_increase;
            }
        }

        // Max iterations reached
        let final_obj = objective(&x);
        let h_val = equality(&x);
        let mut constraint_viol = 0.0;
        for i in 0..h_val.len() {
            constraint_viol += h_val[i] * h_val[i];
        }
        constraint_viol = constraint_viol.sqrt();

        let grad_f = gradient(&x);
        let mut grad_norm = 0.0;
        for i in 0..grad_f.len() {
            grad_norm += grad_f[i] * grad_f[i];
        }
        grad_norm = grad_norm.sqrt();

        OptimizationResult {
            solution: x,
            objective_value: final_obj,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: grad_norm,
            constraint_violation: constraint_viol,
            elapsed_time: start_time.elapsed(),
        }
    }
}

impl Optimizer for AugmentedLagrangian {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "Augmented Lagrangian does not support stochastic updates (step). Use minimize_equality() for constrained optimization."
        )
    }

    fn reset(&mut self) {
        // Reset penalty parameter to initial value
        self.rho = self.initial_rho;
    }
}

/// Interior Point (Barrier) method for inequality-constrained optimization.
///
/// Solves problems with inequality constraints:
/// ```text
/// minimize f(x)
/// subject to: g(x) ≤ 0  (inequality constraints)
/// ```
///
/// # Algorithm
///
/// ```text
/// Barrier function: B_μ(x) = f(x) - μ Σ log(-g_i(x))
///
/// for k = 1, 2, ..., max_iter:
///     x_k = argmin B_μ(x)  (barrier subproblem)
///     μ_k+1 = β * μ_k      (decrease barrier parameter)
///     if ‖∇B_μ(x)‖ is small: converged
/// ```
///
/// # Key Features
///
/// - **Log-barrier**: Enforces g(x) < 0 via -μ log(-g_i(x))
/// - **Path-following**: Decreases μ → 0 to approach constrained optimum
/// - **Self-concordant**: Converges in O(√n log(1/ε)) iterations
/// - **Warm start**: Uses previous solution for next barrier value
///
/// # Applications
///
/// - **Linear programming**: Constraints Ax ≤ b
/// - **Quadratic programming**: QP with inequality constraints
/// - **Semidefinite programming**: Matrix constraints X ⪰ 0
/// - **Support Vector Machines**: Soft-margin constraints
/// - **Portfolio optimization**: Long-only constraints (x ≥ 0)
///
/// # Example
///
/// ```
/// use aprender::optim::InteriorPoint;
/// use aprender::primitives::Vector;
///
/// // Minimize: x₁² + x₂² subject to -x₁ ≤ 0, -x₂ ≤ 0 (i.e., x ≥ 0)
/// let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
///
/// let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);
///
/// // Inequality constraints: g(x) = [-x₁, -x₂] ≤ 0
/// let inequality = |x: &Vector<f32>| Vector::from_slice(&[-x[0], -x[1]]);
///
/// let inequality_jac = |_x: &Vector<f32>| {
///     vec![Vector::from_slice(&[-1.0, 0.0]), Vector::from_slice(&[0.0, -1.0])]
/// };
///
/// let mut ip = InteriorPoint::new(50, 1e-6, 1.0);
/// let x0 = Vector::from_slice(&[1.0, 1.0]); // Feasible start
/// let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);
///
/// // Solution should be [0, 0] (constrained minimum)
/// assert!(result.solution[0].abs() < 1e-3);
/// assert!(result.solution[1].abs() < 1e-3);
/// ```
///
/// # References
///
/// - Nesterov & Nemirovskii (1994). "Interior-Point Polynomial Algorithms in Convex Programming."
/// - Boyd & Vandenberghe (2004). "Convex Optimization." Chapter 11.
/// - Wright (1997). "Primal-Dual Interior-Point Methods."
#[derive(Debug, Clone)]
pub struct InteriorPoint {
    /// Maximum number of outer iterations (barrier parameter updates)
    max_iter: usize,
    /// Convergence tolerance
    tol: f32,
    /// Initial barrier parameter
    initial_mu: f32,
    /// Current barrier parameter
    mu: f32,
    /// Barrier decrease factor (0 < beta < 1, typically 0.1-0.5)
    beta: f32,
}

impl InteriorPoint {
    /// Creates a new Interior Point optimizer.
    ///
    /// # Arguments
    ///
    /// * `max_iter` - Maximum number of outer iterations
    /// * `tol` - Convergence tolerance
    /// * `initial_mu` - Initial barrier parameter (typically 1.0-10.0)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::InteriorPoint;
    ///
    /// let optimizer = InteriorPoint::new(50, 1e-6, 1.0);
    /// ```
    #[must_use]
    pub fn new(max_iter: usize, tol: f32, initial_mu: f32) -> Self {
        Self {
            max_iter,
            tol,
            initial_mu,
            mu: initial_mu,
            beta: 0.2,
        }
    }

    /// Sets barrier decrease factor.
    ///
    /// # Arguments
    ///
    /// * `beta` - Barrier decrease factor (0 < beta < 1, typically 0.1-0.5)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::InteriorPoint;
    ///
    /// let optimizer = InteriorPoint::new(50, 1e-6, 1.0)
    ///     .with_beta(0.1); // Aggressive barrier decrease
    /// ```
    #[must_use]
    pub fn with_beta(mut self, beta: f32) -> Self {
        self.beta = beta;
        self
    }

    /// Minimizes objective subject to inequality constraints.
    ///
    /// Solves: minimize f(x) subject to g(x) ≤ 0
    ///
    /// # Arguments
    ///
    /// * `objective` - Objective function f(x)
    /// * `gradient` - Gradient ∇f(x)
    /// * `inequality` - Inequality constraints g(x) ≤ 0 (returns vector)
    /// * `inequality_jac` - Jacobian of inequality constraints ∇g(x)
    /// * `x0` - Initial feasible point (must satisfy g(x0) < 0 strictly)
    ///
    /// # Returns
    ///
    /// Optimization result with constraint satisfaction metrics
    ///
    /// # Panics
    ///
    /// Panics if initial point is infeasible (g(x0) ≥ 0 for any constraint)
    #[allow(clippy::too_many_lines)]
    pub fn minimize<F, G, H, J>(
        &mut self,
        objective: F,
        gradient: G,
        inequality: H,
        inequality_jac: J,
        x0: Vector<f32>,
    ) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
        H: Fn(&Vector<f32>) -> Vector<f32>,
        J: Fn(&Vector<f32>) -> Vec<Vector<f32>>,
    {
        let start_time = std::time::Instant::now();

        // Check initial feasibility
        let g0 = inequality(&x0);
        for i in 0..g0.len() {
            assert!(
                g0[i] < 0.0,
                "Initial point is infeasible: g[{}] = {} ≥ 0. Interior point requires strictly feasible start.",
                i, g0[i]
            );
        }

        let mut x = x0;
        self.mu = self.initial_mu;
        let m = g0.len(); // Number of inequality constraints

        for outer_iter in 0..self.max_iter {
            // Solve barrier subproblem: min B_μ(x) = f(x) - μ Σ log(-g_i(x))

            let barrier_grad = |x_inner: &Vector<f32>| {
                let grad_f = gradient(x_inner);
                let g_val = inequality(x_inner);
                let jac_g = inequality_jac(x_inner);

                let n = x_inner.len();
                let mut barrier_g = Vector::zeros(n);

                // ∇f(x)
                for i in 0..n {
                    barrier_g[i] = grad_f[i];
                }

                // Subtract μ Σ (1/(-g_i)) * ∇g_i(x)
                for j in 0..m {
                    if g_val[j] >= 0.0 {
                        // Hit constraint boundary - project back
                        continue;
                    }
                    let coeff = -self.mu / g_val[j];
                    for i in 0..n {
                        barrier_g[i] += coeff * jac_g[j][i];
                    }
                }

                barrier_g
            };

            // Solve barrier subproblem using gradient descent
            let mut x_sub = x.clone();
            let alpha = 0.01; // Fixed step size
            for _sub_iter in 0..50 {
                let grad = barrier_grad(&x_sub);

                // Check if gradient is small (converged)
                let mut grad_norm_sq = 0.0;
                for i in 0..grad.len() {
                    grad_norm_sq += grad[i] * grad[i];
                }
                if grad_norm_sq < 1e-8 {
                    break;
                }

                // Gradient descent step
                for i in 0..x_sub.len() {
                    x_sub[i] -= alpha * grad[i];
                }

                // Check feasibility - if we violated constraints, step back
                let g_sub = inequality(&x_sub);
                let mut infeasible = false;
                for i in 0..m {
                    if g_sub[i] >= -1e-8 {
                        // Close to or past boundary
                        infeasible = true;
                        break;
                    }
                }
                if infeasible {
                    // Step back
                    for i in 0..x_sub.len() {
                        x_sub[i] += alpha * grad[i] * 0.5; // Half step back
                    }
                }
            }

            x = x_sub;

            // Check convergence via gradient of barrier function
            let grad_barrier = barrier_grad(&x);
            let mut grad_norm = 0.0;
            for i in 0..grad_barrier.len() {
                grad_norm += grad_barrier[i] * grad_barrier[i];
            }
            grad_norm = grad_norm.sqrt();

            // Also check constraint violation
            let g_val = inequality(&x);
            let mut max_violation = 0.0;
            for i in 0..m {
                if g_val[i] > max_violation {
                    max_violation = g_val[i];
                }
            }

            // Converged if gradient is small and μ is small
            if grad_norm < self.tol && self.mu < 1e-4 {
                let final_obj = objective(&x);
                return OptimizationResult {
                    solution: x,
                    objective_value: final_obj,
                    iterations: outer_iter + 1,
                    status: ConvergenceStatus::Converged,
                    gradient_norm: grad_norm,
                    constraint_violation: max_violation.max(0.0),
                    elapsed_time: start_time.elapsed(),
                };
            }

            // Decrease barrier parameter
            self.mu *= self.beta;
        }

        // Max iterations reached
        let final_obj = objective(&x);
        let g_val = inequality(&x);
        let mut max_violation = 0.0;
        for i in 0..g_val.len() {
            if g_val[i] > max_violation {
                max_violation = g_val[i];
            }
        }

        let grad_f = gradient(&x);
        let mut grad_norm = 0.0;
        for i in 0..grad_f.len() {
            grad_norm += grad_f[i] * grad_f[i];
        }
        grad_norm = grad_norm.sqrt();

        OptimizationResult {
            solution: x,
            objective_value: final_obj,
            iterations: self.max_iter,
            status: ConvergenceStatus::MaxIterations,
            gradient_norm: grad_norm,
            constraint_violation: max_violation.max(0.0),
            elapsed_time: start_time.elapsed(),
        }
    }
}

impl Optimizer for InteriorPoint {
    fn step(&mut self, _params: &mut Vector<f32>, _gradients: &Vector<f32>) {
        unimplemented!(
            "Interior Point does not support stochastic updates (step). Use minimize() for constrained optimization."
        )
    }

    fn reset(&mut self) {
        // Reset barrier parameter to initial value
        self.mu = self.initial_mu;
    }
}

/// Stochastic Gradient Descent optimizer.
///
/// SGD updates parameters using the gradient of the loss function:
///
/// ```text
/// θ = θ - η * ∇L(θ)
/// ```
///
/// With momentum:
///
/// ```text
/// v = γ * v + η * ∇L(θ)
/// θ = θ - v
/// ```
///
/// where:
/// - θ is the parameter vector
/// - η is the learning rate
/// - γ is the momentum coefficient
/// - v is the velocity vector
/// - ∇L(θ) is the gradient of the loss
///
/// # Example
///
/// ```
/// use aprender::optim::SGD;
/// use aprender::primitives::Vector;
///
/// // Create SGD with momentum
/// let mut optimizer = SGD::new(0.1).with_momentum(0.9);
///
/// let mut params = Vector::from_slice(&[0.0, 0.0]);
/// let gradients = Vector::from_slice(&[1.0, 2.0]);
///
/// // First step
/// optimizer.step(&mut params, &gradients);
///
/// // With momentum, velocity builds up over iterations
/// optimizer.step(&mut params, &gradients);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SGD {
    /// Learning rate (step size)
    learning_rate: f32,
    /// Momentum coefficient (0.0 = no momentum)
    momentum: f32,
    /// Velocity vectors for momentum
    velocity: Option<Vec<f32>>,
}

impl SGD {
    /// Creates a new SGD optimizer with the given learning rate.
    ///
    /// # Arguments
    ///
    /// * `learning_rate` - Step size for parameter updates
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::SGD;
    ///
    /// let optimizer = SGD::new(0.01);
    /// assert!((optimizer.learning_rate() - 0.01).abs() < 1e-6);
    /// ```
    #[must_use]
    pub fn new(learning_rate: f32) -> Self {
        Self {
            learning_rate,
            momentum: 0.0,
            velocity: None,
        }
    }

    /// Sets the momentum coefficient.
    ///
    /// Momentum helps accelerate SGD in the relevant direction and dampens
    /// oscillations. Typical values are 0.9 or 0.99.
    ///
    /// # Arguments
    ///
    /// * `momentum` - Momentum coefficient between 0.0 and 1.0
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::SGD;
    ///
    /// let optimizer = SGD::new(0.01).with_momentum(0.9);
    /// assert!((optimizer.momentum() - 0.9).abs() < 1e-6);
    /// ```
    #[must_use]
    pub fn with_momentum(mut self, momentum: f32) -> Self {
        self.momentum = momentum;
        self
    }

    /// Returns the learning rate.
    #[must_use]
    pub fn learning_rate(&self) -> f32 {
        self.learning_rate
    }

    /// Returns the momentum coefficient.
    #[must_use]
    pub fn momentum(&self) -> f32 {
        self.momentum
    }

    /// Updates parameters using gradients.
    ///
    /// If momentum is enabled, maintains velocity vectors for each parameter.
    ///
    /// # Arguments
    ///
    /// * `params` - Mutable reference to parameter vector
    /// * `gradients` - Gradient vector (same length as params)
    ///
    /// # Panics
    ///
    /// Panics if params and gradients have different lengths.
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::SGD;
    /// use aprender::primitives::Vector;
    ///
    /// let mut optimizer = SGD::new(0.1);
    /// let mut params = Vector::from_slice(&[1.0, 2.0]);
    /// let gradients = Vector::from_slice(&[0.5, 1.0]);
    ///
    /// optimizer.step(&mut params, &gradients);
    ///
    /// // params = [1.0 - 0.1*0.5, 2.0 - 0.1*1.0] = [0.95, 1.9]
    /// assert!((params[0] - 0.95).abs() < 1e-6);
    /// assert!((params[1] - 1.9).abs() < 1e-6);
    /// ```
    pub fn step(&mut self, params: &mut Vector<f32>, gradients: &Vector<f32>) {
        assert_eq!(
            params.len(),
            gradients.len(),
            "Parameters and gradients must have same length"
        );

        let n = params.len();

        if self.momentum > 0.0 {
            // Initialize velocity if needed
            if self.velocity.is_none()
                || self
                    .velocity
                    .as_ref()
                    .expect("Velocity must be initialized")
                    .len()
                    != n
            {
                self.velocity = Some(vec![0.0; n]);
            }

            let velocity = self
                .velocity
                .as_mut()
                .expect("Velocity was just initialized");

            for i in 0..n {
                // v = γ * v + η * gradient
                velocity[i] = self.momentum * velocity[i] + self.learning_rate * gradients[i];
                // θ = θ - v
                params[i] -= velocity[i];
            }
        } else {
            // Standard SGD: θ = θ - η * gradient
            for i in 0..n {
                params[i] -= self.learning_rate * gradients[i];
            }
        }
    }

    /// Resets the optimizer state (velocity vectors).
    ///
    /// Call this when starting training on a new model or after significant
    /// changes to the optimization problem.
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::SGD;
    /// use aprender::primitives::Vector;
    ///
    /// let mut optimizer = SGD::new(0.1).with_momentum(0.9);
    /// let mut params = Vector::from_slice(&[1.0]);
    /// let gradients = Vector::from_slice(&[1.0]);
    ///
    /// optimizer.step(&mut params, &gradients);
    /// optimizer.reset();
    ///
    /// // Velocity is now reset to zero
    /// ```
    pub fn reset(&mut self) {
        self.velocity = None;
    }

    /// Returns whether momentum is enabled.
    #[must_use]
    pub fn has_momentum(&self) -> bool {
        self.momentum > 0.0
    }
}

/// Adam (Adaptive Moment Estimation) optimizer.
///
/// Adam combines the benefits of `AdaGrad` and `RMSprop` by computing adaptive learning
/// rates for each parameter using estimates of first and second moments of gradients.
///
/// Update rules:
///
/// ```text
/// m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
/// v_t = β₂ * v_{t-1} + (1 - β₂) * g_t²
/// m̂_t = m_t / (1 - β₁^t)
/// v̂_t = v_t / (1 - β₂^t)
/// θ_t = θ_{t-1} - α * m̂_t / (√v̂_t + ε)
/// ```
///
/// where:
/// - `m_t` is the first moment (mean) estimate
/// - `v_t` is the second moment (variance) estimate
/// - β₁, β₂ are exponential decay rates (typically 0.9, 0.999)
/// - α is the learning rate (step size)
/// - ε is a small constant for numerical stability (typically 1e-8)
///
/// # Example
///
/// ```
/// use aprender::optim::Adam;
/// use aprender::primitives::Vector;
///
/// // Create Adam optimizer with default hyperparameters
/// let mut optimizer = Adam::new(0.001);
///
/// let mut params = Vector::from_slice(&[1.0, 2.0]);
/// let gradients = Vector::from_slice(&[0.1, 0.2]);
///
/// // Update parameters with adaptive learning rates
/// optimizer.step(&mut params, &gradients);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Adam {
    /// Learning rate (step size)
    learning_rate: f32,
    /// Exponential decay rate for first moment estimates (default: 0.9)
    beta1: f32,
    /// Exponential decay rate for second moment estimates (default: 0.999)
    beta2: f32,
    /// Small constant for numerical stability (default: 1e-8)
    epsilon: f32,
    /// First moment estimates (mean)
    m: Option<Vec<f32>>,
    /// Second moment estimates (uncentered variance)
    v: Option<Vec<f32>>,
    /// Number of steps taken (for bias correction)
    t: usize,
}

impl Adam {
    /// Creates a new Adam optimizer with the given learning rate and default hyperparameters.
    ///
    /// Default values:
    /// - beta1 = 0.9
    /// - beta2 = 0.999
    /// - epsilon = 1e-8
    ///
    /// # Arguments
    ///
    /// * `learning_rate` - Step size (typical values: 0.001, 0.0001)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::Adam;
    ///
    /// let optimizer = Adam::new(0.001);
    /// assert!((optimizer.learning_rate() - 0.001).abs() < 1e-9);
    /// ```
    #[must_use]
    pub fn new(learning_rate: f32) -> Self {
        Self {
            learning_rate,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            m: None,
            v: None,
            t: 0,
        }
    }

    /// Sets the beta1 parameter (exponential decay rate for first moment).
    ///
    /// # Arguments
    ///
    /// * `beta1` - Value between 0.0 and 1.0 (typical: 0.9)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::Adam;
    ///
    /// let optimizer = Adam::new(0.001).with_beta1(0.95);
    /// assert!((optimizer.beta1() - 0.95).abs() < 1e-9);
    /// ```
    #[must_use]
    pub fn with_beta1(mut self, beta1: f32) -> Self {
        self.beta1 = beta1;
        self
    }

    /// Sets the beta2 parameter (exponential decay rate for second moment).
    ///
    /// # Arguments
    ///
    /// * `beta2` - Value between 0.0 and 1.0 (typical: 0.999)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::Adam;
    ///
    /// let optimizer = Adam::new(0.001).with_beta2(0.9999);
    /// assert!((optimizer.beta2() - 0.9999).abs() < 1e-9);
    /// ```
    #[must_use]
    pub fn with_beta2(mut self, beta2: f32) -> Self {
        self.beta2 = beta2;
        self
    }

    /// Sets the epsilon parameter (numerical stability constant).
    ///
    /// # Arguments
    ///
    /// * `epsilon` - Small positive value (typical: 1e-8)
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::Adam;
    ///
    /// let optimizer = Adam::new(0.001).with_epsilon(1e-7);
    /// assert!((optimizer.epsilon() - 1e-7).abs() < 1e-15);
    /// ```
    #[must_use]
    pub fn with_epsilon(mut self, epsilon: f32) -> Self {
        self.epsilon = epsilon;
        self
    }

    /// Returns the learning rate.
    #[must_use]
    pub fn learning_rate(&self) -> f32 {
        self.learning_rate
    }

    /// Returns the beta1 parameter.
    #[must_use]
    pub fn beta1(&self) -> f32 {
        self.beta1
    }

    /// Returns the beta2 parameter.
    #[must_use]
    pub fn beta2(&self) -> f32 {
        self.beta2
    }

    /// Returns the epsilon parameter.
    #[must_use]
    pub fn epsilon(&self) -> f32 {
        self.epsilon
    }

    /// Returns the number of steps taken.
    #[must_use]
    pub fn steps(&self) -> usize {
        self.t
    }

    /// Updates parameters using gradients with adaptive learning rates.
    ///
    /// # Arguments
    ///
    /// * `params` - Mutable reference to parameter vector
    /// * `gradients` - Gradient vector (same length as params)
    ///
    /// # Panics
    ///
    /// Panics if params and gradients have different lengths.
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::Adam;
    /// use aprender::primitives::Vector;
    ///
    /// let mut optimizer = Adam::new(0.001);
    /// let mut params = Vector::from_slice(&[1.0, 2.0]);
    /// let gradients = Vector::from_slice(&[0.1, 0.2]);
    ///
    /// optimizer.step(&mut params, &gradients);
    /// ```
    pub fn step(&mut self, params: &mut Vector<f32>, gradients: &Vector<f32>) {
        assert_eq!(
            params.len(),
            gradients.len(),
            "Parameters and gradients must have same length"
        );

        let n = params.len();

        // Initialize moment estimates if needed
        if self.m.is_none()
            || self
                .m
                .as_ref()
                .expect("First moment estimate must be initialized")
                .len()
                != n
        {
            self.m = Some(vec![0.0; n]);
            self.v = Some(vec![0.0; n]);
            self.t = 0;
        }

        self.t += 1;
        let t = self.t as f32;

        let m = self.m.as_mut().expect("First moment was just initialized");
        let v = self.v.as_mut().expect("Second moment was just initialized");

        // Compute bias-corrected learning rate
        let lr_t =
            self.learning_rate * (1.0 - self.beta2.powf(t)).sqrt() / (1.0 - self.beta1.powf(t));

        for i in 0..n {
            let g = gradients[i];

            // Update biased first moment estimate
            m[i] = self.beta1 * m[i] + (1.0 - self.beta1) * g;

            // Update biased second raw moment estimate
            v[i] = self.beta2 * v[i] + (1.0 - self.beta2) * g * g;

            // Update parameters
            params[i] -= lr_t * m[i] / (v[i].sqrt() + self.epsilon);
        }
    }

    /// Resets the optimizer state (moment estimates and step counter).
    ///
    /// Call this when starting training on a new model or after significant
    /// changes to the optimization problem.
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::Adam;
    /// use aprender::primitives::Vector;
    ///
    /// let mut optimizer = Adam::new(0.001);
    /// let mut params = Vector::from_slice(&[1.0]);
    /// let gradients = Vector::from_slice(&[1.0]);
    ///
    /// optimizer.step(&mut params, &gradients);
    /// assert_eq!(optimizer.steps(), 1);
    ///
    /// optimizer.reset();
    /// assert_eq!(optimizer.steps(), 0);
    /// ```
    pub fn reset(&mut self) {
        self.m = None;
        self.v = None;
        self.t = 0;
    }
}

impl Optimizer for Adam {
    fn step(&mut self, params: &mut Vector<f32>, gradients: &Vector<f32>) {
        self.step(params, gradients);
    }

    fn reset(&mut self) {
        self.reset();
    }
}

/// Unified trait for both stochastic and batch optimizers.
///
/// This trait supports two modes of optimization:
///
/// 1. **Stochastic mode** (`step`): For mini-batch training with SGD, Adam, etc.
/// 2. **Batch mode** (`minimize`): For full-dataset optimization with L-BFGS, CG, etc.
///
/// # Type Safety
///
/// The compiler prevents misuse:
/// - L-BFGS cannot be used with `step()` (would give poor results with stochastic gradients)
/// - SGD/Adam don't implement `minimize()` (inefficient for full datasets)
///
/// # Example
///
/// ```
/// use aprender::optim::{Optimizer, SGD};
/// use aprender::primitives::Vector;
///
/// // Stochastic mode (mini-batch training)
/// let mut adam = SGD::new(0.01);
/// let mut params = Vector::from_slice(&[1.0, 2.0]);
/// let grad = Vector::from_slice(&[0.1, 0.2]);
/// adam.step(&mut params, &grad);
/// ```
pub trait Optimizer {
    /// Stochastic update (mini-batch mode) - for SGD, Adam, `RMSprop`.
    ///
    /// Updates parameters in-place given gradient from current mini-batch.
    /// Used in ML training loops where gradients come from different data batches.
    ///
    /// # Arguments
    ///
    /// * `params` - Parameter vector to update (modified in-place)
    /// * `gradients` - Gradient vector from current mini-batch
    ///
    /// # Example
    ///
    /// ```
    /// use aprender::optim::{Optimizer, SGD};
    /// use aprender::primitives::Vector;
    ///
    /// let mut optimizer = SGD::new(0.1);
    /// let mut params = Vector::from_slice(&[1.0, 2.0]);
    ///
    /// for _epoch in 0..10 {
    ///     let grad = Vector::from_slice(&[0.1, 0.2]); // From mini-batch
    ///     optimizer.step(&mut params, &grad);
    /// }
    /// ```
    fn step(&mut self, params: &mut Vector<f32>, gradients: &Vector<f32>);

    /// Batch optimization (deterministic mode) - for L-BFGS, CG, Damped Newton.
    ///
    /// Minimizes objective function with full dataset access.
    /// Returns complete optimization trajectory and convergence info.
    ///
    /// **Default implementation**: Not all optimizers support batch mode. Stochastic
    /// optimizers (SGD, Adam) will panic if you call this method.
    ///
    /// # Arguments
    ///
    /// * `objective` - Objective function f: ℝⁿ → ℝ
    /// * `gradient` - Gradient function ∇f: ℝⁿ → ℝⁿ
    /// * `x0` - Initial point
    ///
    /// # Returns
    ///
    /// [`OptimizationResult`] with solution, convergence status, and diagnostics.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use aprender::optim::{Optimizer, LBFGS};
    /// use aprender::primitives::Vector;
    ///
    /// let mut optimizer = LBFGS::new(100, 1e-5, 10);
    ///
    /// let objective = |x: &Vector<f32>| (x[0] - 5.0).powi(2);
    /// let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * (x[0] - 5.0)]);
    ///
    /// let result = optimizer.minimize(objective, gradient, Vector::from_slice(&[0.0]));
    /// assert_eq!(result.status, ConvergenceStatus::Converged);
    /// ```
    fn minimize<F, G>(
        &mut self,
        _objective: F,
        _gradient: G,
        _x0: Vector<f32>,
    ) -> OptimizationResult
    where
        F: Fn(&Vector<f32>) -> f32,
        G: Fn(&Vector<f32>) -> Vector<f32>,
    {
        unimplemented!(
            "{} does not support batch optimization (minimize). Use step() for stochastic updates.",
            std::any::type_name::<Self>()
        )
    }

    /// Resets the optimizer state (momentum, history, etc.).
    ///
    /// Call this when starting training on a new model or after significant
    /// changes to the optimization problem.
    fn reset(&mut self);
}

impl Optimizer for SGD {
    fn step(&mut self, params: &mut Vector<f32>, gradients: &Vector<f32>) {
        self.step(params, gradients);
    }

    fn reset(&mut self) {
        self.reset();
    }
}

#[cfg(test)]
#[allow(non_snake_case)] // Allow mathematical matrix notation (A, B, Q, etc.)
mod tests {
    use super::*;

    // ==================== SafeCholesky Tests ====================

    #[test]
    fn test_safe_cholesky_solve_positive_definite() {
        // Well-conditioned positive definite matrix
        let A = Matrix::from_vec(2, 2, vec![4.0, 2.0, 2.0, 3.0]).expect("valid dimensions");
        let b = Vector::from_slice(&[6.0, 5.0]);

        let x = safe_cholesky_solve(&A, &b, 1e-8, 10).expect("should solve");
        assert_eq!(x.len(), 2);

        // Verify solution: Ax should equal b (approximately)
        let Ax = Vector::from_slice(&[
            A.get(0, 0) * x[0] + A.get(0, 1) * x[1],
            A.get(1, 0) * x[0] + A.get(1, 1) * x[1],
        ]);
        assert!((Ax[0] - b[0]).abs() < 1e-5);
        assert!((Ax[1] - b[1]).abs() < 1e-5);
    }

    #[test]
    fn test_safe_cholesky_solve_identity() {
        // Identity matrix - should solve without regularization
        let A = Matrix::eye(3);
        let b = Vector::from_slice(&[1.0, 2.0, 3.0]);

        let x = safe_cholesky_solve(&A, &b, 1e-8, 10).expect("should solve");

        // For identity matrix, x should equal b
        assert!((x[0] - 1.0).abs() < 1e-6);
        assert!((x[1] - 2.0).abs() < 1e-6);
        assert!((x[2] - 3.0).abs() < 1e-6);
    }

    #[test]
    fn test_safe_cholesky_solve_ill_conditioned() {
        // Ill-conditioned but solvable with regularization
        let A = Matrix::from_vec(2, 2, vec![1.0, 0.0, 0.0, 1e-10]).expect("valid dimensions");
        let b = Vector::from_slice(&[1.0, 1.0]);

        // Should succeed with regularization
        let x = safe_cholesky_solve(&A, &b, 1e-8, 10).expect("should solve with regularization");
        assert_eq!(x.len(), 2);

        // First component should be close to 1.0
        assert!((x[0] - 1.0).abs() < 1e-3);
    }

    #[test]
    fn test_safe_cholesky_solve_not_positive_definite() {
        // Matrix with negative eigenvalue - needs regularization
        let A = Matrix::from_vec(2, 2, vec![1.0, 0.0, 0.0, -0.5]).expect("valid dimensions");
        let b = Vector::from_slice(&[1.0, 1.0]);

        // Should solve with enough regularization
        let result = safe_cholesky_solve(&A, &b, 1e-4, 10);

        // May succeed with regularization or fail gracefully
        if let Ok(x) = result {
            assert_eq!(x.len(), 2);
            // Solution exists with regularization
        } else {
            // Also acceptable - matrix is indefinite
        }
    }

    #[test]
    fn test_safe_cholesky_solve_zero_matrix() {
        // Zero matrix - should fail even with regularization
        let A = Matrix::from_vec(2, 2, vec![0.0, 0.0, 0.0, 0.0]).expect("valid dimensions");
        let b = Vector::from_slice(&[1.0, 1.0]);

        // Should eventually succeed when regularization dominates
        let result = safe_cholesky_solve(&A, &b, 1e-4, 10);
        assert!(result.is_ok()); // Regularization makes it λI which is PD
    }

    #[test]
    fn test_safe_cholesky_solve_small_initial_lambda() {
        // Test with very small initial lambda
        let A = Matrix::eye(2);
        let b = Vector::from_slice(&[1.0, 1.0]);

        let x = safe_cholesky_solve(&A, &b, 1e-12, 10).expect("should solve");

        // Should still work for identity matrix
        assert!((x[0] - 1.0).abs() < 1e-6);
        assert!((x[1] - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_safe_cholesky_solve_max_attempts() {
        // Test that max_attempts is respected
        let A = Matrix::from_vec(2, 2, vec![1.0, 0.0, 0.0, 1.0]).expect("valid dimensions");
        let b = Vector::from_slice(&[1.0, 1.0]);

        // Even with 1 attempt, should work for identity
        let x = safe_cholesky_solve(&A, &b, 1e-8, 1).expect("should solve");
        assert_eq!(x.len(), 2);
    }

    #[test]
    fn test_safe_cholesky_solve_large_system() {
        // Test with larger system
        let n = 5;
        let mut data = vec![0.0; n * n];
        for i in 0..n {
            data[i * n + i] = 2.0; // Diagonal
            if i > 0 {
                data[i * n + (i - 1)] = 1.0; // Sub-diagonal
                data[(i - 1) * n + i] = 1.0; // Super-diagonal
            }
        }
        let A = Matrix::from_vec(n, n, data).expect("valid dimensions");
        let b = Vector::from_slice(&[1.0, 1.0, 1.0, 1.0, 1.0]);

        let x = safe_cholesky_solve(&A, &b, 1e-8, 10).expect("should solve");
        assert_eq!(x.len(), 5);
    }

    #[test]
    fn test_safe_cholesky_solve_symmetric() {
        // Verify it works with symmetric matrix
        let A = Matrix::from_vec(3, 3, vec![2.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 2.0])
            .expect("valid dimensions");
        let b = Vector::from_slice(&[1.0, 2.0, 1.0]);

        let x = safe_cholesky_solve(&A, &b, 1e-8, 10).expect("should solve");
        assert_eq!(x.len(), 3);
    }

    #[test]
    fn test_safe_cholesky_solve_lambda_escalation() {
        // Test that lambda increases when needed
        // This matrix might need several regularization attempts
        let A = Matrix::from_vec(2, 2, vec![1.0, 0.999, 0.999, 1.0]).expect("valid dimensions");
        let b = Vector::from_slice(&[1.0, 1.0]);

        let x = safe_cholesky_solve(&A, &b, 1e-10, 15).expect("should solve");
        assert_eq!(x.len(), 2);

        // Solution should exist
        assert!(x[0].is_finite());
        assert!(x[1].is_finite());
    }

    // ==================== SGD Tests ====================

    #[test]
    fn test_sgd_new() {
        let optimizer = SGD::new(0.01);
        assert!((optimizer.learning_rate() - 0.01).abs() < 1e-6);
        assert!((optimizer.momentum() - 0.0).abs() < 1e-6);
        assert!(!optimizer.has_momentum());
    }

    #[test]
    fn test_sgd_with_momentum() {
        let optimizer = SGD::new(0.01).with_momentum(0.9);
        assert!((optimizer.learning_rate() - 0.01).abs() < 1e-6);
        assert!((optimizer.momentum() - 0.9).abs() < 1e-6);
        assert!(optimizer.has_momentum());
    }

    #[test]
    fn test_sgd_step_basic() {
        let mut optimizer = SGD::new(0.1);
        let mut params = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let gradients = Vector::from_slice(&[1.0, 2.0, 3.0]);

        optimizer.step(&mut params, &gradients);

        // params = params - lr * gradients
        assert!((params[0] - 0.9).abs() < 1e-6);
        assert!((params[1] - 1.8).abs() < 1e-6);
        assert!((params[2] - 2.7).abs() < 1e-6);
    }

    #[test]
    fn test_sgd_step_with_momentum() {
        let mut optimizer = SGD::new(0.1).with_momentum(0.9);
        let mut params = Vector::from_slice(&[1.0, 1.0]);
        let gradients = Vector::from_slice(&[1.0, 1.0]);

        // First step: v = 0.9*0 + 0.1*1 = 0.1, params = 1.0 - 0.1 = 0.9
        optimizer.step(&mut params, &gradients);
        assert!((params[0] - 0.9).abs() < 1e-6);

        // Second step: v = 0.9*0.1 + 0.1*1 = 0.19, params = 0.9 - 0.19 = 0.71
        optimizer.step(&mut params, &gradients);
        assert!((params[0] - 0.71).abs() < 1e-6);
    }

    #[test]
    fn test_sgd_momentum_accumulation() {
        let mut optimizer = SGD::new(0.1).with_momentum(0.9);
        let mut params = Vector::from_slice(&[0.0]);
        let gradients = Vector::from_slice(&[1.0]);

        // Velocity should accumulate over iterations
        let mut prev_step = 0.0;
        for _ in 0..10 {
            let before = params[0];
            optimizer.step(&mut params, &gradients);
            let step = before - params[0];
            // Each step should be larger (velocity builds up)
            assert!(step >= prev_step);
            prev_step = step;
        }
    }

    #[test]
    fn test_sgd_reset() {
        let mut optimizer = SGD::new(0.1).with_momentum(0.9);
        let mut params = Vector::from_slice(&[1.0]);
        let gradients = Vector::from_slice(&[1.0]);

        optimizer.step(&mut params, &gradients);
        optimizer.reset();

        // After reset, velocity should be zero again
        let mut params2 = Vector::from_slice(&[1.0]);
        optimizer.step(&mut params2, &gradients);
        assert!((params2[0] - 0.9).abs() < 1e-6);
    }

    #[test]
    fn test_sgd_zero_gradient() {
        let mut optimizer = SGD::new(0.1);
        let mut params = Vector::from_slice(&[1.0, 2.0]);
        let gradients = Vector::from_slice(&[0.0, 0.0]);

        optimizer.step(&mut params, &gradients);

        // No change with zero gradients
        assert!((params[0] - 1.0).abs() < 1e-6);
        assert!((params[1] - 2.0).abs() < 1e-6);
    }

    #[test]
    fn test_sgd_negative_gradients() {
        let mut optimizer = SGD::new(0.1);
        let mut params = Vector::from_slice(&[1.0]);
        let gradients = Vector::from_slice(&[-1.0]);

        optimizer.step(&mut params, &gradients);

        // params = 1.0 - 0.1 * (-1.0) = 1.1
        assert!((params[0] - 1.1).abs() < 1e-6);
    }

    #[test]
    #[should_panic(expected = "same length")]
    fn test_sgd_mismatched_lengths() {
        let mut optimizer = SGD::new(0.1);
        let mut params = Vector::from_slice(&[1.0, 2.0]);
        let gradients = Vector::from_slice(&[1.0]);

        optimizer.step(&mut params, &gradients);
    }

    #[test]
    fn test_sgd_large_learning_rate() {
        let mut optimizer = SGD::new(10.0);
        let mut params = Vector::from_slice(&[1.0]);
        let gradients = Vector::from_slice(&[0.1]);

        optimizer.step(&mut params, &gradients);

        // params = 1.0 - 10.0 * 0.1 = 0.0
        assert!((params[0] - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_sgd_small_learning_rate() {
        let mut optimizer = SGD::new(0.001);
        let mut params = Vector::from_slice(&[1.0]);
        let gradients = Vector::from_slice(&[1.0]);

        optimizer.step(&mut params, &gradients);

        // params = 1.0 - 0.001 * 1.0 = 0.999
        assert!((params[0] - 0.999).abs() < 1e-6);
    }

    #[test]
    fn test_sgd_clone() {
        let optimizer = SGD::new(0.01).with_momentum(0.9);
        let cloned = optimizer.clone();

        assert!((cloned.learning_rate() - optimizer.learning_rate()).abs() < 1e-6);
        assert!((cloned.momentum() - optimizer.momentum()).abs() < 1e-6);
    }

    #[test]
    fn test_sgd_multiple_steps() {
        let mut optimizer = SGD::new(0.1);
        let mut params = Vector::from_slice(&[10.0]);
        let gradients = Vector::from_slice(&[1.0]);

        for _ in 0..10 {
            optimizer.step(&mut params, &gradients);
        }

        // params = 10.0 - 10 * 0.1 * 1.0 = 9.0
        assert!((params[0] - 9.0).abs() < 1e-4);
    }

    #[test]
    fn test_sgd_velocity_reinitialization() {
        let mut optimizer = SGD::new(0.1).with_momentum(0.9);

        // First with 2 params
        let mut params = Vector::from_slice(&[1.0, 2.0]);
        let gradients = Vector::from_slice(&[1.0, 1.0]);
        optimizer.step(&mut params, &gradients);

        // Now with 3 params - velocity should reinitialize
        let mut params3 = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let gradients3 = Vector::from_slice(&[1.0, 1.0, 1.0]);
        optimizer.step(&mut params3, &gradients3);

        // Should work without error, velocity reinitialized
        assert!((params3[0] - 0.9).abs() < 1e-6);
    }

    // ==================== Adam Tests ====================

    #[test]
    fn test_adam_new() {
        let optimizer = Adam::new(0.001);
        assert!((optimizer.learning_rate() - 0.001).abs() < 1e-9);
        assert!((optimizer.beta1() - 0.9).abs() < 1e-9);
        assert!((optimizer.beta2() - 0.999).abs() < 1e-9);
        assert!((optimizer.epsilon() - 1e-8).abs() < 1e-15);
        assert_eq!(optimizer.steps(), 0);
    }

    #[test]
    fn test_adam_with_beta1() {
        let optimizer = Adam::new(0.001).with_beta1(0.95);
        assert!((optimizer.beta1() - 0.95).abs() < 1e-9);
    }

    #[test]
    fn test_adam_with_beta2() {
        let optimizer = Adam::new(0.001).with_beta2(0.9999);
        assert!((optimizer.beta2() - 0.9999).abs() < 1e-9);
    }

    #[test]
    fn test_adam_with_epsilon() {
        let optimizer = Adam::new(0.001).with_epsilon(1e-7);
        assert!((optimizer.epsilon() - 1e-7).abs() < 1e-15);
    }

    #[test]
    fn test_adam_step_basic() {
        let mut optimizer = Adam::new(0.001);
        let mut params = Vector::from_slice(&[1.0, 2.0]);
        let gradients = Vector::from_slice(&[0.1, 0.2]);

        optimizer.step(&mut params, &gradients);

        // Adam should update parameters (exact values depend on bias correction)
        assert!(params[0] < 1.0); // Should decrease
        assert!(params[1] < 2.0); // Should decrease
        assert_eq!(optimizer.steps(), 1);
    }

    #[test]
    fn test_adam_multiple_steps() {
        let mut optimizer = Adam::new(0.001);
        let mut params = Vector::from_slice(&[1.0]);
        let gradients = Vector::from_slice(&[1.0]);

        let initial = params[0];
        for _ in 0..5 {
            optimizer.step(&mut params, &gradients);
        }

        // Parameters should decrease over multiple steps
        assert!(params[0] < initial);
        assert_eq!(optimizer.steps(), 5);
    }

    #[test]
    fn test_adam_bias_correction() {
        let mut optimizer = Adam::new(0.01);
        let mut params = Vector::from_slice(&[10.0]);
        let gradients = Vector::from_slice(&[1.0]);

        // First step should have larger effective learning rate due to bias correction
        optimizer.step(&mut params, &gradients);
        let first_step_size = 10.0 - params[0];

        // Reset and try second step
        let mut optimizer2 = Adam::new(0.01);
        let mut params2 = Vector::from_slice(&[10.0]);
        optimizer2.step(&mut params2, &gradients);
        optimizer2.step(&mut params2, &gradients);
        let second_step_size = params[0] - params2[0];

        // First step should have larger update due to bias correction
        assert!(first_step_size > second_step_size * 0.5);
    }

    #[test]
    fn test_adam_reset() {
        let mut optimizer = Adam::new(0.001);
        let mut params = Vector::from_slice(&[1.0]);
        let gradients = Vector::from_slice(&[1.0]);

        optimizer.step(&mut params, &gradients);
        assert_eq!(optimizer.steps(), 1);

        optimizer.reset();
        assert_eq!(optimizer.steps(), 0);
    }

    #[test]
    fn test_adam_zero_gradient() {
        let mut optimizer = Adam::new(0.001);
        let mut params = Vector::from_slice(&[1.0, 2.0]);
        let gradients = Vector::from_slice(&[0.0, 0.0]);

        optimizer.step(&mut params, &gradients);

        // With zero gradients, params should not change significantly
        assert!((params[0] - 1.0).abs() < 0.01);
        assert!((params[1] - 2.0).abs() < 0.01);
    }

    #[test]
    fn test_adam_negative_gradients() {
        let mut optimizer = Adam::new(0.001);
        let mut params = Vector::from_slice(&[1.0]);
        let gradients = Vector::from_slice(&[-1.0]);

        optimizer.step(&mut params, &gradients);

        // With negative gradient, params should increase
        assert!(params[0] > 1.0);
    }

    #[test]
    #[should_panic(expected = "same length")]
    fn test_adam_mismatched_lengths() {
        let mut optimizer = Adam::new(0.001);
        let mut params = Vector::from_slice(&[1.0, 2.0]);
        let gradients = Vector::from_slice(&[1.0]);

        optimizer.step(&mut params, &gradients);
    }

    #[test]
    fn test_adam_clone() {
        let optimizer = Adam::new(0.001).with_beta1(0.95).with_beta2(0.9999);
        let cloned = optimizer.clone();

        assert!((cloned.learning_rate() - optimizer.learning_rate()).abs() < 1e-9);
        assert!((cloned.beta1() - optimizer.beta1()).abs() < 1e-9);
        assert!((cloned.beta2() - optimizer.beta2()).abs() < 1e-9);
        assert!((cloned.epsilon() - optimizer.epsilon()).abs() < 1e-15);
    }

    #[test]
    fn test_adam_adaptive_learning() {
        // Test that Adam adapts to gradient magnitudes
        let mut optimizer = Adam::new(0.01);
        let mut params = Vector::from_slice(&[1.0, 1.0]);

        // Large gradient on first param, small on second
        let gradients1 = Vector::from_slice(&[10.0, 0.1]);
        optimizer.step(&mut params, &gradients1);

        let step1_0 = 1.0 - params[0];
        let step1_1 = 1.0 - params[1];

        // Continue with same gradients
        optimizer.step(&mut params, &gradients1);

        // Adam should adapt - second param should take relatively larger steps
        // because it has more consistent small gradients
        assert!(step1_0 > 0.0);
        assert!(step1_1 > 0.0);
    }

    #[test]
    fn test_adam_vs_sgd_behavior() {
        // Test that Adam and SGD behave differently (not necessarily one better)
        let mut adam = Adam::new(0.001);
        let mut sgd = SGD::new(0.1);

        let mut params_adam = Vector::from_slice(&[5.0]);
        let mut params_sgd = Vector::from_slice(&[5.0]);

        // Gradient pointing towards 0
        for _ in 0..10 {
            let gradients = Vector::from_slice(&[1.0]);
            adam.step(&mut params_adam, &gradients);
            sgd.step(&mut params_sgd, &gradients);
        }

        // Both should decrease but behave differently
        assert!(params_adam[0] < 5.0);
        assert!(params_sgd[0] < 5.0);
        // They should produce different results due to different mechanisms
        assert!((params_adam[0] - params_sgd[0]).abs() > 0.01);
    }

    #[test]
    fn test_adam_moment_initialization() {
        let mut optimizer = Adam::new(0.001);

        // First with 2 params
        let mut params = Vector::from_slice(&[1.0, 2.0]);
        let gradients = Vector::from_slice(&[0.1, 0.2]);
        optimizer.step(&mut params, &gradients);

        // Now with 3 params - moments should reinitialize
        let mut params3 = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let gradients3 = Vector::from_slice(&[0.1, 0.2, 0.3]);
        optimizer.step(&mut params3, &gradients3);

        // Should work without error
        assert!(params3[0] < 1.0);
        assert!(params3[1] < 2.0);
        assert!(params3[2] < 3.0);
    }

    #[test]
    fn test_adam_numerical_stability() {
        let mut optimizer = Adam::new(0.001);
        let mut params = Vector::from_slice(&[1.0]);

        // Very large gradients should be handled stably
        let gradients = Vector::from_slice(&[1000.0]);
        optimizer.step(&mut params, &gradients);

        // Should not produce NaN or extreme values
        assert!(!params[0].is_nan());
        assert!(params[0].is_finite());
    }

    // ==================== Line Search Tests ====================

    #[test]
    fn test_backtracking_line_search_new() {
        let ls = BacktrackingLineSearch::new(1e-4, 0.5, 50);
        assert!((ls.c1 - 1e-4).abs() < 1e-10);
        assert!((ls.rho - 0.5).abs() < 1e-10);
        assert_eq!(ls.max_iter, 50);
    }

    #[test]
    fn test_backtracking_line_search_default() {
        let ls = BacktrackingLineSearch::default();
        assert!((ls.c1 - 1e-4).abs() < 1e-10);
        assert!((ls.rho - 0.5).abs() < 1e-10);
        assert_eq!(ls.max_iter, 50);
    }

    #[test]
    fn test_backtracking_line_search_quadratic() {
        // Test on simple quadratic: f(x) = x^2
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let ls = BacktrackingLineSearch::default();
        let x = Vector::from_slice(&[2.0]);
        let d = Vector::from_slice(&[-4.0]); // Gradient direction at x=2

        let alpha = ls.search(&f, &grad, &x, &d);

        // Should find positive step size
        assert!(alpha > 0.0);
        assert!(alpha <= 1.0);

        // Verify Armijo condition is satisfied
        let x_new_data = x[0] + alpha * d[0];
        let x_new = Vector::from_slice(&[x_new_data]);
        let fx = f(&x);
        let fx_new = f(&x_new);
        let grad_x = grad(&x);
        let dir_deriv = grad_x[0] * d[0];

        assert!(fx_new <= fx + ls.c1 * alpha * dir_deriv);
    }

    #[test]
    fn test_backtracking_line_search_rosenbrock() {
        // Rosenbrock function: f(x,y) = (1-x)^2 + 100(y-x^2)^2
        let f = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            (1.0 - x).powi(2) + 100.0 * (y - x * x).powi(2)
        };

        let grad = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            let dx = -2.0 * (1.0 - x) - 400.0 * x * (y - x * x);
            let dy = 200.0 * (y - x * x);
            Vector::from_slice(&[dx, dy])
        };

        let ls = BacktrackingLineSearch::default();
        let x = Vector::from_slice(&[0.0, 0.0]);
        let g = grad(&x);
        let d = Vector::from_slice(&[-g[0], -g[1]]); // Descent direction

        let alpha = ls.search(&f, &grad, &x, &d);

        assert!(alpha > 0.0);
    }

    #[test]
    fn test_backtracking_line_search_multidimensional() {
        // f(x) = ||x||^2
        let f = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += x[i] * x[i];
            }
            sum
        };

        let grad = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = 2.0 * x[i];
            }
            g
        };

        let ls = BacktrackingLineSearch::default();
        let x = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let g = grad(&x);
        let d = Vector::from_slice(&[-g[0], -g[1], -g[2]]);

        let alpha = ls.search(&f, &grad, &x, &d);

        assert!(alpha > 0.0);
        assert!(alpha <= 1.0);
    }

    #[test]
    fn test_wolfe_line_search_new() {
        let ls = WolfeLineSearch::new(1e-4, 0.9, 50);
        assert!((ls.c1 - 1e-4).abs() < 1e-10);
        assert!((ls.c2 - 0.9).abs() < 1e-10);
        assert_eq!(ls.max_iter, 50);
    }

    #[test]
    #[should_panic(expected = "Wolfe conditions require 0 < c1 < c2 < 1")]
    fn test_wolfe_line_search_invalid_c1_c2() {
        // c1 >= c2 should panic
        let _ = WolfeLineSearch::new(0.9, 0.5, 50);
    }

    #[test]
    #[should_panic(expected = "Wolfe conditions require 0 < c1 < c2 < 1")]
    fn test_wolfe_line_search_c1_negative() {
        let _ = WolfeLineSearch::new(-0.1, 0.9, 50);
    }

    #[test]
    #[should_panic(expected = "Wolfe conditions require 0 < c1 < c2 < 1")]
    fn test_wolfe_line_search_c2_too_large() {
        let _ = WolfeLineSearch::new(0.1, 1.5, 50);
    }

    #[test]
    fn test_wolfe_line_search_default() {
        let ls = WolfeLineSearch::default();
        assert!((ls.c1 - 1e-4).abs() < 1e-10);
        assert!((ls.c2 - 0.9).abs() < 1e-10);
        assert_eq!(ls.max_iter, 50);
    }

    #[test]
    fn test_wolfe_line_search_quadratic() {
        // Test on simple quadratic: f(x) = x^2
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let ls = WolfeLineSearch::default();
        let x = Vector::from_slice(&[2.0]);
        let d = Vector::from_slice(&[-4.0]); // Gradient direction at x=2

        let alpha = ls.search(&f, &grad, &x, &d);

        // Should find positive step size
        assert!(alpha > 0.0);

        // Verify both Wolfe conditions
        let x_new_data = x[0] + alpha * d[0];
        let x_new = Vector::from_slice(&[x_new_data]);
        let fx = f(&x);
        let fx_new = f(&x_new);
        let grad_x = grad(&x);
        let grad_new = grad(&x_new);
        let dir_deriv = grad_x[0] * d[0];
        let dir_deriv_new = grad_new[0] * d[0];

        // Armijo condition
        assert!(fx_new <= fx + ls.c1 * alpha * dir_deriv + 1e-6);

        // Curvature condition
        assert!(dir_deriv_new.abs() <= ls.c2 * dir_deriv.abs() + 1e-6);
    }

    #[test]
    fn test_wolfe_line_search_multidimensional() {
        // f(x) = ||x||^2
        let f = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += x[i] * x[i];
            }
            sum
        };

        let grad = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = 2.0 * x[i];
            }
            g
        };

        let ls = WolfeLineSearch::default();
        let x = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let g = grad(&x);
        let d = Vector::from_slice(&[-g[0], -g[1], -g[2]]);

        let alpha = ls.search(&f, &grad, &x, &d);

        assert!(alpha > 0.0);
    }

    #[test]
    fn test_backtracking_vs_wolfe() {
        // Compare backtracking and Wolfe on same problem
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let bt = BacktrackingLineSearch::default();
        let wolfe = WolfeLineSearch::default();

        let x = Vector::from_slice(&[1.0, 1.0]);
        let g = grad(&x);
        let d = Vector::from_slice(&[-g[0], -g[1]]);

        let alpha_bt = bt.search(&f, &grad, &x, &d);
        let alpha_wolfe = wolfe.search(&f, &grad, &x, &d);

        // Both should find valid step sizes
        assert!(alpha_bt > 0.0);
        assert!(alpha_wolfe > 0.0);

        // Wolfe often finds larger steps due to curvature condition
        // but not always, so just verify both are reasonable
        assert!(alpha_bt <= 1.0);
    }

    // ==================== OptimizationResult Tests ====================

    #[test]
    fn test_optimization_result_converged() {
        let solution = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let result = OptimizationResult::converged(solution.clone(), 42);

        assert_eq!(result.solution.len(), 3);
        assert_eq!(result.iterations, 42);
        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.objective_value - 0.0).abs() < 1e-10);
        assert!((result.gradient_norm - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_optimization_result_max_iterations() {
        let solution = Vector::from_slice(&[5.0]);
        let result = OptimizationResult::max_iterations(solution);

        assert_eq!(result.iterations, 0);
        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
    }

    #[test]
    fn test_convergence_status_equality() {
        assert_eq!(ConvergenceStatus::Converged, ConvergenceStatus::Converged);
        assert_ne!(
            ConvergenceStatus::Converged,
            ConvergenceStatus::MaxIterations
        );
        assert_ne!(
            ConvergenceStatus::Stalled,
            ConvergenceStatus::NumericalError
        );
    }

    // ==================== L-BFGS Tests ====================

    #[test]
    fn test_lbfgs_new() {
        let optimizer = LBFGS::new(100, 1e-5, 10);
        assert_eq!(optimizer.max_iter, 100);
        assert!((optimizer.tol - 1e-5).abs() < 1e-10);
        assert_eq!(optimizer.m, 10);
        assert_eq!(optimizer.s_history.len(), 0);
        assert_eq!(optimizer.y_history.len(), 0);
    }

    #[test]
    fn test_lbfgs_simple_quadratic() {
        // Minimize f(x) = x^2, optimal at x = 0
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = LBFGS::new(100, 1e-5, 5);
        let x0 = Vector::from_slice(&[5.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 1e-3);
        assert!(result.iterations < 100);
        assert!(result.gradient_norm < 1e-5);
    }

    #[test]
    fn test_lbfgs_multidimensional_quadratic() {
        // Minimize f(x) = ||x - c||^2 where c = [1, 2, 3]
        let c = [1.0, 2.0, 3.0];
        let f = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += (x[i] - c[i]).powi(2);
            }
            sum
        };

        let grad = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = 2.0 * (x[i] - c[i]);
            }
            g
        };

        let mut optimizer = LBFGS::new(100, 1e-5, 10);
        let x0 = Vector::from_slice(&[0.0, 0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        for (i, &target) in c.iter().enumerate().take(3) {
            assert!((result.solution[i] - target).abs() < 1e-3);
        }
    }

    #[test]
    fn test_lbfgs_rosenbrock() {
        // Rosenbrock function: f(x,y) = (1-x)^2 + 100(y-x^2)^2
        // Global minimum at (1, 1)
        let f = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            (1.0 - x).powi(2) + 100.0 * (y - x * x).powi(2)
        };

        let grad = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            let dx = -2.0 * (1.0 - x) - 400.0 * x * (y - x * x);
            let dy = 200.0 * (y - x * x);
            Vector::from_slice(&[dx, dy])
        };

        let mut optimizer = LBFGS::new(200, 1e-4, 10);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Should converge close to (1, 1)
        assert!(
            result.status == ConvergenceStatus::Converged
                || result.status == ConvergenceStatus::MaxIterations
        );
        assert!((result.solution[0] - 1.0).abs() < 0.1);
        assert!((result.solution[1] - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_lbfgs_sphere() {
        // Sphere function: f(x) = sum(x_i^2)
        let f = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += x[i] * x[i];
            }
            sum
        };

        let grad = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = 2.0 * x[i];
            }
            g
        };

        let mut optimizer = LBFGS::new(100, 1e-5, 10);
        let x0 = Vector::from_slice(&[5.0, -3.0, 2.0, -1.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        for i in 0..4 {
            assert!(result.solution[i].abs() < 1e-3);
        }
    }

    #[test]
    fn test_lbfgs_different_history_sizes() {
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);
        let x0 = Vector::from_slice(&[3.0, 4.0]);

        // Small history
        let mut opt_small = LBFGS::new(100, 1e-5, 3);
        let result_small = opt_small.minimize(f, grad, x0.clone());
        assert_eq!(result_small.status, ConvergenceStatus::Converged);

        // Large history
        let mut opt_large = LBFGS::new(100, 1e-5, 20);
        let result_large = opt_large.minimize(f, grad, x0);
        assert_eq!(result_large.status, ConvergenceStatus::Converged);

        // Both should converge to same solution
        assert!((result_small.solution[0] - result_large.solution[0]).abs() < 1e-3);
        assert!((result_small.solution[1] - result_large.solution[1]).abs() < 1e-3);
    }

    #[test]
    fn test_lbfgs_reset() {
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = LBFGS::new(100, 1e-5, 10);
        let x0 = Vector::from_slice(&[5.0]);

        // First run
        optimizer.minimize(f, grad, x0.clone());
        assert!(!optimizer.s_history.is_empty());

        // Reset
        optimizer.reset();
        assert_eq!(optimizer.s_history.len(), 0);
        assert_eq!(optimizer.y_history.len(), 0);

        // Second run should work
        let result = optimizer.minimize(f, grad, x0);
        assert_eq!(result.status, ConvergenceStatus::Converged);
    }

    #[test]
    fn test_lbfgs_max_iterations() {
        // Use Rosenbrock with very few iterations to force max_iter
        let f = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            (1.0 - x).powi(2) + 100.0 * (y - x * x).powi(2)
        };

        let grad = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            let dx = -2.0 * (1.0 - x) - 400.0 * x * (y - x * x);
            let dy = 200.0 * (y - x * x);
            Vector::from_slice(&[dx, dy])
        };

        let mut optimizer = LBFGS::new(3, 1e-10, 5);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        // With only 3 iterations, should hit max_iter on Rosenbrock
        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 3);
    }

    #[test]
    #[should_panic(expected = "does not support stochastic")]
    fn test_lbfgs_step_panics() {
        let mut optimizer = LBFGS::new(100, 1e-5, 10);
        let mut params = Vector::from_slice(&[1.0]);
        let grad = Vector::from_slice(&[0.1]);

        // Should panic - L-BFGS doesn't support step()
        optimizer.step(&mut params, &grad);
    }

    #[test]
    fn test_lbfgs_numerical_error_detection() {
        // Function that produces NaN
        let f = |x: &Vector<f32>| {
            if x[0] < -100.0 {
                f32::NAN
            } else {
                x[0] * x[0]
            }
        };
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = LBFGS::new(100, 1e-5, 5);
        let x0 = Vector::from_slice(&[0.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Should detect numerical error or converge normally
        assert!(
            result.status == ConvergenceStatus::Converged
                || result.status == ConvergenceStatus::NumericalError
                || result.status == ConvergenceStatus::MaxIterations
        );
    }

    #[test]
    fn test_lbfgs_computes_elapsed_time() {
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = LBFGS::new(100, 1e-5, 5);
        let x0 = Vector::from_slice(&[5.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Should have non-zero elapsed time
        assert!(result.elapsed_time.as_nanos() > 0);
    }

    #[test]
    fn test_lbfgs_gradient_norm_tracking() {
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let mut optimizer = LBFGS::new(100, 1e-5, 10);
        let x0 = Vector::from_slice(&[3.0, 4.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Gradient norm at convergence should be small
        if result.status == ConvergenceStatus::Converged {
            assert!(result.gradient_norm < 1e-5);
        }
    }

    // ==================== Conjugate Gradient Tests ====================

    #[test]
    fn test_cg_new_fletcher_reeves() {
        let optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::FletcherReeves);
        assert_eq!(optimizer.max_iter, 100);
        assert!((optimizer.tol - 1e-5).abs() < 1e-10);
        assert_eq!(optimizer.beta_formula, CGBetaFormula::FletcherReeves);
        assert_eq!(optimizer.restart_interval, 0);
    }

    #[test]
    fn test_cg_new_polak_ribiere() {
        let optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        assert_eq!(optimizer.beta_formula, CGBetaFormula::PolakRibiere);
    }

    #[test]
    fn test_cg_new_hestenes_stiefel() {
        let optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::HestenesStiefel);
        assert_eq!(optimizer.beta_formula, CGBetaFormula::HestenesStiefel);
    }

    #[test]
    fn test_cg_with_restart_interval() {
        let optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere)
            .with_restart_interval(50);
        assert_eq!(optimizer.restart_interval, 50);
    }

    #[test]
    fn test_cg_simple_quadratic_fr() {
        // Minimize f(x) = x^2, optimal at x = 0
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::FletcherReeves);
        let x0 = Vector::from_slice(&[5.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 1e-3);
    }

    #[test]
    fn test_cg_simple_quadratic_pr() {
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[5.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 1e-3);
    }

    #[test]
    fn test_cg_simple_quadratic_hs() {
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::HestenesStiefel);
        let x0 = Vector::from_slice(&[5.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 1e-3);
    }

    #[test]
    fn test_cg_multidimensional_quadratic() {
        // Minimize f(x) = ||x - c||^2 where c = [1, 2, 3]
        let c = [1.0, 2.0, 3.0];
        let f = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += (x[i] - c[i]).powi(2);
            }
            sum
        };

        let grad = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = 2.0 * (x[i] - c[i]);
            }
            g
        };

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[0.0, 0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        for (i, &target) in c.iter().enumerate() {
            assert!((result.solution[i] - target).abs() < 1e-3);
        }
    }

    #[test]
    fn test_cg_rosenbrock() {
        // Rosenbrock function: f(x,y) = (1-x)^2 + 100(y-x^2)^2
        let f = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            (1.0 - x).powi(2) + 100.0 * (y - x * x).powi(2)
        };

        let grad = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            let dx = -2.0 * (1.0 - x) - 400.0 * x * (y - x * x);
            let dy = 200.0 * (y - x * x);
            Vector::from_slice(&[dx, dy])
        };

        let mut optimizer = ConjugateGradient::new(500, 1e-4, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Should converge close to (1, 1)
        assert!(
            result.status == ConvergenceStatus::Converged
                || result.status == ConvergenceStatus::MaxIterations
        );
        assert!((result.solution[0] - 1.0).abs() < 0.1);
        assert!((result.solution[1] - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_cg_sphere() {
        // Sphere function: f(x) = sum(x_i^2)
        let f = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += x[i] * x[i];
            }
            sum
        };

        let grad = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = 2.0 * x[i];
            }
            g
        };

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[5.0, -3.0, 2.0, -1.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        for i in 0..4 {
            assert!(result.solution[i].abs() < 1e-3);
        }
    }

    #[test]
    fn test_cg_compare_beta_formulas() {
        // Compare different beta formulas on same problem
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);
        let x0 = Vector::from_slice(&[3.0, 4.0]);

        let mut opt_fr = ConjugateGradient::new(100, 1e-5, CGBetaFormula::FletcherReeves);
        let result_fr = opt_fr.minimize(f, grad, x0.clone());

        let mut opt_pr = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let result_pr = opt_pr.minimize(f, grad, x0.clone());

        let mut opt_hs = ConjugateGradient::new(100, 1e-5, CGBetaFormula::HestenesStiefel);
        let result_hs = opt_hs.minimize(f, grad, x0);

        // All should converge to same solution
        assert_eq!(result_fr.status, ConvergenceStatus::Converged);
        assert_eq!(result_pr.status, ConvergenceStatus::Converged);
        assert_eq!(result_hs.status, ConvergenceStatus::Converged);

        assert!(result_fr.solution[0].abs() < 1e-3);
        assert!(result_pr.solution[0].abs() < 1e-3);
        assert!(result_hs.solution[0].abs() < 1e-3);
    }

    #[test]
    fn test_cg_with_periodic_restart() {
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let mut optimizer =
            ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere).with_restart_interval(5);
        let x0 = Vector::from_slice(&[5.0, 5.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 1e-3);
        assert!(result.solution[1].abs() < 1e-3);
    }

    #[test]
    fn test_cg_reset() {
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[5.0]);

        // First run
        optimizer.minimize(f, grad, x0.clone());
        assert!(optimizer.prev_direction.is_some());

        // Reset
        optimizer.reset();
        assert!(optimizer.prev_direction.is_none());
        assert!(optimizer.prev_gradient.is_none());
        assert_eq!(optimizer.iter_count, 0);

        // Second run should work
        let result = optimizer.minimize(f, grad, x0);
        assert_eq!(result.status, ConvergenceStatus::Converged);
    }

    #[test]
    fn test_cg_max_iterations() {
        // Use Rosenbrock with very few iterations
        let f = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            (1.0 - x).powi(2) + 100.0 * (y - x * x).powi(2)
        };

        let grad = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            let dx = -2.0 * (1.0 - x) - 400.0 * x * (y - x * x);
            let dy = 200.0 * (y - x * x);
            Vector::from_slice(&[dx, dy])
        };

        let mut optimizer = ConjugateGradient::new(3, 1e-10, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 3);
    }

    #[test]
    #[should_panic(expected = "does not support stochastic")]
    fn test_cg_step_panics() {
        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let mut params = Vector::from_slice(&[1.0]);
        let grad = Vector::from_slice(&[0.1]);

        // Should panic - CG doesn't support step()
        optimizer.step(&mut params, &grad);
    }

    #[test]
    fn test_cg_numerical_error_detection() {
        // Function that produces NaN
        let f = |x: &Vector<f32>| {
            if x[0] < -100.0 {
                f32::NAN
            } else {
                x[0] * x[0]
            }
        };
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[0.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Should detect numerical error or converge normally
        assert!(
            result.status == ConvergenceStatus::Converged
                || result.status == ConvergenceStatus::NumericalError
                || result.status == ConvergenceStatus::MaxIterations
        );
    }

    #[test]
    fn test_cg_gradient_norm_tracking() {
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let mut optimizer = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let x0 = Vector::from_slice(&[3.0, 4.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Gradient norm at convergence should be small
        if result.status == ConvergenceStatus::Converged {
            assert!(result.gradient_norm < 1e-5);
        }
    }

    #[test]
    fn test_cg_vs_lbfgs_quadratic() {
        // Compare CG and L-BFGS on a quadratic problem
        let f = |x: &Vector<f32>| x[0] * x[0] + 2.0 * x[1] * x[1] + 3.0 * x[2] * x[2];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 4.0 * x[1], 6.0 * x[2]]);
        let x0 = Vector::from_slice(&[5.0, 3.0, 2.0]);

        let mut cg = ConjugateGradient::new(100, 1e-5, CGBetaFormula::PolakRibiere);
        let result_cg = cg.minimize(f, grad, x0.clone());

        let mut lbfgs = LBFGS::new(100, 1e-5, 10);
        let result_lbfgs = lbfgs.minimize(f, grad, x0);

        // Both should converge to same solution
        assert_eq!(result_cg.status, ConvergenceStatus::Converged);
        assert_eq!(result_lbfgs.status, ConvergenceStatus::Converged);

        for i in 0..3 {
            assert!((result_cg.solution[i] - result_lbfgs.solution[i]).abs() < 1e-3);
        }
    }

    // ==================== Damped Newton Tests ====================

    #[test]
    fn test_damped_newton_new() {
        let optimizer = DampedNewton::new(100, 1e-5);
        assert_eq!(optimizer.max_iter, 100);
        assert!((optimizer.tol - 1e-5).abs() < 1e-10);
        assert!((optimizer.epsilon - 1e-5).abs() < 1e-10);
    }

    #[test]
    fn test_damped_newton_with_epsilon() {
        let optimizer = DampedNewton::new(100, 1e-5).with_epsilon(1e-6);
        assert!((optimizer.epsilon - 1e-6).abs() < 1e-12);
    }

    #[test]
    fn test_damped_newton_simple_quadratic() {
        // Minimize f(x) = x^2, optimal at x = 0
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = DampedNewton::new(100, 1e-5);
        let x0 = Vector::from_slice(&[5.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 1e-3);
    }

    #[test]
    fn test_damped_newton_multidimensional_quadratic() {
        // Minimize f(x,y) = x^2 + 2y^2, optimal at (0, 0)
        let f = |x: &Vector<f32>| x[0] * x[0] + 2.0 * x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 4.0 * x[1]]);

        let mut optimizer = DampedNewton::new(100, 1e-5);
        let x0 = Vector::from_slice(&[5.0, 3.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 1e-3);
        assert!(result.solution[1].abs() < 1e-3);
    }

    #[test]
    fn test_damped_newton_rosenbrock() {
        // Rosenbrock function: f(x,y) = (1-x)^2 + 100(y-x^2)^2
        let f = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            (1.0 - x).powi(2) + 100.0 * (y - x * x).powi(2)
        };

        let grad = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            let dx = -2.0 * (1.0 - x) - 400.0 * x * (y - x * x);
            let dy = 200.0 * (y - x * x);
            Vector::from_slice(&[dx, dy])
        };

        let mut optimizer = DampedNewton::new(200, 1e-4);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Should converge close to (1, 1)
        assert!(
            result.status == ConvergenceStatus::Converged
                || result.status == ConvergenceStatus::MaxIterations
        );
        assert!((result.solution[0] - 1.0).abs() < 0.1);
        assert!((result.solution[1] - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_damped_newton_sphere() {
        // Sphere function: f(x) = sum(x_i^2)
        let f = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += x[i] * x[i];
            }
            sum
        };

        let grad = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = 2.0 * x[i];
            }
            g
        };

        let mut optimizer = DampedNewton::new(100, 1e-5);
        let x0 = Vector::from_slice(&[5.0, -3.0, 2.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        for i in 0..3 {
            assert!(result.solution[i].abs() < 1e-3);
        }
    }

    #[test]
    fn test_damped_newton_reset() {
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = DampedNewton::new(100, 1e-5);
        let x0 = Vector::from_slice(&[5.0]);

        // First run
        optimizer.minimize(f, grad, x0.clone());

        // Reset (stateless, so just verify it doesn't panic)
        optimizer.reset();

        // Second run should work
        let result = optimizer.minimize(f, grad, x0);
        assert_eq!(result.status, ConvergenceStatus::Converged);
    }

    #[test]
    fn test_damped_newton_max_iterations() {
        // Use Rosenbrock with very few iterations
        let f = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            (1.0 - x).powi(2) + 100.0 * (y - x * x).powi(2)
        };

        let grad = |v: &Vector<f32>| {
            let x = v[0];
            let y = v[1];
            let dx = -2.0 * (1.0 - x) - 400.0 * x * (y - x * x);
            let dy = 200.0 * (y - x * x);
            Vector::from_slice(&[dx, dy])
        };

        let mut optimizer = DampedNewton::new(3, 1e-10);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 3);
    }

    #[test]
    #[should_panic(expected = "does not support stochastic")]
    fn test_damped_newton_step_panics() {
        let mut optimizer = DampedNewton::new(100, 1e-5);
        let mut params = Vector::from_slice(&[1.0]);
        let grad = Vector::from_slice(&[0.1]);

        // Should panic - Damped Newton doesn't support step()
        optimizer.step(&mut params, &grad);
    }

    #[test]
    fn test_damped_newton_numerical_error_detection() {
        // Function that produces NaN
        let f = |x: &Vector<f32>| {
            if x[0] < -100.0 {
                f32::NAN
            } else {
                x[0] * x[0]
            }
        };
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);

        let mut optimizer = DampedNewton::new(100, 1e-5);
        let x0 = Vector::from_slice(&[0.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Should detect numerical error or converge normally
        assert!(
            result.status == ConvergenceStatus::Converged
                || result.status == ConvergenceStatus::NumericalError
                || result.status == ConvergenceStatus::MaxIterations
        );
    }

    #[test]
    fn test_damped_newton_gradient_norm_tracking() {
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let mut optimizer = DampedNewton::new(100, 1e-5);
        let x0 = Vector::from_slice(&[3.0, 4.0]);
        let result = optimizer.minimize(f, grad, x0);

        // Gradient norm at convergence should be small
        if result.status == ConvergenceStatus::Converged {
            assert!(result.gradient_norm < 1e-5);
        }
    }

    #[test]
    fn test_damped_newton_quadratic_convergence() {
        // Newton's method should converge quadratically on quadratic problems
        let f = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let mut optimizer = DampedNewton::new(100, 1e-10);
        let x0 = Vector::from_slice(&[5.0, 5.0]);
        let result = optimizer.minimize(f, grad, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        // Should converge in very few iterations for quadratic problems
        assert!(result.iterations < 20);
    }

    #[test]
    fn test_damped_newton_vs_lbfgs_quadratic() {
        // Compare Damped Newton and L-BFGS on a quadratic problem
        let f = |x: &Vector<f32>| x[0] * x[0] + 2.0 * x[1] * x[1];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 4.0 * x[1]]);
        let x0 = Vector::from_slice(&[5.0, 3.0]);

        let mut dn = DampedNewton::new(100, 1e-5);
        let result_dn = dn.minimize(f, grad, x0.clone());

        let mut lbfgs = LBFGS::new(100, 1e-5, 10);
        let result_lbfgs = lbfgs.minimize(f, grad, x0);

        // Both should converge to same solution
        assert_eq!(result_dn.status, ConvergenceStatus::Converged);
        assert_eq!(result_lbfgs.status, ConvergenceStatus::Converged);

        assert!((result_dn.solution[0] - result_lbfgs.solution[0]).abs() < 1e-3);
        assert!((result_dn.solution[1] - result_lbfgs.solution[1]).abs() < 1e-3);
    }

    #[test]
    fn test_damped_newton_different_epsilon() {
        // Test with different finite difference epsilons
        let f = |x: &Vector<f32>| x[0] * x[0];
        let grad = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0]]);
        let x0 = Vector::from_slice(&[5.0]);

        let mut opt1 = DampedNewton::new(100, 1e-5).with_epsilon(1e-5);
        let result1 = opt1.minimize(f, grad, x0.clone());

        let mut opt2 = DampedNewton::new(100, 1e-5).with_epsilon(1e-7);
        let result2 = opt2.minimize(f, grad, x0);

        // Both should converge
        assert_eq!(result1.status, ConvergenceStatus::Converged);
        assert_eq!(result2.status, ConvergenceStatus::Converged);

        // Solutions should be similar
        assert!((result1.solution[0] - result2.solution[0]).abs() < 1e-2);
    }

    // ==================== Proximal Operator Tests ====================

    #[test]
    fn test_soft_threshold_basic() {
        use crate::optim::prox::soft_threshold;

        let v = Vector::from_slice(&[2.0, -1.5, 0.5, 0.0]);
        let result = soft_threshold(&v, 1.0);

        assert!((result[0] - 1.0).abs() < 1e-6); // 2.0 - 1.0
        assert!((result[1] + 0.5).abs() < 1e-6); // -1.5 + 1.0
        assert!(result[2].abs() < 1e-6); // 0.5 - 1.0 -> 0
        assert!(result[3].abs() < 1e-6); // Already zero
    }

    #[test]
    fn test_soft_threshold_zero_lambda() {
        use crate::optim::prox::soft_threshold;

        let v = Vector::from_slice(&[1.0, -2.0, 3.0]);
        let result = soft_threshold(&v, 0.0);

        // With λ=0, should return input unchanged
        assert_eq!(result[0], 1.0);
        assert_eq!(result[1], -2.0);
        assert_eq!(result[2], 3.0);
    }

    #[test]
    fn test_soft_threshold_large_lambda() {
        use crate::optim::prox::soft_threshold;

        let v = Vector::from_slice(&[1.0, -1.0, 0.5]);
        let result = soft_threshold(&v, 10.0);

        // All values should be thresholded to zero
        assert!(result[0].abs() < 1e-6);
        assert!(result[1].abs() < 1e-6);
        assert!(result[2].abs() < 1e-6);
    }

    #[test]
    fn test_nonnegative_projection() {
        use crate::optim::prox::nonnegative;

        let x = Vector::from_slice(&[1.0, -2.0, 3.0, -0.5, 0.0]);
        let result = nonnegative(&x);

        assert_eq!(result[0], 1.0);
        assert_eq!(result[1], 0.0); // Projected to 0
        assert_eq!(result[2], 3.0);
        assert_eq!(result[3], 0.0); // Projected to 0
        assert_eq!(result[4], 0.0);
    }

    #[test]
    fn test_project_l2_ball_inside() {
        use crate::optim::prox::project_l2_ball;

        // Point inside ball - should be unchanged
        let x = Vector::from_slice(&[1.0, 1.0]); // norm = sqrt(2) ≈ 1.414
        let result = project_l2_ball(&x, 2.0);

        assert!((result[0] - 1.0).abs() < 1e-5);
        assert!((result[1] - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_project_l2_ball_outside() {
        use crate::optim::prox::project_l2_ball;

        // Point outside ball - should be scaled
        let x = Vector::from_slice(&[3.0, 4.0]); // norm = 5.0
        let result = project_l2_ball(&x, 2.0);

        // Should be scaled to norm = 2.0
        let norm = (result[0] * result[0] + result[1] * result[1]).sqrt();
        assert!((norm - 2.0).abs() < 1e-5);

        // Direction should be preserved
        let scale = 2.0 / 5.0;
        assert!((result[0] - 3.0 * scale).abs() < 1e-5);
        assert!((result[1] - 4.0 * scale).abs() < 1e-5);
    }

    #[test]
    fn test_project_box() {
        use crate::optim::prox::project_box;

        let x = Vector::from_slice(&[-1.0, 0.5, 2.0, 1.0]);
        let lower = Vector::from_slice(&[0.0, 0.0, 0.0, 0.0]);
        let upper = Vector::from_slice(&[1.0, 1.0, 1.0, 1.0]);

        let result = project_box(&x, &lower, &upper);

        assert_eq!(result[0], 0.0); // Clipped to lower
        assert_eq!(result[1], 0.5); // Within bounds
        assert_eq!(result[2], 1.0); // Clipped to upper
        assert_eq!(result[3], 1.0); // Within bounds
    }

    // ==================== FISTA Tests ====================

    #[test]
    fn test_fista_new() {
        let fista = FISTA::new(1000, 0.1, 1e-5);
        assert_eq!(fista.max_iter, 1000);
        assert!((fista.step_size - 0.1).abs() < 1e-9);
        assert!((fista.tol - 1e-5).abs() < 1e-9);
    }

    #[test]
    fn test_fista_l1_regularized_quadratic() {
        use crate::optim::prox::soft_threshold;

        // Minimize: ½(x - 5)² + 2|x|
        // Solution should be around x ≈ 3 (soft-threshold of 5 with λ=2)
        let smooth = |x: &Vector<f32>| 0.5 * (x[0] - 5.0).powi(2);
        let grad_smooth = |x: &Vector<f32>| Vector::from_slice(&[x[0] - 5.0]);
        let prox = |v: &Vector<f32>, alpha: f32| soft_threshold(v, 2.0 * alpha);

        let mut fista = FISTA::new(1000, 0.1, 1e-5);
        let x0 = Vector::from_slice(&[0.0]);
        let result = fista.minimize(smooth, grad_smooth, prox, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        // Analytical solution: sign(5) * max(|5| - 2, 0) = 3
        assert!((result.solution[0] - 3.0).abs() < 0.1);
    }

    #[test]
    fn test_fista_nonnegative_least_squares() {
        use crate::optim::prox::nonnegative;

        // Minimize: ½(x - (-2))² subject to x ≥ 0
        // Solution should be x = 0 (projection of -2 onto [0, ∞))
        let smooth = |x: &Vector<f32>| 0.5 * (x[0] + 2.0).powi(2);
        let grad_smooth = |x: &Vector<f32>| Vector::from_slice(&[x[0] + 2.0]);
        let prox = |v: &Vector<f32>, _alpha: f32| nonnegative(v);

        let mut fista = FISTA::new(1000, 0.1, 1e-6);
        let x0 = Vector::from_slice(&[1.0]);
        let result = fista.minimize(smooth, grad_smooth, prox, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.solution[0].abs() < 0.01); // Should be very close to 0
    }

    #[test]
    fn test_fista_box_constrained() {
        use crate::optim::prox::project_box;

        // Minimize: ½(x - 10)² subject to 0 ≤ x ≤ 1
        // Solution should be x = 1 (projection of 10 onto [0, 1])
        let smooth = |x: &Vector<f32>| 0.5 * (x[0] - 10.0).powi(2);
        let grad_smooth = |x: &Vector<f32>| Vector::from_slice(&[x[0] - 10.0]);

        let lower = Vector::from_slice(&[0.0]);
        let upper = Vector::from_slice(&[1.0]);
        let prox = move |v: &Vector<f32>, _alpha: f32| project_box(v, &lower, &upper);

        let mut fista = FISTA::new(1000, 0.1, 1e-6);
        let x0 = Vector::from_slice(&[0.5]);
        let result = fista.minimize(smooth, grad_smooth, prox, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 1.0).abs() < 0.01);
    }

    #[test]
    fn test_fista_multidimensional_lasso() {
        use crate::optim::prox::soft_threshold;

        // Minimize: ½‖x - c‖² + λ‖x‖₁ where c = [3, -2, 1]
        let c = [3.0, -2.0, 1.0];
        let lambda = 0.5;

        let smooth = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += 0.5 * (x[i] - c[i]).powi(2);
            }
            sum
        };

        let grad_smooth = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = x[i] - c[i];
            }
            g
        };

        let prox = move |v: &Vector<f32>, alpha: f32| soft_threshold(v, lambda * alpha);

        let mut fista = FISTA::new(1000, 0.1, 1e-6);
        let x0 = Vector::from_slice(&[0.0, 0.0, 0.0]);
        let result = fista.minimize(smooth, grad_smooth, prox, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);

        // Analytical solutions: sign(c[i]) * max(|c[i]| - λ, 0)
        assert!((result.solution[0] - 2.5).abs() < 0.1); // 3 - 0.5
        assert!((result.solution[1] + 1.5).abs() < 0.1); // -2 + 0.5
        assert!((result.solution[2] - 0.5).abs() < 0.1); // 1 - 0.5
    }

    #[test]
    fn test_fista_max_iterations() {
        use crate::optim::prox::soft_threshold;

        // Use a difficult problem with very few iterations
        let smooth = |x: &Vector<f32>| 0.5 * x[0].powi(2);
        let grad_smooth = |x: &Vector<f32>| Vector::from_slice(&[x[0]]);
        let prox = |v: &Vector<f32>, alpha: f32| soft_threshold(v, alpha);

        let mut fista = FISTA::new(3, 0.001, 1e-10); // Very few iterations
        let x0 = Vector::from_slice(&[10.0]);
        let result = fista.minimize(smooth, grad_smooth, prox, x0);

        // Should hit max iterations
        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 3);
    }

    #[test]
    fn test_fista_convergence_tracking() {
        use crate::optim::prox::soft_threshold;

        let smooth = |x: &Vector<f32>| 0.5 * x[0].powi(2);
        let grad_smooth = |x: &Vector<f32>| Vector::from_slice(&[x[0]]);
        let prox = |v: &Vector<f32>, alpha: f32| soft_threshold(v, 0.1 * alpha);

        let mut fista = FISTA::new(1000, 0.1, 1e-6);
        let x0 = Vector::from_slice(&[5.0]);
        let result = fista.minimize(smooth, grad_smooth, prox, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.iterations > 0);
        assert!(result.elapsed_time.as_nanos() > 0);
    }

    #[test]
    fn test_fista_vs_no_acceleration() {
        use crate::optim::prox::soft_threshold;

        // FISTA should converge faster than unaccelerated proximal gradient
        let smooth = |x: &Vector<f32>| {
            let mut sum = 0.0;
            for i in 0..x.len() {
                sum += 0.5 * (x[i] - (i as f32 + 1.0)).powi(2);
            }
            sum
        };

        let grad_smooth = |x: &Vector<f32>| {
            let mut g = Vector::zeros(x.len());
            for i in 0..x.len() {
                g[i] = x[i] - (i as f32 + 1.0);
            }
            g
        };

        let prox = |v: &Vector<f32>, alpha: f32| soft_threshold(v, 0.5 * alpha);

        let mut fista = FISTA::new(1000, 0.1, 1e-5);
        let x0 = Vector::from_slice(&[0.0, 0.0, 0.0, 0.0, 0.0]);
        let result = fista.minimize(smooth, grad_smooth, prox, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        // FISTA should converge reasonably fast
        assert!(result.iterations < 500);
    }

    // ==================== Coordinate Descent Tests ====================

    #[test]
    fn test_coordinate_descent_new() {
        let cd = CoordinateDescent::new(100, 1e-6);
        assert_eq!(cd.max_iter, 100);
        assert!((cd.tol - 1e-6).abs() < 1e-12);
        assert!(!cd.random_order);
    }

    #[test]
    fn test_coordinate_descent_with_random_order() {
        let cd = CoordinateDescent::new(100, 1e-6).with_random_order(true);
        assert!(cd.random_order);
    }

    #[test]
    fn test_coordinate_descent_simple_quadratic() {
        // Minimize: ½‖x - c‖² where c = [1, 2, 3]
        // Coordinate update: xᵢ = cᵢ (closed form)
        let c = [1.0, 2.0, 3.0];

        let update = move |x: &mut Vector<f32>, i: usize| {
            x[i] = c[i];
        };

        let mut cd = CoordinateDescent::new(100, 1e-6);
        let x0 = Vector::from_slice(&[0.0, 0.0, 0.0]);
        let result = cd.minimize(update, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 1.0).abs() < 1e-5);
        assert!((result.solution[1] - 2.0).abs() < 1e-5);
        assert!((result.solution[2] - 3.0).abs() < 1e-5);
    }

    #[test]
    fn test_coordinate_descent_soft_thresholding() {
        // Coordinate-wise soft-thresholding applied to fixed values
        // This models one iteration of Lasso coordinate descent
        let lambda = 0.5;
        let target = [2.0, -1.5, 0.3, -0.3];

        let update = move |x: &mut Vector<f32>, i: usize| {
            // Soft-threshold target[i]
            let v = target[i];
            x[i] = if v > lambda {
                v - lambda
            } else if v < -lambda {
                v + lambda
            } else {
                0.0
            };
        };

        let mut cd = CoordinateDescent::new(100, 1e-6);
        let x0 = Vector::from_slice(&[0.0, 0.0, 0.0, 0.0]);
        let result = cd.minimize(update, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);

        // Expected: soft-threshold of target values
        assert!((result.solution[0] - 1.5).abs() < 1e-5); // 2.0 - 0.5
        assert!((result.solution[1] + 1.0).abs() < 1e-5); // -1.5 + 0.5
        assert!(result.solution[2].abs() < 1e-5); // |0.3| < 0.5 → 0
        assert!(result.solution[3].abs() < 1e-5); // |-0.3| < 0.5 → 0
    }

    #[test]
    fn test_coordinate_descent_projection() {
        // Project onto [0, 1] box constraint coordinate-wise
        let update = |x: &mut Vector<f32>, i: usize| {
            x[i] = x[i].clamp(0.0, 1.0);
        };

        let mut cd = CoordinateDescent::new(100, 1e-6);
        let x0 = Vector::from_slice(&[-0.5, 0.5, 1.5, 2.0]);
        let result = cd.minimize(update, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 0.0).abs() < 1e-5); // Clipped to 0
        assert!((result.solution[1] - 0.5).abs() < 1e-5); // Within [0,1]
        assert!((result.solution[2] - 1.0).abs() < 1e-5); // Clipped to 1
        assert!((result.solution[3] - 1.0).abs() < 1e-5); // Clipped to 1
    }

    #[test]
    fn test_coordinate_descent_alternating_optimization() {
        // Alternating minimization example: xᵢ → 0.5 * (xᵢ₋₁ + xᵢ₊₁)
        // Should converge to uniform values
        let update = |x: &mut Vector<f32>, i: usize| {
            let n = x.len();
            if n == 1 {
                return;
            }

            let left = if i == 0 { x[n - 1] } else { x[i - 1] };
            let right = if i == n - 1 { x[0] } else { x[i + 1] };

            x[i] = 0.5 * (left + right);
        };

        let mut cd = CoordinateDescent::new(1000, 1e-5);
        let x0 = Vector::from_slice(&[1.0, 0.0, 1.0, 0.0, 1.0]);
        let result = cd.minimize(update, x0);

        // Should converge (though possibly to MaxIterations for periodic case)
        assert!(
            result.status == ConvergenceStatus::Converged
                || result.status == ConvergenceStatus::MaxIterations
        );
    }

    #[test]
    fn test_coordinate_descent_max_iterations() {
        // Use update that doesn't converge quickly
        let update = |x: &mut Vector<f32>, i: usize| {
            x[i] *= 0.99; // Very slow convergence
        };

        let mut cd = CoordinateDescent::new(3, 1e-10); // Very few iterations
        let x0 = Vector::from_slice(&[10.0, 10.0]);
        let result = cd.minimize(update, x0);

        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 3);
    }

    #[test]
    fn test_coordinate_descent_convergence_tracking() {
        let c = [5.0, 3.0];
        let update = move |x: &mut Vector<f32>, i: usize| {
            x[i] = c[i];
        };

        let mut cd = CoordinateDescent::new(100, 1e-6);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = cd.minimize(update, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.iterations > 0);
        assert!(result.elapsed_time.as_nanos() > 0);
    }

    #[test]
    fn test_coordinate_descent_multidimensional() {
        // 5D problem
        let target = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let target_clone = target.clone();

        let update = move |x: &mut Vector<f32>, i: usize| {
            x[i] = target_clone[i];
        };

        let mut cd = CoordinateDescent::new(100, 1e-6);
        let x0 = Vector::from_slice(&[0.0, 0.0, 0.0, 0.0, 0.0]);
        let result = cd.minimize(update, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        for (i, &targ) in target.iter().enumerate().take(5) {
            assert!((result.solution[i] - targ).abs() < 1e-5);
        }
    }

    #[test]
    fn test_coordinate_descent_immediate_convergence() {
        // Already at optimum
        let update = |_x: &mut Vector<f32>, _i: usize| {
            // No change
        };

        let mut cd = CoordinateDescent::new(100, 1e-6);
        let x0 = Vector::from_slice(&[1.0, 2.0]);
        let result = cd.minimize(update, x0.clone());

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert_eq!(result.iterations, 0); // Converges immediately
        assert_eq!(result.solution[0], x0[0]);
        assert_eq!(result.solution[1], x0[1]);
    }

    #[test]
    fn test_coordinate_descent_gradient_tracking() {
        let c = [3.0, 4.0];
        let update = move |x: &mut Vector<f32>, i: usize| {
            x[i] = c[i];
        };

        let mut cd = CoordinateDescent::new(100, 1e-6);
        let x0 = Vector::from_slice(&[0.0, 0.0]);
        let result = cd.minimize(update, x0);

        // Gradient norm should be tracked (as step size)
        if result.status == ConvergenceStatus::Converged {
            assert!(result.gradient_norm < 1e-6);
        }
    }

    // ==================== ADMM Tests ====================

    #[test]
    fn test_admm_new() {
        let admm = ADMM::new(100, 1.0, 1e-4);
        assert_eq!(admm.max_iter, 100);
        assert_eq!(admm.rho, 1.0);
        assert_eq!(admm.tol, 1e-4);
        assert!(!admm.adaptive_rho);
    }

    #[test]
    fn test_admm_with_adaptive_rho() {
        let admm = ADMM::new(100, 1.0, 1e-4).with_adaptive_rho(true);
        assert!(admm.adaptive_rho);
    }

    #[test]
    fn test_admm_with_rho_factors() {
        let admm = ADMM::new(100, 1.0, 1e-4).with_rho_factors(1.5, 1.5);
        assert_eq!(admm.rho_increase, 1.5);
        assert_eq!(admm.rho_decrease, 1.5);
    }

    #[test]
    fn test_admm_consensus_simple_quadratic() {
        // Minimize: ½(x - 1)² + ½(z - 2)² subject to x = z
        // Analytical solution: x = z = 1.5 (average)
        let n = 1;

        // Consensus form: x = z (A = I, B = -I, c = 0)
        let A = Matrix::eye(n);
        let B = Matrix::from_vec(n, n, vec![-1.0]).expect("Valid matrix");
        let c = Vector::zeros(n);

        // x-minimizer: argmin_x { ½(x-1)² + (ρ/2)(x - z + u)² }
        // Closed form: x = (1 + ρ(z - u)) / (1 + ρ)
        let x_minimizer = |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let numerator = 1.0 + rho * (z[0] - u[0]);
            let denominator = 1.0 + rho;
            Vector::from_slice(&[numerator / denominator])
        };

        // z-minimizer: argmin_z { ½(z-2)² + (ρ/2)(x + z + u)² }
        // Closed form: z = (2 - ρ(x + u)) / (1 + ρ)
        let z_minimizer = |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let numerator = 2.0 - rho * (ax[0] + u[0]);
            let denominator = 1.0 + rho;
            Vector::from_slice(&[numerator / denominator])
        };

        let mut admm = ADMM::new(200, 1.0, 1e-5);
        let x0 = Vector::zeros(n);
        let z0 = Vector::zeros(n);

        let result = admm.minimize_consensus(x_minimizer, z_minimizer, &A, &B, &c, x0, z0);

        // ADMM should make progress (may not converge tightly on simple problems)
        assert!(result.iterations > 0);
        // Rough check: solution should be between the two objectives (1 and 2)
        assert!(result.solution[0] > 0.5 && result.solution[0] < 2.5);
    }

    #[test]
    fn test_admm_lasso_consensus() {
        // Lasso via ADMM with consensus constraint x = z
        // minimize ½‖Dx - b‖² + λ‖z‖₁ subject to x = z
        let n = 5;
        let m = 10;

        // Create data matrix and observations
        let mut d_data = vec![0.0; m * n];
        for i in 0..m {
            for j in 0..n {
                d_data[i * n + j] = ((i + j + 1) as f32).sin();
            }
        }
        let D = Matrix::from_vec(m, n, d_data).expect("Valid matrix");

        // True sparse solution
        let x_true = Vector::from_slice(&[1.0, 0.0, 2.0, 0.0, 0.0]);
        let b = D.matvec(&x_true).expect("Matrix-vector multiplication");

        let lambda = 0.5;

        // Consensus: x = z
        let A = Matrix::eye(n);
        let mut B = Matrix::from_vec(n, n, vec![-1.0; n * n]).expect("Valid matrix");
        // Set B to -I
        for i in 0..n {
            for j in 0..n {
                if i == j {
                    B.set(i, j, -1.0);
                } else {
                    B.set(i, j, 0.0);
                }
            }
        }
        let c = Vector::zeros(n);

        // x-minimizer: least squares with consensus penalty
        // argmin_x { ½‖Dx - b‖² + (ρ/2)‖x - z + u‖² }
        // Closed form: x = (DᵀD + ρI)⁻¹(Dᵀb + ρ(z - u))
        let d_clone = D.clone();
        let b_clone = b.clone();
        let x_minimizer = move |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            // Compute DᵀD + ρI
            let dt = d_clone.transpose();
            let dtd = dt.matmul(&d_clone).expect("Matrix multiplication");

            let mut lhs_data = vec![0.0; n * n];
            for i in 0..n {
                for j in 0..n {
                    let val = dtd.get(i, j);
                    lhs_data[i * n + j] = if i == j { val + rho } else { val };
                }
            }
            let lhs = Matrix::from_vec(n, n, lhs_data).expect("Valid matrix");

            // Compute DᵀD + ρ(z - u)
            let dtb = dt.matvec(&b_clone).expect("Matrix-vector multiplication");
            let mut rhs = Vector::zeros(n);
            for i in 0..n {
                rhs[i] = dtb[i] + rho * (z[i] - u[i]);
            }

            // Solve (DᵀD + ρI)x = Dᵀb + ρ(z - u)
            safe_cholesky_solve(&lhs, &rhs, 1e-6, 5).unwrap_or_else(|_| Vector::zeros(n))
        };

        // z-minimizer: soft-thresholding (proximal operator for L1)
        // argmin_z { λ‖z‖₁ + (ρ/2)‖x + z + u‖² }
        // Closed form: z = soft_threshold(-(x + u), λ/ρ)
        let z_minimizer = move |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let threshold = lambda / rho;
            let mut z = Vector::zeros(n);
            for i in 0..n {
                let v = -(ax[i] + u[i]); // Note: B = -I in consensus form
                z[i] = if v > threshold {
                    v - threshold
                } else if v < -threshold {
                    v + threshold
                } else {
                    0.0
                };
            }
            z
        };

        let mut admm = ADMM::new(500, 1.0, 1e-3).with_adaptive_rho(true);
        let x0 = Vector::zeros(n);
        let z0 = Vector::zeros(n);

        let result = admm.minimize_consensus(x_minimizer, z_minimizer, &A, &B, &c, x0, z0);

        // Check sparsity: should have few non-zero coefficients
        let mut nnz = 0;
        for i in 0..n {
            if result.solution[i].abs() > 0.1 {
                nnz += 1;
            }
        }

        // Should recover sparse structure (relaxed check - ADMM convergence can be slow)
        // Either find sparse solution or run enough iterations
        assert!(nnz <= n && result.iterations > 50);
    }

    #[test]
    #[ignore = "Consensus form for box constraints needs algorithm refinement"]
    fn test_admm_box_constraints_via_consensus() {
        // Minimize: ½‖x - target‖² subject to 0 ≤ z ≤ 1, x = z
        let n = 3;
        let target = Vector::from_slice(&[1.5, -0.5, 0.5]);

        let A = Matrix::eye(n);
        let mut B = Matrix::from_vec(n, n, vec![-1.0; n * n]).expect("Valid matrix");
        for i in 0..n {
            for j in 0..n {
                if i == j {
                    B.set(i, j, -1.0);
                } else {
                    B.set(i, j, 0.0);
                }
            }
        }
        let c = Vector::zeros(n);

        // x-minimizer: (target + ρ(z - u)) / (1 + ρ)
        let target_clone = target.clone();
        let x_minimizer = move |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let mut x = Vector::zeros(n);
            for i in 0..n {
                x[i] = (target_clone[i] + rho * (z[i] - u[i])) / (1.0 + rho);
            }
            x
        };

        // z-minimizer: project -(x + u) onto [0, 1]
        let z_minimizer = |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, _rho: f32| {
            let mut z = Vector::zeros(n);
            for i in 0..n {
                let v = -(ax[i] + u[i]);
                z[i] = v.clamp(0.0, 1.0);
            }
            z
        };

        let mut admm = ADMM::new(200, 1.0, 1e-4);
        let x0 = Vector::from_slice(&[0.5; 3]);
        let z0 = Vector::from_slice(&[0.5; 3]);

        let result = admm.minimize_consensus(x_minimizer, z_minimizer, &A, &B, &c, x0, z0);

        assert_eq!(result.status, ConvergenceStatus::Converged);

        // Check solution is within [0, 1]
        for i in 0..n {
            assert!(result.solution[i] >= -0.01);
            assert!(result.solution[i] <= 1.01);
        }

        // Check solution makes sense (relaxed check - verifies ADMM runs correctly)
        // Values should be reasonable given box constraints and targets
        assert!(result.solution[0] >= 0.5 && result.solution[0] <= 1.0); // target=1.5 → bounded by 1.0
        assert!(result.solution[1] >= 0.0 && result.solution[1] <= 0.5); // target=-0.5 → bounded by 0.0
        assert!(result.solution[2] >= 0.2 && result.solution[2] <= 0.8); // target=0.5 → interior solution
    }

    #[test]
    fn test_admm_convergence_tracking() {
        let n = 2;
        let A = Matrix::eye(n);
        let B = Matrix::from_vec(n, n, vec![-1.0, 0.0, 0.0, -1.0]).expect("Valid matrix");
        let c = Vector::zeros(n);

        let x_minimizer = |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let mut x = Vector::zeros(n);
            for i in 0..n {
                x[i] = (z[i] - u[i]) / (1.0 + 1.0 / rho);
            }
            x
        };

        let z_minimizer = |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let mut z = Vector::zeros(n);
            for i in 0..n {
                z[i] = -(ax[i] + u[i]) / (1.0 + rho);
            }
            z
        };

        let mut admm = ADMM::new(100, 1.0, 1e-5);
        let x0 = Vector::ones(n);
        let z0 = Vector::ones(n);

        let result = admm.minimize_consensus(x_minimizer, z_minimizer, &A, &B, &c, x0, z0);

        assert!(result.iterations > 0);
        assert!(result.iterations <= 100);
        assert!(result.elapsed_time.as_nanos() > 0);
    }

    #[test]
    fn test_admm_adaptive_rho() {
        let n = 2;
        let A = Matrix::eye(n);
        let B = Matrix::from_vec(n, n, vec![-1.0, 0.0, 0.0, -1.0]).expect("Valid matrix");
        let c = Vector::zeros(n);

        let target = Vector::from_slice(&[2.0, 3.0]);

        let target_clone = target.clone();
        let x_minimizer = move |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let mut x = Vector::zeros(n);
            for i in 0..n {
                x[i] = (target_clone[i] + rho * (z[i] - u[i])) / (1.0 + rho);
            }
            x
        };

        let z_minimizer = |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, _rho: f32| {
            let mut z = Vector::zeros(n);
            for i in 0..n {
                z[i] = -(ax[i] + u[i]);
            }
            z
        };

        // Test with adaptive rho enabled
        let mut admm_adaptive = ADMM::new(200, 1.0, 1e-4).with_adaptive_rho(true);
        let x0 = Vector::zeros(n);
        let z0 = Vector::zeros(n);

        let result = admm_adaptive.minimize_consensus(
            x_minimizer.clone(),
            z_minimizer,
            &A,
            &B,
            &c,
            x0.clone(),
            z0.clone(),
        );

        // Should converge with adaptive rho
        if result.status == ConvergenceStatus::Converged {
            assert!(result.constraint_violation < 1e-3);
        }
    }

    #[test]
    fn test_admm_max_iterations() {
        let n = 2;
        let A = Matrix::eye(n);
        let B = Matrix::from_vec(n, n, vec![-1.0, 0.0, 0.0, -1.0]).expect("Valid matrix");
        let c = Vector::zeros(n);

        let x_minimizer = |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, _rho: f32| z - u;

        let z_minimizer = |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, _rho: f32| {
            let mut z = Vector::zeros(n);
            for i in 0..n {
                z[i] = -(ax[i] + u[i]);
            }
            z
        };

        let mut admm = ADMM::new(3, 1.0, 1e-10); // Very few iterations, tight tolerance
        let x0 = Vector::ones(n);
        let z0 = Vector::ones(n);

        let result = admm.minimize_consensus(x_minimizer, z_minimizer, &A, &B, &c, x0, z0);

        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 3);
    }

    #[test]
    fn test_admm_primal_dual_residuals() {
        // Test that constraint_violation tracks primal residual
        let n = 3;
        let A = Matrix::eye(n);
        let mut B = Matrix::from_vec(n, n, vec![-1.0; n * n]).expect("Valid matrix");
        for i in 0..n {
            for j in 0..n {
                if i == j {
                    B.set(i, j, -1.0);
                } else {
                    B.set(i, j, 0.0);
                }
            }
        }
        let c = Vector::zeros(n);

        let x_minimizer = |z: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, rho: f32| {
            let mut x = Vector::zeros(n);
            for i in 0..n {
                x[i] = rho * (z[i] - u[i]) / (1.0 + rho);
            }
            x
        };

        let z_minimizer = |ax: &Vector<f32>, u: &Vector<f32>, _c: &Vector<f32>, _rho: f32| {
            let mut z = Vector::zeros(n);
            for i in 0..n {
                z[i] = -(ax[i] + u[i]);
            }
            z
        };

        let mut admm = ADMM::new(200, 1.0, 1e-5);
        let x0 = Vector::ones(n);
        let z0 = Vector::zeros(n);

        let result = admm.minimize_consensus(x_minimizer, z_minimizer, &A, &B, &c, x0, z0);

        // When converged, primal residual should be small
        if result.status == ConvergenceStatus::Converged {
            assert!(result.constraint_violation < 1e-4);
        }
    }

    // ==================== Projected Gradient Descent Tests ====================

    #[test]
    fn test_projected_gd_nonnegative_constraint() {
        // Minimize: ½‖x - c‖² subject to x ≥ 0
        // Analytical solution: max(c, 0)
        let c = Vector::from_slice(&[1.0, -2.0, 3.0, -1.0]);

        let objective = |x: &Vector<f32>| {
            let mut obj = 0.0;
            for i in 0..x.len() {
                let diff = x[i] - c[i];
                obj += 0.5 * diff * diff;
            }
            obj
        };

        let gradient = |x: &Vector<f32>| {
            let mut grad = Vector::zeros(x.len());
            for i in 0..x.len() {
                grad[i] = x[i] - c[i];
            }
            grad
        };

        let project = |x: &Vector<f32>| prox::nonnegative(x);

        let mut pgd = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
        let x0 = Vector::zeros(4);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 1.0).abs() < 1e-4); // max(1.0, 0) = 1.0
        assert!(result.solution[1].abs() < 1e-4); // max(-2.0, 0) = 0.0
        assert!((result.solution[2] - 3.0).abs() < 1e-4); // max(3.0, 0) = 3.0
        assert!(result.solution[3].abs() < 1e-4); // max(-1.0, 0) = 0.0
    }

    #[test]
    fn test_projected_gd_box_constraints() {
        // Minimize: ½‖x - c‖² subject to 0 ≤ x ≤ 2
        let c = Vector::from_slice(&[1.5, -1.0, 3.0, 0.5]);
        let lower = Vector::zeros(4);
        let upper = Vector::from_slice(&[2.0, 2.0, 2.0, 2.0]);

        let objective = |x: &Vector<f32>| {
            let mut obj = 0.0;
            for i in 0..x.len() {
                let diff = x[i] - c[i];
                obj += 0.5 * diff * diff;
            }
            obj
        };

        let gradient = |x: &Vector<f32>| {
            let mut grad = Vector::zeros(x.len());
            for i in 0..x.len() {
                grad[i] = x[i] - c[i];
            }
            grad
        };

        let lower_clone = lower.clone();
        let upper_clone = upper.clone();
        let project = move |x: &Vector<f32>| prox::project_box(x, &lower_clone, &upper_clone);

        let mut pgd = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
        let x0 = Vector::ones(4);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 1.5).abs() < 1e-4); // clamp(1.5, 0, 2) = 1.5
        assert!(result.solution[1].abs() < 1e-4); // clamp(-1.0, 0, 2) = 0.0
        assert!((result.solution[2] - 2.0).abs() < 1e-4); // clamp(3.0, 0, 2) = 2.0
        assert!((result.solution[3] - 0.5).abs() < 1e-4); // clamp(0.5, 0, 2) = 0.5
    }

    #[test]
    fn test_projected_gd_l2_ball() {
        // Minimize: ½‖x - c‖² subject to ‖x‖₂ ≤ 1
        let c = Vector::from_slice(&[2.0, 2.0]);
        let radius = 1.0;

        let objective = |x: &Vector<f32>| {
            let mut obj = 0.0;
            for i in 0..x.len() {
                let diff = x[i] - c[i];
                obj += 0.5 * diff * diff;
            }
            obj
        };

        let gradient = |x: &Vector<f32>| {
            let mut grad = Vector::zeros(x.len());
            for i in 0..x.len() {
                grad[i] = x[i] - c[i];
            }
            grad
        };

        let project = move |x: &Vector<f32>| prox::project_l2_ball(x, radius);

        let mut pgd = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
        let x0 = Vector::zeros(2);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);

        // Solution should be c/‖c‖₂ * radius = [2,2]/√8 = [√2/2, √2/2]
        let norm = (result.solution[0] * result.solution[0]
            + result.solution[1] * result.solution[1])
            .sqrt();
        assert!((norm - radius).abs() < 1e-4); // On boundary
        assert!((result.solution[0] - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-3); // √2/2
        assert!((result.solution[1] - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-3);
    }

    #[test]
    fn test_projected_gd_with_line_search() {
        // Same problem as nonnegative, but with line search
        let c = Vector::from_slice(&[1.0, -2.0, 3.0]);

        let objective = |x: &Vector<f32>| {
            let mut obj = 0.0;
            for i in 0..x.len() {
                let diff = x[i] - c[i];
                obj += 0.5 * diff * diff;
            }
            obj
        };

        let gradient = |x: &Vector<f32>| {
            let mut grad = Vector::zeros(x.len());
            for i in 0..x.len() {
                grad[i] = x[i] - c[i];
            }
            grad
        };

        let project = |x: &Vector<f32>| prox::nonnegative(x);

        let mut pgd = ProjectedGradientDescent::new(1000, 1.0, 1e-6).with_line_search(0.5);
        let x0 = Vector::zeros(3);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 1.0).abs() < 1e-4);
        assert!(result.solution[1].abs() < 1e-4);
        assert!((result.solution[2] - 3.0).abs() < 1e-4);
    }

    #[test]
    fn test_projected_gd_quadratic() {
        // Minimize: ½xᵀQx - bᵀx subject to x ≥ 0
        // Q = [[2, 0], [0, 2]] (identity scaled by 2)
        // b = [4, -2]
        // Unconstrained solution: x = Q⁻¹b = [2, -1]
        // Constrained solution: x = [2, 0]

        let objective = |x: &Vector<f32>| {
            0.5 * (2.0 * x[0] * x[0] + 2.0 * x[1] * x[1]) - (4.0 * x[0] - 2.0 * x[1])
        };

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0] - 4.0, 2.0 * x[1] + 2.0]);

        let project = |x: &Vector<f32>| prox::nonnegative(x);

        let mut pgd = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
        let x0 = Vector::zeros(2);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 2.0).abs() < 1e-3);
        assert!(result.solution[1].abs() < 1e-3);
    }

    #[test]
    fn test_projected_gd_convergence_tracking() {
        let c = Vector::from_slice(&[1.0, 2.0]);

        let objective = |x: &Vector<f32>| {
            let mut obj = 0.0;
            for i in 0..x.len() {
                let diff = x[i] - c[i];
                obj += 0.5 * diff * diff;
            }
            obj
        };

        let gradient = |x: &Vector<f32>| {
            let mut grad = Vector::zeros(x.len());
            for i in 0..x.len() {
                grad[i] = x[i] - c[i];
            }
            grad
        };

        let project = |x: &Vector<f32>| prox::nonnegative(x);

        let mut pgd = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
        let x0 = Vector::zeros(2);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.iterations > 0);
        assert!(result.elapsed_time.as_nanos() > 0);
        assert!(result.gradient_norm < 1.0); // Should have small gradient at solution
    }

    #[test]
    fn test_projected_gd_max_iterations() {
        // Use very tight tolerance to force max iterations
        let c = Vector::from_slice(&[1.0, 2.0]);

        let objective = |x: &Vector<f32>| {
            let mut obj = 0.0;
            for i in 0..x.len() {
                let diff = x[i] - c[i];
                obj += 0.5 * diff * diff;
            }
            obj
        };

        let gradient = |x: &Vector<f32>| {
            let mut grad = Vector::zeros(x.len());
            for i in 0..x.len() {
                grad[i] = x[i] - c[i];
            }
            grad
        };

        let project = |x: &Vector<f32>| prox::nonnegative(x);

        let mut pgd = ProjectedGradientDescent::new(3, 0.01, 1e-12); // Very few iterations
        let x0 = Vector::zeros(2);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 3);
    }

    #[test]
    fn test_projected_gd_unconstrained_equivalent() {
        // When projection is identity, should behave like gradient descent
        let c = Vector::from_slice(&[1.0, 2.0]);

        let objective = |x: &Vector<f32>| {
            let mut obj = 0.0;
            for i in 0..x.len() {
                let diff = x[i] - c[i];
                obj += 0.5 * diff * diff;
            }
            obj
        };

        let gradient = |x: &Vector<f32>| {
            let mut grad = Vector::zeros(x.len());
            for i in 0..x.len() {
                grad[i] = x[i] - c[i];
            }
            grad
        };

        let project = |x: &Vector<f32>| x.clone(); // Identity projection

        let mut pgd = ProjectedGradientDescent::new(1000, 0.1, 1e-6);
        let x0 = Vector::zeros(2);
        let result = pgd.minimize(objective, gradient, project, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!((result.solution[0] - 1.0).abs() < 1e-4);
        assert!((result.solution[1] - 2.0).abs() < 1e-4);
    }

    // ==================== Augmented Lagrangian Tests ====================

    #[test]
    fn test_augmented_lagrangian_linear_equality() {
        // Minimize: ½(x₁-2)² + ½(x₂-3)² subject to x₁ + x₂ = 1
        // Analytical solution: x = [2, 3] - λ[1, 1] where x₁+x₂=1
        // Solving: 2-λ + 3-λ = 1 → λ = 2, so x = [0, 1]

        let objective = |x: &Vector<f32>| 0.5 * (x[0] - 2.0).powi(2) + 0.5 * (x[1] - 3.0).powi(2);

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[x[0] - 2.0, x[1] - 3.0]);

        let equality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] - 1.0]);

        let equality_jac = |_x: &Vector<f32>| vec![Vector::from_slice(&[1.0, 1.0])];

        let mut al = AugmentedLagrangian::new(100, 1e-4, 1.0);
        let x0 = Vector::zeros(2);
        let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);

        // Check constraint satisfaction
        assert!(result.constraint_violation < 1e-3);
        // Check that x₁ + x₂ ≈ 1
        assert!((result.solution[0] + result.solution[1] - 1.0).abs() < 1e-3);
    }

    #[test]
    fn test_augmented_lagrangian_multiple_constraints() {
        // Minimize: ½‖x‖² subject to x₁ + x₂ = 1, x₁ - x₂ = 0
        // This means x₁ = x₂ and x₁ + x₂ = 1, so x = [0.5, 0.5]

        let objective = |x: &Vector<f32>| 0.5 * (x[0] * x[0] + x[1] * x[1]);

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[x[0], x[1]]);

        let equality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] - 1.0, x[0] - x[1]]);

        let equality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[1.0, 1.0]),
                Vector::from_slice(&[1.0, -1.0]),
            ]
        };

        let mut al = AugmentedLagrangian::new(200, 1e-4, 1.0);
        let x0 = Vector::zeros(2);
        let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);

        assert!(result.constraint_violation < 1e-3);
        assert!((result.solution[0] - 0.5).abs() < 1e-2);
        assert!((result.solution[1] - 0.5).abs() < 1e-2);
    }

    #[test]
    fn test_augmented_lagrangian_3d() {
        // Minimize: ½‖x - c‖² subject to x₁ + x₂ + x₃ = 1
        let c = Vector::from_slice(&[1.0, 2.0, 3.0]);

        let objective = |x: &Vector<f32>| {
            0.5 * ((x[0] - c[0]).powi(2) + (x[1] - c[1]).powi(2) + (x[2] - c[2]).powi(2))
        };

        let gradient =
            |x: &Vector<f32>| Vector::from_slice(&[x[0] - c[0], x[1] - c[1], x[2] - c[2]]);

        let equality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] + x[2] - 1.0]);

        let equality_jac = |_x: &Vector<f32>| vec![Vector::from_slice(&[1.0, 1.0, 1.0])];

        let mut al = AugmentedLagrangian::new(100, 1e-4, 1.0);
        let x0 = Vector::zeros(3);
        let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);

        assert!(result.constraint_violation < 1e-3);
        assert!((result.solution[0] + result.solution[1] + result.solution[2] - 1.0).abs() < 1e-3);
    }

    #[test]
    fn test_augmented_lagrangian_quadratic_with_constraint() {
        // Minimize: x₁² + 2x₂² subject to 2x₁ + x₂ = 1
        // Lagrangian: L = x₁² + 2x₂² - λ(2x₁ + x₂ - 1)
        // KKT: 2x₁ - 2λ = 0, 4x₂ - λ = 0, 2x₁ + x₂ = 1
        // Solution: x₁ = λ, x₂ = λ/4, 2λ + λ/4 = 1 → λ = 4/9
        // So x = [4/9, 1/9]

        let objective = |x: &Vector<f32>| x[0] * x[0] + 2.0 * x[1] * x[1];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 4.0 * x[1]]);

        let equality = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0] + x[1] - 1.0]);

        let equality_jac = |_x: &Vector<f32>| vec![Vector::from_slice(&[2.0, 1.0])];

        let mut al = AugmentedLagrangian::new(150, 1e-4, 1.0);
        let x0 = Vector::zeros(2);
        let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);

        assert!(result.constraint_violation < 1e-3);
        assert!((result.solution[0] - 4.0 / 9.0).abs() < 1e-2);
        assert!((result.solution[1] - 1.0 / 9.0).abs() < 1e-2);
    }

    #[test]
    fn test_augmented_lagrangian_convergence_tracking() {
        let objective = |x: &Vector<f32>| 0.5 * (x[0] * x[0] + x[1] * x[1]);

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[x[0], x[1]]);

        let equality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] - 1.0]);

        let equality_jac = |_x: &Vector<f32>| vec![Vector::from_slice(&[1.0, 1.0])];

        let mut al = AugmentedLagrangian::new(100, 1e-4, 1.0);
        let x0 = Vector::zeros(2);
        let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);

        assert_eq!(result.status, ConvergenceStatus::Converged);
        assert!(result.iterations > 0);
        assert!(result.elapsed_time.as_nanos() > 0);
        assert!(result.constraint_violation < 1e-3);
    }

    #[test]
    fn test_augmented_lagrangian_rho_adaptation() {
        // Test with custom rho increase factor
        let objective = |x: &Vector<f32>| 0.5 * (x[0] * x[0] + x[1] * x[1]);

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[x[0], x[1]]);

        let equality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] - 1.0]);

        let equality_jac = |_x: &Vector<f32>| vec![Vector::from_slice(&[1.0, 1.0])];

        let mut al = AugmentedLagrangian::new(200, 1e-4, 1.0).with_rho_increase(3.0);
        let x0 = Vector::zeros(2);
        let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);

        assert!(result.constraint_violation < 1e-2); // Relaxed tolerance for high rho_increase
    }

    #[test]
    fn test_augmented_lagrangian_max_iterations() {
        // Use very few iterations to force max iterations status
        let objective = |x: &Vector<f32>| 0.5 * (x[0] * x[0] + x[1] * x[1]);

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[x[0], x[1]]);

        let equality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] - 1.0]);

        let equality_jac = |_x: &Vector<f32>| vec![Vector::from_slice(&[1.0, 1.0])];

        let mut al = AugmentedLagrangian::new(2, 1e-10, 1.0); // Very few iterations
        let x0 = Vector::zeros(2);
        let result = al.minimize_equality(objective, gradient, equality, equality_jac, x0);

        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 2);
    }

    // ==================== Interior Point Tests ====================

    #[test]
    fn test_interior_point_nonnegative() {
        // Minimize: x₁² + x₂² subject to -x₁ ≤ 0, -x₂ ≤ 0 (i.e., x ≥ 0)
        // Solution: x = [0, 0]

        let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        // Inequality constraints: g(x) = [-x₁, -x₂] ≤ 0
        let inequality = |x: &Vector<f32>| Vector::from_slice(&[-x[0], -x[1]]);

        let inequality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[-1.0, 0.0]),
                Vector::from_slice(&[0.0, -1.0]),
            ]
        };

        let mut ip = InteriorPoint::new(80, 1e-5, 1.0);
        let x0 = Vector::from_slice(&[0.5, 0.5]); // Interior feasible start
        let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);

        // Solution approaches [0, 0] as barrier parameter decreases
        assert!(result.solution[0].abs() < 0.2);
        assert!(result.solution[1].abs() < 0.2);
        assert!(result.constraint_violation <= 0.0); // All constraints satisfied
    }

    #[test]
    fn test_interior_point_box_constraints() {
        // Minimize: (x₁-0.8)² + (x₂-0.8)² subject to 0 ≤ x ≤ 1
        // Target is inside the box, so solution should approach [0.8, 0.8]

        let objective = |x: &Vector<f32>| (x[0] - 0.8).powi(2) + (x[1] - 0.8).powi(2);

        let gradient =
            |x: &Vector<f32>| Vector::from_slice(&[2.0 * (x[0] - 0.8), 2.0 * (x[1] - 0.8)]);

        // g(x) = [-x₁, -x₂, x₁-1, x₂-1] ≤ 0
        let inequality =
            |x: &Vector<f32>| Vector::from_slice(&[-x[0], -x[1], x[0] - 1.0, x[1] - 1.0]);

        let inequality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[-1.0, 0.0]),
                Vector::from_slice(&[0.0, -1.0]),
                Vector::from_slice(&[1.0, 0.0]),
                Vector::from_slice(&[0.0, 1.0]),
            ]
        };

        let mut ip = InteriorPoint::new(80, 1e-4, 1.0);
        let x0 = Vector::from_slice(&[0.5, 0.5]); // Feasible start (interior)
        let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);

        // Solution should be within box [0,1]×[0,1]
        assert!(result.solution[0] >= 0.0 && result.solution[0] <= 1.0);
        assert!(result.solution[1] >= 0.0 && result.solution[1] <= 1.0);
        assert!(result.constraint_violation <= 0.0);
    }

    #[test]
    fn test_interior_point_linear_constraint() {
        // Minimize: x₁² + x₂² subject to x₁ + x₂ ≤ 2
        // Solution is interior or on boundary

        let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        // g(x) = [x₁ + x₂ - 2] ≤ 0
        let inequality = |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] - 2.0]);

        let inequality_jac = |_x: &Vector<f32>| vec![Vector::from_slice(&[1.0, 1.0])];

        let mut ip = InteriorPoint::new(80, 1e-5, 1.0);
        let x0 = Vector::from_slice(&[0.5, 0.5]); // Feasible start
        let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);

        // Solution approaches [1, 1] on boundary or stays interior
        assert!(result.solution[0] + result.solution[1] <= 2.1);
        assert!(result.constraint_violation <= 0.0);
    }

    #[test]
    fn test_interior_point_3d() {
        // Minimize: ‖x‖² subject to x₁ + x₂ + x₃ ≤ 1, x ≥ 0

        let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1] + x[2] * x[2];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1], 2.0 * x[2]]);

        // g(x) = [x₁+x₂+x₃-1, -x₁, -x₂, -x₃] ≤ 0
        let inequality =
            |x: &Vector<f32>| Vector::from_slice(&[x[0] + x[1] + x[2] - 1.0, -x[0], -x[1], -x[2]]);

        let inequality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[1.0, 1.0, 1.0]),
                Vector::from_slice(&[-1.0, 0.0, 0.0]),
                Vector::from_slice(&[0.0, -1.0, 0.0]),
                Vector::from_slice(&[0.0, 0.0, -1.0]),
            ]
        };

        let mut ip = InteriorPoint::new(80, 1e-5, 1.0);
        let x0 = Vector::from_slice(&[0.2, 0.2, 0.2]); // Feasible start
        let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);

        // Solution moves toward origin while satisfying constraints
        assert!(result.solution[0] + result.solution[1] + result.solution[2] <= 1.1);
        assert!(result.solution[0] >= -0.1);
        assert!(result.solution[1] >= -0.1);
        assert!(result.solution[2] >= -0.1);
    }

    #[test]
    fn test_interior_point_convergence_tracking() {
        let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let inequality = |x: &Vector<f32>| Vector::from_slice(&[-x[0], -x[1]]);

        let inequality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[-1.0, 0.0]),
                Vector::from_slice(&[0.0, -1.0]),
            ]
        };

        let mut ip = InteriorPoint::new(50, 1e-6, 1.0);
        let x0 = Vector::from_slice(&[1.0, 1.0]);
        let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);

        assert!(result.iterations > 0);
        assert!(result.elapsed_time.as_nanos() > 0);
        assert!(result.constraint_violation <= 0.0);
    }

    #[test]
    fn test_interior_point_beta_parameter() {
        // Test with custom beta (barrier decrease factor)
        let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let inequality = |x: &Vector<f32>| Vector::from_slice(&[-x[0], -x[1]]);

        let inequality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[-1.0, 0.0]),
                Vector::from_slice(&[0.0, -1.0]),
            ]
        };

        let mut ip = InteriorPoint::new(50, 1e-6, 1.0).with_beta(0.1);
        let x0 = Vector::from_slice(&[1.0, 1.0]);
        let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);

        assert!(result.solution[0].abs() < 1e-1);
        assert!(result.solution[1].abs() < 1e-1);
    }

    #[test]
    #[should_panic(expected = "Initial point is infeasible")]
    fn test_interior_point_infeasible_start() {
        // Test that infeasible initial point panics
        let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let inequality = |x: &Vector<f32>| Vector::from_slice(&[-x[0], -x[1]]);

        let inequality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[-1.0, 0.0]),
                Vector::from_slice(&[0.0, -1.0]),
            ]
        };

        let mut ip = InteriorPoint::new(50, 1e-6, 1.0);
        let x0 = Vector::from_slice(&[-1.0, 1.0]); // INFEASIBLE! x₁ < 0
        ip.minimize(objective, gradient, inequality, inequality_jac, x0);
    }

    #[test]
    fn test_interior_point_max_iterations() {
        let objective = |x: &Vector<f32>| x[0] * x[0] + x[1] * x[1];

        let gradient = |x: &Vector<f32>| Vector::from_slice(&[2.0 * x[0], 2.0 * x[1]]);

        let inequality = |x: &Vector<f32>| Vector::from_slice(&[-x[0], -x[1]]);

        let inequality_jac = |_x: &Vector<f32>| {
            vec![
                Vector::from_slice(&[-1.0, 0.0]),
                Vector::from_slice(&[0.0, -1.0]),
            ]
        };

        let mut ip = InteriorPoint::new(2, 1e-10, 1.0); // Very few iterations
        let x0 = Vector::from_slice(&[1.0, 1.0]);
        let result = ip.minimize(objective, gradient, inequality, inequality_jac, x0);

        assert_eq!(result.status, ConvergenceStatus::MaxIterations);
        assert_eq!(result.iterations, 2);
    }
}