/// Dense matrix stored row-major for one-dimensional operator factors.
/// Values are i64 fixed-point.
#[derive(Debug, Clone)]
pub struct DenseMatrix {
    pub rows: usize,
    pub cols: usize,
    pub data: Vec<i64>,
}

impl DenseMatrix {
    pub fn zeros(rows: usize, cols: usize) -> Self {
        Self {
            rows,
            cols,
            data: vec![0; rows * cols],
        }
    }

    pub fn identity(n: usize, one: i64) -> Self {
        let mut m = Self::zeros(n, n);
        for i in 0..n {
            m.set(i, i, one);
        }
        m
    }

    #[inline]
    pub fn get(&self, row: usize, col: usize) -> i64 {
        self.data[row * self.cols + col]
    }

    #[inline]
    pub fn set(&mut self, row: usize, col: usize, val: i64) {
        self.data[row * self.cols + col] = val;
    }

    /// LU decomposition in-place. Returns (L, U) as separate matrices.
    /// L has 1s on the diagonal (unit lower triangular, where "1" is
    /// the fixed-point representation of one).
    pub fn lu_decompose(
        &self,
        one: i64,
        fp_mul: impl Fn(i64, i64) -> i64,
        fp_div: impl Fn(i64, i64) -> i64,
    ) -> (DenseMatrix, DenseMatrix) {
        assert_eq!(self.rows, self.cols);
        let n = self.rows;
        let mut lu = self.clone();

        for k in 0..n {
            let pivot = lu.get(k, k);
            for i in (k + 1)..n {
                let factor = fp_div(lu.get(i, k), pivot);
                lu.set(i, k, factor);
                for j in (k + 1)..n {
                    let val = lu.get(i, j) - fp_mul(factor, lu.get(k, j));
                    lu.set(i, j, val);
                }
            }
        }

        let mut lower = DenseMatrix::zeros(n, n);
        let mut upper = DenseMatrix::zeros(n, n);
        for i in 0..n {
            for j in 0..i {
                lower.set(i, j, lu.get(i, j));
            }
            lower.set(i, i, one);
            for j in i..n {
                upper.set(i, j, lu.get(i, j));
            }
        }

        (lower, upper)
    }

    /// Invert a unit lower triangular matrix via forward substitution.
    pub fn invert_unit_lower(&self, one: i64, fp_mul: impl Fn(i64, i64) -> i64) -> DenseMatrix {
        assert_eq!(self.rows, self.cols);
        let n = self.rows;
        let mut inv = DenseMatrix::identity(n, one);

        for col in 0..n {
            for row in (col + 1)..n {
                let mut sum = 0i64;
                for k in col..row {
                    sum = sum.wrapping_add(fp_mul(self.get(row, k), inv.get(k, col)));
                }
                inv.set(row, col, -sum);
            }
        }
        inv
    }

    /// Invert an upper triangular matrix via back substitution.
    pub fn invert_upper(
        &self,
        one: i64,
        fp_mul: impl Fn(i64, i64) -> i64,
        fp_div: impl Fn(i64, i64) -> i64,
    ) -> DenseMatrix {
        assert_eq!(self.rows, self.cols);
        let n = self.rows;
        let mut inv = DenseMatrix::zeros(n, n);

        for col in (0..n).rev() {
            inv.set(col, col, fp_div(one, self.get(col, col)));
            for row in (0..col).rev() {
                let mut sum = 0i64;
                for k in (row + 1)..=col {
                    sum = sum.wrapping_add(fp_mul(self.get(row, k), inv.get(k, col)));
                }
                inv.set(row, col, fp_div(-sum, self.get(row, row)));
            }
        }
        inv
    }
}

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

    #[test]
    fn identity_diag() {
        let fp = FixedPoint::new(16);
        let m = DenseMatrix::identity(3, fp.one());
        for i in 0..3 {
            assert_eq!(m.get(i, i), fp.one());
            for j in 0..3 {
                if i != j {
                    assert_eq!(m.get(i, j), 0);
                }
            }
        }
    }

    #[test]
    fn lu_2x2() {
        let fp = FixedPoint::new(16);
        let mut m = DenseMatrix::zeros(2, 2);
        m.set(0, 0, fp.from_f64(4.0));
        m.set(0, 1, fp.from_f64(3.0));
        m.set(1, 0, fp.from_f64(6.0));
        m.set(1, 1, fp.from_f64(3.0));

        let (l, u) = m.lu_decompose(
            fp.one(),
            |a, b| fp.mul(a, b),
            |a, b| fp.div(a, b),
        );

        // L should be [[1, 0], [1.5, 1]]
        assert_eq!(l.get(0, 0), fp.one());
        assert!((fp.to_f64(l.get(1, 0)) - 1.5).abs() < 0.01);
        assert_eq!(l.get(1, 1), fp.one());

        // U should be [[4, 3], [0, -1.5]]
        assert!((fp.to_f64(u.get(0, 0)) - 4.0).abs() < 0.01);
        assert!((fp.to_f64(u.get(1, 1)) + 1.5).abs() < 0.01);
    }
}
