apex-solver 1.4.0

//! Jacobian linearization — the bridge between the nonlinear factor graph
//! and the linear system solved each iteration.
//!
//! This is the central module for all linearization concerns:
//! - `linearize_block()`: Shared factor evaluation (loss correction, residual accumulation)
//! - [`cpu::sparse`]: Sparse Jacobian assembly using `SparseColMat` and symbolic structure
//! - [`cpu::dense`]: Dense Jacobian assembly using `Mat<f64>`
//! - [`AssemblyBackend`]: Trait bridging linearization with the optimizer's solver types
//!
//! # Architecture
//!
//! ```text
//! Problem (factor graph)
//!     │  AssemblyBackend::assemble()
//!     ▼
//! (r: Mat<f64>, J: M::Jacobian)   ← M: LinearizationMode
//!     │
//!     ▼
//! LinearSolver<M>   (linalg/)
//!     │
//!     ▼
//! dx: Mat<f64>  → manifold update
//! ```

pub mod cpu;
pub mod gpu;

use slotmap::{SecondaryMap, SlotMap};

use faer::Mat;
use faer::sparse::{SparseColMat, Triplet};
use smallvec::SmallVec;
use thiserror::Error;

use crate::core::VarKey;
use crate::core::problem::Problem;
use crate::core::variable::ManifoldVariable;
use crate::{
    core::{corrector::Corrector, residual_block::ResidualBlock},
    linearizer::cpu::{DenseMode, LinearizationMode, SparseMode},
};

pub use cpu::sparse::SymbolicStructure;

// ============================================================================
// Linearizer error types
// ============================================================================

/// Linearizer-specific error types for Jacobian assembly and symbolic structure operations.
///
/// These errors occur during the linearization phase of optimization, where
/// the nonlinear factor graph is converted into a linear system (residual vector
/// and Jacobian matrix).
///
/// # Error Hierarchy
///
/// `LinearizerError` is a Layer C (deep/module) error. It propagates up through:
/// - `CoreError::Linearizer(LinearizerError)` → for core module callers
/// - `ApexSolverError::Linearizer(LinearizerError)` → for direct API callers
///

#[derive(Debug, Clone, Error)]
pub enum LinearizerError {
    /// Symbolic structure construction or usage failed
    #[error("Symbolic structure error: {0}")]
    SymbolicStructure(String),

    /// Parallel computation error (thread/mutex failures during assembly)
    #[error("Parallel computation error: {0}")]
    ParallelComputation(String),

    /// Variable key missing in index mapping during Jacobian scatter
    #[error("Variable error: {0}")]
    Variable(String),

    /// Factor linearization returned no Jacobian when expected
    #[error("Factor linearization failed: {0}")]
    FactorLinearization(String),

    /// Invalid input (e.g., SparseMode requires symbolic structure)
    #[error("Invalid input: {0}")]
    InvalidInput(String),
}

/// Result type for linearizer module operations
pub type LinearizerResult<T> = Result<T, LinearizerError>;

// ============================================================================
// Block linearization (shared by sparse and dense paths)
// ============================================================================

/// Metadata produced by evaluating a single residual block.
///
/// The corrected Jacobian itself lives in a caller-provided buffer (see
/// [`compute_block_into`]) — this struct only describes how to scatter it
/// into the global Jacobian.
pub(crate) struct BlockLinearization {
    /// Maps each variable to (local_col_offset, dof_size) within the block Jacobian.
    /// SmallVec inline storage covers all current factor types (≤ 4 variables per factor).
    pub variable_local_idx_size_list: SmallVec<[(usize, usize); 8]>,
    /// Starting row index in the global residual/Jacobian
    pub residual_row_start_idx: usize,
    /// Residual dimension for this block
    pub residual_dim: usize,
}

/// Split a mutable buffer into non-overlapping slices at the given (start, len) offsets.
///
/// `sorted_offsets_lens` must be sorted ascending by `start` and non-overlapping.
/// Uses chained `split_at_mut` calls — no unsafe code needed.
pub(crate) fn split_by_row_offsets_mut<'a>(
    buf: &'a mut [f64],
    sorted_offsets_lens: &[(usize, usize)],
) -> Vec<&'a mut [f64]> {
    let mut remaining = buf;
    let mut result = Vec::with_capacity(sorted_offsets_lens.len());
    let mut current = 0usize;
    for &(start, len) in sorted_offsets_lens {
        let gap = start - current;
        let (_, rest) = remaining.split_at_mut(gap);
        let (slice, rest2) = rest.split_at_mut(len);
        result.push(slice);
        remaining = rest2;
        current = start + len;
    }
    result
}

/// Evaluate a single residual block: call `factor.linearize()`, apply loss correction,
/// write the corrected residual into the provided slice, return the Jacobian buffer.
///
/// This is the shared core used by both sparse and dense assembly. It:
/// 1. Gathers `&[f64]` parameter slices for each variable (zero copy from manifold storage)
/// 2. Calls `factor.linearize()` — writes into `residual_slice` and a local Jacobian buffer
/// 3. Applies the robust loss function correction (if any)
/// 4. Returns the corrected Jacobian buffer and metadata for the caller to scatter
pub(crate) fn compute_block_into(
    residual_block: &ResidualBlock,
    variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
    residual_slice: &mut [f64],
    jacobian_buf: &mut [f64],
) -> LinearizerResult<BlockLinearization> {
    let mut param_slices: SmallVec<[&[f64]; 8]> = SmallVec::new();
    let mut variable_local_idx_size_list: SmallVec<[(usize, usize); 8]> = SmallVec::new();
    let mut count_variable_local_idx: usize = 0;

    for &var_key in &residual_block.variable_keys {
        if let Some(variable) = variables.get(var_key) {
            param_slices.push(variable.as_param_slice());
            let var_size = variable.dof();
            variable_local_idx_size_list.push((count_variable_local_idx, var_size));
            count_variable_local_idx += var_size;
        }
    }

    let (rows, cols) = residual_block.factor.jacobian_shape();
    debug_assert_eq!(jacobian_buf.len(), rows * cols);
    {
        let jac_mut = faer::mat::MatMut::from_column_major_slice_mut(jacobian_buf, rows, cols);
        residual_block
            .factor
            .linearize(&param_slices, residual_slice, Some(jac_mut));
    }

    // Apply robust-loss correction in-place: no heap allocation, math identical to
    // `Corrector::correct_jacobian` / `correct_residuals` on `DVector`/`DMatrix`.
    if let Some(loss_func) = &residual_block.loss_func {
        let squared_norm: f64 = residual_slice.iter().map(|x| x * x).sum();
        let corrector = Corrector::new(loss_func.as_ref(), squared_norm);
        // Jacobian correction must read the original (un-scaled) residual.
        corrector.correct_jacobian_in_place(residual_slice, jacobian_buf, rows, cols);
        corrector.correct_residual_in_place(residual_slice);
    }

    Ok(BlockLinearization {
        variable_local_idx_size_list,
        residual_row_start_idx: residual_block.residual_row_start_idx,
        residual_dim: rows,
    })
}

// ============================================================================
// AssemblyBackend trait (bridges linearizer output with optimizer solver types)
// ============================================================================

/// Type-level backend for assembling (residuals, Jacobian) and performing
/// matrix operations. Implemented by [`SparseMode`] and [`DenseMode`].
///
/// All methods are static — this trait is used as a compile-time strategy
/// selector, not as an object interface. Extends [`LinearizationMode`] with
/// the five operations an optimizer needs each iteration: building `(r, J)`,
/// scaling `J`, unscaling `dx`, and `H·v`.
///
/// All three optimizers (LM, GN, DogLeg) are generic over `M: AssemblyBackend`,
/// giving zero-cost static dispatch through the entire pipeline.
pub trait AssemblyBackend: LinearizationMode {
    /// Assemble residuals and Jacobian from the problem.
    fn assemble(
        problem: &Problem,
        variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
        variable_index_map: &SecondaryMap<VarKey, usize>,
        symbolic_structure: Option<&SymbolicStructure>,
        total_dof: usize,
    ) -> LinearizerResult<(Mat<f64>, Self::Jacobian)>;

    /// Compute column norms of the Jacobian (for Jacobi scaling).
    fn compute_column_norms(jacobian: &Self::Jacobian) -> Vec<f64>;

    /// Apply diagonal column scaling to the Jacobian.
    /// Returns a new Jacobian with columns scaled by `1 / (1 + norm)`.
    fn apply_column_scaling(jacobian: &Self::Jacobian, scaling: &[f64]) -> Self::Jacobian;

    /// Apply inverse scaling to a step vector: step_i *= scaling_i
    fn apply_inverse_scaling(step: &Mat<f64>, scaling: &[f64]) -> Mat<f64>;

    /// Hessian-vector product: H * v (needed by DogLeg for Cauchy point)
    fn hessian_vec_product(hessian: &Self::Hessian, vec: &Mat<f64>) -> Mat<f64>;
}

impl AssemblyBackend for SparseMode {
    fn assemble(
        problem: &Problem,
        variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
        variable_index_map: &SecondaryMap<VarKey, usize>,
        symbolic_structure: Option<&SymbolicStructure>,
        _total_dof: usize,
    ) -> LinearizerResult<(Mat<f64>, SparseColMat<usize, f64>)> {
        let sym = symbolic_structure.ok_or_else(|| {
            LinearizerError::InvalidInput("SparseMode requires symbolic structure".to_string())
        })?;
        crate::linearizer::cpu::sparse::assemble_sparse(problem, variables, variable_index_map, sym)
    }

    fn compute_column_norms(jacobian: &SparseColMat<usize, f64>) -> Vec<f64> {
        let ncols = jacobian.ncols();
        let sparse_ref = jacobian.as_ref();
        (0..ncols)
            .map(|c| {
                let col_norm_squared: f64 =
                    sparse_ref.val_of_col(c).iter().map(|&val| val * val).sum();
                col_norm_squared.sqrt()
            })
            .collect()
    }

    fn apply_column_scaling(
        jacobian: &SparseColMat<usize, f64>,
        scaling: &[f64],
    ) -> SparseColMat<usize, f64> {
        let ncols = jacobian.ncols();
        let triplets: Vec<Triplet<usize, usize, f64>> =
            (0..ncols).map(|c| Triplet::new(c, c, scaling[c])).collect();
        let scaling_mat = match SparseColMat::try_new_from_triplets(ncols, ncols, &triplets) {
            Ok(mat) => mat,
            Err(_) => return jacobian.clone(),
        };
        jacobian * &scaling_mat
    }

    fn apply_inverse_scaling(step: &Mat<f64>, scaling: &[f64]) -> Mat<f64> {
        let mut result = step.clone();
        for i in 0..step.nrows() {
            result[(i, 0)] *= scaling[i];
        }
        result
    }

    fn hessian_vec_product(hessian: &SparseColMat<usize, f64>, vec: &Mat<f64>) -> Mat<f64> {
        use std::ops::Mul;
        hessian.as_ref().mul(vec)
    }
}

impl AssemblyBackend for DenseMode {
    fn assemble(
        problem: &Problem,
        variables: &SlotMap<VarKey, Box<dyn ManifoldVariable>>,
        variable_index_map: &SecondaryMap<VarKey, usize>,
        _symbolic_structure: Option<&SymbolicStructure>,
        total_dof: usize,
    ) -> LinearizerResult<(Mat<f64>, Mat<f64>)> {
        crate::linearizer::cpu::dense::assemble_dense(
            problem,
            variables,
            variable_index_map,
            total_dof,
        )
    }

    fn compute_column_norms(jacobian: &Mat<f64>) -> Vec<f64> {
        let ncols = jacobian.ncols();
        (0..ncols)
            .map(|c| {
                let mut norm_sq = 0.0;
                for r in 0..jacobian.nrows() {
                    let v = jacobian[(r, c)];
                    norm_sq += v * v;
                }
                norm_sq.sqrt()
            })
            .collect()
    }

    fn apply_column_scaling(jacobian: &Mat<f64>, scaling: &[f64]) -> Mat<f64> {
        let mut result = jacobian.clone();
        for c in 0..jacobian.ncols() {
            for r in 0..jacobian.nrows() {
                result[(r, c)] *= scaling[c];
            }
        }
        result
    }

    fn apply_inverse_scaling(step: &Mat<f64>, scaling: &[f64]) -> Mat<f64> {
        let mut result = step.clone();
        for i in 0..step.nrows() {
            result[(i, 0)] *= scaling[i];
        }
        result
    }

    fn hessian_vec_product(hessian: &Mat<f64>, vec: &Mat<f64>) -> Mat<f64> {
        hessian * vec
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::{VarKey, problem::Problem},
        factors,
        linalg::JacobianMode,
    };
    use apex_manifolds::ManifoldType;
    use faer::prelude::ReborrowMut;
    use nalgebra::dvector;
    use slotmap::{SecondaryMap, SlotMap};

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

    struct LinearFactor {
        target: f64,
    }

    impl factors::Factor for LinearFactor {
        fn linearize(
            &self,
            params: &[&[f64]],
            residual: &mut [f64],
            jacobian: Option<faer::mat::MatMut<'_, f64>>,
        ) {
            residual[0] = params[0][0] - self.target;
            if let Some(mut jac) = jacobian {
                *jac.rb_mut().get_mut(0, 0) = 1.0;
            }
        }
        fn residual_dim(&self) -> usize {
            1
        }
        fn jacobian_shape(&self) -> (usize, usize) {
            (1, 1)
        }
    }

    fn one_var_problem() -> (Problem, VarKey) {
        let mut problem = Problem::new(JacobianMode::Sparse);
        let k = problem.add_variable(ManifoldType::RN, dvector![5.0]);
        problem.add_residual_block(&[k], Box::new(LinearFactor { target: 0.0 }), None);
        (problem, k)
    }

    #[allow(clippy::type_complexity)]
    fn make_index_map(
        problem: &Problem,
    ) -> (
        SlotMap<VarKey, Box<dyn crate::core::variable::ManifoldVariable>>,
        SecondaryMap<VarKey, usize>,
        usize,
    ) {
        let variables = problem.variables.clone();
        let mut index_map: SecondaryMap<VarKey, usize> = SecondaryMap::new();
        let mut offset = 0;
        for (k, v) in &variables {
            index_map.insert(k, offset);
            offset += v.dof();
        }
        (variables, index_map, offset)
    }

    // -------------------------------------------------------------------------
    // compute_block_into
    // -------------------------------------------------------------------------

    #[test]
    fn test_compute_block_into_residual_value() -> TestResult {
        let (problem, _k) = one_var_problem();
        let (variables, _, _) = make_index_map(&problem);
        let block = problem
            .residual_blocks()
            .values()
            .next()
            .ok_or("no blocks")?;
        let mut residual_slice = vec![0.0f64; 1];
        let mut jac_buf = vec![0.0f64; 1];
        compute_block_into(block, &variables, &mut residual_slice, &mut jac_buf)?;
        assert!((residual_slice[0] - 5.0).abs() < 1e-12);
        Ok(())
    }

    #[test]
    fn test_compute_block_into_jacobian_shape() -> TestResult {
        let (problem, _k) = one_var_problem();
        let (variables, _, _) = make_index_map(&problem);
        let block = problem
            .residual_blocks()
            .values()
            .next()
            .ok_or("no blocks")?;
        let mut residual_slice = vec![0.0f64; 1];
        let mut jac_buf = vec![0.0f64; 1];
        let result = compute_block_into(block, &variables, &mut residual_slice, &mut jac_buf)?;
        assert_eq!(result.residual_dim, 1);
        assert_eq!(jac_buf.len(), 1); // 1×1
        Ok(())
    }

    #[test]
    fn test_compute_block_into_variable_local_idx() -> TestResult {
        let (problem, _k) = one_var_problem();
        let (variables, _, _) = make_index_map(&problem);
        let block = problem
            .residual_blocks()
            .values()
            .next()
            .ok_or("no blocks")?;
        let mut residual_slice = vec![0.0f64; 1];
        let mut jac_buf = vec![0.0f64; 1];
        let result = compute_block_into(block, &variables, &mut residual_slice, &mut jac_buf)?;
        assert_eq!(result.variable_local_idx_size_list.len(), 1);
        let (local_idx, size) = result.variable_local_idx_size_list[0];
        assert_eq!(local_idx, 0);
        assert_eq!(size, 1);
        Ok(())
    }

    // -------------------------------------------------------------------------
    // split_by_row_offsets_mut
    // -------------------------------------------------------------------------

    #[test]
    fn test_split_by_row_offsets_mut_basic() {
        let mut buf = vec![1.0f64, 2.0, 3.0, 4.0, 5.0];
        let offsets = vec![(0, 2), (3, 2)];
        let slices = split_by_row_offsets_mut(&mut buf, &offsets);
        assert_eq!(slices.len(), 2);
        assert_eq!(slices[0], &[1.0, 2.0]);
        assert_eq!(slices[1], &[4.0, 5.0]);
    }

    #[test]
    fn test_split_by_row_offsets_mut_write() {
        let mut buf = vec![0.0f64; 4];
        let offsets = vec![(0, 2), (2, 2)];
        {
            let mut slices = split_by_row_offsets_mut(&mut buf, &offsets);
            slices[0][0] = 1.0;
            slices[0][1] = 2.0;
            slices[1][0] = 3.0;
            slices[1][1] = 4.0;
        }
        assert_eq!(buf, vec![1.0, 2.0, 3.0, 4.0]);
    }

    // -------------------------------------------------------------------------
    // SparseMode AssemblyBackend
    // -------------------------------------------------------------------------

    #[test]
    fn test_sparse_backend_assemble() -> TestResult {
        let (problem, _k) = one_var_problem();
        let (variables, index_map, total_dof) = make_index_map(&problem);
        let sym = crate::linearizer::cpu::sparse::build_symbolic_structure(
            &problem, &variables, &index_map, total_dof,
        )?;
        let (residual, _) =
            SparseMode::assemble(&problem, &variables, &index_map, Some(&sym), total_dof)?;
        assert!((residual[(0, 0)] - 5.0).abs() < 1e-12);
        Ok(())
    }

    #[test]
    fn test_sparse_backend_assemble_no_symbolic_returns_error() -> TestResult {
        let (problem, _k) = one_var_problem();
        let (variables, index_map, total_dof) = make_index_map(&problem);
        let result = SparseMode::assemble(&problem, &variables, &index_map, None, total_dof);
        assert!(result.is_err());
        Ok(())
    }

    #[test]
    fn test_sparse_backend_compute_column_norms() -> TestResult {
        let (problem, _k) = one_var_problem();
        let (variables, index_map, total_dof) = make_index_map(&problem);
        let sym = crate::linearizer::cpu::sparse::build_symbolic_structure(
            &problem, &variables, &index_map, total_dof,
        )?;
        let (_, jacobian) =
            SparseMode::assemble(&problem, &variables, &index_map, Some(&sym), total_dof)?;
        let norms = SparseMode::compute_column_norms(&jacobian);
        assert_eq!(norms.len(), 1);
        assert!((norms[0] - 1.0).abs() < 1e-12);
        Ok(())
    }

    #[test]
    fn test_sparse_backend_apply_column_scaling() -> TestResult {
        let (problem, _k) = one_var_problem();
        let (variables, index_map, total_dof) = make_index_map(&problem);
        let sym = crate::linearizer::cpu::sparse::build_symbolic_structure(
            &problem, &variables, &index_map, total_dof,
        )?;
        let (_, jacobian) =
            SparseMode::assemble(&problem, &variables, &index_map, Some(&sym), total_dof)?;
        let scaling = vec![0.5_f64];
        let scaled = SparseMode::apply_column_scaling(&jacobian, &scaling);
        let val = scaled.as_ref().val_of_col(0)[0];
        assert!((val - 0.5).abs() < 1e-12);
        Ok(())
    }

    #[test]
    fn test_sparse_backend_apply_inverse_scaling() {
        let step = Mat::from_fn(1, 1, |_, _| 1.0_f64);
        let scaling = vec![2.0_f64];
        let result = SparseMode::apply_inverse_scaling(&step, &scaling);
        assert!((result[(0, 0)] - 2.0).abs() < 1e-12);
    }

    #[test]
    fn test_sparse_backend_hessian_vec_product() -> TestResult {
        let triplets = vec![faer::sparse::Triplet::new(0usize, 0usize, 4.0_f64)];
        let h = SparseColMat::try_new_from_triplets(1, 1, &triplets)?;
        let v = Mat::from_fn(1, 1, |_, _| 2.0_f64);
        let result = SparseMode::hessian_vec_product(&h, &v);
        assert!((result[(0, 0)] - 8.0).abs() < 1e-12);
        Ok(())
    }

    // -------------------------------------------------------------------------
    // DenseMode AssemblyBackend
    // -------------------------------------------------------------------------

    #[test]
    fn test_dense_backend_assemble() -> TestResult {
        let mut problem = Problem::new(JacobianMode::Dense);
        let k = problem.add_variable(ManifoldType::RN, dvector![5.0]);
        problem.add_residual_block(&[k], Box::new(LinearFactor { target: 0.0 }), None);
        let (variables, index_map, total_dof) = make_index_map(&problem);
        let (residual, _) = DenseMode::assemble(&problem, &variables, &index_map, None, total_dof)?;
        assert!((residual[(0, 0)] - 5.0).abs() < 1e-12);
        Ok(())
    }

    #[test]
    fn test_dense_backend_compute_column_norms() {
        let jacobian = Mat::from_fn(1, 1, |_, _| 1.0_f64);
        let norms = DenseMode::compute_column_norms(&jacobian);
        assert_eq!(norms.len(), 1);
        assert!((norms[0] - 1.0).abs() < 1e-12);
    }

    #[test]
    fn test_dense_backend_apply_column_scaling() {
        let jacobian = Mat::from_fn(1, 1, |_, _| 1.0_f64);
        let scaling = vec![0.5_f64];
        let scaled = DenseMode::apply_column_scaling(&jacobian, &scaling);
        assert!((scaled[(0, 0)] - 0.5).abs() < 1e-12);
    }

    #[test]
    fn test_dense_backend_apply_inverse_scaling() {
        let step = Mat::from_fn(1, 1, |_, _| 1.0_f64);
        let scaling = vec![2.0_f64];
        let result = DenseMode::apply_inverse_scaling(&step, &scaling);
        assert!((result[(0, 0)] - 2.0).abs() < 1e-12);
    }

    #[test]
    fn test_dense_backend_hessian_vec_product() {
        let h = Mat::from_fn(1, 1, |_, _| 4.0_f64);
        let v = Mat::from_fn(1, 1, |_, _| 2.0_f64);
        let result = DenseMode::hessian_vec_product(&h, &v);
        assert!((result[(0, 0)] - 8.0).abs() < 1e-12);
    }
}