rosa 0.1.0

use crate::Matrix;

/// Which axis to operate along.
#[derive(Clone, Copy)]
pub enum Axis {
    /// Axis 0: reduce across rows → result has `cols` elements.
    Row,
    /// Axis 1: reduce across columns → result has `rows` elements.
    Col,
}

// --- Axis reductions (Step 4) ---

impl Matrix {
    /// Sum along axis.
    pub fn sum_axis(&self, axis: Axis) -> Vec<f64> {
        match axis {
            Axis::Col => {
                // Reduce across columns → one value per row
                (0..self.rows()).map(|r| self.row(r).iter().sum()).collect()
            }
            Axis::Row => {
                // Reduce across rows → one value per column
                let mut result = vec![0.0; self.cols()];
                for r in 0..self.rows() {
                    let row = self.row(r);
                    for (c, val) in row.iter().enumerate() {
                        result[c] += val;
                    }
                }
                result
            }
        }
    }

    /// Mean along axis.
    pub fn mean_axis(&self, axis: Axis) -> Vec<f64> {
        let sums = self.sum_axis(axis);
        let divisor = match axis {
            Axis::Col => self.cols() as f64,
            Axis::Row => self.rows() as f64,
        };
        sums.into_iter().map(|s| s / divisor).collect()
    }

    /// Max along axis.
    pub fn max_axis(&self, axis: Axis) -> Vec<f64> {
        match axis {
            Axis::Col => (0..self.rows())
                .map(|r| {
                    self.row(r)
                        .iter()
                        .copied()
                        .fold(f64::NEG_INFINITY, f64::max)
                })
                .collect(),
            Axis::Row => {
                let mut result = vec![f64::NEG_INFINITY; self.cols()];
                for r in 0..self.rows() {
                    let row = self.row(r);
                    for (c, &val) in row.iter().enumerate() {
                        result[c] = result[c].max(val);
                    }
                }
                result
            }
        }
    }

    /// Min along axis.
    pub fn min_axis(&self, axis: Axis) -> Vec<f64> {
        match axis {
            Axis::Col => (0..self.rows())
                .map(|r| self.row(r).iter().copied().fold(f64::INFINITY, f64::min))
                .collect(),
            Axis::Row => {
                let mut result = vec![f64::INFINITY; self.cols()];
                for r in 0..self.rows() {
                    let row = self.row(r);
                    for (c, &val) in row.iter().enumerate() {
                        result[c] = result[c].min(val);
                    }
                }
                result
            }
        }
    }

    /// Argmin along axis. Returns indices of minimum values.
    /// First occurrence wins (matches numpy behavior).
    pub fn argmin_axis(&self, axis: Axis) -> Vec<usize> {
        match axis {
            Axis::Col => (0..self.rows())
                .map(|r| {
                    let row = self.row(r);
                    let mut min_idx = 0;
                    let mut min_val = row[0];
                    for (i, &val) in row.iter().enumerate().skip(1) {
                        if val < min_val {
                            min_val = val;
                            min_idx = i;
                        }
                    }
                    min_idx
                })
                .collect(),
            Axis::Row => {
                let mut result = vec![0usize; self.cols()];
                let mut min_vals: Vec<f64> = self.row(0).to_vec();
                for r in 1..self.rows() {
                    let row = self.row(r);
                    for (c, &val) in row.iter().enumerate() {
                        if val < min_vals[c] {
                            min_vals[c] = val;
                            result[c] = r;
                        }
                    }
                }
                result
            }
        }
    }

    /// Argmax along axis. Returns indices of maximum values.
    /// First occurrence wins (matches numpy behavior).
    pub fn argmax_axis(&self, axis: Axis) -> Vec<usize> {
        match axis {
            Axis::Col => (0..self.rows())
                .map(|r| {
                    let row = self.row(r);
                    let mut max_idx = 0;
                    let mut max_val = row[0];
                    for (i, &val) in row.iter().enumerate().skip(1) {
                        if val > max_val {
                            max_val = val;
                            max_idx = i;
                        }
                    }
                    max_idx
                })
                .collect(),
            Axis::Row => {
                let mut result = vec![0usize; self.cols()];
                let mut max_vals: Vec<f64> = self.row(0).to_vec();
                for r in 1..self.rows() {
                    let row = self.row(r);
                    for (c, &val) in row.iter().enumerate() {
                        if val > max_vals[c] {
                            max_vals[c] = val;
                            result[c] = r;
                        }
                    }
                }
                result
            }
        }
    }

    /// Cumulative sum along axis. Returns Matrix of same shape.
    pub fn cumsum_axis(&self, axis: Axis) -> Matrix {
        let mut result = self.clone();
        match axis {
            Axis::Col => {
                // Cumsum across columns within each row
                for r in 0..self.rows() {
                    let row = result.row_mut(r);
                    for c in 1..self.cols() {
                        row[c] += row[c - 1];
                    }
                }
            }
            Axis::Row => {
                // Cumsum across rows within each column
                for r in 1..self.rows() {
                    for c in 0..self.cols() {
                        let prev = result.get(r - 1, c);
                        *result.get_mut(r, c) += prev;
                    }
                }
            }
        }
        result
    }
}

// --- Normalize (Step 5) ---

impl Matrix {
    /// Lp-normalize along axis.
    /// - `norm == 1.0`: L1 (divide by sum of abs values)
    /// - `norm == 2.0`: L2 (divide by sqrt of sum of squares)
    /// - `norm == f64::INFINITY`: max-norm (divide by max abs value)
    ///
    /// If the norm is zero for a row/col, leave it as zeros.
    pub fn normalize_axis(&self, axis: Axis, norm: f64) -> Matrix {
        let mut result = self.clone();

        match axis {
            Axis::Col => {
                // Normalize each row
                for r in 0..self.rows() {
                    let row = self.row(r);
                    let divisor = compute_norm(row.iter().copied(), norm);
                    if divisor > 0.0 {
                        let out_row = result.row_mut(r);
                        for val in out_row.iter_mut() {
                            *val /= divisor;
                        }
                    }
                }
            }
            Axis::Row => {
                // Normalize each column
                for c in 0..self.cols() {
                    let col_vals: Vec<f64> = self.col(c).copied().collect();
                    let divisor = compute_norm(col_vals.iter().copied(), norm);
                    if divisor > 0.0 {
                        for r in 0..self.rows() {
                            *result.get_mut(r, c) /= divisor;
                        }
                    }
                }
            }
        }
        result
    }
}

fn compute_norm(values: impl Iterator<Item = f64>, norm: f64) -> f64 {
    if norm == 1.0 {
        values.map(|x| x.abs()).sum()
    } else if norm == 2.0 {
        values.map(|x| x * x).sum::<f64>().sqrt()
    } else if norm.is_infinite() {
        values.map(|x| x.abs()).fold(0.0_f64, f64::max)
    } else {
        values
            .map(|x| x.abs().powf(norm))
            .sum::<f64>()
            .powf(1.0 / norm)
    }
}

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

    // --- Step 4 tests: Axis reductions ---

    fn test_matrix() -> Matrix {
        // 3×4 matrix with known values:
        // [ 1  2  3  4 ]
        // [ 5  6  7  8 ]
        // [ 9 10 11 12 ]
        Matrix::from_fn(3, 4, |r, c| (r * 4 + c + 1) as f64)
    }

    #[test]
    fn test_sum_axis_col() {
        let m = test_matrix();
        let sums = m.sum_axis(Axis::Col);
        assert_eq!(sums, vec![10.0, 26.0, 42.0]);
    }

    #[test]
    fn test_sum_axis_row() {
        let m = test_matrix();
        let sums = m.sum_axis(Axis::Row);
        assert_eq!(sums, vec![15.0, 18.0, 21.0, 24.0]);
    }

    #[test]
    fn test_mean_axis_col() {
        let m = test_matrix();
        let means = m.mean_axis(Axis::Col);
        assert_eq!(means, vec![2.5, 6.5, 10.5]);
    }

    #[test]
    fn test_mean_axis_row() {
        let m = test_matrix();
        let means = m.mean_axis(Axis::Row);
        assert_eq!(means, vec![5.0, 6.0, 7.0, 8.0]);
    }

    #[test]
    fn test_max_axis_col() {
        let m = test_matrix();
        assert_eq!(m.max_axis(Axis::Col), vec![4.0, 8.0, 12.0]);
    }

    #[test]
    fn test_max_axis_row() {
        let m = test_matrix();
        assert_eq!(m.max_axis(Axis::Row), vec![9.0, 10.0, 11.0, 12.0]);
    }

    #[test]
    fn test_min_axis_col() {
        let m = test_matrix();
        assert_eq!(m.min_axis(Axis::Col), vec![1.0, 5.0, 9.0]);
    }

    #[test]
    fn test_min_axis_row() {
        let m = test_matrix();
        assert_eq!(m.min_axis(Axis::Row), vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn test_argmax_axis_col() {
        let m = test_matrix();
        assert_eq!(m.argmax_axis(Axis::Col), vec![3, 3, 3]); // max at last col
    }

    #[test]
    fn test_argmax_axis_row() {
        let m = test_matrix();
        assert_eq!(m.argmax_axis(Axis::Row), vec![2, 2, 2, 2]); // max at last row
    }

    #[test]
    fn test_argmin_axis_col() {
        let m = test_matrix();
        assert_eq!(m.argmin_axis(Axis::Col), vec![0, 0, 0]); // min at first col
    }

    #[test]
    fn test_argmin_axis_row() {
        let m = test_matrix();
        assert_eq!(m.argmin_axis(Axis::Row), vec![0, 0, 0, 0]); // min at first row
    }

    #[test]
    fn test_argmax_ties_first_wins() {
        let m = Matrix::from_rows(&[&[1.0, 3.0, 3.0, 2.0]]);
        assert_eq!(m.argmax_axis(Axis::Col), vec![1]); // first occurrence of 3.0
    }

    #[test]
    fn test_cumsum_axis_col() {
        let m = Matrix::from_rows(&[&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0]]);
        let cs = m.cumsum_axis(Axis::Col);
        assert_eq!(cs.row(0), &[1.0, 3.0, 6.0]);
        assert_eq!(cs.row(1), &[4.0, 9.0, 15.0]);
    }

    #[test]
    fn test_cumsum_axis_row() {
        let m = Matrix::from_rows(&[&[1.0, 2.0], &[3.0, 4.0], &[5.0, 6.0]]);
        let cs = m.cumsum_axis(Axis::Row);
        assert_eq!(cs.row(0), &[1.0, 2.0]);
        assert_eq!(cs.row(1), &[4.0, 6.0]);
        assert_eq!(cs.row(2), &[9.0, 12.0]);
    }

    #[test]
    fn test_reductions_1x1() {
        let m = Matrix::from_vec(vec![42.0], 1, 1);
        assert_eq!(m.sum_axis(Axis::Col), vec![42.0]);
        assert_eq!(m.sum_axis(Axis::Row), vec![42.0]);
        assert_eq!(m.mean_axis(Axis::Col), vec![42.0]);
        assert_eq!(m.max_axis(Axis::Col), vec![42.0]);
        assert_eq!(m.min_axis(Axis::Row), vec![42.0]);
        assert_eq!(m.argmax_axis(Axis::Col), vec![0]);
        assert_eq!(m.argmin_axis(Axis::Row), vec![0]);
    }

    // --- Step 5 tests: normalize_axis ---

    #[test]
    fn test_normalize_l1_col() {
        let m = Matrix::from_rows(&[&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0]]);
        let n = m.normalize_axis(Axis::Col, 1.0);
        // Each row should sum to 1.0
        for r in 0..n.rows() {
            let sum: f64 = n.row(r).iter().sum();
            assert_relative_eq!(sum, 1.0, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_normalize_l1_row() {
        let m = Matrix::from_rows(&[&[1.0, 2.0], &[3.0, 4.0], &[5.0, 6.0]]);
        let n = m.normalize_axis(Axis::Row, 1.0);
        // Each column should sum to 1.0
        for c in 0..n.cols() {
            let sum: f64 = n.col(c).sum();
            assert_relative_eq!(sum, 1.0, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_normalize_l2_col() {
        let m = Matrix::from_rows(&[&[3.0, 4.0]]);
        let n = m.normalize_axis(Axis::Col, 2.0);
        // L2 norm of [3, 4] = 5, so [0.6, 0.8]
        assert_relative_eq!(n.get(0, 0), 0.6, epsilon = 1e-10);
        assert_relative_eq!(n.get(0, 1), 0.8, epsilon = 1e-10);
        let norm: f64 = n.row(0).iter().map(|x| x * x).sum::<f64>().sqrt();
        assert_relative_eq!(norm, 1.0, epsilon = 1e-10);
    }

    #[test]
    fn test_normalize_linf_col() {
        let m = Matrix::from_rows(&[&[1.0, -5.0, 3.0]]);
        let n = m.normalize_axis(Axis::Col, f64::INFINITY);
        // Max abs value is 5, so divide by 5
        assert_relative_eq!(n.get(0, 0), 0.2, epsilon = 1e-10);
        assert_relative_eq!(n.get(0, 1), -1.0, epsilon = 1e-10);
        assert_relative_eq!(n.get(0, 2), 0.6, epsilon = 1e-10);
    }

    #[test]
    fn test_normalize_zero_row() {
        let m = Matrix::from_rows(&[&[0.0, 0.0, 0.0], &[1.0, 2.0, 3.0]]);
        let n = m.normalize_axis(Axis::Col, 1.0);
        // All-zero row stays zero
        assert_eq!(n.row(0), &[0.0, 0.0, 0.0]);
        // Non-zero row normalizes
        let sum: f64 = n.row(1).iter().sum();
        assert_relative_eq!(sum, 1.0, epsilon = 1e-10);
    }

    #[test]
    fn test_normalize_zero_col() {
        let m = Matrix::from_rows(&[&[0.0, 1.0], &[0.0, 2.0]]);
        let n = m.normalize_axis(Axis::Row, 2.0);
        // First column is all zeros → stays zero
        assert_eq!(n.get(0, 0), 0.0);
        assert_eq!(n.get(1, 0), 0.0);
        // Second column normalizes
        let norm: f64 = n.col(1).map(|x| x * x).sum::<f64>().sqrt();
        assert_relative_eq!(norm, 1.0, epsilon = 1e-10);
    }
}