apex-solver 1.5.0

use faer::{
    Mat, Side,
    linalg::solvers::Solve,
    sparse::SparseColMat,
    sparse::linalg::solvers::{Llt, SymbolicLlt},
};

use crate::error::ErrorLogging;
use crate::linalg::sparse::normal_eq::{LazyNormalEquations, NormalEquations};
use crate::linalg::{Damping, LinAlgError, LinAlgResult, LinearSolver, SparseMode};

#[derive(Debug, Clone)]
pub struct SparseCholeskySolver {
    /// Cached symbolic factorization for reuse across iterations.
    ///
    /// This is computed once and reused when the sparsity pattern doesn't change,
    /// providing a 10-15% performance improvement for iterative optimization.
    symbolic_factorization: Option<SymbolicLlt<usize>>,

    /// Cached symbolic machinery for forming `JᵀJ` and `Jᵀr` in parallel.
    ne_cache: LazyNormalEquations,

    /// The Hessian matrix, computed as `(J^T *  J)`.
    ///
    /// This is `None` if the Hessian could not be computed.
    hessian: Option<SparseColMat<usize, f64>>,

    /// The gradient vector, computed as `J^T *  r`.
    ///
    /// This is `None` if the gradient could not be computed.
    gradient: Option<Mat<f64>>,
}

impl SparseCholeskySolver {
    pub fn new() -> Self {
        SparseCholeskySolver {
            symbolic_factorization: None,
            ne_cache: LazyNormalEquations::default(),
            hessian: None,
            gradient: None,
        }
    }

    pub fn hessian(&self) -> Option<&SparseColMat<usize, f64>> {
        self.hessian.as_ref()
    }

    pub fn gradient(&self) -> Option<&Mat<f64>> {
        self.gradient.as_ref()
    }
}

impl Default for SparseCholeskySolver {
    fn default() -> Self {
        Self::new()
    }
}
impl LinearSolver<SparseMode> for SparseCholeskySolver {
    fn solve_normal_equation(
        &mut self,
        residuals: &Mat<f64>,
        jacobians: &SparseColMat<usize, f64>,
    ) -> LinAlgResult<Mat<f64>> {
        // Form the normal equations: H = JᵀJ, g = Jᵀr (parallel faer kernels,
        // symbolic product cached across iterations).
        let NormalEquations { hessian, gradient } = self.ne_cache.compute(residuals, jacobians)?;

        let sym = if let Some(ref cached_sym) = self.symbolic_factorization {
            // Reuse cached symbolic factorization
            // Note: SymbolicLlt is reference-counted, so clone() is cheap (O(1))
            // We assume the sparsity pattern is constant across iterations
            // which is typical in iterative optimization
            cached_sym.clone()
        } else {
            // Create new symbolic factorization and cache it
            let new_sym = SymbolicLlt::try_new(hessian.symbolic(), Side::Lower).map_err(|e| {
                LinAlgError::FactorizationFailed(
                    "Symbolic Cholesky decomposition failed".to_string(),
                )
                .log_with_source(e)
            })?;
            // Cache it (clone is cheap due to reference counting)
            self.symbolic_factorization = Some(new_sym.clone());
            new_sym
        };

        // Perform numeric factorization using the symbolic structure
        let cholesky =
            Llt::try_new_with_symbolic(sym, hessian.as_ref(), Side::Lower).map_err(|e| {
                LinAlgError::SingularMatrix(
                    "Cholesky factorization failed (matrix may be singular)".to_string(),
                )
                .log_with_source(e)
            })?;

        let dx = cholesky.solve(-&gradient);
        self.hessian = Some(hessian);
        self.gradient = Some(gradient);

        Ok(dx)
    }

    fn solve_augmented_equation(
        &mut self,
        residuals: &Mat<f64>,
        jacobians: &SparseColMat<usize, f64>,
        damping: &Damping,
    ) -> LinAlgResult<Mat<f64>> {
        // H = JᵀJ, g = Jᵀr (parallel faer kernels)
        let NormalEquations { hessian, gradient } = self.ne_cache.compute(residuals, jacobians)?;

        // H_aug = H + λ·D — diagonal edit on the cached product pattern.
        let augmented_hessian = self.ne_cache.damped_hessian(damping)?;

        let sym = if let Some(ref cached_sym) = self.symbolic_factorization {
            // Reuse cached symbolic factorization
            // Note: SymbolicLlt is reference-counted, so clone() is cheap (O(1))
            // We assume the sparsity pattern is constant across iterations
            // which is typical in iterative optimization
            cached_sym.clone()
        } else {
            // Create new symbolic factorization and cache it
            let new_sym =
                SymbolicLlt::try_new(augmented_hessian.symbolic(), Side::Lower).map_err(|e| {
                    LinAlgError::FactorizationFailed(
                        "Symbolic Cholesky decomposition failed for augmented system".to_string(),
                    )
                    .log_with_source(e)
                })?;
            // Cache it (clone is cheap due to reference counting)
            self.symbolic_factorization = Some(new_sym.clone());
            new_sym
        };

        // Perform numeric factorization
        let cholesky = Llt::try_new_with_symbolic(sym, augmented_hessian.as_ref(), Side::Lower)
            .map_err(|e| {
                LinAlgError::SingularMatrix(
                    "Cholesky factorization failed (matrix may be singular)".to_string(),
                )
                .log_with_source(e)
            })?;

        let dx = cholesky.solve(-&gradient);
        self.hessian = Some(hessian);
        self.gradient = Some(gradient);

        Ok(dx)
    }

    fn hessian_vec_product(&self, v: &Mat<f64>) -> Option<Mat<f64>> {
        Some(
            <SparseMode as crate::linearizer::AssemblyBackend>::hessian_vec_product(
                self.hessian.as_ref()?,
                v,
            ),
        )
    }

    fn get_hessian(&self) -> Option<&SparseColMat<usize, f64>> {
        self.hessian.as_ref()
    }

    fn get_gradient(&self) -> Option<&Mat<f64>> {
        self.gradient.as_ref()
    }
}

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

    const TOLERANCE: f64 = 1e-10;

    type TestResult = Result<(), Box<dyn std::error::Error>>;

    /// Helper function to create a simple test matrix and vectors
    fn create_test_data()
    -> Result<(SparseColMat<usize, f64>, Mat<f64>), faer::sparse::CreationError> {
        // Create an overdetermined system (4x3) so that weights have an effect
        let triplets = vec![
            Triplet::new(0, 0, 2.0),
            Triplet::new(0, 1, 1.0),
            Triplet::new(1, 0, 1.0),
            Triplet::new(1, 1, 3.0),
            Triplet::new(1, 2, 1.0),
            Triplet::new(2, 1, 1.0),
            Triplet::new(2, 2, 2.0),
            Triplet::new(3, 0, 1.5), // Add a 4th row for overdetermined system
            Triplet::new(3, 2, 0.5),
        ];
        let jacobian = SparseColMat::try_new_from_triplets(4, 3, &triplets)?;

        let residuals = Mat::from_fn(4, 1, |i, _| match i {
            0 => 1.0,
            1 => -2.0,
            2 => 0.5,
            3 => 1.2,
            _ => 0.0,
        });

        Ok((jacobian, residuals))
    }

    /// Test basic solver creation and default implementation
    #[test]
    fn test_solver_creation() {
        let solver = SparseCholeskySolver::new();
        assert!(solver.hessian.is_none());
        assert!(solver.gradient.is_none());

        let default_solver = SparseCholeskySolver::default();
        assert!(default_solver.hessian.is_none());
        assert!(default_solver.gradient.is_none());
    }

    /// Test normal equation solving with well-conditioned matrix
    #[test]
    fn test_solve_normal_equation_well_conditioned() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        let (jacobian, residuals) = create_test_data()?;

        let solution =
            LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
        assert_eq!(solution.nrows(), 3);
        assert_eq!(solution.ncols(), 1);

        // Verify the symbolic pattern was cached
        Ok(())
    }

    /// Test that symbolic pattern is reused on subsequent calls
    #[test]
    fn test_symbolic_pattern_caching() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        let (jacobian, residuals) = create_test_data()?;

        // First solve
        let sol1 =
            LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;

        // Second solve should reuse pattern
        let sol2 =
            LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;

        // Results should be identical
        for i in 0..sol1.nrows() {
            assert!((sol1[(i, 0)] - sol2[(i, 0)]).abs() < TOLERANCE);
        }
        Ok(())
    }

    /// Test augmented equation solving
    #[test]
    fn test_solve_augmented_equation() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        let (jacobian, residuals) = create_test_data()?;
        let lambda = 0.1;

        let solution = LinearSolver::<SparseMode>::solve_augmented_equation(
            &mut solver,
            &residuals,
            &jacobian,
            &Damping::identity(lambda),
        )?;
        assert_eq!(solution.nrows(), 3);
        assert_eq!(solution.ncols(), 1);
        Ok(())
    }

    /// Test with different lambda values in augmented system
    #[test]
    fn test_augmented_equation_different_lambdas() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        let (jacobian, residuals) = create_test_data()?;

        let lambda1 = 0.01;
        let lambda2 = 1.0;

        let sol1 = LinearSolver::<SparseMode>::solve_augmented_equation(
            &mut solver,
            &residuals,
            &jacobian,
            &Damping::identity(lambda1),
        )?;
        let sol2 = LinearSolver::<SparseMode>::solve_augmented_equation(
            &mut solver,
            &residuals,
            &jacobian,
            &Damping::identity(lambda2),
        )?;

        // Solutions should be different due to different regularization
        let mut different = false;
        for i in 0..sol1.nrows() {
            if (sol1[(i, 0)] - sol2[(i, 0)]).abs() > TOLERANCE {
                different = true;
                break;
            }
        }
        assert!(
            different,
            "Solutions should differ with different lambda values"
        );
        Ok(())
    }

    /// Test with singular matrix (should return None)
    #[test]
    fn test_singular_matrix() -> TestResult {
        let mut solver = SparseCholeskySolver::new();

        // Create a singular matrix
        let triplets = vec![
            Triplet::new(0, 0, 1.0),
            Triplet::new(0, 1, 2.0),
            Triplet::new(1, 0, 2.0),
            Triplet::new(1, 1, 4.0), // Second row is 2x first row
        ];
        let singular_jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
        let residuals = Mat::from_fn(2, 1, |i, _| i as f64);

        let result = LinearSolver::<SparseMode>::solve_normal_equation(
            &mut solver,
            &residuals,
            &singular_jacobian,
        );
        // Without regularization, singular matrices should fail
        assert!(result.is_err(), "Singular matrix should return Err");
        Ok(())
    }

    /// Test with empty matrix (edge case)
    #[test]
    fn test_empty_matrix() -> TestResult {
        let mut solver = SparseCholeskySolver::new();

        let empty_jacobian = SparseColMat::try_new_from_triplets(0, 0, &[])?;
        let empty_residuals = Mat::zeros(0, 1);

        let result = LinearSolver::<SparseMode>::solve_normal_equation(
            &mut solver,
            &empty_residuals,
            &empty_jacobian,
        );
        if let Ok(solution) = result {
            assert_eq!(solution.nrows(), 0);
        }
        Ok(())
    }

    /// Test numerical accuracy with known solution
    #[test]
    fn test_numerical_accuracy() -> TestResult {
        let mut solver = SparseCholeskySolver::new();

        // Create a simple 2x2 system with known solution
        let triplets = vec![
            Triplet::new(0, 0, 1.0),
            Triplet::new(0, 1, 0.0),
            Triplet::new(1, 0, 0.0),
            Triplet::new(1, 1, 1.0),
        ];
        let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
        let residuals = Mat::from_fn(2, 1, |i, _| -((i + 1) as f64)); // [-1, -2]

        let solution =
            LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
        // Expected solution should be [1, 2] since J^T * J = I and J^T * (-r) = [1, 2]
        assert!((solution[(0, 0)] - 1.0).abs() < TOLERANCE);
        assert!((solution[(1, 0)] - 2.0).abs() < TOLERANCE);
        Ok(())
    }

    /// Test clone functionality
    #[test]
    fn test_solver_clone() {
        let solver1 = SparseCholeskySolver::new();
        let solver2 = solver1.clone();
        assert!(solver2.hessian.is_none());
        assert!(solver2.gradient.is_none());
    }

    /// Test Cholesky decomposition properties
    #[test]
    fn test_cholesky_decomposition_properties() -> TestResult {
        let mut solver = SparseCholeskySolver::new();

        // Create a simple positive definite system
        let triplets = vec![Triplet::new(0, 0, 2.0), Triplet::new(1, 1, 3.0)];
        let jacobian = SparseColMat::try_new_from_triplets(2, 2, &triplets)?;
        let residuals = Mat::from_fn(2, 1, |i, _| (i + 1) as f64);

        LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;

        // Verify that we have a factorizer and hessian
        assert!(solver.hessian.is_some());

        // The hessian should be positive definite for Cholesky to work
        if let Some(hessian) = &solver.hessian {
            assert_eq!(hessian.nrows(), 2);
            assert_eq!(hessian.ncols(), 2);
        }
        Ok(())
    }

    /// Test numerical stability with different condition numbers
    #[test]
    fn test_cholesky_numerical_stability() -> TestResult {
        let mut solver = SparseCholeskySolver::new();

        // Create a well-conditioned system
        let triplets = vec![
            Triplet::new(0, 0, 1.0),
            Triplet::new(1, 1, 1.0),
            Triplet::new(2, 2, 1.0),
        ];
        let jacobian = SparseColMat::try_new_from_triplets(3, 3, &triplets)?;
        let residuals = Mat::from_fn(3, 1, |i, _| -((i + 1) as f64)); // [-1, -2, -3]

        let solution =
            LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;
        // Expected solution should be [1, 2, 3] since H = I and g = [1, 2, 3]
        for i in 0..3 {
            let expected = (i + 1) as f64;
            assert!(
                (solution[(i, 0)] - expected).abs() < TOLERANCE,
                "Expected {}, got {}",
                expected,
                solution[(i, 0)]
            );
        }

        Ok(())
    }

    /// Test hessian() getter returns None before solve and Some after
    #[test]
    fn test_cholesky_hessian_getter() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        assert!(solver.hessian().is_none());

        let (jacobian, residuals) = create_test_data()?;
        LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;

        assert!(solver.hessian().is_some());
        Ok(())
    }

    /// Test gradient() getter returns None before solve and Some after
    #[test]
    fn test_cholesky_gradient_getter() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        assert!(solver.gradient().is_none());

        let (jacobian, residuals) = create_test_data()?;
        LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;

        assert!(solver.gradient().is_some());
        Ok(())
    }

    /// Test get_hessian() trait method returns Some after solve
    #[test]
    fn test_cholesky_get_hessian_trait() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_none());

        let (jacobian, residuals) = create_test_data()?;
        LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;

        assert!(LinearSolver::<SparseMode>::get_hessian(&solver).is_some());
        Ok(())
    }

    /// Test get_gradient() trait method returns Some after solve
    #[test]
    fn test_cholesky_get_gradient_trait() -> TestResult {
        let mut solver = SparseCholeskySolver::new();
        assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_none());

        let (jacobian, residuals) = create_test_data()?;
        LinearSolver::<SparseMode>::solve_normal_equation(&mut solver, &residuals, &jacobian)?;

        assert!(LinearSolver::<SparseMode>::get_gradient(&solver).is_some());
        Ok(())
    }
}