rosa 0.1.0

use std::fmt;
use std::ops::Index;

use super::matrix_view::ColIter;

/// Dense 2D matrix, row-major layout.
///
/// Convention for spectrograms: rows = frequency bins, cols = time frames.
/// Element at (row, col) is stored at `data[row * cols + col]`.
#[derive(Clone)]
pub struct Matrix {
    data: Vec<f64>,
    rows: usize,
    cols: usize,
}

// --- Construction ---

impl Matrix {
    /// Zero-filled matrix.
    pub fn zeros(rows: usize, cols: usize) -> Self {
        Self {
            data: vec![0.0; rows * cols],
            rows,
            cols,
        }
    }

    /// From flat data + shape (row-major order).
    /// Panics if `data.len() != rows * cols`.
    pub fn from_vec(data: Vec<f64>, rows: usize, cols: usize) -> Self {
        assert_eq!(
            data.len(),
            rows * cols,
            "data length {} does not match shape {}x{}",
            data.len(),
            rows,
            cols
        );
        Self { data, rows, cols }
    }

    /// From a generator function: `f(row, col) -> value`.
    pub fn from_fn(rows: usize, cols: usize, f: impl Fn(usize, usize) -> f64) -> Self {
        let mut data = Vec::with_capacity(rows * cols);
        for r in 0..rows {
            for c in 0..cols {
                data.push(f(r, c));
            }
        }
        Self { data, rows, cols }
    }

    /// From row slices. Each slice is one row. Direct copy.
    /// Panics if rows have different lengths.
    pub fn from_rows(rows: &[&[f64]]) -> Self {
        if rows.is_empty() {
            return Self::zeros(0, 0);
        }
        let n_cols = rows[0].len();
        let n_rows = rows.len();
        let mut data = Vec::with_capacity(n_rows * n_cols);
        for (i, row) in rows.iter().enumerate() {
            assert_eq!(
                row.len(),
                n_cols,
                "row {} has length {}, expected {}",
                i,
                row.len(),
                n_cols
            );
            data.extend_from_slice(row);
        }
        Self {
            data,
            rows: n_rows,
            cols: n_cols,
        }
    }

    /// From column slices. Each slice is one column.
    /// Transposes into row-major internally.
    /// Panics if columns have different lengths.
    pub fn from_columns(cols: &[&[f64]]) -> Self {
        if cols.is_empty() {
            return Self::zeros(0, 0);
        }
        let n_rows = cols[0].len();
        let n_cols = cols.len();
        let mut data = vec![0.0; n_rows * n_cols];
        for (j, col) in cols.iter().enumerate() {
            assert_eq!(
                col.len(),
                n_rows,
                "column {} has length {}, expected {}",
                j,
                col.len(),
                n_rows
            );
            for (i, &val) in col.iter().enumerate() {
                data[i * n_cols + j] = val;
            }
        }
        Self {
            data,
            rows: n_rows,
            cols: n_cols,
        }
    }
}

// --- Access ---

impl Matrix {
    /// Number of rows.
    pub fn rows(&self) -> usize {
        self.rows
    }

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

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

    /// Element access. Panics on out-of-bounds.
    pub fn get(&self, row: usize, col: usize) -> f64 {
        assert!(
            row < self.rows && col < self.cols,
            "index ({}, {}) out of bounds for {}x{} matrix",
            row,
            col,
            self.rows,
            self.cols
        );
        self.data[row * self.cols + col]
    }

    /// Mutable element access. Panics on out-of-bounds.
    pub fn get_mut(&mut self, row: usize, col: usize) -> &mut f64 {
        assert!(
            row < self.rows && col < self.cols,
            "index ({}, {}) out of bounds for {}x{} matrix",
            row,
            col,
            self.rows,
            self.cols
        );
        &mut self.data[row * self.cols + col]
    }

    /// Row as a contiguous slice. Zero-cost.
    pub fn row(&self, r: usize) -> &[f64] {
        assert!(
            r < self.rows,
            "row {} out of bounds for {} rows",
            r,
            self.rows
        );
        let start = r * self.cols;
        &self.data[start..start + self.cols]
    }

    /// Mutable row slice.
    pub fn row_mut(&mut self, r: usize) -> &mut [f64] {
        assert!(
            r < self.rows,
            "row {} out of bounds for {} rows",
            r,
            self.rows
        );
        let start = r * self.cols;
        &mut self.data[start..start + self.cols]
    }

    /// Sub-row slice for a range of columns.
    pub fn row_range(&self, r: usize, col_start: usize, col_end: usize) -> &[f64] {
        assert!(r < self.rows);
        assert!(col_start <= col_end && col_end <= self.cols);
        let start = r * self.cols + col_start;
        &self.data[start..start + (col_end - col_start)]
    }

    /// Column iterator (strided access).
    pub fn col(&self, c: usize) -> ColIter<'_> {
        assert!(
            c < self.cols,
            "col {} out of bounds for {} cols",
            c,
            self.cols
        );
        ColIter::new(&self.data, c, self.cols, self.rows)
    }

    /// Raw data as a flat slice (row-major).
    pub fn as_slice(&self) -> &[f64] {
        &self.data
    }

    /// Raw data as a mutable flat slice.
    pub fn as_mut_slice(&mut self) -> &mut [f64] {
        &mut self.data
    }
}

// --- Element-wise operations (Step 3) ---

impl Matrix {
    /// Apply `f` to every element, return new Matrix.
    pub fn map(&self, f: impl Fn(f64) -> f64) -> Matrix {
        Matrix {
            data: self.data.iter().map(|&x| f(x)).collect(),
            rows: self.rows,
            cols: self.cols,
        }
    }

    /// Apply `f` to every element in-place.
    pub fn map_inplace(&mut self, f: impl Fn(f64) -> f64) {
        for x in &mut self.data {
            *x = f(*x);
        }
    }

    /// Element-wise binary operation with another Matrix of same shape.
    /// Panics if shapes differ.
    pub fn zip_map(&self, other: &Matrix, f: impl Fn(f64, f64) -> f64) -> Matrix {
        assert_eq!(
            self.shape(),
            other.shape(),
            "shape mismatch: {:?} vs {:?}",
            self.shape(),
            other.shape()
        );
        Matrix {
            data: self
                .data
                .iter()
                .zip(other.data.iter())
                .map(|(&a, &b)| f(a, b))
                .collect(),
            rows: self.rows,
            cols: self.cols,
        }
    }

    /// Multiply every element by a scalar.
    pub fn scale(&self, scalar: f64) -> Matrix {
        self.map(|x| x * scalar)
    }

    /// Add a scalar to every element.
    pub fn add_scalar(&self, scalar: f64) -> Matrix {
        self.map(|x| x + scalar)
    }
}

// --- Trait impls (Step 10) ---

impl PartialEq for Matrix {
    fn eq(&self, other: &Self) -> bool {
        self.rows == other.rows && self.cols == other.cols && self.data == other.data
    }
}

impl Index<(usize, usize)> for Matrix {
    type Output = f64;

    fn index(&self, (row, col): (usize, usize)) -> &f64 {
        assert!(
            row < self.rows && col < self.cols,
            "index ({}, {}) out of bounds for {}x{} matrix",
            row,
            col,
            self.rows,
            self.cols
        );
        &self.data[row * self.cols + col]
    }
}

impl fmt::Debug for Matrix {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Matrix({}x{}, [", self.rows, self.cols)?;
        let max_elements = 20;
        let total = self.data.len();
        if total <= max_elements {
            for (i, &v) in self.data.iter().enumerate() {
                if i > 0 {
                    write!(f, ", ")?;
                }
                write!(f, "{v}")?;
            }
        } else {
            for (i, &v) in self.data[..10].iter().enumerate() {
                if i > 0 {
                    write!(f, ", ")?;
                }
                write!(f, "{v}")?;
            }
            write!(f, ", ... ")?;
            for (i, &v) in self.data[total - 10..].iter().enumerate() {
                if i > 0 {
                    write!(f, ", ")?;
                }
                write!(f, "{v}")?;
            }
        }
        write!(f, "])")
    }
}

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

    // --- Step 1 tests: Construction and Access ---

    #[test]
    fn test_zeros() {
        let m = Matrix::zeros(3, 4);
        assert_eq!(m.shape(), (3, 4));
        for r in 0..3 {
            for c in 0..4 {
                assert_eq!(m.get(r, c), 0.0);
            }
        }
    }

    #[test]
    fn test_from_vec() {
        let m = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 2, 3);
        assert_eq!(m.shape(), (2, 3));
        assert_eq!(m.get(0, 0), 1.0);
        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]
    #[should_panic(expected = "data length")]
    fn test_from_vec_wrong_size() {
        Matrix::from_vec(vec![1.0, 2.0, 3.0], 2, 3);
    }

    #[test]
    fn test_from_fn() {
        let m = Matrix::from_fn(3, 4, |r, c| (r * 4 + c) as f64);
        assert_eq!(m.shape(), (3, 4));
        assert_eq!(m.get(0, 0), 0.0);
        assert_eq!(m.get(1, 2), 6.0);
        assert_eq!(m.get(2, 3), 11.0);
    }

    #[test]
    fn test_from_rows() {
        let m = Matrix::from_rows(&[&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0]]);
        assert_eq!(m.shape(), (2, 3));
        assert_eq!(m.row(0), &[1.0, 2.0, 3.0]);
        assert_eq!(m.row(1), &[4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_from_columns() {
        let m = Matrix::from_columns(&[&[1.0, 4.0], &[2.0, 5.0], &[3.0, 6.0]]);
        assert_eq!(m.shape(), (2, 3));
        assert_eq!(m.row(0), &[1.0, 2.0, 3.0]);
        assert_eq!(m.row(1), &[4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_from_columns_roundtrip() {
        // Build from columns, read back rows
        let col0 = &[10.0, 40.0, 70.0][..];
        let col1 = &[20.0, 50.0, 80.0][..];
        let col2 = &[30.0, 60.0, 90.0][..];
        let m = Matrix::from_columns(&[col0, col1, col2]);
        assert_eq!(m.shape(), (3, 3));
        assert_eq!(m.row(0), &[10.0, 20.0, 30.0]);
        assert_eq!(m.row(1), &[40.0, 50.0, 60.0]);
        assert_eq!(m.row(2), &[70.0, 80.0, 90.0]);
    }

    #[test]
    fn test_access() {
        let mut m = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        assert_eq!(m.rows(), 2);
        assert_eq!(m.cols(), 2);
        assert_eq!(m.get(0, 0), 1.0);

        *m.get_mut(0, 0) = 99.0;
        assert_eq!(m.get(0, 0), 99.0);

        assert_eq!(m.row(1), &[3.0, 4.0]);
        m.row_mut(1)[0] = 33.0;
        assert_eq!(m.get(1, 0), 33.0);

        assert_eq!(m.row_range(0, 0, 2), &[99.0, 2.0]);
        assert_eq!(m.as_slice().len(), 4);
    }

    #[test]
    #[should_panic(expected = "out of bounds")]
    fn test_get_out_of_bounds() {
        let m = Matrix::zeros(2, 3);
        m.get(2, 0);
    }

    #[test]
    #[should_panic(expected = "out of bounds")]
    fn test_get_col_out_of_bounds() {
        let m = Matrix::zeros(2, 3);
        m.get(0, 3);
    }

    // --- Step 3 tests: Element-wise operations ---

    #[test]
    fn test_map() {
        let m = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        let doubled = m.map(|x| x * 2.0);
        assert_eq!(doubled.get(0, 0), 2.0);
        assert_eq!(doubled.get(0, 1), 4.0);
        assert_eq!(doubled.get(1, 0), 6.0);
        assert_eq!(doubled.get(1, 1), 8.0);
    }

    #[test]
    fn test_map_inplace() {
        let mut m = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        m.map_inplace(|x| x + 10.0);
        assert_eq!(m.get(0, 0), 11.0);
        assert_eq!(m.get(1, 1), 14.0);
    }

    #[test]
    fn test_zip_map() {
        let a = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        let b = Matrix::from_vec(vec![10.0, 20.0, 30.0, 40.0], 2, 2);
        let sum = a.zip_map(&b, |x, y| x + y);
        assert_eq!(sum.get(0, 0), 11.0);
        assert_eq!(sum.get(1, 1), 44.0);
    }

    #[test]
    #[should_panic(expected = "shape mismatch")]
    fn test_zip_map_shape_mismatch() {
        let a = Matrix::zeros(2, 3);
        let b = Matrix::zeros(3, 2);
        a.zip_map(&b, |x, y| x + y);
    }

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

    #[test]
    fn test_add_scalar() {
        let m = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        let result = m.add_scalar(10.0);
        assert_eq!(result.get(0, 0), 11.0);
        assert_eq!(result.get(1, 1), 14.0);
    }

    // --- Step 10 tests: Trait impls ---

    #[test]
    fn test_clone_independent() {
        let m = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        let mut c = m.clone();
        *c.get_mut(0, 0) = 99.0;
        assert_eq!(m.get(0, 0), 1.0);
        assert_eq!(c.get(0, 0), 99.0);
    }

    #[test]
    fn test_partial_eq() {
        let a = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        let b = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        let c = Matrix::from_vec(vec![1.0, 2.0, 3.0, 5.0], 2, 2);
        let d = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 1, 4);
        assert_eq!(a, b);
        assert_ne!(a, c);
        assert_ne!(a, d);
    }

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

    #[test]
    fn test_debug_small() {
        let m = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
        let s = format!("{:?}", m);
        assert!(s.contains("2x2"));
        assert!(s.contains("1"));
    }

    #[test]
    fn test_debug_large() {
        let m = Matrix::zeros(10, 10);
        let s = format!("{:?}", m);
        assert!(s.contains("10x10"));
        assert!(s.contains("..."));
    }
}