rosa 0.1.0

use std::ops::Range;

use crate::Matrix;

/// Iterator over a column (strided access through row-major data).
pub struct ColIter<'a> {
    data: &'a [f64],
    pos: usize,
    stride: usize,
    remaining: usize,
}

impl<'a> ColIter<'a> {
    pub(crate) fn new(data: &'a [f64], col: usize, stride: usize, rows: usize) -> Self {
        Self {
            data,
            pos: col,
            stride,
            remaining: rows,
        }
    }
}

impl<'a> Iterator for ColIter<'a> {
    type Item = &'a f64;

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining == 0 {
            return None;
        }
        let val = &self.data[self.pos];
        self.pos += self.stride;
        self.remaining -= 1;
        Some(val)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.remaining, Some(self.remaining))
    }
}

impl ExactSizeIterator for ColIter<'_> {}

/// Immutable view into a rectangular region of a Matrix.
pub struct MatrixView<'a> {
    data: &'a [f64],
    row_start: usize,
    col_start: usize,
    rows: usize,
    cols: usize,
    stride: usize, // parent cols
}

impl<'a> MatrixView<'a> {
    /// Number of rows in the view.
    pub fn rows(&self) -> usize {
        self.rows
    }

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

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

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

    /// Row slice. Only zero-cost when col_start == 0 and cols == stride.
    pub fn row(&self, r: usize) -> &[f64] {
        assert!(r < self.rows);
        let start = (self.row_start + r) * self.stride + self.col_start;
        &self.data[start..start + self.cols]
    }

    /// Column iterator (strided).
    pub fn col(&self, c: usize) -> ColIter<'a> {
        assert!(c < self.cols);
        let start = self.row_start * self.stride + self.col_start + c;
        ColIter {
            data: self.data,
            pos: start,
            stride: self.stride,
            remaining: self.rows,
        }
    }
}

/// Mutable view into a rectangular region of a Matrix.
pub struct MatrixViewMut<'a> {
    data: &'a mut [f64],
    row_start: usize,
    col_start: usize,
    rows: usize,
    cols: usize,
    stride: usize,
}

impl MatrixViewMut<'_> {
    /// Number of rows in the view.
    pub fn rows(&self) -> usize {
        self.rows
    }

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

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

    /// Element access.
    pub fn get(&self, row: usize, col: usize) -> f64 {
        assert!(row < self.rows && col < self.cols);
        self.data[(self.row_start + row) * self.stride + self.col_start + col]
    }

    /// Mutable element access.
    pub fn get_mut(&mut self, row: usize, col: usize) -> &mut f64 {
        assert!(row < self.rows && col < self.cols);
        &mut self.data[(self.row_start + row) * self.stride + self.col_start + col]
    }

    /// Set element value.
    pub fn set(&mut self, row: usize, col: usize, val: f64) {
        *self.get_mut(row, col) = val;
    }

    /// Add values element-wise from a MatrixView into this mutable view.
    /// Used for ISTFT overlap-add.
    /// Panics if shapes differ.
    pub fn add_assign(&mut self, other: &MatrixView<'_>) {
        assert_eq!(
            (self.rows, self.cols),
            (other.rows, other.cols),
            "add_assign shape mismatch: {:?} vs {:?}",
            (self.rows, self.cols),
            (other.rows, other.cols)
        );
        for r in 0..self.rows {
            for c in 0..self.cols {
                let idx = (self.row_start + r) * self.stride + self.col_start + c;
                self.data[idx] += other.get(r, c);
            }
        }
    }
}

// --- View construction on Matrix ---

impl Matrix {
    /// View of rows `[row_start..row_end]`.
    pub fn rows_range(&self, row_start: usize, row_end: usize) -> MatrixView<'_> {
        assert!(row_start <= row_end && row_end <= self.rows());
        MatrixView {
            data: self.as_slice(),
            row_start,
            col_start: 0,
            rows: row_end - row_start,
            cols: self.cols(),
            stride: self.cols(),
        }
    }

    /// View of a rectangular subregion.
    pub fn view(&self, row_range: Range<usize>, col_range: Range<usize>) -> MatrixView<'_> {
        assert!(row_range.end <= self.rows() && col_range.end <= self.cols());
        MatrixView {
            data: self.as_slice(),
            row_start: row_range.start,
            col_start: col_range.start,
            rows: row_range.end - row_range.start,
            cols: col_range.end - col_range.start,
            stride: self.cols(),
        }
    }

    /// Full immutable view.
    pub fn as_view(&self) -> MatrixView<'_> {
        MatrixView {
            data: self.as_slice(),
            row_start: 0,
            col_start: 0,
            rows: self.rows(),
            cols: self.cols(),
            stride: self.cols(),
        }
    }

    /// Mutable view of a rectangular subregion.
    pub fn view_mut(
        &mut self,
        row_range: Range<usize>,
        col_range: Range<usize>,
    ) -> MatrixViewMut<'_> {
        assert!(row_range.end <= self.rows() && col_range.end <= self.cols());
        let cols = self.cols();
        MatrixViewMut {
            data: self.as_mut_slice(),
            row_start: row_range.start,
            col_start: col_range.start,
            rows: row_range.end - row_range.start,
            cols: col_range.end - col_range.start,
            stride: cols,
        }
    }
}

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

    // --- Step 2 tests: ColIter ---

    #[test]
    fn test_col_iter() {
        let m = Matrix::from_rows(&[&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0], &[7.0, 8.0, 9.0]]);
        let col0: Vec<f64> = m.col(0).copied().collect();
        assert_eq!(col0, vec![1.0, 4.0, 7.0]);
        let col2: Vec<f64> = m.col(2).copied().collect();
        assert_eq!(col2, vec![3.0, 6.0, 9.0]);
    }

    #[test]
    fn test_col_iter_size_hint() {
        let m = Matrix::from_rows(&[&[1.0, 2.0], &[3.0, 4.0], &[5.0, 6.0]]);
        let iter = m.col(0);
        assert_eq!(iter.len(), 3);
        assert_eq!(iter.size_hint(), (3, Some(3)));
    }

    #[test]
    fn test_col_single_column() {
        let m = Matrix::from_vec(vec![1.0, 2.0, 3.0], 3, 1);
        let col: Vec<f64> = m.col(0).copied().collect();
        assert_eq!(col, vec![1.0, 2.0, 3.0]);
    }

    // --- Step 8 tests: Views ---

    #[test]
    fn test_view_subregion() {
        let m = Matrix::from_fn(4, 4, |r, c| (r * 4 + c) as f64);
        let v = m.view(1..3, 1..3);
        assert_eq!(v.shape(), (2, 2));
        assert_eq!(v.get(0, 0), 5.0); // m[1][1]
        assert_eq!(v.get(0, 1), 6.0); // m[1][2]
        assert_eq!(v.get(1, 0), 9.0); // m[2][1]
        assert_eq!(v.get(1, 1), 10.0); // m[2][2]
    }

    #[test]
    fn test_view_row() {
        let m = Matrix::from_fn(3, 4, |r, c| (r * 4 + c) as f64);
        let v = m.view(1..2, 0..4);
        assert_eq!(v.row(0), &[4.0, 5.0, 6.0, 7.0]);
    }

    #[test]
    fn test_view_col() {
        let m = Matrix::from_fn(3, 4, |r, c| (r * 4 + c) as f64);
        let v = m.view(0..3, 1..3);
        let col: Vec<f64> = v.col(0).copied().collect();
        assert_eq!(col, vec![1.0, 5.0, 9.0]);
    }

    #[test]
    fn test_view_mut_add_assign() {
        let mut m = Matrix::zeros(4, 4);
        let vals = Matrix::from_fn(2, 2, |r, c| (r * 2 + c + 1) as f64);
        let vals_view = vals.as_view();

        {
            let mut vm = m.view_mut(1..3, 1..3);
            vm.add_assign(&vals_view);
        }

        assert_eq!(m.get(0, 0), 0.0);
        assert_eq!(m.get(1, 1), 1.0);
        assert_eq!(m.get(1, 2), 2.0);
        assert_eq!(m.get(2, 1), 3.0);
        assert_eq!(m.get(2, 2), 4.0);
        assert_eq!(m.get(3, 3), 0.0);
    }

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

    #[test]
    fn test_rows_range() {
        let m = Matrix::from_fn(4, 3, |r, c| (r * 3 + c) as f64);
        let v = m.rows_range(1, 3);
        assert_eq!(v.shape(), (2, 3));
        assert_eq!(v.get(0, 0), 3.0); // m[1][0]
        assert_eq!(v.get(1, 2), 8.0); // m[2][2]
    }

    #[test]
    fn test_overlapping_rows_range() {
        let m = Matrix::from_fn(4, 2, |r, c| (r * 2 + c) as f64);
        let v1 = m.rows_range(0, 3);
        let v2 = m.rows_range(1, 4);
        // Row 1 and 2 overlap
        assert_eq!(v1.get(1, 0), v2.get(0, 0)); // both are m[1][0]
        assert_eq!(v1.get(2, 1), v2.get(1, 1)); // both are m[2][1]
    }

    #[test]
    fn test_view_mut_set() {
        let mut m = Matrix::zeros(3, 3);
        {
            let mut vm = m.view_mut(0..2, 0..2);
            vm.set(0, 0, 42.0);
            vm.set(1, 1, 99.0);
        }
        assert_eq!(m.get(0, 0), 42.0);
        assert_eq!(m.get(1, 1), 99.0);
        assert_eq!(m.get(2, 2), 0.0);
    }
}