datarust/

matrix.rs

//! Dense and sparse matrix containers used throughout the crate.

use crate::error::{DatarustError, Result};

/// Row-major dense matrix of `f64` backed by a single contiguous `Vec<f64>`.
///
/// The flat layout (one allocation, stride-1 row traversal) keeps every numeric
/// hot loop cache-friendly and auto-vectorizable — a substantial speedup over
/// the previous `Vec<Vec<f64>>` representation on large dense inputs.
#[derive(Debug, Clone, PartialEq)]
pub struct Matrix {
    /// Flat row-major data, length `rows * cols`.
    data: Vec<f64>,
    rows: usize,
    cols: usize,
}

impl Matrix {
    /// Creates a matrix from a nested vector, validating a rectangular shape.
    pub fn new(data: Vec<Vec<f64>>) -> Result<Self> {
        if data.is_empty() {
            return Err(DatarustError::EmptyInput("matrix has no rows".into()));
        }
        let cols = data[0].len();
        if cols == 0 {
            return Err(DatarustError::EmptyInput("matrix has no columns".into()));
        }
        for (i, row) in data.iter().enumerate() {
            if row.len() != cols {
                return Err(DatarustError::ShapeMismatch {
                    expected: format!("{} columns", cols),
                    actual: format!("{} columns at row {}", row.len(), i),
                });
            }
        }
        let rows = data.len();
        let mut flat = Vec::with_capacity(rows * cols);
        for row in data {
            flat.extend(row);
        }
        Ok(Self {
            data: flat,
            rows,
            cols,
        })
    }

    /// Creates a matrix from a nested vector of rows.
    pub fn from_rows(rows: Vec<Vec<f64>>) -> Result<Self> {
        Self::new(rows)
    }

    /// Creates a matrix from row-major flat data of the given shape.
    ///
    /// With the flat internal representation this stores the buffer directly,
    /// with no per-row chunking or copy.
    pub fn from_flat(rows: usize, cols: usize, flat: Vec<f64>) -> Result<Self> {
        if rows == 0 || cols == 0 {
            return Err(DatarustError::EmptyInput("zero dimension".into()));
        }
        let expected = rows
            .checked_mul(cols)
            .ok_or_else(|| DatarustError::ShapeMismatch {
                expected: "rows * cols within usize range".into(),
                actual: format!("{} rows × {} cols overflows usize", rows, cols),
            })?;
        if flat.len() != expected {
            return Err(DatarustError::ShapeMismatch {
                expected: format!("{} elements", expected),
                actual: format!("{} elements", flat.len()),
            });
        }
        Ok(Self {
            data: flat,
            rows,
            cols,
        })
    }

    /// Creates a matrix filled with zeros of the given shape.
    pub fn zeros(rows: usize, cols: usize) -> Result<Self> {
        if rows == 0 || cols == 0 {
            return Err(DatarustError::EmptyInput("zero dimension".into()));
        }
        Ok(Self {
            data: vec![0.0; rows * cols],
            rows,
            cols,
        })
    }

    /// Creates an `n` by `n` identity matrix.
    pub fn identity(n: usize) -> Result<Self> {
        if n == 0 {
            return Err(DatarustError::EmptyInput("zero dimension".into()));
        }
        let mut data = vec![0.0; n * n];
        for i in 0..n {
            data[i * n + i] = 1.0;
        }
        Ok(Self {
            data,
            rows: n,
            cols: n,
        })
    }

    /// Returns the number of rows.
    #[inline]
    pub fn nrows(&self) -> usize {
        self.rows
    }

    /// Returns the number of columns.
    #[inline]
    pub fn ncols(&self) -> usize {
        self.cols
    }

    /// Returns the underlying flat row-major data as a slice.
    ///
    /// Elements are laid out so that element `(i, j)` is at index `i * ncols + j`.
    /// Use this for cache-friendly, auto-vectorizable numeric loops in preference
    /// to per-element [`get`](Self::get) calls.
    #[inline]
    pub fn as_slice(&self) -> &[f64] {
        &self.data
    }

    /// Returns the underlying flat row-major data as a mutable slice.
    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [f64] {
        &mut self.data
    }

    /// Returns the element at row `i`, column `j`.
    ///
    /// # Panics
    ///
    /// Panics if `i >= nrows()` or `j >= ncols()`. Use [`checked_get`](Self::checked_get)
    /// for a safe alternative that returns `None` on out-of-bounds access.
    #[inline]
    pub fn get(&self, i: usize, j: usize) -> f64 {
        debug_assert!(
            i < self.rows && j < self.cols,
            "Matrix::get: index ({}, {}) out of bounds for {}×{}",
            i,
            j,
            self.rows,
            self.cols
        );
        // SAFETY: debug_assert above + the struct invariant (data.len() == rows*cols).
        unsafe { *self.data.get_unchecked(i * self.cols + j) }
    }

    /// Returns `Some(element)` at row `i`, column `j`, or `None` if the indices
    /// are out of bounds.
    #[inline]
    pub fn checked_get(&self, i: usize, j: usize) -> Option<f64> {
        if i < self.rows && j < self.cols {
            Some(self.data[i * self.cols + j])
        } else {
            None
        }
    }

    /// Sets the element at row `i`, column `j`.
    #[inline]
    pub fn set(&mut self, i: usize, j: usize, v: f64) {
        self.data[i * self.cols + j] = v;
    }

    /// Returns `Ok(())` if the matrix contains no NaN values.
    pub fn validate_no_nan(&self) -> Result<()> {
        for (flat_idx, &v) in self.data.iter().enumerate() {
            if v.is_nan() {
                let i = flat_idx / self.cols;
                let j = flat_idx % self.cols;
                return Err(DatarustError::InvalidInput(format!(
                    "NaN value at position ({}, {})",
                    i, j
                )));
            }
        }
        Ok(())
    }

    /// Returns the row at index `i` as a contiguous slice.
    #[inline]
    pub fn row(&self, i: usize) -> &[f64] {
        let start = i * self.cols;
        &self.data[start..start + self.cols]
    }

    /// Returns column `j` as a new vector.
    pub fn col(&self, j: usize) -> Vec<f64> {
        (0..self.rows)
            .map(|i| self.data[i * self.cols + j])
            .collect()
    }

    /// Iterates over the rows as contiguous slices.
    pub fn iter_rows(&self) -> impl Iterator<Item = &[f64]> {
        self.data.chunks_exact(self.cols)
    }

    /// Returns the rows as a nested `Vec<Vec<f64>>`.
    ///
    /// This allocates and copies — prefer [`as_slice`](Self::as_slice) or
    /// [`iter_rows`](Self::iter_rows) in hot loops. Kept for callers that
    /// need the nested shape (e.g. passing to `stats` functions that still
    /// take `&[Vec<f64>]`).
    #[doc(hidden)]
    pub fn rows_ref(&self) -> Vec<Vec<f64>> {
        self.data
            .chunks_exact(self.cols)
            .map(|chunk| chunk.to_vec())
            .collect()
    }

    /// Consumes the matrix and returns the rows as a nested `Vec<Vec<f64>>`.
    #[doc(hidden)]
    pub fn into_rows(self) -> Vec<Vec<f64>> {
        let cols = self.cols;
        self.data.chunks(cols).map(|chunk| chunk.to_vec()).collect()
    }

    /// Returns the transpose of the matrix.
    pub fn transpose(&self) -> Matrix {
        let rows = self.rows;
        let cols = self.cols;
        let mut out = vec![0.0; rows * cols];
        for i in 0..rows {
            for j in 0..cols {
                out[j * rows + i] = self.data[i * cols + j];
            }
        }
        Matrix {
            data: out,
            rows: cols,
            cols: rows,
        }
    }

    /// Multiplies two matrices and returns the product.
    #[allow(clippy::needless_range_loop)]
    pub fn matmul(&self, other: &Matrix) -> Result<Matrix> {
        if self.cols != other.rows {
            return Err(DatarustError::ShapeMismatch {
                expected: format!("second operand with {} rows", self.cols),
                actual: format!("{} rows", other.rows),
            });
        }
        let m = self.rows;
        let k = self.cols;
        let n = other.cols;
        let mut out = vec![0.0; m * n];

        #[cfg(feature = "matrixmultiply")]
        {
            // Tuned pure-Rust GEMM: C(m×n) = 1.0*A(m×k)·B(k×n) + 0.0*C, row-major.
            // dgemm takes isize strides.
            unsafe {
                matrixmultiply::dgemm(
                    m,
                    k,
                    n,
                    1.0,
                    self.data.as_ptr(),
                    k as isize,
                    1,
                    other.data.as_ptr(),
                    n as isize,
                    1,
                    0.0,
                    out.as_mut_ptr(),
                    n as isize,
                    1,
                );
            }
        }

        #[cfg(not(feature = "matrixmultiply"))]
        {
            // Indexed inner loops retained for cache-friendly row-major access.
            for i in 0..m {
                let out_base = i * n;
                let self_base = i * k;
                for l in 0..k {
                    let a = self.data[self_base + l];
                    if a == 0.0 {
                        continue;
                    }
                    let other_base = l * n;
                    for j in 0..n {
                        out[out_base + j] += a * other.data[other_base + j];
                    }
                }
            }
        }
        Ok(Matrix {
            data: out,
            rows: m,
            cols: n,
        })
    }

    /// Returns the mean of each column.
    pub fn column_mean(&self) -> Vec<f64> {
        crate::stats::column_mean_flat(&self.data, self.rows, self.cols)
    }

    /// Creates a matrix from a vector of columns.
    pub fn from_columns(cols: Vec<Vec<f64>>) -> Result<Self> {
        if cols.is_empty() || cols[0].is_empty() {
            return Err(DatarustError::EmptyInput("no columns".into()));
        }
        let rows = cols[0].len();
        for c in &cols {
            if c.len() != rows {
                return Err(DatarustError::ShapeMismatch {
                    expected: format!("{} rows", rows),
                    actual: format!("{} rows", c.len()),
                });
            }
        }
        let ncols = cols.len();
        let mut data = vec![0.0; rows * ncols];
        for (j, col) in cols.iter().enumerate() {
            for (i, &v) in col.iter().enumerate() {
                data[i * ncols + j] = v;
            }
        }
        Ok(Self {
            data,
            rows,
            cols: ncols,
        })
    }

    /// Select a subset of columns by index (0-based).
    ///
    /// ```rust
    /// use datarust::Matrix;
    ///
    /// let m = Matrix::new(vec![
    ///     vec![1.0, 2.0, 3.0],
    ///     vec![4.0, 5.0, 6.0],
    /// ])?;
    /// let sub = m.select_columns(&[0, 2])?;
    /// assert_eq!(sub.ncols(), 2);
    /// assert_eq!(sub.get(0, 0), 1.0);
    /// assert_eq!(sub.get(0, 1), 3.0);
    /// # Ok::<_, Box<dyn std::error::Error>>(())
    /// ```
    pub fn select_columns(&self, indices: &[usize]) -> Result<Self> {
        if indices.is_empty() {
            return Err(DatarustError::EmptyInput("no columns selected".into()));
        }
        let ncols = self.cols;
        for &c in indices {
            if c >= ncols {
                return Err(DatarustError::InvalidInput(format!(
                    "column index {} out of range (ncols {})",
                    c, ncols
                )));
            }
        }
        let out_cols = indices.len();
        let mut out = Vec::with_capacity(self.rows * out_cols);
        for i in 0..self.rows {
            let base = i * ncols;
            for &c in indices {
                out.push(self.data[base + c]);
            }
        }
        Ok(Self {
            data: out,
            rows: self.rows,
            cols: out_cols,
        })
    }

    /// Select a subset of rows by index (0-based).
    ///
    /// ```rust
    /// use datarust::Matrix;
    ///
    /// let m = Matrix::new(vec![
    ///     vec![10.0],
    ///     vec![20.0],
    ///     vec![30.0],
    /// ])?;
    /// let sub = m.select_rows(&[0, 2])?;
    /// assert_eq!(sub.nrows(), 2);
    /// assert_eq!(sub.get(0, 0), 10.0);
    /// assert_eq!(sub.get(1, 0), 30.0);
    /// # Ok::<_, Box<dyn std::error::Error>>(())
    /// ```
    pub fn select_rows(&self, indices: &[usize]) -> Result<Self> {
        if indices.is_empty() {
            return Err(DatarustError::EmptyInput("no rows selected".into()));
        }
        let nrows = self.rows;
        for &r in indices {
            if r >= nrows {
                return Err(DatarustError::InvalidInput(format!(
                    "row index {} out of range (nrows {})",
                    r, nrows
                )));
            }
        }
        let mut out = Vec::with_capacity(indices.len() * self.cols);
        for &r in indices {
            let start = r * self.cols;
            out.extend_from_slice(&self.data[start..start + self.cols]);
        }
        Ok(Self {
            data: out,
            rows: indices.len(),
            cols: self.cols,
        })
    }
}

#[cfg(feature = "serde")]
impl serde::Serialize for Matrix {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::SerializeStruct;
        // Preserve the nested `{"data":[[...]]}` wire format even though the
        // internal storage is now flat row-major.
        let rows: Vec<&[f64]> = self.data.chunks_exact(self.cols).collect();
        let mut s = serializer.serialize_struct("Matrix", 1)?;
        s.serialize_field("data", &rows)?;
        s.end()
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for Matrix {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(serde::Deserialize)]
        struct Raw {
            data: Vec<Vec<f64>>,
        }
        let raw = Raw::deserialize(deserializer)?;
        Matrix::new(raw.data).map_err(serde::de::Error::custom)
    }
}

/// Row-major matrix of strings used by the categorical encoders.
#[derive(Debug, Clone, PartialEq)]
pub struct StrMatrix {
    pub(crate) data: Vec<Vec<String>>,
}

impl StrMatrix {
    /// Creates a string matrix from a nested vector, validating a rectangular shape.
    pub fn new(data: Vec<Vec<String>>) -> Result<Self> {
        if data.is_empty() {
            return Err(DatarustError::EmptyInput("matrix has no rows".into()));
        }
        let cols = data[0].len();
        if cols == 0 {
            return Err(DatarustError::EmptyInput("matrix has no columns".into()));
        }
        for (i, row) in data.iter().enumerate() {
            if row.len() != cols {
                return Err(DatarustError::ShapeMismatch {
                    expected: format!("{} columns", cols),
                    actual: format!("{} columns at row {}", row.len(), i),
                });
            }
        }
        Ok(Self { data })
    }

    /// Creates a single-column string matrix from an iterator of values.
    pub fn from_column<I, S>(col: I) -> Result<Self>
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let data: Vec<Vec<String>> = col.into_iter().map(|s| vec![s.into()]).collect::<Vec<_>>();
        if data.is_empty() {
            return Err(DatarustError::EmptyInput("column has no rows".into()));
        }
        Self::new(data)
    }

    /// Creates a string matrix from an iterator of rows.
    pub fn from_strings<I, S>(rows: I) -> Result<Self>
    where
        I: IntoIterator<Item = Vec<S>>,
        S: Into<String>,
    {
        let data: Vec<Vec<String>> = rows
            .into_iter()
            .map(|r| r.into_iter().map(|s| s.into()).collect())
            .collect();
        Self::new(data)
    }

    /// Returns the number of rows.
    #[inline]
    pub fn nrows(&self) -> usize {
        self.data.len()
    }

    /// Returns the number of columns.
    #[inline]
    pub fn ncols(&self) -> usize {
        self.data[0].len()
    }

    /// Returns the string at row `i`, column `j`.
    ///
    /// # Panics
    ///
    /// Panics if `i >= nrows()` or `j >= ncols()`. Use
    /// [`checked_get`](Self::checked_get) for a safe alternative.
    #[inline]
    pub fn get(&self, i: usize, j: usize) -> &str {
        debug_assert!(
            i < self.data.len() && j < self.data[i].len(),
            "StrMatrix::get: index ({}, {}) out of bounds for {}×{}",
            i,
            j,
            self.data.len(),
            self.data[0].len()
        );
        &self.data[i][j]
    }

    /// Returns `Some(element)` at row `i`, column `j`, or `None` if the indices
    /// are out of bounds.
    #[inline]
    pub fn checked_get(&self, i: usize, j: usize) -> Option<&str> {
        self.data.get(i)?.get(j).map(|s| s.as_str())
    }

    /// Returns column `j` as a new vector of strings.
    pub fn column(&self, j: usize) -> Vec<String> {
        self.data.iter().map(|r| r[j].clone()).collect()
    }

    /// Returns the row at index `i` as a slice.
    pub fn row(&self, i: usize) -> &[String] {
        &self.data[i]
    }
}

#[cfg(feature = "serde")]
impl serde::Serialize for StrMatrix {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::SerializeStruct;
        let mut s = serializer.serialize_struct("StrMatrix", 1)?;
        s.serialize_field("data", &self.data)?;
        s.end()
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for StrMatrix {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(serde::Deserialize)]
        struct Raw {
            data: Vec<Vec<String>>,
        }
        let raw = Raw::deserialize(deserializer)?;
        StrMatrix::new(raw.data).map_err(serde::de::Error::custom)
    }
}

impl TryFrom<Vec<Vec<f64>>> for Matrix {
    type Error = DatarustError;

    /// Fallibly construct a [`Matrix`] from a nested vector.
    ///
    /// Returns an error if the rows are empty or jagged. This replaces the
    /// previous panicking `From` impl to keep validation consistent with
    /// [`Matrix::new`].
    fn try_from(data: Vec<Vec<f64>>) -> Result<Self> {
        Matrix::new(data)
    }
}

/// Compressed Sparse Row (CSR) matrix for memory-efficient storage of
/// mostly-zero 2-D data, mirroring scipy.sparse `csr_matrix`.
///
/// Three arrays define the non-zero entries:
/// - `indptr` (length `nrows + 1`): row `i` occupies
///   `indptr[i]..indptr[i+1]` in `indices`/`data`.
/// - `indices`: column index of each non-zero.
/// - `data`: value of each non-zero.
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct SparseMatrix {
    nrows: usize,
    ncols: usize,
    indptr: Vec<usize>,
    indices: Vec<usize>,
    data: Vec<f64>,
}

impl SparseMatrix {
    /// Build a CSR matrix from raw CSR arrays.
    pub fn new(
        nrows: usize,
        ncols: usize,
        indptr: Vec<usize>,
        indices: Vec<usize>,
        data: Vec<f64>,
    ) -> Result<Self> {
        if nrows == 0 || ncols == 0 {
            return Err(DatarustError::EmptyInput("zero dimension".into()));
        }
        if indptr.len() != nrows + 1 {
            return Err(DatarustError::ShapeMismatch {
                expected: format!("{} indptr entries", nrows + 1),
                actual: format!("{} indptr entries", indptr.len()),
            });
        }
        if indices.len() != data.len() {
            return Err(DatarustError::ShapeMismatch {
                expected: format!("{} indices", data.len()),
                actual: format!("{} indices", indices.len()),
            });
        }
        let nnz = data.len();
        if indptr[0] != 0 || indptr[nrows] != nnz {
            return Err(DatarustError::InvalidInput(
                "indptr must start at 0 and end at nnz".into(),
            ));
        }
        for &c in &indices {
            if c >= ncols {
                return Err(DatarustError::InvalidInput(format!(
                    "column index {} out of range (ncols {})",
                    c, ncols
                )));
            }
        }
        Ok(Self {
            nrows,
            ncols,
            indptr,
            indices,
            data,
        })
    }

    /// Build from `(row, col, value)` triplets. Zero-valued triplets are
    /// dropped automatically. Within a row, entries are sorted by column.
    pub fn from_triplets(
        nrows: usize,
        ncols: usize,
        triplets: &[(usize, usize, f64)],
    ) -> Result<Self> {
        if nrows == 0 || ncols == 0 {
            return Err(DatarustError::EmptyInput("zero dimension".into()));
        }
        let mut per_row: Vec<Vec<(usize, f64)>> = vec![vec![]; nrows];
        for &(r, c, v) in triplets {
            if r >= nrows {
                return Err(DatarustError::InvalidInput(format!(
                    "row {} out of range (nrows {})",
                    r, nrows
                )));
            }
            if c >= ncols {
                return Err(DatarustError::InvalidInput(format!(
                    "col {} out of range (ncols {})",
                    c, ncols
                )));
            }
            if v != 0.0 {
                per_row[r].push((c, v));
            }
        }
        let mut indptr = Vec::with_capacity(nrows + 1);
        let mut indices = Vec::new();
        let mut data = Vec::new();
        indptr.push(0);
        for row_entries in &mut per_row {
            row_entries.sort_by_key(|(c, _)| *c);
            for &(c, v) in row_entries.iter() {
                indices.push(c);
                data.push(v);
            }
            indptr.push(indices.len());
        }
        Ok(Self {
            nrows,
            ncols,
            indptr,
            indices,
            data,
        })
    }

    /// Create an all-zeros sparse matrix with the given shape.
    pub fn zeros(nrows: usize, ncols: usize) -> Result<Self> {
        if nrows == 0 || ncols == 0 {
            return Err(DatarustError::EmptyInput("zero dimension".into()));
        }
        Ok(Self {
            nrows,
            ncols,
            indptr: vec![0; nrows + 1],
            indices: vec![],
            data: vec![],
        })
    }

    /// Returns the number of rows.
    #[inline]
    pub fn nrows(&self) -> usize {
        self.nrows
    }

    /// Returns the number of columns.
    #[inline]
    pub fn ncols(&self) -> usize {
        self.ncols
    }

    /// Number of stored non-zero entries.
    #[inline]
    pub fn nnz(&self) -> usize {
        self.data.len()
    }

    /// Density (fraction of non-zero entries).
    pub fn density(&self) -> f64 {
        let total = self.nrows.saturating_mul(self.ncols);
        if total == 0 {
            return 0.0;
        }
        self.nnz() as f64 / total as f64
    }

    /// Get element at `(i, j)`. Returns 0.0 if not stored.
    pub fn get(&self, i: usize, j: usize) -> f64 {
        let start = self.indptr[i];
        let end = self.indptr[i + 1];
        // Binary search within the row's column indices.
        let slice = &self.indices[start..end];
        match slice.binary_search(&j) {
            Ok(local) => self.data[start + local],
            Err(_) => 0.0,
        }
    }

    /// Returns `Some(element)` at row `i`, column `j`, or `None` if `i` is
    /// out of bounds or the value is not stored (equivalent to 0.0).
    pub fn checked_get(&self, i: usize, j: usize) -> Option<f64> {
        let start = *self.indptr.get(i)?;
        let end = *self.indptr.get(i + 1)?;
        let slice = self.indices.get(start..end)?;
        match slice.binary_search(&j) {
            Ok(local) => Some(self.data[start + local]),
            Err(_) => Some(0.0),
        }
    }

    /// Iterate over `(col, value)` non-zero entries in row `i`.
    pub fn row_nz(&self, i: usize) -> impl Iterator<Item = (usize, f64)> + '_ {
        let start = self.indptr[i];
        let end = self.indptr[i + 1];
        self.indices[start..end]
            .iter()
            .zip(self.data[start..end].iter())
            .map(|(&c, &v)| (c, v))
    }

    /// Convert to a dense [`Matrix`].
    pub fn to_dense(&self) -> Result<Matrix> {
        let mut rows = vec![vec![0.0; self.ncols]; self.nrows];
        for (i, row) in rows.iter_mut().enumerate() {
            for (c, v) in self.row_nz(i) {
                row[c] = v;
            }
        }
        Matrix::new(rows)
    }
}

#[cfg(test)]
#[allow(dead_code)]
pub(crate) fn approx_eq_matrices(a: &Matrix, b: &Matrix, tol: f64) -> bool {
    if a.nrows() != b.nrows() || a.ncols() != b.ncols() {
        return false;
    }
    for i in 0..a.nrows() {
        for j in 0..a.ncols() {
            if (a.get(i, j) - b.get(i, j)).abs() > tol {
                return false;
            }
        }
    }
    true
}

#[cfg(test)]
#[allow(dead_code)]
pub(crate) fn approx_eq_vecs(a: &[f64], b: &[f64], tol: f64) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.iter().zip(b.iter()).all(|(x, y)| (x - y).abs() <= tol)
}

#[cfg(test)]
#[macro_use]
mod assert_macros {
    #[allow(unused_macros)]
    macro_rules! assert_mat_eq {
        ($a:expr, $b:expr, $tol:expr) => {{
            assert!(
                $crate::matrix::approx_eq_matrices(&$a, &$b, $tol),
                "matrices not equal within tolerance {}\n left: {:?}\nright: {:?}",
                $tol,
                $a.as_slice(),
                $b.as_slice()
            );
        }};
    }
}

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

    #[test]
    fn new_valid() {
        let m = Matrix::new(vec![vec![1.0, 2.0], vec![3.0, 4.0]]).unwrap();
        assert_eq!(m.nrows(), 2);
        assert_eq!(m.ncols(), 2);
    }

    #[test]
    fn new_jagged_rejected() {
        let err = Matrix::new(vec![vec![1.0, 2.0], vec![3.0]]).unwrap_err();
        assert!(matches!(err, DatarustError::ShapeMismatch { .. }));
    }

    #[test]
    fn new_empty_rejected() {
        assert!(Matrix::new(vec![]).is_err());
        assert!(Matrix::new(vec![vec![]]).is_err());
    }

    #[test]
    fn from_flat() {
        let m = Matrix::from_flat(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        assert_eq!(m.get(0, 2), 3.0);
        assert_eq!(m.get(1, 0), 4.0);
        assert_eq!(m.get(1, 2), 6.0);
    }

    #[test]
    fn from_flat_bad_count() {
        let err = Matrix::from_flat(2, 2, vec![1.0, 2.0, 3.0]).unwrap_err();
        assert!(matches!(err, DatarustError::ShapeMismatch { .. }));
    }

    #[test]
    fn zeros_and_identity() {
        let z = Matrix::zeros(2, 3).unwrap();
        assert_eq!(z.get(1, 2), 0.0);
        let id = Matrix::identity(3).unwrap();
        assert_eq!(id.get(0, 0), 1.0);
        assert_eq!(id.get(0, 1), 0.0);
        assert_eq!(id.get(1, 1), 1.0);
    }

    #[test]
    fn transpose() {
        let m = Matrix::new(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]).unwrap();
        let t = m.transpose();
        assert_eq!(t.nrows(), 3);
        assert_eq!(t.ncols(), 2);
        assert_eq!(t.get(2, 0), 3.0);
        assert_eq!(t.get(1, 1), 5.0);
    }

    #[test]
    fn matmul() {
        let a = Matrix::new(vec![vec![1.0, 2.0], vec![3.0, 4.0]]).unwrap();
        let b = Matrix::new(vec![vec![5.0, 6.0], vec![7.0, 8.0]]).unwrap();
        let c = a.matmul(&b).unwrap();
        // [[1*5+2*7, 1*6+2*8],[3*5+4*7, 3*6+4*8]] = [[19,22],[43,50]]
        assert_eq!(c.get(0, 0), 19.0);
        assert_eq!(c.get(0, 1), 22.0);
        assert_eq!(c.get(1, 0), 43.0);
        assert_eq!(c.get(1, 1), 50.0);
    }

    #[test]
    fn matmul_shape_mismatch() {
        let a = Matrix::new(vec![vec![1.0, 2.0, 3.0]]).unwrap();
        let b = Matrix::new(vec![vec![1.0, 2.0]]).unwrap();
        assert!(a.matmul(&b).is_err());
    }

    #[test]
    fn from_columns() {
        let m = Matrix::from_columns(vec![vec![1.0, 2.0], vec![10.0, 20.0]]).unwrap();
        assert_eq!(m.nrows(), 2);
        assert_eq!(m.ncols(), 2);
        assert_eq!(m.get(0, 0), 1.0);
        assert_eq!(m.get(1, 1), 20.0);
    }

    #[test]
    fn strmatrix_from_column() {
        let s = StrMatrix::from_column(["a", "b", "a"]).unwrap();
        assert_eq!(s.nrows(), 3);
        assert_eq!(s.ncols(), 1);
        assert_eq!(s.get(2, 0), "a");
    }

    #[test]
    fn strmatrix_from_strings() {
        let s = StrMatrix::from_strings(vec![vec!["x", "y"], vec!["x", "z"]]).unwrap();
        assert_eq!(s.ncols(), 2);
        assert_eq!(s.get(1, 1), "z");
    }

    #[test]
    fn sparse_from_triplets_basic() {
        let sp = SparseMatrix::from_triplets(
            3,
            4,
            &[(0, 0, 1.0), (1, 2, 3.0), (2, 3, 5.0), (0, 3, 7.0)],
        )
        .unwrap();
        assert_eq!(sp.nrows(), 3);
        assert_eq!(sp.ncols(), 4);
        assert_eq!(sp.nnz(), 4);
        assert_eq!(sp.get(0, 0), 1.0);
        assert_eq!(sp.get(1, 2), 3.0);
        assert_eq!(sp.get(0, 3), 7.0);
        assert_eq!(sp.get(1, 0), 0.0);
    }

    #[test]
    fn sparse_zero_triplets_dropped() {
        let sp =
            SparseMatrix::from_triplets(2, 2, &[(0, 0, 0.0), (0, 1, 5.0), (1, 0, 0.0)]).unwrap();
        assert_eq!(sp.nnz(), 1);
        assert_eq!(sp.get(0, 1), 5.0);
    }

    #[test]
    fn sparse_to_dense() {
        let sp = SparseMatrix::from_triplets(2, 3, &[(0, 1, 2.0), (1, 0, 4.0)]).unwrap();
        let dense = sp.to_dense().unwrap();
        assert_eq!(dense.row(0), [0.0, 2.0, 0.0]);
        assert_eq!(dense.row(1), [4.0, 0.0, 0.0]);
    }

    #[test]
    fn sparse_zeros() {
        let sp = SparseMatrix::zeros(2, 3).unwrap();
        assert_eq!(sp.nnz(), 0);
        assert_eq!(sp.density(), 0.0);
        assert_eq!(sp.get(0, 1), 0.0);
    }

    #[test]
    fn sparse_density() {
        let sp = SparseMatrix::from_triplets(2, 4, &[(0, 0, 1.0), (1, 3, 1.0)]).unwrap();
        assert!((sp.density() - 0.25).abs() < 1e-12);
    }

    #[test]
    fn sparse_row_nz() {
        let sp =
            SparseMatrix::from_triplets(2, 3, &[(0, 1, 2.0), (0, 2, 9.0), (1, 0, 4.0)]).unwrap();
        let row0: Vec<(usize, f64)> = sp.row_nz(0).collect();
        assert_eq!(row0, vec![(1, 2.0), (2, 9.0)]);
        let row1: Vec<(usize, f64)> = sp.row_nz(1).collect();
        assert_eq!(row1, vec![(0, 4.0)]);
    }

    #[test]
    fn sparse_bad_indptr_rejected() {
        let err = SparseMatrix::new(2, 2, vec![0, 0, 5], vec![0], vec![1.0]).unwrap_err();
        assert!(matches!(err, DatarustError::InvalidInput(_)));
    }

    #[test]
    fn sparse_col_out_of_range_rejected() {
        let err = SparseMatrix::from_triplets(2, 2, &[(0, 5, 1.0)]).unwrap_err();
        assert!(matches!(err, DatarustError::InvalidInput(_)));
    }

    #[test]
    fn select_columns_basic() {
        let m = Matrix::new(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]).unwrap();
        let sub = m.select_columns(&[0, 2]).unwrap();
        assert_eq!(sub.ncols(), 2);
        assert_eq!(sub.get(0, 0), 1.0);
        assert_eq!(sub.get(0, 1), 3.0);
        assert_eq!(sub.get(1, 0), 4.0);
        assert_eq!(sub.get(1, 1), 6.0);
    }

    #[test]
    fn select_columns_out_of_range() {
        let m = Matrix::new(vec![vec![1.0, 2.0]]).unwrap();
        assert!(m.select_columns(&[0, 5]).is_err());
    }

    #[test]
    fn select_columns_empty() {
        let m = Matrix::new(vec![vec![1.0]]).unwrap();
        assert!(m.select_columns(&[]).is_err());
    }

    #[test]
    fn select_rows_basic() {
        let m = Matrix::new(vec![vec![10.0], vec![20.0], vec![30.0]]).unwrap();
        let sub = m.select_rows(&[0, 2]).unwrap();
        assert_eq!(sub.nrows(), 2);
        assert_eq!(sub.get(0, 0), 10.0);
        assert_eq!(sub.get(1, 0), 30.0);
    }

    #[test]
    fn select_rows_out_of_range() {
        let m = Matrix::new(vec![vec![1.0]]).unwrap();
        assert!(m.select_rows(&[0, 5]).is_err());
    }

    #[test]
    fn select_rows_empty() {
        let m = Matrix::new(vec![vec![1.0]]).unwrap();
        assert!(m.select_rows(&[]).is_err());
    }

    #[test]
    fn select_columns_reordered() {
        let m = Matrix::new(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]).unwrap();
        let sub = m.select_columns(&[2, 0]).unwrap();
        assert_eq!(sub.get(0, 0), 3.0);
        assert_eq!(sub.get(0, 1), 1.0);
    }

    #[test]
    fn select_rows_duplicates() {
        let m = Matrix::new(vec![vec![10.0], vec![20.0]]).unwrap();
        let sub = m.select_rows(&[0, 0, 1]).unwrap();
        assert_eq!(sub.nrows(), 3);
        assert_eq!(sub.get(0, 0), 10.0);
        assert_eq!(sub.get(1, 0), 10.0);
        assert_eq!(sub.get(2, 0), 20.0);
    }

    #[test]
    fn dense_accessors_and_row_conversions_preserve_layout() {
        let mut m = Matrix::from_rows(vec![vec![1.0, 2.0], vec![3.0, 4.0]]).unwrap();

        assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
        assert_eq!(m.checked_get(1, 0), Some(3.0));
        assert_eq!(m.checked_get(2, 0), None);
        assert_eq!(m.checked_get(0, 2), None);
        m.as_mut_slice()[1] = 20.0;
        m.set(1, 0, 30.0);
        assert_eq!(m.row(0), [1.0, 20.0]);
        assert_eq!(m.col(0), vec![1.0, 30.0]);
        assert_eq!(
            m.iter_rows().collect::<Vec<_>>(),
            vec![&[1.0, 20.0][..], &[30.0, 4.0][..]]
        );
        assert_eq!(m.rows_ref(), vec![vec![1.0, 20.0], vec![30.0, 4.0]]);
        assert_eq!(
            m.clone().into_rows(),
            vec![vec![1.0, 20.0], vec![30.0, 4.0]]
        );
        assert!(m.validate_no_nan().is_ok());

        m.set(1, 1, f64::NAN);
        assert!(matches!(
            m.validate_no_nan(),
            Err(DatarustError::InvalidInput(_))
        ));
    }

    #[test]
    fn dense_constructor_edge_cases_are_rejected() {
        assert!(Matrix::zeros(0, 1).is_err());
        assert!(Matrix::identity(0).is_err());
        assert!(Matrix::from_flat(0, 1, vec![]).is_err());
        assert!(Matrix::from_flat(usize::MAX, 2, vec![]).is_err());
        assert!(Matrix::from_columns(vec![]).is_err());
        assert!(Matrix::from_columns(vec![vec![]]).is_err());
        assert!(Matrix::from_columns(vec![vec![1.0], vec![2.0, 3.0]]).is_err());
        assert!(Matrix::try_from(vec![vec![1.0], vec![]]).is_err());
    }

    #[test]
    fn strmatrix_accessors_and_validation_cover_edge_cases() {
        let strings = StrMatrix::new(vec![
            vec!["a".into(), "b".into()],
            vec!["c".into(), "d".into()],
        ])
        .unwrap();
        assert_eq!(strings.checked_get(1, 1), Some("d"));
        assert_eq!(strings.checked_get(2, 0), None);
        assert_eq!(strings.checked_get(0, 2), None);
        assert_eq!(strings.column(1), vec!["b".to_string(), "d".to_string()]);
        assert_eq!(strings.row(0), ["a".to_string(), "b".to_string()]);
        assert!(StrMatrix::new(vec![]).is_err());
        assert!(StrMatrix::new(vec![vec![]]).is_err());
        assert!(StrMatrix::new(vec![vec!["a".into()], vec!["b".into(), "c".into()]]).is_err());
        assert!(StrMatrix::from_column(Vec::<String>::new()).is_err());
    }

    #[test]
    fn raw_sparse_construction_and_bounds_checks() {
        let sparse = SparseMatrix::new(2, 3, vec![0, 1, 2], vec![0, 2], vec![1.0, 3.0]).unwrap();
        assert_eq!(sparse.checked_get(0, 0), Some(1.0));
        assert_eq!(sparse.checked_get(0, 2), Some(0.0));
        assert_eq!(sparse.checked_get(2, 0), None);
        assert_eq!(
            sparse.to_dense().unwrap().rows_ref(),
            vec![vec![1.0, 0.0, 0.0], vec![0.0, 0.0, 3.0]]
        );

        assert!(SparseMatrix::new(0, 1, vec![0], vec![], vec![]).is_err());
        assert!(SparseMatrix::new(2, 2, vec![0, 0], vec![], vec![]).is_err());
        assert!(SparseMatrix::new(1, 2, vec![0, 1], vec![], vec![1.0]).is_err());
        assert!(SparseMatrix::new(1, 2, vec![1, 1], vec![0], vec![1.0]).is_err());
        assert!(SparseMatrix::new(1, 2, vec![0, 1], vec![2], vec![1.0]).is_err());
        assert!(SparseMatrix::from_triplets(0, 1, &[]).is_err());
        assert!(SparseMatrix::from_triplets(1, 1, &[(1, 0, 1.0)]).is_err());
    }

    #[cfg(feature = "serde")]
    #[test]
    fn serde_round_trips_dense_and_string_matrices_and_rejects_jagged_data() {
        let matrix = Matrix::new(vec![vec![1.0, 2.0], vec![3.0, 4.0]]).unwrap();
        let encoded = serde_json::to_string(&matrix).unwrap();
        assert_eq!(serde_json::from_str::<Matrix>(&encoded).unwrap(), matrix);
        assert!(serde_json::from_str::<Matrix>(r#"{"data":[[1.0],[2.0,3.0]]}"#).is_err());

        let strings = StrMatrix::from_strings(vec![vec!["north", "south"]]).unwrap();
        let encoded = serde_json::to_string(&strings).unwrap();
        assert_eq!(
            serde_json::from_str::<StrMatrix>(&encoded).unwrap(),
            strings
        );
        assert!(
            serde_json::from_str::<StrMatrix>(r#"{"data":[["north"],["south","east"]]}"#).is_err()
        );
    }
}