//! Sparse 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.
//!
//! # Why this is sparse (#52)
//!
//! A time-expanded joint is *extremely* sparse and gets sparser as the history
//! grows: a row couples only to its own previous and next appearance through
//! the drift link, and to whoever co-appeared in its slice. Measured on a
//! 76-slice, 988-duel, 200-competitor history: `n = 1976`, `nnz = 7504`,
//! **0.19% dense**.
//!
//! This used to store all `n^2` entries and run a dense `O(n^3)` factorisation
//! over them. Two measurements decided the replacement:
//!
//! - **Ordering alone does nothing to a dense factorisation.** Its inner loops
//!   run over every `k` whether or not the entry is zero. A 700x700 banded
//!   matrix at 0.43% density factorised in 30.196 ms in band order and
//!   29.544 ms under a scramble that destroyed the band — identical, as the
//!   flop count says it must be. Fill-reducing order is worth nothing until
//!   the factorisation skips zeros.
//! - **Together they are worth four orders of magnitude.** On that `n = 1976`
//!   fixture, against `n^3/3 = 2.572e9` flops dense: sparse in the natural
//!   order needs `5.597e7` (46x better), and sparse under an AMD fill-reducing
//!   order needs `8.656e4` — **29,710x**. AMD is worth 646x *on top of*
//!   sparsity and nothing without it.
//!
//! Natural ordering fills in badly here for the reason #52 predicted: a
//! competitor who appears in slice 0 and not again until slice 75 creates a
//! drift link spanning nearly the whole matrix. `nnz(L)` is 292,437 under the
//! natural order against 11,583 under AMD, from an `A` with 7,504.
//!
//! The ordering comes from `feral-amd`. The factorisation is the up-looking
//! sparse Cholesky of Davis's *Direct Methods for Sparse Linear Systems*,
//! written here rather than taken from a crate: the sparse solvers on
//! crates.io either pull SIMD dispatch (`faer`, and `feral` itself, both
//! through `pulp`), which would make results differ between an AVX-512 host
//! and an AVX2 one — the same class of drift the `libm`-over-`std` decision
//! was made to avoid — or are LGPL, or disclaim fill-reduction in their own
//! docs.

use std::collections::BTreeMap;

/// A symmetric matrix accumulated entry by entry, before factorisation.
///
/// A `BTreeMap` rather than a hash map because the iteration order becomes the
/// factorisation's summation order, and a hash map's order varies per process.
/// `tests/cross_process_determinism.rs` exists because that has bitten before.
#[derive(Default)]
pub(crate) struct SymmetricBuilder {
    entries: BTreeMap<(usize, usize), f64>,
}

impl SymmetricBuilder {
    pub(crate) fn new() -> Self {
        Self::default()
    }

    /// Add `value` to entry `(row, col)`. Both triangles must be supplied.
    pub(crate) fn add(&mut self, row: usize, col: usize, value: f64) {
        *self.entries.entry((row, col)).or_insert(0.0) += value;
    }

    /// The `(row, col)` positions that hold a nonzero. For the #52 measurement.
    #[cfg(feature = "measure-sparsity")]
    pub(crate) fn pattern(&self) -> impl Iterator<Item = (usize, usize)> + '_ {
        self.entries
            .iter()
            .filter(|(_, v)| **v != 0.0)
            .map(|(&rc, _)| rc)
    }
}

/// A factorised symmetric positive-definite matrix, reusable across queries.
pub(crate) struct Cholesky {
    n: usize,
    /// `inv[old] = new`: where each original row sits after the AMD reorder.
    inv: Vec<usize>,
    /// `L` in compressed-column form, permuted. Within a column the diagonal
    /// is first and the rest ascend by row.
    col_ptr: Vec<usize>,
    row_idx: Vec<usize>,
    val: Vec<f64>,
}

impl Cholesky {
    /// Factorise the accumulated matrix into `L L^T`, under a fill-reducing
    /// permutation.
    ///
    /// 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 — or if the ordering fails.
    pub(crate) fn factor(built: SymmetricBuilder, n: usize) -> Option<Self> {
        if n == 0 {
            return Some(Self {
                n: 0,
                inv: Vec::new(),
                col_ptr: vec![0],
                row_idx: Vec::new(),
                val: Vec::new(),
            });
        }

        let inv = Self::amd_permutation(n, &built)?;

        // Upper triangle of the permuted matrix, column-major: column `c`
        // holds the rows `r <= c`. Exactly one of a symmetric pair survives
        // the `r <= c` filter, so nothing is double-counted.
        let mut cols: Vec<Vec<(usize, f64)>> = vec![Vec::new(); n];
        for (&(old_r, old_c), &v) in &built.entries {
            if v == 0.0 {
                continue;
            }
            let (r, c) = (inv[old_r], inv[old_c]);
            if r <= c {
                cols[c].push((r, v));
            }
        }
        let mut a_ptr = Vec::with_capacity(n + 1);
        let mut a_row = Vec::new();
        let mut a_val = Vec::new();
        a_ptr.push(0usize);
        for col in &mut cols {
            col.sort_unstable_by_key(|&(r, _)| r);
            for &(r, v) in col.iter() {
                a_row.push(r);
                a_val.push(v);
            }
            a_ptr.push(a_row.len());
        }

        let parent = Self::etree(n, &a_ptr, &a_row);

        // Symbolic pass: how many entries each column of L will hold. Running
        // `ereach` per column costs O(nnz(L)) in total, which is the same order
        // as the numeric pass it sizes.
        let mut counts = vec![0usize; n];
        let mut stack = vec![0usize; n];
        let mut mark = vec![false; n];
        for k in 0..n {
            let top = Self::ereach(k, &a_ptr, &a_row, &parent, &mut stack, &mut mark);
            for &i in &stack[top..] {
                counts[i] += 1;
            }
            counts[k] += 1; // the diagonal
        }

        let mut col_ptr = Vec::with_capacity(n + 1);
        col_ptr.push(0usize);
        for &c in &counts {
            col_ptr.push(col_ptr[col_ptr.len() - 1] + c);
        }
        let nnz = col_ptr[n];
        let mut row_idx = vec![0usize; nnz];
        let mut val = vec![0.0f64; nnz];

        // `next[i]` is the slot column `i` will fill next. Column `i`'s
        // diagonal lands first, at `col_ptr[i]`, because nothing is written to
        // a column before its own iteration.
        let mut next: Vec<usize> = col_ptr[..n].to_vec();
        let mut x = vec![0.0f64; n];

        for k in 0..n {
            let top = Self::ereach(k, &a_ptr, &a_row, &parent, &mut stack, &mut mark);

            for p in a_ptr[k]..a_ptr[k + 1] {
                if a_row[p] <= k {
                    x[a_row[p]] = a_val[p];
                }
            }

            let mut d = x[k];
            x[k] = 0.0;

            for &i in &stack[top..] {
                let lki = x[i] / val[col_ptr[i]];
                x[i] = 0.0;
                for p in col_ptr[i] + 1..next[i] {
                    x[row_idx[p]] -= val[p] * lki;
                }
                d -= lki * lki;
                let p = next[i];
                next[i] += 1;
                row_idx[p] = k;
                val[p] = lki;
            }

            // 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 p = next[k];
            next[k] += 1;
            row_idx[p] = k;
            val[p] = d.sqrt();
        }

        Some(Self {
            n,
            inv,
            col_ptr,
            row_idx,
            val,
        })
    }

    /// AMD fill-reducing order, as `inv[old] = new`.
    fn amd_permutation(n: usize, built: &SymmetricBuilder) -> Option<Vec<usize>> {
        let mut cols: Vec<Vec<i32>> = vec![Vec::new(); n];
        for (&(r, c), &v) in &built.entries {
            if v != 0.0 {
                cols[c].push(i32::try_from(r).ok()?);
            }
        }
        let mut col_ptr = Vec::with_capacity(n + 1);
        let mut row_idx = Vec::new();
        col_ptr.push(0i32);
        for (j, col) in cols.iter_mut().enumerate() {
            col.push(i32::try_from(j).ok()?);
            col.sort_unstable();
            col.dedup();
            row_idx.extend_from_slice(col);
            col_ptr.push(i32::try_from(row_idx.len()).ok()?);
        }

        let pattern = feral_amd::CscPattern::new(n, &col_ptr, &row_idx)?;
        // `perm[new] = old`; we want the inverse.
        let perm = feral_amd::amd_order(&pattern).ok()?;
        let mut inv = vec![0usize; n];
        for (new, &old) in perm.iter().enumerate() {
            inv[usize::try_from(old).ok()?] = new;
        }
        Some(inv)
    }

    /// Elimination tree of the upper-triangular pattern. `usize::MAX` is "no
    /// parent", i.e. a root.
    fn etree(n: usize, col_ptr: &[usize], row_idx: &[usize]) -> Vec<usize> {
        let mut parent = vec![usize::MAX; n];
        let mut ancestor = vec![usize::MAX; n];
        for k in 0..n {
            for &row in &row_idx[col_ptr[k]..col_ptr[k + 1]] {
                let mut i = row;
                while i != usize::MAX && i < k {
                    let next = ancestor[i];
                    ancestor[i] = k;
                    if next == usize::MAX {
                        parent[i] = k;
                    }
                    i = next;
                }
            }
        }
        parent
    }

    /// Nonzero pattern of row `k` of `L`, written into `stack[top..n]` in
    /// topological order. Returns `top`.
    ///
    /// `stack` is used from both ends — a scratch region from `0` while walking
    /// each path up the tree, and the result from `n` downwards. They cannot
    /// collide because every node is pushed at most once across the whole call.
    fn ereach(
        k: usize,
        col_ptr: &[usize],
        row_idx: &[usize],
        parent: &[usize],
        stack: &mut [usize],
        mark: &mut [bool],
    ) -> usize {
        let n = mark.len();
        let mut top = n;
        mark[k] = true;
        for &row in &row_idx[col_ptr[k]..col_ptr[k + 1]] {
            let mut i = row;
            if i > k {
                continue;
            }
            let mut len = 0usize;
            while i != usize::MAX && !mark[i] {
                stack[len] = i;
                len += 1;
                mark[i] = true;
                i = parent[i];
            }
            // Reverse the path onto the output end, so the result stays in
            // topological order overall.
            while len > 0 {
                len -= 1;
                top -= 1;
                stack[top] = stack[len];
            }
        }
        for &i in &stack[top..] {
            mark[i] = false;
        }
        mark[k] = false;
        top
    }

    /// Whiten a contrast: `y = L^-1 P 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.
    ///
    /// The result is in the permuted order, and stays there — a dot product
    /// does not care, as long as both operands were permuted the same way.
    pub(crate) fn whiten(&self, b: &[f64]) -> Vec<f64> {
        debug_assert_eq!(b.len(), self.n);
        let n = self.n;
        let mut y = vec![0.0f64; n];
        for (old, &v) in b.iter().enumerate() {
            y[self.inv[old]] = v;
        }
        for j in 0..n {
            y[j] /= self.val[self.col_ptr[j]];
            let yj = y[j];
            for p in self.col_ptr[j] + 1..self.col_ptr[j + 1] {
                y[self.row_idx[p]] -= self.val[p] * yj;
            }
        }
        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::*;

    /// Factorise a dense row-major matrix, for the goldens below.
    fn dense(a: &[f64], n: usize) -> Option<Cholesky> {
        let mut b = SymmetricBuilder::new();
        for i in 0..n {
            for j in 0..n {
                if a[i * n + j] != 0.0 {
                    b.add(i, j, a[i * n + j]);
                }
            }
        }
        Cholesky::factor(b, n)
    }

    /// `[[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 = dense(&[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 = [2.0, -1.0, 0.0, -1.0, 2.0, -1.0, 0.0, -1.0, 2.0];
        let c = dense(&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 = [2.0, -1.0, 0.0, -1.0, 2.0, -1.0, 0.0, -1.0, 2.0];
        let c = dense(&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 = [1e12, 1e12 - 1.0, 1e12 - 1.0, 1e12];
        let c = dense(&a, 2).unwrap();
        let y = c.whiten(&[1.0, -1.0]);
        assert!(bilinear(&y, &y) >= 0.0);
    }

    /// Against an independent dense reference, on random sparse SPD matrices.
    ///
    /// The goldens above are 2x2 and 3x3 — small enough that AMD does nothing
    /// and no fill-in occurs, so they cannot catch a symbolic-pass bug. This
    /// builds matrices big enough to permute and fill in, and checks every
    /// bilinear form against a textbook dense factorisation of the *same*
    /// matrix in its original order.
    #[test]
    fn agrees_with_a_dense_reference_on_random_sparse_systems() {
        /// Dense Cholesky and quadratic form, deliberately naive: this is the
        /// reference, so it must not share code with what it is checking.
        fn dense_quadratic_form(a: &[f64], n: usize, b: &[f64], c: &[f64]) -> f64 {
            let mut l = a.to_vec();
            for j in 0..n {
                let mut d = l[j * n + j];
                for k in 0..j {
                    d -= l[j * n + k] * l[j * n + k];
                }
                let d = d.sqrt();
                l[j * n + j] = d;
                for i in j + 1..n {
                    let mut sum = l[i * n + j];
                    for k in 0..j {
                        sum -= l[i * n + k] * l[j * n + k];
                    }
                    l[i * n + j] = sum / d;
                }
            }
            let solve = |rhs: &[f64]| -> Vec<f64> {
                let mut y = rhs.to_vec();
                for i in 0..n {
                    for k in 0..i {
                        y[i] -= l[i * n + k] * y[k];
                    }
                    y[i] /= l[i * n + i];
                }
                y
            };
            let (yb, yc) = (solve(b), solve(c));
            yb.iter().zip(&yc).map(|(x, y)| x * y).sum()
        }

        // A cheap deterministic generator; no dependency, and reproducible.
        let mut seed = 0x2545_F491_4F6C_DD1Du64;
        let mut rand = move || {
            seed ^= seed << 13;
            seed ^= seed >> 7;
            seed ^= seed << 17;
            (seed >> 11) as f64 / (1u64 << 53) as f64
        };

        for n in [7usize, 23, 60] {
            let mut a = vec![0.0f64; n * n];
            // A chain plus scattered long-range couplings: the shape of a
            // time-expanded joint, where a competitor's drift link can span
            // the whole matrix.
            for i in 0..n {
                a[i * n + i] = 4.0 + rand();
                if i + 1 < n {
                    let v = -(0.5 + rand() * 0.5);
                    a[i * n + i + 1] = v;
                    a[(i + 1) * n + i] = v;
                }
            }
            for step in 0..n / 3 {
                let i = (step * 7) % n;
                let j = (step * 29 + 3) % n;
                if i != j {
                    let v = -(0.1 + rand() * 0.2);
                    a[i * n + j] = v;
                    a[j * n + i] = v;
                    // Keep it diagonally dominant, hence positive-definite.
                    a[i * n + i] += 0.6;
                    a[j * n + j] += 0.6;
                }
            }

            let sparse = dense(&a, n).expect("spd");

            for trial in 0..8 {
                let b: Vec<f64> = (0..n).map(|_| rand() * 2.0 - 1.0).collect();
                let c: Vec<f64> = (0..n).map(|_| rand() * 2.0 - 1.0).collect();
                let got = bilinear(&sparse.whiten(&b), &sparse.whiten(&c));
                let want = dense_quadratic_form(&a, n, &b, &c);
                assert!(
                    (got - want).abs() <= 1e-10 * want.abs().max(1.0),
                    "n={n} trial={trial}: sparse {got} vs dense {want}"
                );
            }
        }
    }

    /// A permutation must not change which matrices are rejected.
    #[test]
    fn rejects_a_non_positive_definite_matrix() {
        // Singular: the second row is a multiple of the first.
        assert!(dense(&[1.0, 2.0, 2.0, 4.0], 2).is_none());
    }
}