use std::default::Default;
use std::ops::{Add, AddAssign, Div, Mul};

#[cfg(test)]
mod tests;

/// A representation of a standard Matrix of size `r * c`
///
/// # Examples
///
/// ```
/// # // no idea why I need this but eh, well..
/// # use math_concept::lalg::matrix::Matrix;
/// let mut m = Matrix::new(2, 2, vec![0; 4]);
///
/// assert_eq!(m.at(0, 0), 0);
///
/// *m.at_mut(1, 1) = 5;
/// assert_eq!(m.at(1, 1), 5);
/// ```
/// ```
/// # use math_concept::lalg::matrix::Matrix;
/// let m = Matrix::from_vec(1, 2, &vec![9, 8]);
///
/// assert_eq!(&vec![9, 8], m.data());
/// ```
pub struct Matrix<T> {
    r: usize,
    c: usize,
    buf: Vec<T>,
}

impl<T> Matrix<T>
where
    T: Copy + Add<Output = T> + Div<Output = T>,
{
    /// Constructs a new `Matrix<T>`, taking ownership of `v` and using it
    /// as the buffer.
    pub fn new(r: usize, c: usize, v: Vec<T>) -> Matrix<T> {
        assert_eq!(r * c, v.len(), "Matrix dimensions and data size differ");
        Matrix { r, c, buf: v }
    }

    /// Constructs a new `Matrix<T>`, using `v`
    /// as the buffer without taking ownership of it.
    pub fn from_vec(r: usize, c: usize, v: &Vec<T>) -> Matrix<T> {
        assert_eq!(r * c, v.len(), "Matrix dimensions and data size differ");
        Matrix::new(r, c, v.clone())
    }

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

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

    /// Returns the element at the `r`-th row and `c`-th column,
    /// provided those are within bounds.
    pub fn at(&self, r: usize, c: usize) -> T {
        assert!(r < self.r && c < self.c, "Out of bounds");
        self.buf[r * self.c + c]
    }

    /// Returns a mutable reference to the element at the `r`-th row
    /// and `c`-th column, provided those are within bounds.
    pub fn at_mut(&mut self, r: usize, c: usize) -> &mut T {
        assert!(r < self.r && c < self.c, "Out of bounds");
        &mut self.buf[r * self.c + c]
    }

    // TODO: implement this
    // pub fn determinant(&self) -> T;

    /// Returns a reference to the matrix data vector
    pub fn data(&self) -> &Vec<T> {
        &self.buf
    }

    /// Maps over the elements of the matrix.
    ///
    /// ## Example
    /// ```
    /// # // no idea why I need this but eh, well..
    /// # use math_concept::lalg::matrix::Matrix;
    /// let m = Matrix::new(2, 2, vec![5; 4]);
    ///
    /// let mapped = m.map(|&x| x * 8);
    /// assert_eq!(mapped.data(), &vec![40; 4]);
    /// ```
    pub fn map<U, F>(&self, f: F) -> Matrix<U>
    where
        U: Copy + Add<Output = U> + Div<Output = U>,
        F: Fn(&T) -> U,
    {
        let new_v = self.buf.iter().map(f).collect();

        Matrix::new(self.r, self.c, new_v)
    }

    /// Produces the diagonal of an 'm x m' matrix.
    /// Panics otherwise.
    ///
    /// ## Example
    /// ```
    /// # // no idea why I need this but eh, well..
    /// # use math_concept::lalg::matrix::Matrix;
    /// let m = Matrix::new(2, 2, vec![1, 2, 3, 4]);
    ///
    /// let diag = m.diagonal();
    /// assert_eq!(diag, vec![1, 4]);
    /// ```
    pub fn diagonal(&self) -> Vec<T> {
        self.assert_square();

        self.buf
            .iter()
            .zip(0..self.buf.len())
            .filter(|(_, i)| i % self.c == i / self.r)
            .map(|(&elem, _)| elem)
            .collect()
    }

    fn assert_square(&self) {
        assert_eq!(self.r, self.c, "Matrix is not square");
    }
}

impl<'a, T> Add<&'a Matrix<T>> for &'a Matrix<T>
where
    T: Copy + Add<Output = T> + Div<Output = T>,
{
    type Output = Matrix<T>;

    fn add(self, other: Self) -> Self::Output {
        assert!(
            self.r == other.rows() && self.c == other.cols(),
            "Matrices are not of the same size"
        );

        let new_v = self
            .buf
            .iter()
            .zip(other.data())
            .map(|(&x, &y)| x + y)
            .collect();

        Matrix::new(self.r, self.c, new_v)
    }
}

impl<'a, T> Mul<&'a Matrix<T>> for &'a Matrix<T>
where
    T: Copy + Add<Output = T> + AddAssign<T> + Mul<Output = T> + Div<Output = T> + Default,
{
    type Output = Matrix<T>;

    fn mul(self, other: Self) -> Self::Output {
        assert_eq!(self.c, other.r);

        let mut v: Vec<T> = vec![Default::default(); self.r * other.c];

        for i in 0..self.r {
            for j in 0..other.c {
                for k in 0..self.c {
                    v[i * self.r + j] += self.buf[i * self.c + k] * other.buf[k * other.c + j];
                }
            }
        }

        Matrix::new(self.r, other.c, v)
    }
}