//! Cholesky factorisation of a joint precision matrix.
//!
//! Every question the joint answers is a *bilinear form* in the precision
//! matrix's inverse — the variance of a contrast is `c^T L^-1 c`, and the
//! covariance of two contrasts is `c^T L^-1 a`. None of them wants `L^-1 c`
//! itself, which is what makes the shape here worth stating explicitly.
//!
//! Writing the precision as `A = L L^T`,
//!
//! ```text
//! c^T A^-1 a = c^T L^-T L^-1 a = (L^-1 c) . (L^-1 a)
//! ```
//!
//! so a single forward substitution per contrast answers everything, and the
//! back substitution a general solve would do is wasted work. That halves the
//! cost of a query, and it removes a failure mode: a variance computed as
//! `c . (A^-1 c)` is a difference of products that can round to a small
//! negative number, where the same quantity as `|L^-1 c|^2` is a sum of
//! squares and cannot.
//!
//! Factorising is `O(n^3)` and whitening is `O(n^2)`, so the split also
//! matters structurally: the expensive half depends only on the fit, and is
//! shared across every query a [`Joint`](crate::Joint) answers.

/// A factorised symmetric positive-definite matrix, reusable across queries.
pub(crate) struct Cholesky {
    /// Lower triangle of `L`, row-major `n * n`. The upper triangle is
    /// leftover scratch from the factorisation and is never read.
    l: Vec<f64>,
    n: usize,
}

impl Cholesky {
    /// Factorise `a` (row-major, `n * n`, symmetric) into `L L^T`.
    ///
    /// `a` is consumed as scratch.
    ///
    /// Returns `None` if the matrix is not positive-definite, which for a
    /// precision matrix means the model is improper — a competitor with
    /// neither a proper prior nor any evidence.
    pub(crate) fn factor(mut a: Vec<f64>, n: usize) -> Option<Self> {
        debug_assert_eq!(a.len(), n * n);

        for j in 0..n {
            let mut d = a[j * n + j];
            for k in 0..j {
                d -= a[j * n + k] * a[j * n + k];
            }
            // Explicit rather than `!(d > 0.0)`: a NaN pivot must fail here
            // too, and a negated comparison would let it through as "not
            // positive".
            if d.is_nan() || d <= 0.0 {
                return None;
            }
            let d = d.sqrt();
            a[j * n + j] = d;

            for i in j + 1..n {
                let mut s = a[i * n + j];
                for k in 0..j {
                    s -= a[i * n + k] * a[j * n + k];
                }
                a[i * n + j] = s / d;
            }
        }

        Some(Self { l: a, n })
    }

    /// Whiten a contrast: `y = L^-1 b`.
    ///
    /// The point of the result is the dot product, not the vector: for two
    /// contrasts `b` and `b'`, `y . y'` is `b^T A^-1 b'`. See the module docs.
    pub(crate) fn whiten(&self, b: &[f64]) -> Vec<f64> {
        debug_assert_eq!(b.len(), self.n);
        let n = self.n;
        let mut y = b.to_vec();
        for i in 0..n {
            // Folded from `y[i]` rather than summed and subtracted once, so the
            // accumulation order matches a plain substitution loop exactly.
            let row = &self.l[i * n..i * n + i];
            let s = row
                .iter()
                .zip(&y[..i])
                .fold(y[i], |acc, (l, v)| acc - l * v);
            y[i] = s / self.l[i * n + i];
        }
        y
    }
}

/// `b^T A^-1 b'`, given the two whitened contrasts.
pub(crate) fn bilinear(y: &[f64], y_prime: &[f64]) -> f64 {
    y.iter().zip(y_prime).map(|(a, b)| a * b).sum()
}

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

    /// `[[4, 1], [1, 3]] z = [1, 2]` has `z = [1/11, 7/11]`, so the quadratic
    /// form `b^T A^-1 b` is `1 * 1/11 + 2 * 7/11 = 15/11`.
    #[test]
    fn reproduces_a_known_quadratic_form() {
        let c = Cholesky::factor(vec![4.0, 1.0, 1.0, 3.0], 2).unwrap();
        let y = c.whiten(&[1.0, 2.0]);
        assert!((bilinear(&y, &y) - 15.0 / 11.0).abs() < 1e-12);
    }

    /// Whitening `e_i` recovers the inverse's diagonal, which is the variance
    /// of a single variable.
    #[test]
    fn recovers_the_inverse_diagonal() {
        // A = [[2, -1, 0], [-1, 2, -1], [0, -1, 2]]; inverse diagonal is
        // [0.75, 1.0, 0.75].
        let a = vec![2.0, -1.0, 0.0, -1.0, 2.0, -1.0, 0.0, -1.0, 2.0];
        let c = Cholesky::factor(a, 3).unwrap();
        for (i, expected) in [0.75, 1.0, 0.75].into_iter().enumerate() {
            let mut e = vec![0.0; 3];
            e[i] = 1.0;
            let y = c.whiten(&e);
            assert!((bilinear(&y, &y) - expected).abs() < 1e-12, "row {i}");
        }
    }

    /// The off-diagonal bilinear form is symmetric and matches the inverse.
    #[test]
    fn recovers_an_off_diagonal_covariance() {
        // Same A; (A^-1)_{0,1} = 0.5.
        let a = vec![2.0, -1.0, 0.0, -1.0, 2.0, -1.0, 0.0, -1.0, 2.0];
        let c = Cholesky::factor(a, 3).unwrap();
        let y0 = c.whiten(&[1.0, 0.0, 0.0]);
        let y1 = c.whiten(&[0.0, 1.0, 0.0]);
        assert!((bilinear(&y0, &y1) - 0.5).abs() < 1e-12);
        assert!((bilinear(&y1, &y0) - 0.5).abs() < 1e-12);
    }

    /// A variance can never come out negative, because it is a sum of squares.
    #[test]
    fn a_quadratic_form_is_never_negative() {
        let a = vec![1e12, 1e12 - 1.0, 1e12 - 1.0, 1e12];
        let c = Cholesky::factor(a, 2).unwrap();
        let y = c.whiten(&[1.0, -1.0]);
        assert!(bilinear(&y, &y) >= 0.0);
    }

    #[test]
    fn rejects_a_non_positive_definite_matrix() {
        // Singular: the second row is a multiple of the first.
        assert!(Cholesky::factor(vec![1.0, 2.0, 2.0, 4.0], 2).is_none());
    }
}