rustebra 0.4.0

use alloc::vec;
use alloc::vec::Vec;

use crate::scalar::Scalar;

use super::{CscError, CscMatrix, CsrError, CsrMatrix};

/// Removes stored entries from `m` whose absolute value does not exceed `tolerance`,
/// returning a new `CsrMatrix<T>` with a reduced sparsity pattern and unchanged shape.
///
/// An entry `v` is **kept** when its absolute value strictly exceeds `tolerance`.
/// Setting `tolerance` to `T::zero()` removes only entries that are exactly `0`.
/// Negative `tolerance` is clamped to `T::zero()`. NaN entries are always kept.
///
/// This is useful after repeated sparse additions or Krylov-solver updates, where
/// accumulated fill-in or exact cancellations leave stored zeros that inflate `nnz`
/// without contributing to the matrix value.
///
/// # Errors
///
/// Returns [`CsrError::DimensionOverflow`] if `rows + 1` overflows `usize`, or if the
/// pruned non-zero count exceeds `u32::MAX`.
///
/// # Examples
///
/// ```
/// use rustebra::sparse::{CsrMatrix, prune_csr};
///
/// // Row 0: three entries — two negligible, one significant.
/// let m = CsrMatrix::new(
///     1, 3,
///     vec![0, 3],
///     vec![0, 1, 2],
///     vec![1e-15_f64, 2.0, -1e-15],
/// )
/// .unwrap();
///
/// let pruned = prune_csr(m, 1e-10).unwrap();
/// assert_eq!(pruned.nnz(), 1);
/// assert_eq!(pruned.col_indices(), &[1]);
/// assert_eq!(pruned.values(),      &[2.0]);
/// ```
pub fn prune_csr<T: Scalar + PartialOrd>(
    m: CsrMatrix<T>,
    tolerance: T,
) -> Result<CsrMatrix<T>, CsrError> {
    let (rows, cols, old_row_ptr, old_col, old_val) = m.into_raw_parts();
    let tolerance = if tolerance < T::zero() {
        T::zero()
    } else {
        tolerance
    };
    let neg_tol = T::zero().sub(tolerance);

    let expected_len = rows.checked_add(1).ok_or(CsrError::DimensionOverflow)?;
    let mut new_row_ptr = vec![0u32; expected_len];
    let mut new_col: Vec<u32> = Vec::new();
    let mut new_val: Vec<T> = Vec::new();

    for r in 0..rows {
        for k in old_row_ptr[r] as usize..old_row_ptr[r + 1] as usize {
            let v = old_val[k];
            if !(v <= tolerance && v >= neg_tol) {
                new_col.push(old_col[k]);
                new_val.push(v);
            }
        }
        new_row_ptr[r + 1] =
            u32::try_from(new_col.len()).map_err(|_| CsrError::DimensionOverflow)?;
    }

    Ok(CsrMatrix::new_raw(
        rows,
        cols,
        new_row_ptr,
        new_col,
        new_val,
    ))
}

/// Removes stored entries from `m` whose absolute value does not exceed `tolerance`,
/// returning a new `CscMatrix<T>` with a reduced sparsity pattern and unchanged shape.
///
/// An entry `v` is **kept** when its absolute value strictly exceeds `tolerance`.
/// Setting `tolerance` to `T::zero()` removes only entries that are exactly `0`.
/// Negative `tolerance` is clamped to `T::zero()`. NaN entries are always kept.
///
/// This is the CSC counterpart to [`prune_csr`], providing symmetric support for both
/// sparse matrix formats.
///
/// # Errors
///
/// Returns [`CscError::DimensionOverflow`] if `cols + 1` overflows `usize`, or if the
/// pruned non-zero count exceeds `u32::MAX`.
///
/// # Examples
///
/// ```
/// use rustebra::sparse::{CscMatrix, prune_csc};
///
/// // Column 0: three entries — two negligible, one significant.
/// let m = CscMatrix::new(
///     3, 1,
///     vec![0, 3],
///     vec![0, 1, 2],
///     vec![1e-15_f64, 2.0, -1e-15],
/// )
/// .unwrap();
///
/// let pruned = prune_csc(m, 1e-10).unwrap();
/// assert_eq!(pruned.nnz(), 1);
/// assert_eq!(pruned.row_indices(), &[1]);
/// assert_eq!(pruned.values(),      &[2.0]);
/// ```
pub fn prune_csc<T: Scalar + PartialOrd>(
    m: CscMatrix<T>,
    tolerance: T,
) -> Result<CscMatrix<T>, CscError> {
    let (rows, cols, old_col_ptr, old_row, old_val) = m.into_raw_parts();
    let tolerance = if tolerance < T::zero() {
        T::zero()
    } else {
        tolerance
    };
    let neg_tol = T::zero().sub(tolerance);

    let expected_len = cols.checked_add(1).ok_or(CscError::DimensionOverflow)?;
    let mut new_col_ptr = vec![0u32; expected_len];
    let mut new_row: Vec<u32> = Vec::new();
    let mut new_val: Vec<T> = Vec::new();

    for c in 0..cols {
        for k in old_col_ptr[c] as usize..old_col_ptr[c + 1] as usize {
            let v = old_val[k];
            if !(v <= tolerance && v >= neg_tol) {
                new_row.push(old_row[k]);
                new_val.push(v);
            }
        }
        new_col_ptr[c + 1] =
            u32::try_from(new_row.len()).map_err(|_| CscError::DimensionOverflow)?;
    }

    Ok(CscMatrix::new_raw(
        rows,
        cols,
        new_col_ptr,
        new_row,
        new_val,
    ))
}

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

    #[test]
    fn keeps_entries_outside_tolerance_band() {
        let m = CsrMatrix::new(
            1,
            3,
            vec![0, 3],
            vec![0, 1, 2],
            vec![1e-15_f64, 2.0, -1e-15],
        )
        .unwrap();
        let p = prune_csr(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.col_indices(), &[1]);
        assert_eq!(p.values(), &[2.0]);
    }

    #[test]
    fn zero_tolerance_removes_exact_zeros() {
        let m = CsrMatrix::new(1, 3, vec![0, 3], vec![0, 1, 2], vec![0.0_f64, 5.0, 0.0]).unwrap();
        let p = prune_csr(m, 0.0).unwrap();
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.col_indices(), &[1]);
    }

    #[test]
    fn keeps_negative_values_outside_band() {
        let m = CsrMatrix::new(1, 2, vec![0, 2], vec![0, 1], vec![-3.0_f64, 0.5]).unwrap();
        let p = prune_csr(m, 1.0).unwrap();
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.col_indices(), &[0]);
        assert_eq!(p.values(), &[-3.0]);
    }

    #[test]
    fn preserves_row_structure_with_multi_row_matrix() {
        // Row 0 keeps one entry; row 1 is fully pruned; row 2 keeps both.
        let m = CsrMatrix::new(
            3,
            3,
            vec![0, 2, 4, 6],
            vec![0, 1, 0, 1, 0, 1],
            vec![1e-15_f64, 5.0, 1e-14, 1e-13, 7.0, 8.0],
        )
        .unwrap();
        let p = prune_csr(m, 1e-10).unwrap();
        assert_eq!(p.row_ptr(), &[0, 1, 1, 3]);
        assert_eq!(p.col_indices(), &[1, 0, 1]);
        assert_eq!(p.values(), &[5.0, 7.0, 8.0]);
    }

    #[test]
    fn all_entries_below_tolerance_gives_empty_matrix() {
        let m = CsrMatrix::new(2, 2, vec![0, 1, 2], vec![0, 1], vec![1e-20_f64, 1e-20]).unwrap();
        let p = prune_csr(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 0);
        assert_eq!(p.row_ptr(), &[0, 0, 0]);
    }

    #[test]
    fn empty_matrix_unchanged() {
        let m = CsrMatrix::<f64>::new(3, 3, vec![0, 0, 0, 0], vec![], vec![]).unwrap();
        let p = prune_csr(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 0);
        assert_eq!(p.rows(), 3);
    }

    #[test]
    fn negative_tolerance_clamped_to_zero() {
        // Negative tolerance must not invert the predicate and keep everything.
        let m = CsrMatrix::new(1, 3, vec![0, 3], vec![0, 1, 2], vec![0.0_f64, 5.0, 0.0]).unwrap();
        let p = prune_csr(m, -1.0).unwrap();
        // Clamped to 0.0: only exact zeros removed, 5.0 kept.
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.col_indices(), &[1]);
        assert_eq!(p.values(), &[5.0]);
    }

    #[test]
    fn nan_entry_is_kept_not_silently_dropped() {
        let m = CsrMatrix::new(1, 2, vec![0, 2], vec![0, 1], vec![f64::NAN, 2.0]).unwrap();
        let p = prune_csr(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 2);
        assert!(p.values()[0].is_nan());
        assert_eq!(p.values()[1], 2.0);
    }

    #[test]
    fn rows_plus_one_overflow_is_an_error_not_a_panic() {
        let m: CsrMatrix<f64> = CsrMatrix::new_raw(usize::MAX, 1, vec![], vec![], vec![]);
        let err = prune_csr(m, 1e-10);
        assert_eq!(err, Err(CsrError::DimensionOverflow));
    }
}

#[cfg(test)]
mod tests_csc {
    use super::*;
    use crate::sparse::CscMatrix;

    #[test]
    fn csc_keeps_entries_outside_tolerance_band() {
        let m = CscMatrix::new(
            3,
            1,
            vec![0, 3],
            vec![0, 1, 2],
            vec![1e-15_f64, 2.0, -1e-15],
        )
        .unwrap();
        let p = prune_csc(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.row_indices(), &[1]);
        assert_eq!(p.values(), &[2.0]);
    }

    #[test]
    fn csc_zero_tolerance_removes_exact_zeros() {
        let m = CscMatrix::new(3, 1, vec![0, 3], vec![0, 1, 2], vec![0.0_f64, 5.0, 0.0]).unwrap();
        let p = prune_csc(m, 0.0).unwrap();
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.row_indices(), &[1]);
    }

    #[test]
    fn csc_keeps_negative_values_outside_band() {
        let m = CscMatrix::new(2, 1, vec![0, 2], vec![0, 1], vec![-3.0_f64, 0.5]).unwrap();
        let p = prune_csc(m, 1.0).unwrap();
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.row_indices(), &[0]);
        assert_eq!(p.values(), &[-3.0]);
    }

    #[test]
    fn csc_preserves_column_structure_with_multi_column_matrix() {
        // Column 0 keeps one entry; column 1 is fully pruned; column 2 keeps both.
        let m = CscMatrix::new(
            3,
            3,
            vec![0, 2, 4, 6],
            vec![0, 1, 0, 1, 0, 1],
            vec![1e-15_f64, 5.0, 1e-14, 1e-13, 7.0, 8.0],
        )
        .unwrap();
        let p = prune_csc(m, 1e-10).unwrap();
        assert_eq!(p.col_ptr(), &[0, 1, 1, 3]);
        assert_eq!(p.row_indices(), &[1, 0, 1]);
        assert_eq!(p.values(), &[5.0, 7.0, 8.0]);
    }

    #[test]
    fn csc_all_entries_below_tolerance_gives_empty_matrix() {
        let m = CscMatrix::new(2, 2, vec![0, 1, 2], vec![0, 1], vec![1e-20_f64, 1e-20]).unwrap();
        let p = prune_csc(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 0);
        assert_eq!(p.col_ptr(), &[0, 0, 0]);
    }

    #[test]
    fn csc_empty_matrix_unchanged() {
        let m = CscMatrix::<f64>::new(3, 3, vec![0, 0, 0, 0], vec![], vec![]).unwrap();
        let p = prune_csc(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 0);
        assert_eq!(p.cols(), 3);
    }

    #[test]
    fn csc_negative_tolerance_clamped_to_zero() {
        let m = CscMatrix::new(3, 1, vec![0, 3], vec![0, 1, 2], vec![0.0_f64, 5.0, 0.0]).unwrap();
        let p = prune_csc(m, -1.0).unwrap();
        assert_eq!(p.nnz(), 1);
        assert_eq!(p.row_indices(), &[1]);
        assert_eq!(p.values(), &[5.0]);
    }

    #[test]
    fn csc_nan_entry_is_kept_not_silently_dropped() {
        let m = CscMatrix::new(2, 1, vec![0, 2], vec![0, 1], vec![f64::NAN, 2.0]).unwrap();
        let p = prune_csc(m, 1e-10).unwrap();
        assert_eq!(p.nnz(), 2);
        assert!(p.values()[0].is_nan());
        assert_eq!(p.values()[1], 2.0);
    }

    #[test]
    fn cols_plus_one_overflow_is_an_error_not_a_panic() {
        let m: CscMatrix<f64> = CscMatrix::new_raw(1, usize::MAX, vec![], vec![], vec![]);
        let err = prune_csc(m, 1e-10);
        assert_eq!(err, Err(CscError::DimensionOverflow));
    }
}