bixverse-rs 0.3.5

//! Sparse matrix formats, sparse operations and helpers to transform different
//! formats into each other.

use faer::{Mat, MatRef};
use num_traits::ToPrimitive;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rayon::prelude::*;
use std::ops::{Add, AddAssign, Mul};

use crate::core::math::pca_svd::SvdResults;
use crate::prelude::*;

/////////////
// Helpers //
/////////////

/// Counts the zeroes in a given faer matrix
///
/// ### Params
///
/// * `mat` - The respective faer matrix
///
/// ### Returns
///
/// A tuple with the first being the total zeroes, the second the zeroes per
/// row and the last element being the column zeroes.
pub fn count_zeroes<T>(mat: &MatRef<T>) -> (usize, Vec<usize>, Vec<usize>)
where
    T: BixverseFloat,
{
    let (nrow, ncol) = mat.shape();
    let mut total_zeroes = 0_usize;
    let mut row_zeroes = vec![0_usize; nrow];
    let mut col_zeroes = vec![0_usize; ncol];

    let zero = T::zero();

    for j in 0..ncol {
        for i in 0..nrow {
            let val = unsafe { mat.get_unchecked(i, j) };
            if *val == zero {
                total_zeroes += 1;
                row_zeroes[i] += 1;
                col_zeroes[j] += 1;
            }
        }
    }

    (total_zeroes, row_zeroes, col_zeroes)
}

///////////////////////
// Sparse structures //
///////////////////////

///////////
// Enums //
///////////

/// Type to describe the CompressedSparseFormat
#[derive(Debug, Clone, Copy)]
pub enum CompressedSparseFormat {
    /// CSC-formatted data
    Csc,
    /// CSR-formatted data
    Csr,
}

/// Display implementation for [CompressedSparseFormat]
impl std::fmt::Display for CompressedSparseFormat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CompressedSparseFormat::Csc => write!(f, "CSC"),
            CompressedSparseFormat::Csr => write!(f, "CSR"),
        }
    }
}

impl CompressedSparseFormat {
    /// Returns boolean if it's CSC
    ///
    /// ### Returns
    ///
    /// Boolean indicating if CSC
    #[inline(always)]
    pub fn is_csc(&self) -> bool {
        matches!(self, CompressedSparseFormat::Csc)
    }
    /// Returns boolean if it's CSR
    ///
    /// ### Returns
    ///
    /// Boolean indicating if CSR
    #[inline(always)]
    pub fn is_csr(&self) -> bool {
        matches!(self, CompressedSparseFormat::Csr)
    }
}

/// Helper function to parse compressed sparse format
///
/// ### Params
///
/// * `s` - String specifying the type
///
/// ### Return
///
/// Returns an `Option<CompressedSparseFormat>`
pub fn parse_compressed_sparse_format(s: &str) -> Option<CompressedSparseFormat> {
    match s.to_lowercase().as_str() {
        "csr" => Some(CompressedSparseFormat::Csr),
        "csc" => Some(CompressedSparseFormat::Csc),
        _ => None,
    }
}

////////////////
// SparseAxis //
////////////////

/// Generate structure to store sparse rows or columns
pub struct SparseAxis<T, U = T> {
    /// The indices of the values/non-zero positions
    pub indices: Vec<usize>,
    /// The values in the row/columns
    pub data: Vec<T>,
    /// An optional second data layer
    pub data_2: Option<Vec<U>>,
    /// Is the data stored in `Csr` or `Csc`. `Csr` -> sparse row; `Csc` ->
    /// sparse column
    pub cs_type: CompressedSparseFormat,
    /// Total values in that dimension
    pub len: usize,
}

impl<T, U> SparseAxis<T, U>
where
    T: BixverseNumeric,
    U: BixverseNumeric,
{
    /// Generate a new `SparseAxis` in CSC format
    ///
    /// ### Params
    ///
    /// * `indices` - The indices of the values/non-zero positions
    /// * `data` - The values in the row/columns
    /// * `data_2` - An optional second data layer
    /// * `len` - Number of rows in this sparse column
    pub fn new_csc(indices: Vec<usize>, data: Vec<T>, data_2: Option<Vec<U>>, len: usize) -> Self {
        SparseAxis {
            indices,
            data,
            data_2,
            cs_type: CompressedSparseFormat::Csc,
            len,
        }
    }

    /// Get references to the indices and second layer
    ///
    /// ### Returns
    ///
    /// A tuple of `(indices, data_2)`
    pub fn get_indices_data_2(&self) -> Result<(&[usize], &[U]), BixverseErrors> {
        let indices = &self.indices;
        let data_2 = self
            .data_2
            .as_ref()
            .ok_or(BixverseErrors::Data2NotAvailable)?;

        Ok((indices, data_2))
    }

    /// Construct a `SparseAxis` in CSC format from index and value vectors.
    ///
    /// Stores `indices` (converted to `usize`) in `self.indices`, leaves
    /// `data` empty, and places `values` into `data_2`. Useful for building
    /// single-target test fixtures and lightweight sparse columns where only
    /// the float layer is needed.
    ///
    /// ### Params
    ///
    /// * `indices` - Non-zero positions (must be convertible to `usize`).
    /// * `values` - Corresponding values, stored in `data_2`.
    ///
    /// ### Panics
    ///
    /// Panics if `indices` and `values` differ in length.
    pub fn from_vecs_to_csc(indices: Vec<T>, values: Vec<U>) -> Result<Self, BixverseErrors>
    where
        T: Into<usize>,
    {
        if indices.len() != values.len() {
            return Err(BixverseErrors::DimensionMisMatchSparse {
                indices_len: indices.len(),
                data_len: values.len(),
            });
        }

        let usize_indices: Vec<usize> = indices.into_iter().map(Into::into).collect();
        let len = usize_indices.last().map_or(0, |&i| i + 1);

        Ok(Self {
            indices: usize_indices,
            data: Vec::new(),
            data_2: Some(values),
            cs_type: CompressedSparseFormat::Csc,
            len,
        })
    }
}

///////////////////////////
// CompressedSparseData2 //
///////////////////////////

/// Structure to store compressed sparse data with two potential data layers.
/// Think for example in single cell raw counts and normalised counts.
#[derive(Debug, Clone)]
pub struct CompressedSparseData2<T, U = T>
where
    T: Clone,
    U: Clone,
{
    /// The first data slot for this compressed sparse data format
    pub data: Vec<T>,
    /// The indices of the data points
    pub indices: Vec<u32>,
    /// The indptr of the data points
    pub indptr: Vec<u32>,
    /// Enum defining if the data is stored in CSC or CSR, see
    /// [CompressedSparseFormat]
    pub cs_type: CompressedSparseFormat,
    /// Optional second data slot for a different layer of the data (for
    /// example raw and normalised counts)
    pub data_2: Option<Vec<U>>,
    /// Shape of the data (rows, cols)
    pub shape: (usize, usize),
}

impl<T, U> CompressedSparseData2<T, U>
where
    T: BixverseNumeric,
    U: BixverseNumeric,
{
    /// Generate a nes CSC version of the matrix
    ///
    /// ### Params
    ///
    /// * `data` - The underlying data
    /// * `indices` - The index positions (in this case row indices)
    /// * `indptr` - The index pointer (in this case the column index pointers)
    /// * `data2` - An optional second layer
    #[allow(dead_code)]
    pub fn new_csc(
        data: &[T],
        indices: &[u32],
        indptr: &[u32],
        data2: Option<&[U]>,
        shape: (usize, usize),
    ) -> Self {
        Self {
            data: data.to_vec(),
            indices: indices.to_vec(),
            indptr: indptr.to_vec(), // Fixed: was using indices instead of indptr
            cs_type: CompressedSparseFormat::Csc,
            data_2: data2.map(|d| d.to_vec()),
            shape,
        }
    }

    /// Generate a nes CSR version of the matrix
    ///
    /// ### Params
    ///
    /// * `data` - The underlying data
    /// * `indices` - The index positions (in this case row indices)
    /// * `indptr` - The index pointer (in this case the column index pointers)
    /// * `data2` - An optional second layer
    pub fn new_csr(
        data: &[T],
        indices: &[u32],
        indptr: &[u32],
        data2: Option<&[U]>,
        shape: (usize, usize),
    ) -> Self {
        Self {
            data: data.to_vec(),
            indices: indices.to_vec(),
            indptr: indptr.to_vec(), // Fixed: was using indices instead of indptr
            cs_type: CompressedSparseFormat::Csr,
            data_2: data2.map(|d| d.to_vec()),
            shape,
        }
    }

    /// Transform from CSC to CSR or vice versa
    ///
    /// ### Returns
    ///
    /// The transformed/transposed version
    pub fn transform(&self) -> Self {
        transpose_sparse(self)
    }

    /// Transform from CSC to CSR or vice versa, keeping only one data layer.
    ///
    /// ### Params
    ///
    /// * `use_second_layer` - If `true`, the output keeps `data_2` and `data`
    ///   is empty. If `false`, the output keeps `data` and `data_2` is `None`.
    ///
    /// ### Returns
    ///
    /// The transformed/transposed version with only the requested layer.
    pub fn transform_single_layer(&self, use_second_layer: bool) -> Result<Self, BixverseErrors> {
        transpose_sparse_single_layer(self, use_second_layer)
    }

    /// Transpose and convert
    ///
    /// ### Returns
    ///
    /// The transformed/transposed version
    pub fn transpose_and_convert(&self) -> Self {
        match self.cs_type {
            CompressedSparseFormat::Csr => {
                // convert first and then switch around
                let csc_version = transpose_sparse(self);
                CompressedSparseData2 {
                    data: csc_version.data,
                    indices: csc_version.indices,
                    indptr: csc_version.indptr,
                    cs_type: CompressedSparseFormat::Csr, // relabel as CSR
                    data_2: csc_version.data_2,
                    shape: (self.shape.1, self.shape.0), // swap dimensions
                }
            }
            CompressedSparseFormat::Csc => {
                // no conversion needed here! simple transpose is enough...
                CompressedSparseData2 {
                    data: self.data.clone(),
                    indices: self.indices.clone(),
                    indptr: self.indptr.clone(),
                    cs_type: CompressedSparseFormat::Csr,
                    data_2: self.data_2.clone(),
                    shape: (self.shape.1, self.shape.0),
                }
            }
        }
    }

    /// Generates a sparse matrix from a dense matrix
    ///
    /// ### Params
    ///
    /// * `mat`: The dense matrix to convert to a sparse matrix
    /// * `format`: The format of the sparse matrix to generate
    ///
    /// ### Returns
    ///
    /// * `Self`: The sparse matrix generated from the dense matrix
    pub fn from_dense_matrix(mat: faer::MatRef<T>, format: CompressedSparseFormat) -> Self
    where
        T: BixverseFloat,
    {
        let (nrows, ncols) = (mat.nrows(), mat.ncols());
        let mut data = Vec::new();
        let mut indices: Vec<u32> = Vec::new();
        let mut indptr: Vec<u32> = Vec::new();

        match format {
            CompressedSparseFormat::Csr => {
                indptr.push(0);
                for i in 0..nrows {
                    for j in 0..ncols {
                        let val = mat[(i, j)];
                        if val != T::zero() {
                            data.push(val);
                            indices.push(j as u32);
                        }
                    }
                    indptr.push(data.len() as u32);
                }
            }
            CompressedSparseFormat::Csc => {
                indptr.push(0);
                for j in 0..ncols {
                    for i in 0..nrows {
                        let val = mat[(i, j)];
                        if val != T::zero() {
                            data.push(val);
                            indices.push(i as u32);
                        }
                    }
                    indptr.push(data.len() as u32);
                }
            }
        }

        Self {
            data,
            indices,
            indptr,
            cs_type: format,
            data_2: None,
            shape: (nrows, ncols),
        }
    }

    /// Create a sparse matrix from an upper triangular matrix.
    ///
    /// ### Params
    ///
    /// * `upper_triangle` - The upper triangular matrix.
    /// * `n` - The number of rows and columns in the matrix.
    /// * `include_diagonal` - Whether to include the diagonal elements.
    /// * `format` - The format of the sparse matrix.
    ///
    /// ### Returns
    ///
    /// A sparse matrix.
    pub fn from_upper_triangle_sym(
        upper_triangle: &[T],
        n: usize,
        include_diagonal: bool,
        format: CompressedSparseFormat,
    ) -> Self
    where
        T: BixverseFloat,
    {
        // lambda function in Rust style...
        let get_value = |row: usize, col: usize| -> T {
            if row == col {
                if include_diagonal {
                    let offset = row * n - row * (row + 1) / 2 + col;
                    upper_triangle[offset]
                } else {
                    T::one()
                }
            } else if row < col {
                let offset = if include_diagonal {
                    row * n - row * (row + 1) / 2 + col
                } else {
                    row * (n - 1) - row * (row + 1) / 2 + col - 1
                };
                upper_triangle[offset]
            } else {
                let offset = if include_diagonal {
                    col * n - col * (col + 1) / 2 + row
                } else {
                    col * (n - 1) - col * (col + 1) / 2 + row - 1
                };
                upper_triangle[offset]
            }
        };

        let mut data = Vec::new();
        let mut indices: Vec<u32> = Vec::new();
        let mut indptr: Vec<u32> = Vec::new();

        match format {
            CompressedSparseFormat::Csr => {
                indptr.push(0);
                for row in 0..n {
                    for col in 0..n {
                        let value = get_value(row, col);
                        if value != T::zero() {
                            data.push(value);
                            indices.push(col as u32);
                        }
                    }
                    indptr.push(data.len() as u32);
                }
            }
            CompressedSparseFormat::Csc => {
                indptr.push(0);
                for col in 0..n {
                    for row in 0..n {
                        let value = get_value(row, col);
                        if value != T::zero() {
                            data.push(value);
                            indices.push(row as u32);
                        }
                    }
                    indptr.push(data.len() as u32);
                }
            }
        }

        Self {
            data,
            indices,
            indptr,
            cs_type: format,
            data_2: None,
            shape: (n, n),
        }
    }

    /// Returns the shape of the matrix
    ///
    /// ### Returns
    ///
    /// A tuple of `(nrow, ncol)`
    pub fn shape(&self) -> (usize, usize) {
        self.shape
    }

    /// Return the number of rows
    ///
    /// ### Returns
    ///
    /// Number of rows
    pub fn nrows(&self) -> usize {
        self.shape.0
    }

    /// Return the number of columns
    ///
    /// ### Returns
    ///
    /// Number of columns
    pub fn ncols(&self) -> usize {
        self.shape.1
    }

    /// Returns the NNZ
    ///
    /// ### Returns
    ///
    /// The number of NNZ
    pub fn get_nnz(&self) -> usize {
        self.indices.len()
    }

    /// Return the second layer
    ///
    /// If this does not exist, the function will panic
    ///
    /// ### Returns
    ///
    /// Vector of the second layer
    pub fn get_data2_unsafe(&self) -> Vec<U> {
        self.data_2.clone().unwrap()
    }

    /////////////
    // Slicing //
    /////////////

    /// Gather along the major axis (rows for CSR, cols for CSC).
    ///
    /// Arbitrary order allowed; duplicates and out-of-range indices rejected.
    ///
    /// ### Params
    ///
    /// * `keep` - The indices on the major dimensions to keep
    ///
    /// ### Returns
    ///
    /// Self sliced along the major axis
    fn slice_major(&self, keep: &[usize]) -> Result<Self, BixverseErrors> {
        let major = self.major_dim();
        let mut seen = vec![false; major];

        let mut indptr = Vec::with_capacity(keep.len() + 1);
        indptr.push(0u32);
        let mut indices = Vec::new();
        let mut data = Vec::new();
        let mut data_2 = self.data_2.as_ref().map(|_| Vec::new());

        for &m in keep {
            if m >= major {
                return Err(BixverseErrors::SliceIndexOutOfBounds {
                    index: m,
                    len: major,
                });
            }
            if seen[m] {
                return Err(BixverseErrors::DuplicateSliceIndex(m));
            }
            seen[m] = true;

            let s = self.indptr[m] as usize;
            let e = self.indptr[m + 1] as usize;
            indices.extend_from_slice(&self.indices[s..e]);
            data.extend_from_slice(&self.data[s..e]);
            if let (Some(dst), Some(src)) = (data_2.as_mut(), self.data_2.as_ref()) {
                dst.extend_from_slice(&src[s..e]);
            }
            indptr.push(indices.len() as u32);
        }

        let shape = match self.cs_type {
            CompressedSparseFormat::Csr => (keep.len(), self.shape.1),
            CompressedSparseFormat::Csc => (self.shape.0, keep.len()),
        };

        let out = Self {
            data,
            indices,
            indptr,
            cs_type: self.cs_type,
            data_2,
            shape,
        };
        out.assert_invariants();
        Ok(out)
    }

    /// Filter + remap the minor axis (cols for CSR, rows for CSC).
    ///
    /// Arbitrary order allowed; duplicates and out-of-range indices rejected.
    /// Minor indices within each major run are kept ascending in the new space.
    ///
    /// ### Params
    ///
    /// * `keep` - The indices on the major dimensions to keep
    ///
    /// ### Returns
    ///
    /// Self sliced along the minor axis (with remapping)
    fn slice_minor(&self, keep: &[usize]) -> Result<Self, BixverseErrors> {
        let minor = self.minor_dim();
        let mut map: Vec<Option<usize>> = vec![None; minor];
        for (new_idx, &old) in keep.iter().enumerate() {
            if old >= minor {
                return Err(BixverseErrors::SliceIndexOutOfBounds {
                    index: old,
                    len: minor,
                });
            }
            if map[old].is_some() {
                return Err(BixverseErrors::DuplicateSliceIndex(old));
            }
            map[old] = Some(new_idx);
        }

        let major = self.major_dim();
        let mut indptr = Vec::with_capacity(major + 1);
        indptr.push(0u32);
        let mut indices = Vec::new();
        let mut data = Vec::new();
        let mut data_2 = self.data_2.as_ref().map(|_| Vec::new());

        let mut kept: Vec<(usize, usize)> = Vec::new(); // (new_minor, src_pos)
        for m in 0..major {
            let s = self.indptr[m];
            let e = self.indptr[m + 1];
            kept.clear();
            for p in s..e {
                if let Some(nm) = map[self.indices[p as usize] as usize] {
                    kept.push((nm, p as usize));
                }
            }
            kept.sort_unstable_by_key(|&(nm, _)| nm);
            for &(nm, p) in &kept {
                indices.push(nm as u32);
                data.push(self.data[p]);
                if let (Some(dst), Some(src)) = (data_2.as_mut(), self.data_2.as_ref()) {
                    dst.push(src[p]);
                }
            }
            indptr.push(indices.len() as u32);
        }

        let shape = match self.cs_type {
            CompressedSparseFormat::Csr => (self.shape.0, keep.len()),
            CompressedSparseFormat::Csc => (keep.len(), self.shape.1),
        };

        let out = Self {
            data,
            indices,
            indptr,
            cs_type: self.cs_type,
            data_2,
            shape,
        };
        out.assert_invariants();
        Ok(out)
    }

    /// Slice the rows of the matrix
    ///
    /// ### Params
    ///
    /// * `rows` - The row indices to keep
    ///
    /// ### Returns
    ///
    /// Self with the rows to keep
    pub fn slice_rows(&self, rows: &[usize]) -> Result<Self, BixverseErrors> {
        match self.cs_type {
            CompressedSparseFormat::Csr => self.slice_major(rows),
            CompressedSparseFormat::Csc => self.slice_minor(rows),
        }
    }

    /// Slice the columns of the matrix
    ///
    /// ### Params
    ///
    /// * `cols` - The column indices to keep
    ///
    /// ### Returns
    ///
    /// Self with the cols to keep
    pub fn slice_cols(&self, cols: &[usize]) -> Result<Self, BixverseErrors> {
        match self.cs_type {
            CompressedSparseFormat::Csr => self.slice_minor(cols),
            CompressedSparseFormat::Csc => self.slice_major(cols),
        }
    }

    /// Slice across rows and columns
    ///
    /// ### Params
    ///
    /// * `rows` - The row indices to keep
    /// * `cols` - The column indices to keep
    ///
    /// ### Returns
    ///
    /// Self with the rows and columns to keep
    pub fn slice(&self, rows: &[usize], cols: &[usize]) -> Result<Self, BixverseErrors> {
        self.slice_rows(rows)?.slice_cols(cols)
    }

    /////////////
    // Helpers //
    /////////////

    /// Return the major dimensions
    ///
    /// ### Returns
    ///
    /// Major dimension
    fn major_dim(&self) -> usize {
        self.indptr.len() - 1
    }

    /// Returns the minor dimension
    ///
    /// ### Returns
    ///
    /// Minor dimension
    fn minor_dim(&self) -> usize {
        match self.cs_type {
            CompressedSparseFormat::Csr => self.shape.1,
            CompressedSparseFormat::Csc => self.shape.0,
        }
    }

    /// Assertion helper for invariants
    pub fn assert_invariants(&self) {
        let expected_major = match self.cs_type {
            CompressedSparseFormat::Csr => self.shape.0,
            CompressedSparseFormat::Csc => self.shape.1,
        };
        debug_assert_eq!(self.indices.len(), self.data.len());
        debug_assert_eq!(self.indptr.len(), expected_major + 1);
        debug_assert_eq!(*self.indptr.last().unwrap() as usize, self.data.len());
        if let Some(d2) = &self.data_2 {
            debug_assert_eq!(d2.len(), self.data.len());
        }
    }
}

////////////////////////
// Format conversions //
////////////////////////

/// Transpose a compressed sparse matrix (CSC→CSR or CSR→CSC).
///
/// This is the standard two-pass sparse transpose in O(nnz) time.
///
/// ### Params
///
/// * `sparse_data`: The input compressed sparse matrix to be transformed.
///
/// ### Returns
///
/// The transposed compressed sparse matrix.
pub fn transpose_sparse<T, U>(
    sparse_data: &CompressedSparseData2<T, U>,
) -> CompressedSparseData2<T, U>
where
    T: BixverseNumeric,
    U: BixverseNumeric,
{
    let nnz = sparse_data.get_nnz();
    let (nrow, ncol) = sparse_data.shape();

    // the "minor" dimension is what becomes the new indptr axis.
    let (new_major, new_type) = match sparse_data.cs_type {
        CompressedSparseFormat::Csc => (nrow, CompressedSparseFormat::Csr),
        CompressedSparseFormat::Csr => (ncol, CompressedSparseFormat::Csc),
    };

    // first pass: count entries per new-major index
    let mut new_indptr = vec![0u32; new_major + 1];
    for &idx in &sparse_data.indices {
        new_indptr[(idx + 1) as usize] += 1;
    }
    for i in 0..new_major {
        new_indptr[i + 1] += new_indptr[i];
    }

    // second pass: scatter data
    let mut new_data: Vec<T> = vec![T::default(); nnz];
    let mut new_indices: Vec<u32> = vec![0u32; nnz];
    let mut new_data2: Option<Vec<U>> =
        sparse_data.data_2.as_ref().map(|_| vec![U::default(); nnz]);
    unsafe {
        new_data.set_len(nnz);
        new_indices.set_len(nnz);
    }

    let old_major_len = sparse_data.indptr.len() - 1;
    for major in 0..old_major_len {
        for idx in sparse_data.indptr[major]..sparse_data.indptr[major + 1] {
            let minor = sparse_data.indices[idx as usize];
            let pos = new_indptr[minor as usize];

            let pos_usize = pos as usize;

            // SAFETY: pos < nnz guaranteed by the counting pass
            unsafe {
                *new_data.get_unchecked_mut(pos_usize) = sparse_data.data[idx as usize];
                *new_indices.get_unchecked_mut(pos_usize) = major as u32;
            }

            if let (Some(src), Some(dst)) = (&sparse_data.data_2, &mut new_data2) {
                unsafe {
                    *dst.get_unchecked_mut(pos_usize) = src[idx as usize];
                }
            }

            new_indptr[minor as usize] += 1;
        }
    }

    // restore new_indptr: the scatter pass shifted every entry forward by its
    // count, so we shift the whole array right by one position.
    for i in (1..=new_major).rev() {
        new_indptr[i] = new_indptr[i - 1];
    }
    new_indptr[0] = 0;

    CompressedSparseData2 {
        data: new_data,
        indices: new_indices,
        indptr: new_indptr,
        cs_type: new_type,
        data_2: new_data2,
        shape: (nrow, ncol),
    }
}

/// Transpose a compressed sparse matrix (CSC→CSR or CSR→CSC), keeping only
/// one of the two data layers.
///
/// Useful when the caller knows the other layer is dead weight (e.g. raw
/// counts when only normalised counts are consumed downstream). Skips the
/// scatter for the unused layer entirely.
///
/// ### Params
///
/// * `sparse_data` - The input compressed sparse matrix to be transposed.
/// * `use_second_layer` - If `true`, only `data_2` is transposed and the
///   output's `data` is empty. If `false`, only `data` is transposed and the
///   output's `data_2` is `None`.
///
/// ### Returns
///
/// The transposed compressed sparse matrix with only the requested layer
/// populated.
pub fn transpose_sparse_single_layer<T, U>(
    sparse_data: &CompressedSparseData2<T, U>,
    use_second_layer: bool,
) -> Result<CompressedSparseData2<T, U>, BixverseErrors>
where
    T: BixverseNumeric,
    U: BixverseNumeric,
{
    let nnz = sparse_data.get_nnz();
    let (nrow, ncol) = sparse_data.shape();

    let (new_major, new_type) = match sparse_data.cs_type {
        CompressedSparseFormat::Csc => (nrow, CompressedSparseFormat::Csr),
        CompressedSparseFormat::Csr => (ncol, CompressedSparseFormat::Csc),
    };

    // first pass: count entries per new-major index
    let mut new_indptr = vec![0u32; new_major + 1];
    for &idx in &sparse_data.indices {
        new_indptr[(idx + 1) as usize] += 1;
    }
    for i in 0..new_major {
        new_indptr[i + 1] += new_indptr[i];
    }

    let mut new_indices: Vec<u32> = vec![0u32; nnz];

    // allocate only for the kept layer
    let (mut new_data, mut new_data2): (Vec<T>, Option<Vec<U>>) = if use_second_layer {
        (Vec::new(), Some(vec![U::default(); nnz]))
    } else {
        (vec![T::default(); nnz], None)
    };

    let old_major_len = sparse_data.indptr.len() - 1;

    // second pass: scatter. Two branches to keep the inner loop tight.
    if use_second_layer {
        let src = sparse_data
            .data_2
            .as_ref()
            .ok_or(BixverseErrors::Data2NotAvailable)?
            .as_slice();
        let dst = new_data2.as_mut().unwrap();
        for major in 0..old_major_len {
            for idx in sparse_data.indptr[major]..sparse_data.indptr[major + 1] {
                let idx_usize = idx as usize;

                let minor = sparse_data.indices[idx_usize];
                let minor_usize = minor as usize;

                let pos = new_indptr[minor_usize];
                let pos_usize = pos as usize;
                // SAFETY: pos < nnz guaranteed by the counting pass
                unsafe {
                    *new_indices.get_unchecked_mut(pos_usize) = major as u32;
                    *dst.get_unchecked_mut(pos_usize) = src[idx_usize];
                }
                new_indptr[minor_usize] += 1;
            }
        }
    } else {
        for major in 0..old_major_len {
            for idx in sparse_data.indptr[major]..sparse_data.indptr[major + 1] {
                let idx_usize = idx as usize;

                let minor = sparse_data.indices[idx_usize];
                let minor_usize = minor as usize;

                let pos = new_indptr[minor_usize];
                let pos_usize = pos as usize;
                // SAFETY: pos < nnz guaranteed by the counting pass
                unsafe {
                    *new_indices.get_unchecked_mut(pos_usize) = major as u32;
                    *new_data.get_unchecked_mut(pos_usize) = sparse_data.data[idx_usize];
                }
                new_indptr[minor_usize] += 1;
            }
        }
    }

    // restore new_indptr: the scatter pass shifted every entry forward by its
    // count, so we shift the whole array right by one position.
    for i in (1..=new_major).rev() {
        new_indptr[i] = new_indptr[i - 1];
    }
    new_indptr[0] = 0;

    Ok(CompressedSparseData2 {
        data: new_data,
        indices: new_indices,
        indptr: new_indptr,
        cs_type: new_type,
        data_2: new_data2,
        shape: (nrow, ncol),
    })
}

/// Transform COO stored data into CSR
///
/// ### Params
///
/// * `rows` - Row indices
/// * `cols` - Col indices
/// * `vals` - The values to store in the matrix
///
/// ### Returns
///
/// `CompressedSparseData2` in CSR format
pub fn coo_to_csr<T>(
    rows: &[u32],
    cols: &[u32],
    vals: &[T],
    shape: (usize, usize),
) -> CompressedSparseData2<T>
where
    T: BixverseNumeric,
{
    let n_rows = shape.0;

    // sort by (row, col) and merge duplicates
    let mut entries: Vec<(u32, u32, T)> = rows
        .iter()
        .zip(cols.iter())
        .zip(vals.iter())
        .map(|((&r, &c), &v)| (r, c, v))
        .collect();

    entries.sort_unstable_by_key(|&(r, c, _)| (r, c));

    // merge duplicates; can happen during additions
    let mut merged_entries = Vec::new();
    if !entries.is_empty() {
        let mut current = entries[0];

        for &(r, c, v) in &entries[1..] {
            if r == current.0 && c == current.1 {
                current.2 += v;
            } else {
                if current.2 != T::default() {
                    merged_entries.push(current);
                }
                current = (r, c, v);
            }
        }
        if current.2 != T::default() {
            merged_entries.push(current);
        }
    }

    // build CSR from merged entries
    let final_nnz = merged_entries.len();
    let mut data = Vec::with_capacity(final_nnz);
    let mut indices = Vec::with_capacity(final_nnz);
    let mut indptr = vec![0u32; n_rows + 1];

    for &(row, col, val) in &merged_entries {
        data.push(val);
        indices.push(col);
        indptr[(row + 1) as usize] += 1;
    }

    // Convert counts to cumulative offsets
    for i in 0..n_rows {
        indptr[i + 1] += indptr[i];
    }

    CompressedSparseData2::new_csr(&data, &indices, &indptr, None, shape)
}

/// Optimised COO to CSR - assumes input is already sorted by (row, col)
///
/// ### Params
///
/// * `rows` - Row indices (must be sorted by row first, then col)
/// * `cols` - Col indices
/// * `vals` - Values
/// * `shape` - Matrix dimensions
/// * `is_sorted` - If true, skips sorting step
///
/// ### Returns
///
/// CSR matrix
pub fn coo_to_csr_presorted<T>(
    rows: &[u32],
    cols: &[u32],
    vals: &[T],
    shape: (usize, usize),
) -> CompressedSparseData2<T>
where
    T: BixverseNumeric,
{
    let n_rows = shape.0;
    let nnz = rows.len();

    let mut data = Vec::with_capacity(nnz);
    let mut indices: Vec<u32> = Vec::with_capacity(nnz);
    let mut indptr = vec![0u32; n_rows + 1];

    // unsafe to squeeze out performance...
    unsafe {
        data.set_len(nnz);
        indices.set_len(nnz);

        let data_ptr: *mut T = data.as_mut_ptr();
        let indices_ptr: *mut u32 = indices.as_mut_ptr();
        let indptr_ptr: *mut u32 = indptr.as_mut_ptr();

        for i in 0..nnz {
            *data_ptr.add(i) = *vals.get_unchecked(i);
            *indices_ptr.add(i) = *cols.get_unchecked(i);
            let row = *rows.get_unchecked(i) as usize;
            *indptr_ptr.add(row + 1) += 1;
        }

        for i in 0..n_rows {
            *indptr_ptr.add(i + 1) += *indptr_ptr.add(i);
        }
    }

    CompressedSparseData2::new_csr(&data, &indices, &indptr, None, shape)
}

///////////////////////
// Sparse operations //
///////////////////////

/// Add two CSR matrices together
///
/// ### Params
///
/// * `a` - Reference to the first CompressedSparseData2 (in CSR format!)
/// * `b` - Reference to the second CompressedSparseData2 (in CSR format!)
///
/// ### Returns
///
/// `CompressedSparseData2` with added values between the two.
pub fn sparse_add_csr<T>(
    a: &CompressedSparseData2<T>,
    b: &CompressedSparseData2<T>,
) -> Result<CompressedSparseData2<T>, BixverseErrors>
where
    T: BixverseNumeric + Into<f64> + Add<Output = T>,
{
    if !a.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if !b.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if a.shape != b.shape {
        return Err(BixverseErrors::ShapeMismatchSparse);
    }

    const EPSILON: f32 = 1e-9;
    let n_rows = a.shape.0;

    let mut rows: Vec<u32> = Vec::new();
    let mut cols: Vec<u32> = Vec::new();
    let mut vals = Vec::new();

    for i in 0..n_rows {
        let a_start = a.indptr[i] as usize;
        let a_end = a.indptr[i + 1] as usize;
        let b_start = b.indptr[i] as usize;
        let b_end = b.indptr[i + 1] as usize;

        let mut a_idx = a_start;
        let mut b_idx = b_start;

        while a_idx < a_end || b_idx < b_end {
            if a_idx < a_end && (b_idx >= b_end || a.indices[a_idx] < b.indices[b_idx]) {
                rows.push(i as u32);
                cols.push(a.indices[a_idx]);
                vals.push(a.data[a_idx]);
                a_idx += 1;
            } else if b_idx < b_end && (a_idx >= a_end || b.indices[b_idx] < a.indices[a_idx]) {
                rows.push(i as u32);
                cols.push(b.indices[b_idx]);
                vals.push(b.data[b_idx]);
                b_idx += 1;
            } else {
                let val = a.data[a_idx] + b.data[b_idx];
                if val.into().abs() > EPSILON as f64 {
                    rows.push(i as u32);
                    cols.push(a.indices[a_idx]);
                    vals.push(val);
                }
                a_idx += 1;
                b_idx += 1;
            }
        }
    }

    Ok(coo_to_csr_presorted(&rows, &cols, &vals, a.shape))
}

/// Scalar multiplication of CSR matrix
///
/// ### Params
///
/// * `a` - Reference to the first CompressedSparseData2 (in CSR format!)
/// * `scalar` - The scalar value to multiply with
///
/// ### Returns
///
/// `CompressedSparseData2` with the data multiplied by the scalar.
pub fn sparse_scalar_multiply_csr<T>(
    a: &CompressedSparseData2<T>,
    scalar: T,
) -> CompressedSparseData2<T>
where
    T: BixverseNumeric,
    <T as Mul>::Output: Send,
    Vec<T>: FromParallelIterator<<T as Mul>::Output>,
{
    let data: Vec<T> = a.data.par_iter().map(|&v| v * scalar).collect();
    CompressedSparseData2::new_csr(&data, &a.indices, &a.indptr, None, a.shape)
}

/// Sparse matrix subtraction
///
/// ### Params
///
/// * `a` - Reference to the first CompressedSparseData2 (in CSR format!)
/// * `b` - Reference to the second CompressedSparseData2 (in CSR format!)
///
/// ### Returns
///
/// The subtracted new matrix
pub fn sparse_subtract_csr<T>(
    a: &CompressedSparseData2<T>,
    b: &CompressedSparseData2<T>,
) -> Result<CompressedSparseData2<T>, BixverseErrors>
where
    T: BixverseNumeric + Into<f64>,
{
    if !a.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if !b.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if a.shape != b.shape {
        return Err(BixverseErrors::ShapeMismatchSparse);
    }

    const EPSILON: f32 = 1e-9;
    let n_rows = a.shape.0;

    let mut rows: Vec<u32> = Vec::new();
    let mut cols: Vec<u32> = Vec::new();
    let mut vals = Vec::new();

    for i in 0..n_rows {
        let a_start = a.indptr[i] as usize;
        let a_end = a.indptr[i + 1] as usize;
        let b_start = b.indptr[i] as usize;
        let b_end = b.indptr[i + 1] as usize;

        let mut a_idx = a_start;
        let mut b_idx = b_start;

        while a_idx < a_end || b_idx < b_end {
            if a_idx < a_end && (b_idx >= b_end || a.indices[a_idx] < b.indices[b_idx]) {
                rows.push(i as u32);
                cols.push(a.indices[a_idx]);
                vals.push(a.data[a_idx]);
                a_idx += 1;
            } else if b_idx < b_end && (a_idx >= a_end || b.indices[b_idx] < a.indices[a_idx]) {
                rows.push(i as u32);
                cols.push(b.indices[b_idx]);
                vals.push(T::default() - b.data[b_idx]);
                b_idx += 1;
            } else {
                let val = a.data[a_idx] - b.data[b_idx];
                if val.into().abs() > EPSILON as f64 {
                    rows.push(i as u32);
                    cols.push(a.indices[a_idx]);
                    vals.push(val);
                }
                a_idx += 1;
                b_idx += 1;
            }
        }
    }

    Ok(coo_to_csr_presorted(&rows, &cols, &vals, a.shape))
}

/// Element-wise sparse multiplication
///
/// ### Params
///
/// * `a` - Reference to the first CompressedSparseData2 (in CSR format!)
/// * `b` - Reference to the second CompressedSparseData2 (in CSR format!)
///
/// ### Returns
///
/// The multiplied matrix.
pub fn sparse_multiply_elementwise_csr<T>(
    a: &CompressedSparseData2<T>,
    b: &CompressedSparseData2<T>,
) -> Result<CompressedSparseData2<T>, BixverseErrors>
where
    T: BixverseNumeric,
    <T as std::ops::Add>::Output: std::cmp::PartialEq<T>,
{
    if !a.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if !b.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if a.shape != b.shape {
        return Err(BixverseErrors::ShapeMismatchSparse);
    }

    let n_rows = a.shape.0;
    let mut rows: Vec<u32> = Vec::new();
    let mut cols: Vec<u32> = Vec::new();
    let mut vals = Vec::new();
    for i in 0..n_rows {
        let a_start = a.indptr[i] as usize;
        let a_end = a.indptr[i + 1] as usize;
        let b_start = b.indptr[i] as usize;
        let b_end = b.indptr[i + 1] as usize;
        let mut a_idx = a_start;
        let mut b_idx = b_start;
        while a_idx < a_end && b_idx < b_end {
            match a.indices[a_idx].cmp(&b.indices[b_idx]) {
                std::cmp::Ordering::Less => {
                    a_idx += 1;
                }
                std::cmp::Ordering::Greater => {
                    b_idx += 1;
                }
                std::cmp::Ordering::Equal => {
                    // Same column - multiply
                    let val = a.data[a_idx] * b.data[b_idx];
                    if val != T::default() {
                        rows.push(i as u32);
                        cols.push(a.indices[a_idx]);
                        vals.push(val);
                    }
                    a_idx += 1;
                    b_idx += 1;
                }
            }
        }
    }

    Ok(coo_to_csr(&rows, &cols, &vals, a.shape))
}

/// Normalises the columns of a CSR matrix to a sum of 1 (L1 norm)
///
/// ### Params
///
/// * `csr` - Mutable reference to the CSR matrix (modified in-place)
pub fn normalise_csr_columns_l1<T>(csr: &mut CompressedSparseData2<T>)
where
    T: BixverseNumeric + Into<f64>,
    <T as std::ops::Add>::Output: std::cmp::PartialEq<T>,
{
    assert!(csr.cs_type.is_csr(), "Matrix must be in CSR format");

    let ncols = csr.shape.1;

    let mut col_sums = vec![T::default(); ncols];

    for (idx, &col) in csr.indices.iter().enumerate() {
        col_sums[col as usize] += csr.data[idx]
    }

    for (idx, &col) in csr.indices.iter().enumerate() {
        let sum = col_sums[col as usize];
        if sum.into() > 1e-15 {
            csr.data[idx] /= sum;
        }
    }
}

/// Normalises the rows of a CSR matrix to a sum of 1 (L1 norm)
///
/// ### Params
///
/// * `csr` - Mutable reference to the CSR matrix (modified in-place)
pub fn normalise_csr_rows_l1<T>(csr: &mut CompressedSparseData2<T>) -> Result<(), BixverseErrors>
where
    T: BixverseNumeric + Into<f64>,
    T: std::iter::Sum<T>,
{
    if !csr.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }

    let nrows = csr.shape.0;

    for i in 0..nrows {
        let start = csr.indptr[i] as usize;
        let end = csr.indptr[i + 1] as usize;
        let row_data_slice = &mut csr.data[start..end];

        let row_sum: T = row_data_slice.iter().copied().sum();

        if row_sum.into() > 1e-15 {
            for val in row_data_slice.iter_mut() {
                *val /= row_sum;
            }
        } else {
            panic!(
                "Row {} has sum {}, indicating isolated node",
                i,
                row_sum.into()
            );
        }
    }

    Ok(())
}

/// Compute Frobenius norm of sparse matrix
///
/// ### Params
///
/// * `mat` - Sparse matrix in CSR or CSC format
///
/// ### Returns
///
/// Frobenius norm ||A||_F = sqrt(sum(A_ij^2))
pub fn frobenius_norm<T>(mat: &CompressedSparseData2<T>) -> f32
where
    T: BixverseNumeric + Into<f32>,
{
    mat.data
        .par_iter()
        .with_min_len(10000)
        .map(|&v| {
            let val: f32 = v.into();
            val * val
        })
        .sum::<f32>()
        .sqrt()
}

/// Remove zeros from sparse matrix
///
/// ### Params
///
/// * `mat` - Matrix from which to remove the zeroes
///
/// ### Returns
///
/// The Matrix with 0's removed.
pub fn eliminate_zeros_csr<T>(mat: CompressedSparseData2<T>) -> CompressedSparseData2<T>
where
    T: BixverseNumeric,
    <T as std::ops::Add>::Output: std::cmp::PartialEq<T>,
{
    let mut rows: Vec<u32> = Vec::new();
    let mut cols: Vec<u32> = Vec::new();
    let mut vals = Vec::new();

    let n_rows = mat.shape.0;
    for i in 0..n_rows {
        let start = mat.indptr[i] as usize;
        let end = mat.indptr[i + 1] as usize;

        for j in start..end {
            if mat.data[j] != T::default() {
                rows.push(i as u32);
                cols.push(mat.indices[j]);
                vals.push(mat.data[j]);
            }
        }
    }

    coo_to_csr(&rows, &cols, &vals, mat.shape)
}

/// Sparse matrix @ dense vector: `M @ v`.
///
/// Computes `result[i] = sum_j M[i, j] * v[j]` by iterating CSR rows.
///
/// ### Params
///
/// * `mat` - CSR matrix of shape `(m, n)`.
/// * `vec` - Dense vector of length `n`.
///
/// ### Returns
///
/// Dense vector of length `m`. Errors if `mat` is not CSR.
pub fn csr_matvec<T>(mat: &CompressedSparseData2<T>, vec: &[T]) -> Result<Vec<T>, BixverseErrors>
where
    T: BixverseNumeric,
    <T as std::ops::Add>::Output: std::cmp::PartialEq<T>,
{
    if !mat.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }

    let mut result = vec![T::default(); mat.shape.0];
    for i in 0..mat.shape.0 {
        let row_start = mat.indptr[i] as usize;
        let row_end = mat.indptr[i + 1] as usize;
        let mut sum = T::default();
        for idx in row_start..row_end {
            sum += mat.data[idx] * vec[mat.indices[idx] as usize];
        }
        result[i] = sum;
    }

    Ok(result)
}

/// Dense vector @ sparse matrix: `v @ M`.
///
/// Computes `result[j] = sum_i v[i] * M[i, j]` by iterating CSR rows of `M`
/// and scattering scaled row contributions into the output. Treats `v` as a
/// row vector.
///
/// ### Params
///
/// * `vec` - Dense vector of length `m`.
/// * `mat` - CSR matrix of shape `(m, n)`.
///
/// ### Returns
///
/// Dense vector of length `n`. Errors if `mat` is not CSR.
pub fn csr_vecmat<T>(vec: &[T], mat: &CompressedSparseData2<T>) -> Result<Vec<T>, BixverseErrors>
where
    T: BixverseNumeric,
{
    if !mat.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }

    let n_cols = mat.ncols();
    let mut out = vec![T::default(); n_cols];
    for (j, &vj) in vec.iter().enumerate() {
        if vj == T::default() {
            continue;
        }
        let gs = mat.indptr[j] as usize;
        let ge = mat.indptr[j + 1] as usize;
        for q in gs..ge {
            let t = mat.indices[q] as usize;
            out[t] += vj * mat.data[q];
        }
    }

    Ok(out)
}

/// Sparse accumulator for efficient sparse matrix multiplication
///
/// ### Fields
///
/// * `values` - Vector storing accumulated values for each index
/// * `indices` - Vector of active (non-zero) indices
/// * `flags` - Boolean flags indicating which indices are active
struct SparseAccumulator<T>
where
    T: Copy + Default + AddAssign,
{
    values: Vec<T>,
    indices: Vec<usize>,
    flags: Vec<bool>,
}

impl<T> SparseAccumulator<T>
where
    T: Copy + Default + AddAssign,
{
    /// Create a new sparse accumulator
    ///
    /// ### Params
    ///
    /// * `size` - Maximum number of indices to accumulate
    fn new(size: usize) -> Self {
        Self {
            values: vec![T::default(); size],
            indices: Vec::with_capacity(size / 10),
            flags: vec![false; size],
        }
    }

    /// Add a value to the accumulator at the given index
    ///
    /// ### Params
    ///
    /// * `idx` - Index to accumulate at
    /// * `val` - Value to add
    ///
    /// ### Safety
    ///
    /// `idx` must be less than the size specified during construction
    #[inline]
    unsafe fn add(&mut self, idx: usize, val: T) {
        unsafe {
            if !*self.flags.get_unchecked(idx) {
                *self.flags.get_unchecked_mut(idx) = true;
                self.indices.push(idx);
                *self.values.get_unchecked_mut(idx) = val;
            } else {
                *self.values.get_unchecked_mut(idx) += val;
            }
        }
    }

    /// Drain accumulated values into caller-provided buffers and reset the
    /// accumulator.
    ///
    /// ### Params
    ///
    /// * `indices_out` - Buffer to append active indices into
    /// * `data_out` - Buffer to append accumulated values into
    ///
    /// ### Returns
    ///
    /// Number of entries written
    #[inline]
    fn extract_into(&mut self, indices_out: &mut Vec<usize>, data_out: &mut Vec<T>) -> usize {
        self.indices.sort_unstable();
        let n = self.indices.len();
        unsafe {
            for &i in &self.indices {
                indices_out.push(i);
                data_out.push(*self.values.get_unchecked(i));
                *self.flags.get_unchecked_mut(i) = false;
                *self.values.get_unchecked_mut(i) = T::default();
            }
        }
        self.indices.clear();
        n
    }
}

/// Multiply two CSR matrices using sparse accumulators and parallel processing
///
/// ### Params
///
/// * `a` - Left CSR matrix
/// * `b` - Right CSR matrix
///
/// ### Returns
///
/// Product matrix in CSR format
pub fn csr_matmul_csr<T>(
    a: &CompressedSparseData2<T>,
    b: &CompressedSparseData2<T>,
) -> Result<CompressedSparseData2<T>, BixverseErrors>
where
    T: BixverseNumeric,
{
    let ncol_a = a.shape().1;
    let nrow_b = b.shape().0;

    if !a.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if !b.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if ncol_a != nrow_b {
        return Err(BixverseErrors::SparseMatrixMultiplication {
            n_col_a: ncol_a,
            n_row_b: nrow_b,
        });
    }

    let nrows = a.shape.0;
    let ncols = b.shape.1;

    const CHUNK: usize = 256;
    let n_chunks = nrows.div_ceil(CHUNK);

    let chunks: Vec<(Vec<usize>, Vec<T>, Vec<usize>)> = (0..n_chunks)
        .into_par_iter()
        .map_init(
            || SparseAccumulator::new(ncols),
            |acc, chunk_idx| {
                let start = chunk_idx * CHUNK;
                let end = ((chunk_idx + 1) * CHUNK).min(nrows);

                let mut idx_buf = Vec::new();
                let mut data_buf = Vec::new();
                let mut lengths = Vec::with_capacity(end - start);

                for i in start..end {
                    for a_idx in a.indptr[i] as usize..a.indptr[i + 1] as usize {
                        let k = a.indices[a_idx] as usize;
                        let a_val = a.data[a_idx];
                        for b_idx in b.indptr[k] as usize..b.indptr[k + 1] as usize {
                            unsafe {
                                acc.add(b.indices[b_idx] as usize, a_val * b.data[b_idx]);
                            }
                        }
                    }
                    lengths.push(acc.extract_into(&mut idx_buf, &mut data_buf));
                }

                (idx_buf, data_buf, lengths)
            },
        )
        .collect();

    let total_nnz: usize = chunks.iter().map(|(i, _, _)| i.len()).sum();

    let mut indptr: Vec<u32> = Vec::with_capacity(nrows + 1);
    indptr.push(0);
    let mut running = 0usize;
    for (_, _, lengths) in &chunks {
        for &len in lengths {
            running += len;
            indptr.push(running as u32);
        }
    }

    let mut indices: Vec<u32> = Vec::with_capacity(total_nnz);
    let mut data = Vec::with_capacity(total_nnz);
    for (idx_buf, data_buf, _) in chunks {
        indices.extend(idx_buf.into_iter().map(|x| x as u32));
        data.extend(data_buf);
    }

    // build directly rather than via new_csr, which would .to_vec() the lot.
    Ok(CompressedSparseData2 {
        data,
        indices,
        indptr,
        cs_type: CompressedSparseFormat::Csr,
        data_2: None,
        shape: (nrows, ncols),
    })
}

/////////////////////////////
// Sparse dense operations //
/////////////////////////////

/// Sparse CSR @ sparse CSR -> dense
///
/// Parallel over output rows. For cases where you assume that the resulting
/// matrix will become dense.
///
/// ### Params
///
/// * `a` - First matrix
/// * `b` - Second matrix
///
/// ### Returns
///
/// Dense matrix product of both CSR matrices
pub fn csr_sparse_matmul_dense<T>(
    a: &CompressedSparseData2<T>,
    b: &CompressedSparseData2<T>,
) -> Result<Mat<T>, BixverseErrors>
where
    T: BixverseFloat + Send + Sync + Default,
{
    let ncol_a = a.shape().1;
    let nrow_b = b.shape().0;

    if !a.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if !b.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsr);
    }
    if ncol_a != nrow_b {
        return Err(BixverseErrors::SparseMatrixMultiplication {
            n_col_a: ncol_a,
            n_row_b: nrow_b,
        });
    }

    let n_rows = a.nrows();
    let n_cols = b.ncols();
    let m_indptr = &a.indptr;
    let m_indices = &a.indices;
    let m_data = &a.data;
    let g_indptr = &b.indptr;
    let g_indices = &b.indices;
    let g_data = &b.data;

    let dense_rows: Vec<Vec<T>> = (0..n_rows)
        .into_par_iter()
        .map(|i| {
            let mut row = vec![T::zero(); n_cols];
            let rs = m_indptr[i] as usize;
            let re = m_indptr[i + 1] as usize;
            for p in rs..re {
                let j = m_indices[p] as usize;
                let w = m_data[p];
                let gs = g_indptr[j] as usize;
                let ge = g_indptr[j + 1] as usize;
                for q in gs..ge {
                    let t = g_indices[q] as usize;
                    row[t] += w * g_data[q];
                }
            }
            row
        })
        .collect();

    // this part can be sequential...
    let mut out = Mat::zeros(n_rows, n_cols);
    for (i, row) in dense_rows.into_iter().enumerate() {
        for (j, v) in row.into_iter().enumerate() {
            out[(i, j)] = v;
        }
    }

    Ok(out)
}

///////////////////////
// Sparse statistics //
///////////////////////

/// Calculate the column means for CSC [CompressedSparseData2]
///
/// ### Params
///
/// * `csc` - The [CompressedSparseData2] (needs to have floats)
/// * `use_second_layer` - Use the second data layer
///
/// ### Returns
///
/// The column means
pub fn sparse_col_means_csc<T>(
    csc: &CompressedSparseData2<T>,
    use_second_layer: bool,
) -> Result<Vec<T>, BixverseErrors>
where
    T: BixverseFloat + BixverseSimd,
{
    if csc.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsc);
    }

    let active_data: &[T] = if use_second_layer {
        csc.data_2
            .as_ref()
            .ok_or(BixverseErrors::Data2NotAvailable)?
            .as_slice()
    } else {
        csc.data.as_slice()
    };

    let (nrows, ncols) = csc.shape();
    let nrows_t = T::from_usize(nrows).unwrap();
    let mut col_means: Vec<T> = Vec::with_capacity(ncols);

    for j in 0..ncols {
        let start = csc.indptr[j] as usize;
        let end = csc.indptr[j + 1] as usize;
        let sum = T::bxv_sum(&active_data[start..end]);
        col_means.push(sum / nrows_t);
    }

    Ok(col_means)
}

/// Calculate the column standard deviations for CSC [CompressedSparseData2]
///
/// ### Params
///
/// * `csc` - The [CompressedSparseData2] (needs to have floats)
/// * `use_second_layer` - Use the second data layer
///
/// ### Returns
///
/// The column standard deviations
pub fn sparse_col_sds_csc<T>(
    csc: &CompressedSparseData2<T>,
    use_second_layer: bool,
) -> Result<Vec<T>, BixverseErrors>
where
    T: BixverseFloat + BixverseSimd,
{
    if csc.cs_type.is_csr() {
        return Err(BixverseErrors::SparseMatrixMustBeCsc);
    }

    let active_data: &[T] = if use_second_layer {
        csc.data_2
            .as_ref()
            .ok_or(BixverseErrors::Data2NotAvailable)?
            .as_slice()
    } else {
        csc.data.as_slice()
    };

    let (nrows, ncols) = csc.shape();
    let nrows_t = T::from_usize(nrows).unwrap();
    let denom = nrows_t - T::one();
    let mut col_sds: Vec<T> = Vec::with_capacity(ncols);

    for j in 0..ncols {
        let start = csc.indptr[j] as usize;
        let end = csc.indptr[j + 1] as usize;
        let col_slice = &active_data[start..end];
        let implicit_zeros = T::from_usize(nrows - (end - start)).unwrap();

        let mean = T::bxv_sum(col_slice) / nrows_t;
        let ssd_nonzero = T::bxv_sum_squared_deviation(col_slice, mean);
        let ssd_total = ssd_nonzero + implicit_zeros * mean * mean;
        col_sds.push((ssd_total / denom).sqrt());
    }

    Ok(col_sds)
}

////////////////////////
// Lanczos Eigenvalue //
////////////////////////

/// Helper function for dot product of two vectors
///
/// ### Params
///
/// * `a` - Vector a
/// * `b` - Vector b
///
/// ### Returns
///
/// Dot product of the two vectors
fn dot<T>(a: &[T], b: &[T]) -> T
where
    T: BixverseSimd,
{
    assert_same_len!(a, b);
    T::bxv_dot_simd(a, b)
}

/// Helper function to normalise a vector
///
/// ### Params
///
/// * `v` - Initial vector
///
/// ### Returns
///
/// Normalised dot product of the vector `v`
fn norm<T>(v: &[T]) -> T
where
    T: BixverseSimd + BixverseFloat,
{
    let dot = dot(v, v);
    dot.sqrt()
}

/// Helper function to normalise a vector
///
/// ### Params
///
/// * `v` - Mutable reference of the vector to normalise
fn normalise<T>(v: &mut [T])
where
    T: BixverseSimd + BixverseFloat,
{
    let n = norm(v);
    v.par_iter_mut().for_each(|x| *x /= n);
}

/// Helper function to calculate eigenvalues
///
/// ### Params
///
/// * `alpha` - alpha vector
/// * `beta` - beta vector
///
/// ### Returns
///
/// Tuple of `(eigenvectors, eigenvalues)`
fn tridiag_eig<T>(alpha: &[T], beta: &[T]) -> Result<(Vec<T>, Mat<T>), BixverseErrors>
where
    T: BixverseFloat,
{
    let n = alpha.len();
    let mut t = Mat::<T>::zeros(n, n);

    for i in 0..n {
        t[(i, i)] = alpha[i];
        if i < n - 1 {
            t[(i, i + 1)] = beta[i];
            t[(i + 1, i)] = beta[i];
        }
    }

    let eig = t
        .self_adjoint_eigen(faer::Side::Lower)
        .map_err(|_| BixverseErrors::FaerEigenError)?;
    let evals = eig.S().column_vector().iter().copied().collect();
    let evecs = eig.U().to_owned();

    Ok((evals, evecs))
}

/// Compute largest eigenvalues and eigenvectors using Lanczos
///
/// ### Params
///
/// * `matrix` - Sparse matrix in CSR format
/// * `n_components` - Number of eigenpairs to compute
/// * `seed` - For reproducibility
///
/// ### Returns
///
/// (eigenvalues, eigenvectors) where `eigenvectors[i][j]` is element j of
/// eigenvector i
pub fn compute_largest_eigenpairs_lanczos<T>(
    matrix: &CompressedSparseData2<T>,
    n_components: usize,
    seed: u64,
) -> Result<(Vec<f32>, Vec<Vec<f32>>), BixverseErrors>
where
    T: BixverseNumeric + BixverseSimd + Into<f64>,
{
    let n = matrix.shape.0;
    let n_iter = (n_components * 2 + 10).max(n_components).min(n);

    // Convert to CSR for efficient row access
    let csr = match matrix.cs_type {
        CompressedSparseFormat::Csr => matrix.clone(),
        CompressedSparseFormat::Csc => matrix.transform(),
    };

    let data_f64: Vec<f64> = csr.data.iter().map(|&v| v.into()).collect();

    // Parallelised matvec: y = A * x
    let matvec = |x: &[f64], y: &mut [f64]| {
        y.par_iter_mut().enumerate().for_each(|(i, yi)| {
            let mut sum = 0.0;
            for idx in csr.indptr[i] as usize..csr.indptr[i + 1] as usize {
                let j = csr.indices[idx] as usize;
                sum += data_f64[idx] * x[j];
            }
            *yi = sum;
        });
    };

    // Lanczos iteration
    let mut v = vec![0.0; n];
    let mut v_old = vec![0.0; n];
    let mut w = vec![0.0; n];
    let mut v_matrix = vec![vec![0.0; n]; n_iter];

    let mut rng = StdRng::seed_from_u64(seed);

    for i in 0..n {
        v[i] = rng.random::<f64>() - 0.5;
    }
    normalise(&mut v);

    let mut alpha = vec![0.0; n_iter];
    let mut beta = vec![0.0; n_iter];

    for j in 0..n_iter {
        v_matrix[j].copy_from_slice(&v);

        matvec(&v, &mut w);
        alpha[j] = dot(&w, &v);

        // w = w - alpha[j]*v - beta[j-1]*v_old
        for i in 0..n {
            w[i] -= alpha[j] * v[i];
            if j > 0 {
                w[i] -= beta[j - 1] * v_old[i];
            }
        }

        // full reorthogonalisation
        for k in 0..=j {
            let coeff = dot(&w, &v_matrix[k]);
            for i in 0..n {
                w[i] -= coeff * v_matrix[k][i];
            }
        }

        beta[j] = norm(&w);
        if beta[j] < 1e-12 {
            break;
        }

        v_old.copy_from_slice(&v);
        v.copy_from_slice(&w);
        normalise(&mut v);
    }

    let (evals, evecs) = tridiag_eig(&alpha[..n_iter], &beta[..n_iter - 1])?;

    let mut indices: Vec<usize> = (0..evals.len()).collect();
    indices.sort_by(|&i, &j| evals[j].partial_cmp(&evals[i]).unwrap());

    let mut largest_evals: Vec<f32> = Vec::with_capacity(n_components);
    let mut largest_evecs: Vec<Vec<f32>> = Vec::with_capacity(n_components);

    for &idx in indices.iter().take(n_components) {
        // Transform eigenvector back to original space: v_original = V * v_tridiag
        let mut evec = vec![0.0; n];
        for i in 0..n {
            for j in 0..n_iter {
                evec[i] += v_matrix[j][i] * evecs[(j, idx)].to_f64().unwrap();
            }
        }

        // Normalise the transformed eigenvector
        let norm: f64 = evec.iter().map(|x| x * x).sum::<f64>().sqrt();
        for x in &mut evec {
            *x /= norm;
        }

        largest_evals.push(evals[idx].to_f64().unwrap() as f32);
        largest_evecs.push(evec.iter().map(|&x| x as f32).collect());
    }

    let mut transposed = vec![vec![0.0f32; n_components]; n];
    for comp_idx in 0..n_components {
        for point_idx in 0..n {
            transposed[point_idx][comp_idx] = largest_evecs[comp_idx][point_idx];
        }
    }

    Ok((largest_evals, transposed))
}

/////////////////
// Lanczos SVD //
/////////////////

/// Compute sparse SVD using Lanczos on A^T A or AA^T
///
/// ### Params
///
/// * `matrix` - Sparse matrix (CSR or CSC)
/// * `n_components` - Number of singular values/vectors to compute
/// * `seed` - For reproducibility
/// * `use_second_layer` - If true, use data_2 instead of data
/// * `col_means` - Optional column means for implicit mean centering
/// * `col_stds` - Optional column sds for implicit variance normalising
/// * `row_offsets` - Additional offsets (for example for CLR-type PCA in single
///   cell).
///
/// ### Returns
///
/// `SvdResults` containing U (n×k), S (length k), and V (m×k)
pub fn sparse_svd_lanczos<T, U, F>(
    matrix: &CompressedSparseData2<T, U>,
    n_components: usize,
    seed: u64,
    use_second_layer: bool,
    col_means: Option<&[F]>,
    col_stds: Option<&[F]>,
    row_offsets: Option<&[F]>,
) -> Result<SvdResults<F>, BixverseErrors>
where
    T: BixverseNumeric + BixverseSimd + Into<F> + Clone,
    U: BixverseNumeric + Into<F> + Clone,
    F: BixverseFloat + BixverseSimd + std::iter::Sum,
{
    let (n, m) = matrix.shape;
    let use_ata = n > m;
    let krylov_dim = if use_ata { m } else { n };
    let n_iter = (n_components * 2 + 10).max(n_components).min(krylov_dim);

    // keep both representations to make matvec operations fast
    let (csr, csc);
    let csr_owned;
    let csc_owned;

    match matrix.cs_type {
        CompressedSparseFormat::Csr => {
            csr = matrix;
            csc_owned = matrix.transform_single_layer(use_second_layer)?;
            csc = &csc_owned;
        }
        CompressedSparseFormat::Csc => {
            csc = matrix;
            csr_owned = matrix.transform_single_layer(use_second_layer)?;
            csr = &csr_owned;
        }
    };

    // helper to extract the right data layer and cast to F
    let extract_data = |mat: &CompressedSparseData2<T, U>| -> Vec<F> {
        if use_second_layer {
            mat.data_2
                .as_ref()
                .expect("data_2 is None but use_second_layer is true")
                .iter()
                .copied()
                .map(|v| v.into())
                .collect()
        } else {
            mat.data.iter().copied().map(|v| v.into()).collect()
        }
    };

    let data_csr_f = extract_data(csr);
    let data_csc_f = extract_data(csc);

    // matrix-vector product for A (using CSR)
    let matvec_a = |x: &[F], y: &mut [F]| {
        let x_scaled: Vec<F> = if let Some(sd) = col_stds {
            x.iter().enumerate().map(|(j, &v)| v / sd[j]).collect()
        } else {
            x.to_vec()
        };
        let mean_dot: F = if let Some(mu) = col_means {
            x_scaled.iter().enumerate().map(|(j, &v)| mu[j] * v).sum()
        } else {
            F::zero()
        };
        let x_sum: F = if row_offsets.is_some() {
            x_scaled.iter().copied().sum()
        } else {
            F::zero()
        };

        y.par_iter_mut().enumerate().for_each(|(i, yi)| {
            let mut sum = F::zero();
            for idx in csr.indptr[i] as usize..csr.indptr[i + 1] as usize {
                let j = csr.indices[idx] as usize;
                sum += data_csr_f[idx] * x_scaled[j];
            }
            if col_means.is_some() {
                sum -= mean_dot;
            }
            if let Some(off) = row_offsets {
                sum -= off[i] * x_sum;
            }
            *yi = sum;
        });
    };

    let matvec_at = |x: &[F], y: &mut [F]| {
        let x_sum: F = x.iter().copied().sum();
        let m_dot_x: F = if let Some(off) = row_offsets {
            off.iter().zip(x.iter()).map(|(&m_i, &xi)| m_i * xi).sum()
        } else {
            F::zero()
        };

        y.par_iter_mut().enumerate().for_each(|(j, yj)| {
            let mut sum = F::zero();
            for idx in csc.indptr[j] as usize..csc.indptr[j + 1] as usize {
                let i = csc.indices[idx] as usize;
                sum += data_csc_f[idx] * x[i];
            }
            if let Some(mu) = col_means {
                sum -= mu[j] * x_sum;
            }
            if row_offsets.is_some() {
                sum -= m_dot_x;
            }
            if let Some(sd) = col_stds {
                sum /= sd[j];
            }
            *yj = sum;
        });
    };

    // select Gram matrix operator
    #[allow(clippy::type_complexity)]
    let matvec_gram: Box<dyn Fn(&[F], &mut [F]) + Sync> = if use_ata {
        Box::new(|x: &[F], y: &mut [F]| {
            let mut temp = vec![F::zero(); n];
            matvec_a(x, &mut temp);
            matvec_at(&temp, y);
        })
    } else {
        Box::new(|x: &[F], y: &mut [F]| {
            let mut temp = vec![F::zero(); m];
            matvec_at(x, &mut temp);
            matvec_a(&temp, y);
        })
    };

    // lanczos iteration
    let mut v_matrix = Mat::<F>::zeros(krylov_dim, n_iter);
    let mut v = vec![F::zero(); krylov_dim];
    let mut v_old = vec![F::zero(); krylov_dim];
    let mut w = vec![F::zero(); krylov_dim];
    let mut w_faer = faer::Col::<F>::zeros(krylov_dim);

    let mut rng = StdRng::seed_from_u64(seed);
    for i in 0..krylov_dim {
        v[i] = F::from(rng.random::<f64>() - 0.5).unwrap();
    }
    normalise(&mut v);

    let mut alpha = vec![F::zero(); n_iter];
    let mut beta = vec![F::zero(); n_iter];

    for j in 0..n_iter {
        for i in 0..krylov_dim {
            v_matrix[(i, j)] = v[i];
        }

        matvec_gram(&v, &mut w);
        alpha[j] = dot(&w, &v);

        for i in 0..krylov_dim {
            w[i] -= alpha[j] * v[i];
            if j > 0 {
                w[i] -= beta[j - 1] * v_old[i];
            }
        }

        // Gram-Schmidt / Orthogonalisation
        // w -= Vj * (Vj^T * w)
        for i in 0..krylov_dim {
            w_faer[i] = w[i];
        }
        let vj = v_matrix.as_ref().subcols(0, j + 1);
        let coeffs = vj.transpose() * w_faer.as_ref();
        let proj = vj * coeffs.as_ref();
        for i in 0..krylov_dim {
            w[i] -= proj[i];
        }

        beta[j] = norm(&w);
        if beta[j] < F::from(1e-12).unwrap() {
            break;
        }

        v_old.copy_from_slice(&v);
        v.copy_from_slice(&w);
        normalise(&mut v);
    }

    // eigendecomposition and reconstruction
    let (evals, evecs) = tridiag_eig(&alpha[..n_iter], &beta[..n_iter - 1])?;

    let mut indices: Vec<usize> = (0..evals.len()).collect();
    indices.sort_by(|&i, &j| evals[j].partial_cmp(&evals[i]).unwrap());

    let mut singular_values: Vec<F> = Vec::with_capacity(n_components);
    let mut u_vecs: Vec<Vec<F>> = Vec::with_capacity(n_components);
    let mut v_vecs: Vec<Vec<F>> = Vec::with_capacity(n_components);

    for &idx in indices.iter().take(n_components) {
        let eval = evals[idx];
        if eval <= F::zero() {
            continue;
        }

        let sigma = eval.sqrt();
        singular_values.push(sigma);

        // gram_evec = v_matrix * evecs.col(idx) via faer operator overloading
        let gram_col = &v_matrix * evecs.col(idx);
        let norm_val: F = gram_col.iter().map(|x| *x * *x).sum::<F>().sqrt();
        let gram_evec: Vec<F> = gram_col.iter().map(|x| *x / norm_val).collect();

        if use_ata {
            let mut u_vec = vec![F::zero(); n];
            matvec_a(&gram_evec, &mut u_vec);
            for x in &mut u_vec {
                *x /= sigma;
            }
            u_vecs.push(u_vec);
            v_vecs.push(gram_evec);
        } else {
            let mut v_vec = vec![F::zero(); m];
            matvec_at(&gram_evec, &mut v_vec);
            for x in &mut v_vec {
                *x /= sigma;
            }
            u_vecs.push(gram_evec);
            v_vecs.push(v_vec);
        }
    }

    let u = Mat::from_fn(n, singular_values.len(), |i, j| u_vecs[j][i]);
    let v = Mat::from_fn(m, singular_values.len(), |i, j| v_vecs[j][i]);

    Ok(SvdResults {
        u,
        s: singular_values,
        v,
    })
}

///////////
// Tests //
///////////

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

    #[test]
    fn test_parse_sparse_format() {
        assert!(parse_compressed_sparse_format("csr").unwrap().is_csr());
        assert!(parse_compressed_sparse_format("CSC").unwrap().is_csc());
        assert!(parse_compressed_sparse_format("dense").is_none());
    }

    #[test]
    fn test_from_dense_and_count_zeroes() {
        let mat: Mat<f64> = Mat::from_fn(3, 2, |i, j| if i == j { (i + 1) as f64 } else { 0.0 });
        let (total_zeroes, row_zeroes, col_zeroes) = count_zeroes(&mat.as_ref());

        assert_eq!(total_zeroes, 4);
        assert_eq!(row_zeroes, vec![1, 1, 2]);
        assert_eq!(col_zeroes, vec![2, 2]);

        let csr = CompressedSparseData2::<f64, f64>::from_dense_matrix(
            mat.as_ref(),
            CompressedSparseFormat::Csr,
        );
        assert_eq!(csr.shape, (3, 2));
        assert_eq!(csr.data, vec![1.0, 2.0]);
    }

    #[test]
    fn test_sparse_add_csr() {
        let shape = (2, 2);
        let a = CompressedSparseData2::<f64, f64>::new_csr(
            &[1.0, 2.0],
            &[0, 1],
            &[0, 1, 2],
            None,
            shape,
        );
        let b = CompressedSparseData2::<f64, f64>::new_csr(
            &[3.0, 4.0],
            &[1, 1],
            &[0, 1, 2],
            None,
            shape,
        );
        let c = sparse_add_csr(&a, &b).unwrap();

        assert_eq!(c.data, vec![1.0, 3.0, 6.0]);
        assert_eq!(c.indices, vec![0, 1, 1]);
        assert_eq!(c.indptr, vec![0, 2, 3]);
    }

    #[test]
    fn test_csr_matvec() {
        let a = CompressedSparseData2::<f64, f64>::new_csr(
            &[1.0, 2.0, 3.0],
            &[0, 1, 1],
            &[0, 2, 3],
            None,
            (2, 2),
        );
        let vec = vec![2.0, 1.0];
        let result = csr_matvec(&a, &vec).unwrap();
        assert_eq!(result, vec![4.0, 3.0]);
    }

    #[test]
    fn test_lanczos_eigenpairs_logic() {
        // Symmetric rank-1 matrix M = x * x^T
        // Let x = [1.0, 0.0, 2.0, 0.0]^T
        let data = vec![1.0, 2.0, 2.0, 4.0];
        let indices = vec![0, 2, 0, 2];
        let indptr = vec![0, 2, 2, 4, 4]; // Rows 1 and 3 are empty
        let shape = (4, 4);

        let csr = CompressedSparseData2::<f64, f64>::new_csr(&data, &indices, &indptr, None, shape);

        // Lanczos expects symmetric matrix, this one is symmetric
        let (evals, evecs) = compute_largest_eigenpairs_lanczos(&csr, 1, 42).unwrap();

        // True top eigenvalue should be exactly sum(x_i^2) = 1.0 + 4.0 = 5.0
        assert!((evals[0] - 5.0).abs() < 1e-3);

        // Eigenvectors are returned transposed: evecs[point_idx][comp_idx]
        // So the first principal component is [evecs[0][0], evecs[1][0], evecs[2][0], evecs[3][0]]
        let x_norm = 5.0_f32.sqrt();
        let dot_x = (evecs[0][0] * 1.0 + evecs[2][0] * 2.0) / x_norm;

        assert!(dot_x.abs() > 0.999);
    }

    #[test]
    fn test_sparse_svd_lanczos_logic() {
        // Sparse rank-1 matrix A = x * y^T
        // x = [0.0, 2.0, 0.0, 4.0]^T
        // y = [1.0, 0.0, 0.5]^T
        let data = vec![2.0, 1.0, 4.0, 2.0];
        let indices = vec![0, 2, 0, 2];
        let indptr = vec![0, 0, 2, 2, 4];
        let shape = (4, 3);

        let csr = CompressedSparseData2::<f64, f64>::new_csr(&data, &indices, &indptr, None, shape);
        let no_params: Option<&[f64]> = None;

        let svd = sparse_svd_lanczos(&csr, 1, 42, false, no_params, no_params, None).unwrap();

        // Test correlation with theoretical U
        let u_col = svd.u.col(0);
        let x_norm = 20.0_f64.sqrt();
        let dot_u = (u_col[1] * 2.0 + u_col[3] * 4.0) / x_norm;
        assert!(dot_u.abs() > 0.999);

        // Test correlation with theoretical V
        let v_col = svd.v.col(0);
        let y_norm = 1.25_f64.sqrt();
        let dot_v = (v_col[0] * 1.0 + v_col[2] * 0.5) / y_norm;
        assert!(dot_v.abs() > 0.999);
    }
}