trueno 0.17.2

//! Matrix storage and construction operations
//!
//! This module provides storage-related operations for Matrix:
//! - Constructors: `new()`, `from_vec()`, `from_slice()`, `zeros()`, `identity()`
//! - Accessors: `rows()`, `cols()`, `shape()`, `get()`, `get_mut()`, `as_slice()`
//!
//! ## Domain Separation (PMAT-018)
//!
//! Storage is separate from Algebra (arithmetic operations).
//! A matrix's memory layout is independent of its mathematical operations.

use crate::{Backend, TruenoError};

use super::super::Matrix;

impl std::ops::Index<(usize, usize)> for Matrix<f32> {
    type Output = f32;

    fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
        &self.data[row * self.cols + col]
    }
}

impl Matrix<f32> {
    // =========================================================================
    // Constructors
    // =========================================================================

    /// Creates a new matrix with uninitialized values
    ///
    /// # Arguments
    ///
    /// * `rows` - Number of rows
    /// * `cols` - Number of columns
    ///
    /// # Returns
    ///
    /// A new matrix with dimensions `rows x cols` containing uninitialized values
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::new(3, 4);
    /// assert_eq!(m.rows(), 3);
    /// assert_eq!(m.cols(), 4);
    /// ```
    pub fn new(rows: usize, cols: usize) -> Self {
        let backend = Backend::select_best();
        Matrix { rows, cols, data: vec![0.0; rows * cols], backend }
    }

    /// Creates a matrix from a vector of data
    ///
    /// # Arguments
    ///
    /// * `rows` - Number of rows
    /// * `cols` - Number of columns
    /// * `data` - Vector containing matrix elements in row-major order
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if `data.len() != rows * cols`
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    /// assert_eq!(m.rows(), 2);
    /// assert_eq!(m.cols(), 2);
    /// ```
    pub fn from_vec(rows: usize, cols: usize, data: Vec<f32>) -> Result<Self, TruenoError> {
        if data.len() != rows * cols {
            return Err(TruenoError::InvalidInput(format!(
                "Data length {} does not match matrix dimensions {}x{} (expected {})",
                data.len(),
                rows,
                cols,
                rows * cols
            )));
        }

        let backend = Backend::select_best();
        Ok(Matrix { rows, cols, data, backend })
    }

    /// Creates a matrix from a vector with a specific backend
    ///
    /// This is useful for testing specific SIMD code paths.
    pub fn from_vec_with_backend(
        rows: usize,
        cols: usize,
        data: Vec<f32>,
        backend: Backend,
    ) -> Self {
        assert_eq!(
            data.len(),
            rows * cols,
            "Data length {} does not match matrix dimensions {}x{}",
            data.len(),
            rows,
            cols
        );
        Matrix { rows, cols, data, backend }
    }

    /// Creates a matrix from a slice by copying the data
    ///
    /// This is a convenience method that copies the slice into an owned vector.
    /// For zero-copy scenarios, consider using the data directly with `from_vec`
    /// if you already have an owned `Vec`.
    ///
    /// # Arguments
    ///
    /// * `rows` - Number of rows
    /// * `cols` - Number of columns
    /// * `data` - Slice containing matrix elements in row-major order
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if `data.len() != rows * cols`
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let data = [1.0, 2.0, 3.0, 4.0];
    /// let m = Matrix::from_slice(2, 2, &data).unwrap();
    /// assert_eq!(m.get(0, 0), Some(&1.0));
    /// ```
    pub fn from_slice(rows: usize, cols: usize, data: &[f32]) -> Result<Self, TruenoError> {
        Self::from_vec(rows, cols, data.to_vec())
    }

    /// Creates a matrix filled with zeros
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::zeros(3, 3);
    /// assert_eq!(m.get(1, 1), Some(&0.0));
    /// ```
    pub fn zeros(rows: usize, cols: usize) -> Self {
        Matrix::new(rows, cols)
    }

    /// Creates a matrix filled with zeros using a specific backend
    /// (Internal use only - reuses backend from parent matrix)
    pub(crate) fn zeros_with_backend(rows: usize, cols: usize, backend: Backend) -> Self {
        Matrix { rows, cols, data: vec![0.0; rows * cols], backend }
    }

    /// Creates an identity matrix (square matrix with 1s on diagonal)
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::identity(3);
    /// assert_eq!(m.get(0, 0), Some(&1.0));
    /// assert_eq!(m.get(0, 1), Some(&0.0));
    /// assert_eq!(m.get(1, 1), Some(&1.0));
    /// ```
    pub fn identity(n: usize) -> Self {
        let mut data = vec![0.0; n * n];
        for i in 0..n {
            data[i * n + i] = 1.0;
        }
        let backend = Backend::select_best();
        Matrix { rows: n, cols: n, data, backend }
    }

    // =========================================================================
    // Accessors
    // =========================================================================

    /// Returns the number of rows
    pub fn rows(&self) -> usize {
        self.rows
    }

    /// Returns the number of columns
    pub fn cols(&self) -> usize {
        self.cols
    }

    /// Returns the shape as (rows, cols)
    pub fn shape(&self) -> (usize, usize) {
        (self.rows, self.cols)
    }

    /// Gets a reference to an element at (row, col)
    ///
    /// Returns `None` if indices are out of bounds
    pub fn get(&self, row: usize, col: usize) -> Option<&f32> {
        if row >= self.rows || col >= self.cols {
            None
        } else {
            self.data.get(row * self.cols + col)
        }
    }

    /// Gets a mutable reference to an element at (row, col)
    ///
    /// Returns `None` if indices are out of bounds
    pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut f32> {
        if row >= self.rows || col >= self.cols {
            None
        } else {
            let idx = row * self.cols + col;
            self.data.get_mut(idx)
        }
    }

    /// Returns a reference to the underlying data
    pub fn as_slice(&self) -> &[f32] {
        &self.data
    }

    /// Returns the backend used by this matrix
    pub fn backend(&self) -> Backend {
        self.backend
    }
}

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

    #[test]
    fn test_new_creates_zero_matrix() {
        let m = Matrix::new(3, 4);
        assert_eq!(m.rows(), 3);
        assert_eq!(m.cols(), 4);
        assert!(m.as_slice().iter().all(|&x| x == 0.0));
    }

    #[test]
    fn test_from_vec_success() {
        let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        assert_eq!(m.get(0, 0), Some(&1.0));
        assert_eq!(m.get(1, 1), Some(&4.0));
    }

    #[test]
    fn test_from_vec_dimension_mismatch() {
        let result = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0]);
        assert!(result.is_err());
    }

    #[test]
    fn test_identity() {
        let m = Matrix::identity(3);
        assert_eq!(m.get(0, 0), Some(&1.0));
        assert_eq!(m.get(1, 1), Some(&1.0));
        assert_eq!(m.get(2, 2), Some(&1.0));
        assert_eq!(m.get(0, 1), Some(&0.0));
    }

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

    #[test]
    fn test_get_out_of_bounds() {
        let m = Matrix::new(2, 2);
        assert_eq!(m.get(2, 0), None);
        assert_eq!(m.get(0, 2), None);
    }

    #[test]
    fn test_get_mut() {
        let mut m = Matrix::new(2, 2);
        if let Some(val) = m.get_mut(1, 1) {
            *val = 42.0;
        }
        assert_eq!(m.get(1, 1), Some(&42.0));
    }
}