slain-num 0.4.0

use crate::{simd, Mat};

#[must_use]
pub fn solve(a: &Mat, b: &[f64]) -> Option<Vec<f64>> {
    let n = a.rows();
    assert!(a.cols() == n && b.len() == n, "square system");
    let scale = a.as_slice().iter().fold(0f64, |m, v| m.max(v.abs())).max(f64::MIN_POSITIVE);
    let mut m = a.clone();
    let mut x = b.to_vec();
    for i in 0..n {
        let p = (i..n).max_by(|&r, &s| m[(r, i)].abs().total_cmp(&m[(s, i)].abs()))?;
        if m[(p, i)].abs() < 1e-12 * scale {
            return None;
        }
        if p != i {
            for c in 0..n {
                let t = m[(i, c)];
                m[(i, c)] = m[(p, c)];
                m[(p, c)] = t;
            }
            x.swap(i, p);
        }
        let pivot_row: Vec<f64> = m.row(i)[i..].to_vec();
        for r in i + 1..n {
            let f = m[(r, i)] / pivot_row[0];
            if f != 0.0 {
                simd::axpy(-f, &pivot_row, &mut m.row_mut(r)[i..]);
                x[r] -= f * x[i];
            }
        }
    }
    for i in (0..n).rev() {
        let s = simd::dot(&m.row(i)[i + 1..], &x[i + 1..]);
        x[i] = (x[i] - s) / m[(i, i)];
    }
    Some(x)
}

#[must_use]
pub fn solve_spd(a: &Mat, b: &[f64]) -> Option<Vec<f64>> {
    let n = a.rows();
    assert!(a.cols() == n && b.len() == n, "square system");
    let mut l = Mat::zeros(n, n);
    for i in 0..n {
        for j in 0..=i {
            let s = a[(i, j)] - simd::dot(&l.row(i)[..j], &l.row(j)[..j]);
            if i == j {
                if s <= 0.0 {
                    return None;
                }
                l[(i, i)] = s.sqrt();
            } else {
                l[(i, j)] = s / l[(j, j)];
            }
        }
    }
    let mut y = vec![0.0; n];
    for i in 0..n {
        y[i] = (b[i] - simd::dot(&l.row(i)[..i], &y[..i])) / l[(i, i)];
    }
    let lt = l.transpose();
    let mut x = vec![0.0; n];
    for i in (0..n).rev() {
        x[i] = (y[i] - simd::dot(&lt.row(i)[i + 1..], &x[i + 1..])) / l[(i, i)];
    }
    Some(x)
}

#[must_use]
pub fn lstsq(a: &Mat, b: &[f64], ridge: f64) -> Option<Vec<f64>> {
    let mut g = a.gram();
    let n = g.rows();
    let trace: f64 = (0..n).map(|i| g[(i, i)]).sum();
    let bump = ridge * trace / n.max(1) as f64 + f64::MIN_POSITIVE;
    for i in 0..n {
        g[(i, i)] += bump;
    }
    let rhs = a.t_mul_vec(b);
    solve_spd(&g, &rhs).or_else(|| solve(&g, &rhs))
}

#[must_use]
pub fn lstsq_columns(feats: &[&[f64]], target: &[f64], ridge: f64) -> Option<Vec<f64>> {
    let n = feats.len();
    let mut gram = vec![0.0; n * n];
    let mut rhs = vec![0.0; n];
    simd::gram_rhs(feats, target, &mut gram, &mut rhs);
    lstsq_gram(gram, &rhs, ridge)
}

#[must_use]
pub fn lstsq_gram(gram: Vec<f64>, rhs: &[f64], ridge: f64) -> Option<Vec<f64>> {
    let n = rhs.len();
    let mut g = Mat::from_vec(n, n, gram);
    let trace: f64 = (0..n).map(|i| g[(i, i)]).sum();
    let bump = ridge * trace / n.max(1) as f64 + f64::MIN_POSITIVE;
    for i in 0..n {
        g[(i, i)] += bump;
    }
    solve_spd(&g, rhs).or_else(|| solve(&g, rhs))
}

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

    #[test]
    fn solves_recover_the_answer() {
        let a = Mat::from_vec(3, 3, vec![4.0, 1.0, 2.0, 1.0, 5.0, 3.0, 2.0, 3.0, 6.0]);
        let want = [1.0, -2.0, 0.5];
        let b = a.mul_vec(&want);
        for x in [solve(&a, &b).unwrap(), solve_spd(&a, &b).unwrap()] {
            assert!(x.iter().zip(want).all(|(g, w)| (g - w).abs() < 1e-10));
        }
        assert!(solve(&Mat::zeros(2, 2), &[1.0, 1.0]).is_none());
    }

    #[test]
    fn least_squares_fits_a_line() {
        let xs: Vec<f64> = (0..50).map(f64::from).collect();
        let ys: Vec<f64> = xs.iter().map(|x| 3.0 * x - 7.0).collect();
        let a = Mat::from_vec(50, 2, xs.iter().flat_map(|&x| [x, 1.0]).collect());
        let c = lstsq(&a, &ys, 0.0).unwrap();
        assert!((c[0] - 3.0).abs() < 1e-8 && (c[1] + 7.0).abs() < 1e-6);
        let ones = vec![1.0; 50];
        let c = lstsq_columns(&[&xs, &ones], &ys, 0.0).unwrap();
        assert!((c[0] - 3.0).abs() < 1e-8 && (c[1] + 7.0).abs() < 1e-6);
    }
}