slain-num 0.4.0

use crate::simd;

#[derive(Clone, Debug, PartialEq)]
pub struct Mat {
    rows: usize,
    cols: usize,
    data: Vec<f64>,
}

impl Mat {
    #[must_use]
    pub fn zeros(rows: usize, cols: usize) -> Self {
        Self { rows, cols, data: vec![0.0; rows * cols] }
    }

    #[must_use]
    pub fn identity(n: usize) -> Self {
        let mut m = Self::zeros(n, n);
        for i in 0..n {
            m[(i, i)] = 1.0;
        }
        m
    }

    #[must_use]
    pub fn from_vec(rows: usize, cols: usize, data: Vec<f64>) -> Self {
        assert_eq!(data.len(), rows * cols, "matrix data length");
        Self { rows, cols, data }
    }

    #[must_use]
    pub fn rows(&self) -> usize {
        self.rows
    }

    #[must_use]
    pub fn cols(&self) -> usize {
        self.cols
    }

    #[must_use]
    pub fn row(&self, r: usize) -> &[f64] {
        &self.data[r * self.cols..(r + 1) * self.cols]
    }

    pub fn row_mut(&mut self, r: usize) -> &mut [f64] {
        &mut self.data[r * self.cols..(r + 1) * self.cols]
    }

    #[must_use]
    pub fn as_slice(&self) -> &[f64] {
        &self.data
    }

    #[must_use]
    pub fn transpose(&self) -> Self {
        let mut t = Self::zeros(self.cols, self.rows);
        for r in 0..self.rows {
            for c in 0..self.cols {
                t.data[c * self.rows + r] = self.data[r * self.cols + c];
            }
        }
        t
    }

    #[must_use]
    pub fn matmul(&self, b: &Self) -> Self {
        assert_eq!(self.cols, b.rows, "inner dimensions");
        let mut out = Self::zeros(self.rows, b.cols);
        for r in 0..self.rows {
            let orow = &mut out.data[r * b.cols..(r + 1) * b.cols];
            for (k, &a) in self.row(r).iter().enumerate() {
                if a != 0.0 {
                    simd::axpy(a, b.row(k), orow);
                }
            }
        }
        out
    }

    #[must_use]
    pub fn mul_vec(&self, v: &[f64]) -> Vec<f64> {
        (0..self.rows).map(|r| simd::dot(self.row(r), v)).collect()
    }

    #[must_use]
    pub fn gram(&self) -> Self {
        let t = self.transpose();
        let n = self.cols;
        let mut g = Self::zeros(n, n);
        for i in 0..n {
            for j in i..n {
                let v = simd::dot(t.row(i), t.row(j));
                g.data[i * n + j] = v;
                g.data[j * n + i] = v;
            }
        }
        g
    }

    #[must_use]
    pub fn t_mul_vec(&self, v: &[f64]) -> Vec<f64> {
        let mut out = vec![0.0; self.cols];
        for (r, &k) in v.iter().enumerate().take(self.rows) {
            simd::axpy(k, self.row(r), &mut out);
        }
        out
    }
}

impl std::ops::Index<(usize, usize)> for Mat {
    type Output = f64;
    fn index(&self, (r, c): (usize, usize)) -> &f64 {
        &self.data[r * self.cols + c]
    }
}

impl std::ops::IndexMut<(usize, usize)> for Mat {
    fn index_mut(&mut self, (r, c): (usize, usize)) -> &mut f64 {
        &mut self.data[r * self.cols + c]
    }
}

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

    fn filled(rows: usize, cols: usize, seed: f64) -> Mat {
        Mat::from_vec(rows, cols, (0..rows * cols).map(|i| ((i as f64 + seed) * 0.713).sin()).collect())
    }

    #[test]
    fn products_match_naive() {
        let (a, b) = (filled(7, 13, 1.0), filled(13, 9, 2.0));
        let c = a.matmul(&b);
        for r in 0..7 {
            for k in 0..9 {
                let want: f64 = (0..13).map(|i| a[(r, i)] * b[(i, k)]).sum();
                assert!((c[(r, k)] - want).abs() < 1e-12);
            }
        }
        let g = a.gram();
        let want = a.transpose().matmul(&a);
        assert!(g.as_slice().iter().zip(want.as_slice()).all(|(x, y)| (x - y).abs() < 1e-12));
        let v: Vec<f64> = (0..7).map(|i| i as f64).collect();
        let tv = a.t_mul_vec(&v);
        let want = a.transpose().mul_vec(&v);
        assert!(tv.iter().zip(&want).all(|(x, y)| (x - y).abs() < 1e-12));
    }
}