edge-rs 0.3.2

//! limma's `lmFit`: one weighted least squares fit per gene.
//!
//! The linear-model half of voom. Genes are fitted independently, as in
//! [`crate::glm::levenberg`].
//!
//! ### One code path
//!
//! limma splits `lmFit` across `lm.series` and `gls.series`, each with a fast
//! all-observed branch and a slow per-gene branch. They are one computation.
//! With working covariance `V = D^-1/2 R D^-1/2` (`D` the weights, `R` the
//! compound-symmetry correlation from `block`), whitening by the Cholesky
//! factor of `V` turns every case into ordinary least squares. Plain weights
//! are `R = I`, where whitening is a row scaling by `sqrt(w)`.
//!
//! The factor is computed once, not per gene: `V` factors as `L = D^-1/2 L_R`
//! (scaling a lower triangular matrix by a positive diagonal keeps it lower
//! triangular, and Cholesky factors are unique), so `L_R` depends only on
//! `block` and the correlation. Genes that keep all their observations reuse
//! it; only genes with dropped observations refactorise, on their own
//! submatrix.
//!
//! ### Shared fast path
//!
//! When nothing is dropped anywhere and there is no block correlation, every
//! gene has the same design and its QR is the shared one, so all genes are
//! fitted by two GEMMs (`lm.series`' all-observed branch). The per-gene path
//! above stays for everything else.
//!
//! ### Missing observations
//!
//! A gene with non-finite responses or zero weights is refitted on what
//! survives, with its own residual df and rank.
//!
//! ### Rank deficiency
//!
//! The per-gene solve is a Householder QR that accepts design columns left to
//! right and rejects any whose residual norm has collapsed (LINPACK `dqrdc2`,
//! as in R's `lm.fit`). Rejected columns report `NaN` for their coefficient and
//! unscaled standard deviation, as R reports `NA`.
//! ### References
//!
//! Smyth, Statistical Applications in Genetics and Molecular Biology 3(1), 2004
//! Law, Chen, Shi and Smyth, Genome Biology 15, R29, 2014

use faer::linalg::matmul::matmul;
use faer::{Accum, MatMut, MatRef, Par};
use rayon::prelude::*;

use crate::errors::EdgeErrors;
use crate::utils::design::non_estimable;
use crate::utils::linalg::{
    cholesky_lower, cross_inverse, forward_substitute, invert_upper_triangular,
};
use crate::utils::recycled::{Recycled, RecycledRow};

/// Relative tolerance for declaring a design column linearly dependent.
/// R's `lm.fit` (LINPACK `dqrdc2`) cycles a column to the back when its
/// remaining norm falls below `1e-7` times its original norm. limma inherits the
/// threshold through `lm.fit`.
const RANK_TOL: f64 = 1e-7;

/// Weights at or below this drop the observation rather than scaling it.
///
/// limma's `gls.series` sets `M[weights < 1e-15] <- NA`; `lm.series` sets
/// `weights[weights <= 0] <- NA`. One threshold covers both: they differ only on
/// weights in `(0, 1e-15]`, which carry no information.
const MIN_WEIGHT: f64 = 1e-15;

/////////////////
// LmFitResult //
/////////////////

/// Everything `lmFit` produces, per gene.
#[derive(Clone, Debug)]
pub struct LmFitResult {
    /// Coefficients, row-major `n_genes * n_coef`.
    ///
    /// `NaN` for a coefficient that is not estimable (aliased design column, or
    /// the gene lost too many observations), where R reports `NA`.
    pub coefficients: Vec<f64>,
    /// Unscaled standard deviations, row-major `n_genes * n_coef`.
    ///
    /// The square root of the diagonal of `(X' V^-1 X)^-1`. The residual scale
    /// is kept in [`LmFitResult::sigma`] so `eBayes` can moderate it.
    pub stdev_unscaled: Vec<f64>,
    /// Residual standard deviation per gene.
    ///
    /// `sqrt(rss / df_residual)` with `rss` the whitened residual sum of
    /// squares. Zero where `df_residual` is zero, where limma reports `NA`;
    /// this is the one deliberate departure here, harmless because `squeezeVar`
    /// drops zero-df genes.
    pub sigma: Vec<f64>,
    /// Residual degrees of freedom per gene.
    ///
    /// Observations kept minus the rank achieved on them.
    pub df_residual: Vec<f64>,
    /// Fitted values, row-major `n_genes * n_samples`.
    ///
    /// `X * beta` on the original scale, aliased coefficients contributing
    /// nothing, at every sample including dropped ones. `NaN` throughout for a
    /// gene with no usable observation.
    pub fitted: Vec<f64>,
    /// Unscaled covariance of the estimable coefficients, row-major
    /// `rank * rank`.
    ///
    /// `(X' V^-1 X)^-1` for the design as a whole, in the order of
    /// [`LmFitResult::pivot`]. One matrix for all genes; see [`lm_fit`].
    /// probe weights.
    pub cov_coefficients: Vec<f64>,
    /// Design column indices, accepted ones first, then the rejected ones.
    ///
    /// Length `n_coef`, zero-based. `pivot[..rank]` names the estimable columns
    /// in the order [`LmFitResult::cov_coefficients`] uses (R's `qr()$pivot`).
    pub pivot: Vec<usize>,
    /// Rank of the design, shared by every fully observed gene.
    ///
    /// A gene that lost observations can have a lower rank of its own; that
    /// shows as extra `NaN` coefficients and a smaller
    /// [`LmFitResult::df_residual`].
    pub rank: usize,
}

///////////////////////
// Block correlation //
///////////////////////

/// The compound-symmetry correlation shared by every gene.
struct BlockCorrelation<'a> {
    /// Block label per sample. Samples sharing a label are correlated.
    block: &'a [usize],
    /// Within-block correlation.
    correlation: f64,
    /// Lower Cholesky factor of the full correlation matrix, row-major
    /// `n_samples * n_samples`. Reused by every fully observed gene.
    chol: Vec<f64>,
}

/// Writes the compound-symmetry correlation matrix for a subset of samples.
///
/// ### Params
///
/// * `out` - Destination, row-major with row stride `stride`
/// * `stride` - Row stride of `out`
/// * `obs` - Sample indices to include, in increasing order
/// * `block` - Block label per sample
/// * `correlation` - Within-block correlation
fn fill_correlation(
    out: &mut [f64],
    stride: usize,
    obs: &[usize],
    block: &[usize],
    correlation: f64,
) {
    for (i, &si) in obs.iter().enumerate() {
        for (j, &sj) in obs.iter().enumerate() {
            out[i * stride + j] = if i == j {
                1.0
            } else if block[si] == block[sj] {
                correlation
            } else {
                0.0
            };
        }
    }
}

/////////////
// Scratch //
/////////////

/// Per-thread buffers, allocated once per rayon worker.
struct Scratch {
    /// Sample indices this gene actually contributes, in increasing order.
    obs: Vec<usize>,
    /// Whitened design, column-major: column `j` starts at `j * n_samples`.
    ///
    /// Column-major: the Householder sweep reads one column at a time.
    a: Vec<f64>,
    /// Whitened response, overwritten by the Householder-transformed one. The
    /// tail below the rank is the residual, exactly R's `effects`.
    z: Vec<f64>,
    /// Householder reflectors, column-major with the same stride as `a`.
    hh: Vec<f64>,
    /// Reciprocal of `u'u / 2` for each stored reflector.
    tau: Vec<f64>,
    /// Original norm of each whitened design column, for the rank test.
    col_norm: Vec<f64>,
    /// Upper triangular factor of the accepted columns, row-major `n_coef`.
    r: Vec<f64>,
    /// Inverse of `r`, row-major `n_coef`.
    r_inv: Vec<f64>,
    /// Original indices of the accepted design columns, increasing.
    est: Vec<usize>,
    /// Solution in the accepted-column basis.
    beta: Vec<f64>,
    /// Correlation matrix for this gene's observed subset, then its Cholesky
    /// factor. Empty unless a block correlation was supplied.
    chol: Vec<f64>,
}

impl Scratch {
    /// Allocates scratch for one worker.
    ///
    /// ### Params
    ///
    /// * `n_samples` - Number of samples
    /// * `n_coef` - Number of coefficients
    /// * `blocked` - Whether a block correlation is in play, which is the only
    ///   case needing an `n_samples` by `n_samples` buffer
    ///
    /// ### Returns
    ///
    /// Buffers sized for a single gene's fit.
    fn new(n_samples: usize, n_coef: usize, blocked: bool) -> Self {
        Self {
            obs: Vec::with_capacity(n_samples),
            a: vec![0.0; n_samples * n_coef],
            z: vec![0.0; n_samples],
            hh: vec![0.0; n_samples * n_coef],
            tau: vec![0.0; n_coef],
            col_norm: vec![0.0; n_coef],
            r: vec![0.0; n_coef * n_coef],
            r_inv: vec![0.0; n_coef * n_coef],
            est: Vec::with_capacity(n_coef),
            beta: vec![0.0; n_coef],
            chol: if blocked {
                vec![0.0; n_samples * n_samples]
            } else {
                Vec::new()
            },
        }
    }
}

/// Householder QR that accepts design columns left to right and rejects any
/// whose residual norm has collapsed.
///
/// This is LINPACK's `dqrdc2` as used by R's `lm.fit`: a rejected column is set
/// aside, not cycled to the end, so `est` is an increasing subset. The per-gene
/// fit and the shared design factorisation both use it, so their ranks agree.
///
/// ### Params
///
/// * `a` - Whitened design, column-major with column stride `stride`,
///   overwritten by the transformed columns
/// * `z` - Whitened response to carry along, overwritten by `Q' z`. `None` when
///   only the factorisation itself is wanted.
/// * `hh` - Scratch for the reflectors, column-major with stride `stride`
/// * `tau` - Scratch for the reflector scales, length `n_coef`
/// * `col_norm` - Original norm of each column, read only when `pivot` is set
/// * `r` - Destination for the upper triangular factor, row-major `n_coef`
/// * `est` - Destination for the accepted column indices, cleared on entry
/// * `m` - Number of observations in play
/// * `stride` - Column stride of `a` and `hh`
/// * `n_coef` - Number of design columns
/// * `pivot` - Whether to apply the rank test. A design full rank on all
///   samples stays full rank after whitening, so a fully observed gene can skip
///   it.
///
/// ### Returns
///
/// The rank achieved, which is `est.len()`.
#[allow(clippy::too_many_arguments)]
fn rank_revealing_qr(
    a: &mut [f64],
    mut z: Option<&mut [f64]>,
    hh: &mut [f64],
    tau: &mut [f64],
    col_norm: &[f64],
    r: &mut [f64],
    est: &mut Vec<usize>,
    m: usize,
    stride: usize,
    n_coef: usize,
    pivot: bool,
) -> usize {
    est.clear();
    let mut rank = 0usize;
    for j in 0..n_coef {
        if rank == m {
            break;
        }
        let col = &a[j * stride..j * stride + m];
        let norm = col[rank..].iter().map(|v| v * v).sum::<f64>().sqrt();
        if norm == 0.0 || (pivot && norm <= RANK_TOL * col_norm[j]) {
            continue;
        }

        // Record the column of R above and on the diagonal before the column is
        // consumed by its own reflector.
        for i in 0..rank {
            r[i * n_coef + rank] = a[j * stride + i];
        }
        let head = a[j * stride + rank];
        let alpha = if head >= 0.0 { -norm } else { norm };
        r[rank * n_coef + rank] = alpha;

        // u = x - alpha e_1, so u'u = 2 * norm * (norm + |head|) > 0.
        let base = rank * stride;
        for i in rank..m {
            hh[base + i] = a[j * stride + i];
        }
        hh[base + rank] = head - alpha;
        tau[rank] = 1.0 / (norm * (norm + head.abs()));

        // Apply the reflector to every column still to come, and to the
        // response. Columns already passed over are never revisited.
        for jj in (j + 1)..n_coef {
            let cbase = jj * stride;
            let mut dot = 0.0;
            for i in rank..m {
                dot += hh[base + i] * a[cbase + i];
            }
            let factor = dot * tau[rank];
            for i in rank..m {
                a[cbase + i] -= factor * hh[base + i];
            }
        }
        if let Some(z) = z.as_deref_mut() {
            let mut dot = 0.0;
            for i in rank..m {
                dot += hh[base + i] * z[i];
            }
            let factor = dot * tau[rank];
            for i in rank..m {
                z[i] -= factor * hh[base + i];
            }
        }

        est.push(j);
        rank += 1;
    }
    rank
}

//////////////////
// Per-gene fit //
//////////////////

/// Outputs of one gene's fit, borrowed from the result buffers.
struct GeneOut<'a> {
    /// Coefficients for this gene, length `n_coef`.
    coef: &'a mut [f64],
    /// Unscaled standard deviations, length `n_coef`.
    stdev: &'a mut [f64],
    /// Fitted values, length `n_samples`.
    fitted: &'a mut [f64],
    /// Residual standard deviation.
    sigma: &'a mut f64,
    /// Residual degrees of freedom.
    df: &'a mut f64,
}

/// Fits one gene by whitening, then a rank-revealing Householder QR.
///
/// Whitening turns GLS against `V = D^-1/2 R D^-1/2` into OLS, so the rest is
/// the same for weights, block correlation, both or neither.
///
/// ### Params
///
/// * `scratch` - Per-thread buffers
/// * `y_row` - This gene's response, length `n_samples`
/// * `design` - Row-major design, `n_samples * n_coef`
/// * `n_samples` - Number of samples
/// * `n_coef` - Number of coefficients
/// * `weights` - Optional weight row for this gene
/// * `gls` - Optional block correlation
/// * `design_full_rank` - Whether the design is full rank on all samples, which
///   lets a fully observed gene skip the rank test
/// * `out` - Destination slices for this gene
#[allow(clippy::too_many_arguments)]
fn fit_gene(
    scratch: &mut Scratch,
    y_row: &[f64],
    design: &[f64],
    n_samples: usize,
    n_coef: usize,
    weights: Option<RecycledRow<'_, f64>>,
    gls: Option<&BlockCorrelation<'_>>,
    design_full_rank: bool,
    out: GeneOut<'_>,
) {
    let Scratch {
        obs,
        a,
        z,
        hh,
        tau,
        col_norm,
        r,
        r_inv,
        est,
        beta,
        chol,
    } = scratch;

    // -------------------------- //
    // which observations survive //
    // -------------------------- //
    obs.clear();
    for (sample, &value) in y_row.iter().enumerate() {
        let usable = value.is_finite()
            && match weights {
                Some(w) => {
                    let wj = w.get(sample);
                    wj.is_finite() && wj > MIN_WEIGHT
                }
                None => true,
            };
        if usable {
            obs.push(sample);
        }
    }
    let m = obs.len();

    if m == 0 {
        out.coef.fill(f64::NAN);
        out.stdev.fill(f64::NAN);
        out.fitted.fill(f64::NAN);
        *out.sigma = 0.0;
        *out.df = 0.0;
        return;
    }

    // --------------------------------------------------------------------- //
    // whiten: scale rows by sqrt(w), then apply the inverse Cholesky factor //
    // --------------------------------------------------------------------- //
    for (i, &sample) in obs.iter().enumerate() {
        let scale = match weights {
            Some(w) => w.get(sample).sqrt(),
            None => 1.0,
        };
        let row = &design[sample * n_coef..(sample + 1) * n_coef];
        for (j, &x) in row.iter().enumerate() {
            a[j * n_samples + i] = scale * x;
        }
        z[i] = scale * y_row[sample];
    }

    if let Some(g) = gls {
        // A principal submatrix of a positive definite matrix is positive
        // definite, so refactorising cannot fail once the full matrix passed.
        let factor: &[f64] = if m < n_samples {
            fill_correlation(chol, n_samples, obs, g.block, g.correlation);
            let factored = cholesky_lower(chol, m, n_samples);
            debug_assert!(factored, "principal submatrix must be positive definite");
            &chol[..]
        } else {
            &g.chol
        };
        for j in 0..n_coef {
            forward_substitute(
                factor,
                n_samples,
                m,
                &mut a[j * n_samples..j * n_samples + m],
            );
        }
        forward_substitute(factor, n_samples, m, &mut z[..m]);
    }

    // Whitening is invertible, so a full rank design stays full rank. Only a
    // gene that lost observations can lose rank and needs the reference norms.
    let pivot = !design_full_rank || m < n_samples;
    if pivot {
        for j in 0..n_coef {
            let col = &a[j * n_samples..j * n_samples + m];
            col_norm[j] = col.iter().map(|v| v * v).sum::<f64>().sqrt();
        }
    }

    // ------------------------------------------------------------- //
    // rank-revealing Householder QR, columns accepted left to right //
    // ------------------------------------------------------------- //
    let rank = rank_revealing_qr(
        a,
        Some(z.as_mut_slice()),
        hh,
        tau,
        col_norm,
        r,
        est,
        m,
        n_samples,
        n_coef,
        pivot,
    );

    // -------------- //
    // residual scale //
    // -------------- //
    let df = m - rank;
    let rss: f64 = z[rank..m].iter().map(|v| v * v).sum();
    *out.df = df as f64;
    *out.sigma = if df > 0 {
        (rss / df as f64).sqrt()
    } else {
        0.0
    };

    // ------------------------------------------------------------------- //
    // coefficients: back substitution against the leading rank block of R //
    // ------------------------------------------------------------------- //
    for i in (0..rank).rev() {
        let mut sum = z[i];
        for t in (i + 1)..rank {
            sum -= r[i * n_coef + t] * beta[t];
        }
        beta[i] = sum / r[i * n_coef + i];
    }

    out.coef.fill(f64::NAN);
    for (t, &j) in est.iter().enumerate() {
        out.coef[j] = beta[t];
    }

    // ----------------------------------------------- //
    // unscaled standard deviations: row norms of R^-1 //
    // ----------------------------------------------- //
    //
    // diag((R'R)^-1) = diag(R^-1 R^-T), so entry i is the sum of squares of row
    // i of R^-1. Inverting a tiny triangular matrix beats a solve per coefficient.
    invert_upper_triangular(r, r_inv, rank, n_coef);

    out.stdev.fill(f64::NAN);
    for (t, &j) in est.iter().enumerate() {
        let row: f64 = (t..rank).map(|s| r_inv[t * n_coef + s].powi(2)).sum();
        out.stdev[j] = row.sqrt();
    }

    // ---------------------------------------------------- //
    // fitted values on the original scale, at every sample //
    // ---------------------------------------------------- //
    for (sample, slot) in out.fitted.iter_mut().enumerate() {
        let row = &design[sample * n_coef..(sample + 1) * n_coef];
        let mut acc = 0.0;
        for (t, &j) in est.iter().enumerate() {
            acc += row[j] * beta[t];
        }
        *slot = acc;
    }
}

//////////////////////////
// Shared design factor //
//////////////////////////

/// The design factorised once for the whole fit.
struct DesignFactor {
    /// Unscaled covariance of the estimable columns, row-major `rank * rank`.
    cov: Vec<f64>,
    /// Design column indices, accepted first, then rejected.
    pivot: Vec<usize>,
    /// Rank of the (weighted, whitened) design.
    rank: usize,
    /// Least-squares solve matrix, row-major `n_samples * n_coef`:
    /// `beta = y' P` for any fully observed gene. `W^1/2 Q R^-T` scattered to
    /// the accepted columns; rejected columns are zero.
    solve: Vec<f64>,
    /// Unscaled standard deviation per design column, `NaN` for rejected ones.
    stdev: Vec<f64>,
}

/// Factors the design once for the whole fit, giving the unscaled covariance.
///
/// **This deliberately ignores per-gene weights.** limma computes
/// `cov.coefficients` from `qr(design)` on the shared design even when every
/// gene is fitted under probe weights (`R/lmfit.R:180-182`, and `R/lmfit.R:359`
/// for the blocked form, where only the correlation whitening is applied).
/// `contrasts.fit` consumes it regardless (`R/contrasts.R:5-6`), so reproducing
/// it keeps the contrast standard errors matching.
///
/// A weight row shared by every gene is folded in, as limma does when no
/// response is missing either (`R/lmfit.R:147`, `R/lmfit.R:316-318`).
///
/// ### Params
///
/// * `design` - Row-major design, `n_samples * n_coef`
/// * `n_samples` - Number of samples
/// * `n_coef` - Number of coefficients
/// * `shared_weights` - One weight per sample when every gene shares the same
///   row and no response is missing, otherwise `None`
/// * `gls` - Optional block correlation, whose Cholesky factor whitens the
///   design
///
/// ### Returns
///
/// The unscaled covariance, pivot and rank, plus the solve matrix and
/// standard deviations the shared fast path in [`lm_fit`] uses. The latter two
/// are only meaningful without `gls`, the only case that path is taken.
fn design_factorisation(
    design: &[f64],
    n_samples: usize,
    n_coef: usize,
    shared_weights: Option<&[f64]>,
    gls: Option<&BlockCorrelation<'_>>,
) -> DesignFactor {
    let mut a = vec![0.0; n_samples * n_coef];
    for sample in 0..n_samples {
        let scale = match shared_weights {
            Some(w) => w[sample].sqrt(),
            None => 1.0,
        };
        let row = &design[sample * n_coef..(sample + 1) * n_coef];
        for (j, &x) in row.iter().enumerate() {
            a[j * n_samples + sample] = scale * x;
        }
    }

    if let Some(g) = gls {
        for j in 0..n_coef {
            forward_substitute(
                &g.chol,
                n_samples,
                n_samples,
                &mut a[j * n_samples..(j + 1) * n_samples],
            );
        }
    }

    let col_norm: Vec<f64> = (0..n_coef)
        .map(|j| {
            a[j * n_samples..(j + 1) * n_samples]
                .iter()
                .map(|v| v * v)
                .sum::<f64>()
                .sqrt()
        })
        .collect();

    let mut hh = vec![0.0; n_samples * n_coef];
    let mut tau = vec![0.0; n_coef];
    let mut r = vec![0.0; n_coef * n_coef];
    let mut est = Vec::with_capacity(n_coef);
    let rank = rank_revealing_qr(
        &mut a, None, &mut hh, &mut tau, &col_norm, &mut r, &mut est, n_samples, n_samples, n_coef,
        true,
    );

    let mut r_inv = vec![0.0; n_coef * n_coef];
    invert_upper_triangular(&r, &mut r_inv, rank, n_coef);
    let cov = cross_inverse(&r_inv, rank, n_coef);

    // Accepted columns first, then the rejected ones in their original order,
    // which is what `dqrdc2` leaves behind and what `contrasts.fit` indexes.
    let mut pivot = est.clone();
    pivot.extend((0..n_coef).filter(|j| !est.contains(j)));

    // Thin Q, column-major `n_samples * rank`: the reflectors applied in
    // reverse to the leading identity columns. Reflector i touches rows i..,
    // so columns before it are still unit vectors and are skipped.
    let mut q = vec![0.0; n_samples * rank];
    for j in 0..rank {
        q[j * n_samples + j] = 1.0;
    }
    for i in (0..rank).rev() {
        let u = &hh[i * n_samples + i..(i + 1) * n_samples];
        for j in i..rank {
            let col = &mut q[j * n_samples + i..(j + 1) * n_samples];
            let dot: f64 = u.iter().zip(col.iter()).map(|(a, b)| a * b).sum();
            let factor = dot * tau[i];
            for (c, &uv) in col.iter_mut().zip(u) {
                *c -= factor * uv;
            }
        }
    }

    // beta_t = sum_l R^-1[t, l] (Q' W^1/2 y)_l, so P[s, est[t]] is
    // sqrt(w_s) sum_{l >= t} R^-1[t, l] Q[s, l].
    let mut solve = vec![0.0; n_samples * n_coef];
    for s in 0..n_samples {
        let scale = shared_weights.map_or(1.0, |w| w[s].sqrt());
        for (t, &j) in est.iter().enumerate() {
            let acc: f64 = (t..rank)
                .map(|l| r_inv[t * n_coef + l] * q[l * n_samples + s])
                .sum();
            solve[s * n_coef + j] = scale * acc;
        }
    }

    // Row norms of R^-1, as `fit_gene` forms them.
    let mut stdev = vec![f64::NAN; n_coef];
    for (t, &j) in est.iter().enumerate() {
        let row: f64 = (t..rank).map(|s| r_inv[t * n_coef + s].powi(2)).sum();
        stdev[j] = row.sqrt();
    }

    DesignFactor {
        cov,
        pivot,
        rank,
        solve,
        stdev,
    }
}

//////////////////////
// Shared fast path //
//////////////////////

/// Fits every gene at once against the shared design factorisation.
///
/// limma's `lm.series` all-observed branch (`lm.fit(design, t(M))`). Valid
/// when every response is finite, every gene shares one weight row with no
/// dropped observation, and there is no block correlation: then every gene
/// has the same whitened design, so the per-gene QR is the same QR. Two GEMMs
/// replace it, both through faer at the top level:
///
/// * `coefficients = Y P` (`n_genes * n_samples` by `n_samples * n_coef`);
/// * `fitted = coefficients X'`, needed by voom regardless.
///
/// The residual sum of squares comes from `y - fitted` in one fused pass over
/// genes, not from `||y||^2 - ||Q'y||^2`, which cancels for well-fitted genes
/// at high expression.
///
/// ### Params
///
/// * `y` - Response, row-major `n_genes * n_samples`, all finite
/// * `n_genes` - Number of genes
/// * `n_samples` - Number of samples
/// * `design` - Row-major design, `n_samples * n_coef`
/// * `n_coef` - Number of coefficients
/// * `factor` - Shared factorisation of the (weighted) design
/// * `shared_weights` - The one weight row, or `None` for unit weights
/// * `out` - Result buffers to fill, sized for all genes
#[allow(clippy::too_many_arguments)]
fn fit_shared(
    y: &[f64],
    n_genes: usize,
    n_samples: usize,
    design: &[f64],
    n_coef: usize,
    factor: &DesignFactor,
    shared_weights: Option<&[f64]>,
    out: &mut LmFitResult,
) {
    let y_mat = MatRef::from_row_major_slice(y, n_genes, n_samples);
    let x_mat = MatRef::from_row_major_slice(design, n_samples, n_coef);
    let p_mat = MatRef::from_row_major_slice(&factor.solve, n_samples, n_coef);

    // Rejected columns of P are zero, so their coefficients come out zero and
    // drop out of the fitted values. They become NaN below.
    matmul(
        MatMut::from_row_major_slice_mut(&mut out.coefficients, n_genes, n_coef),
        Accum::Replace,
        y_mat,
        p_mat,
        1.0,
        Par::rayon(0),
    );
    matmul(
        MatMut::from_row_major_slice_mut(&mut out.fitted, n_genes, n_samples),
        Accum::Replace,
        MatRef::from_row_major_slice(&out.coefficients, n_genes, n_coef),
        x_mat.transpose(),
        1.0,
        Par::rayon(0),
    );

    let df = n_samples - factor.rank;
    out.df_residual.fill(df as f64);
    out.sigma
        .par_iter_mut()
        .zip(y.par_chunks(n_samples))
        .zip(out.fitted.par_chunks(n_samples))
        .for_each(|((sig, y_row), fit_row)| {
            let rss: f64 = match shared_weights {
                Some(w) => y_row
                    .iter()
                    .zip(fit_row)
                    .zip(w)
                    .map(|((a, b), wj)| wj * (a - b).powi(2))
                    .sum(),
                None => y_row
                    .iter()
                    .zip(fit_row)
                    .map(|(a, b)| (a - b).powi(2))
                    .sum(),
            };
            *sig = if df > 0 {
                (rss / df as f64).sqrt()
            } else {
                0.0
            };
        });

    let rejected: Vec<usize> = factor.pivot[factor.rank..].to_vec();
    out.coefficients
        .par_chunks_mut(n_coef)
        .zip(out.stdev_unscaled.par_chunks_mut(n_coef))
        .for_each(|(coef, stdev)| {
            for &j in &rejected {
                coef[j] = f64::NAN;
            }
            stdev.copy_from_slice(&factor.stdev);
        });
}

////////////
// lm_fit //
////////////

/// Fits a linear model to every gene by weighted or generalised least squares.
///
/// Port of limma's `lmFit`, covering the `lm.series` and `gls.series` paths.
/// Each gene is fitted on its own finite, positively weighted observations, so
/// genes with missing values get their own residual df.
///
/// `block` plus `correlation` gives the mixed-model form: samples sharing a
/// label have the stated correlation, and the fit is GLS against that
/// compound-symmetry covariance. `correlation` without `block` is ignored, as in
/// `lmFit`.
///
/// When every response is finite, the weights are absent or one row shared by
/// all genes (all positive), and there is no block correlation, every gene has
/// the same design and the fit is two GEMMs against one factorisation. This is
/// limma's own all-observed branch. Anything else goes gene by gene.
///
/// ### Params
///
/// * `y` - Response, row-major `n_genes * n_samples`. `f64` throughout (log-CPM
///   for voom). Non-finite entries are treated as missing.
/// * `n_genes` - Number of genes
/// * `n_samples` - Number of samples
/// * `design` - Design matrix, row-major `n_samples * n_coef`
/// * `n_coef` - Number of coefficients
/// * `weights` - Optional observation weights, recycled over genes and samples.
///   A weight at or below `1e-15`, or a non-finite one, drops the observation.
/// * `block` - Optional block label per sample
/// * `correlation` - Within-block correlation, required when `block` is given
///
/// ### Returns
///
/// Coefficients, unscaled standard deviations, residual standard deviations,
/// residual degrees of freedom, fitted values, and the shared design
/// factorisation `eBayes` and `contrasts.fit` need. [`EdgeErrors`] if the
/// shapes disagree, the correlation is out of range, or the implied correlation
/// matrix is not positive definite.
/// [`LmFitResult::cov_coefficients`] comes from the design as a whole, not each
/// gene's weighted fit. That is limma's choice; see `design_factorisation`.
///
/// ### References
///
/// Smyth, Michaud and Scott, Bioinformatics 21(9), 2005 (the blocked form)
#[allow(clippy::too_many_arguments)]
pub fn lm_fit(
    y: &[f64],
    n_genes: usize,
    n_samples: usize,
    design: &[f64],
    n_coef: usize,
    weights: Option<&Recycled<f64>>,
    block: Option<&[usize]>,
    correlation: Option<f64>,
) -> Result<LmFitResult, EdgeErrors> {
    lm_fit_impl(
        y,
        n_genes,
        n_samples,
        design,
        n_coef,
        weights,
        block,
        correlation,
        true,
    )
}

/// Body of [`lm_fit`], with the shared fast path optional.
///
/// ### Params
///
/// * `y`, `n_genes`, `n_samples`, `design`, `n_coef`, `weights`, `block`,
///   `correlation` - As for [`lm_fit`]
/// * `allow_shared` - Whether the shared fast path may be taken. `false` forces
///   the per-gene path, so tests can compare the two.
///
/// ### Returns
///
/// As for [`lm_fit`].
#[allow(clippy::too_many_arguments)]
fn lm_fit_impl(
    y: &[f64],
    n_genes: usize,
    n_samples: usize,
    design: &[f64],
    n_coef: usize,
    weights: Option<&Recycled<f64>>,
    block: Option<&[usize]>,
    correlation: Option<f64>,
    allow_shared: bool,
) -> Result<LmFitResult, EdgeErrors> {
    if n_genes == 0 || n_samples == 0 {
        return Err(EdgeErrors::EmptyCounts { n_genes, n_samples });
    }
    if n_coef == 0 {
        return Err(EdgeErrors::MustBePositive("n_coef".to_string()));
    }
    if y.len() != n_genes * n_samples {
        return Err(EdgeErrors::LengthMismatch {
            name: "y",
            expected: n_genes * n_samples,
            got: y.len(),
        });
    }
    if design.len() != n_samples * n_coef {
        return Err(EdgeErrors::LengthMismatch {
            name: "design",
            expected: n_samples * n_coef,
            got: design.len(),
        });
    }
    if let Some(w) = weights {
        w.validate(n_genes, n_samples)?;
    }

    // ------------------------------------------------- //
    // block correlation set-up, done once for all genes //
    // ------------------------------------------------- //
    let gls = match block {
        None => None,
        Some(b) => {
            if b.len() != n_samples {
                return Err(EdgeErrors::LengthMismatch {
                    name: "block",
                    expected: n_samples,
                    got: b.len(),
                });
            }
            let rho = correlation.ok_or_else(|| {
                EdgeErrors::InvalidArgument(
                    "block requires a correlation; limma would estimate one with duplicateCorrelation"
                        .to_string(),
                )
            })?;
            if !rho.is_finite() || rho.abs() >= 1.0 {
                return Err(EdgeErrors::InvalidArgument(format!(
                    "correlation must be finite and strictly inside (-1, 1); got {rho}"
                )));
            }
            if rho == 0.0 {
                // Uncorrelated blocks give the identity; skip whitening.
                None
            } else {
                let obs: Vec<usize> = (0..n_samples).collect();
                let mut chol = vec![0.0; n_samples * n_samples];
                fill_correlation(&mut chol, n_samples, &obs, b, rho);
                if !cholesky_lower(&mut chol, n_samples, n_samples) {
                    return Err(EdgeErrors::CholeskyFailed(format!(
                        "correlation {rho} makes the block covariance indefinite; \
                         a block of size k needs correlation above -1/(k - 1)"
                    )));
                }
                Some(BlockCorrelation {
                    block: b,
                    correlation: rho,
                    chol,
                })
            }
        }
    };

    // Columns no gene can estimate. A design full rank here is full rank for
    // every gene that keeps all its observations, which skip the rank test.
    let design_full_rank = non_estimable(design, n_samples, n_coef)?.is_none();

    // limma folds weights into the shared factorisation only when one row serves
    // every gene and no response is missing, where `lm.wfit` on the shared design
    // is exact (`R/lmfit.R:142`). Otherwise it uses the bare design.
    let shared_weights: Option<Vec<f64>> = if y.par_iter().all(|v| v.is_finite()) {
        match weights {
            Some(Recycled::Scalar(w)) => Some(vec![*w; n_samples]),
            Some(Recycled::BySample(v)) => Some(v.clone()),
            _ => None,
        }
    } else {
        None
    };
    let factor = design_factorisation(
        design,
        n_samples,
        n_coef,
        shared_weights.as_deref(),
        gls.as_ref(),
    );

    let mut out = LmFitResult {
        coefficients: vec![0.0; n_genes * n_coef],
        stdev_unscaled: vec![0.0; n_genes * n_coef],
        sigma: vec![0.0; n_genes],
        df_residual: vec![0.0; n_genes],
        fitted: vec![0.0; n_genes * n_samples],
        cov_coefficients: Vec::new(),
        pivot: Vec::new(),
        rank: factor.rank,
    };

    // No observation dropped anywhere, so every gene sees the same whitened
    // design: all responses finite (`shared_weights` is only set then), and
    // either no weights or a shared row that drops nothing.
    let all_observed = y.par_iter().all(|v| v.is_finite())
        && match weights {
            None => true,
            Some(_) => shared_weights
                .as_deref()
                .is_some_and(|w| w.iter().all(|&v| v.is_finite() && v > MIN_WEIGHT)),
        };
    if allow_shared && gls.is_none() && all_observed {
        fit_shared(
            y,
            n_genes,
            n_samples,
            design,
            n_coef,
            &factor,
            shared_weights.as_deref(),
            &mut out,
        );
        out.cov_coefficients = factor.cov;
        out.pivot = factor.pivot;
        return Ok(out);
    }

    let blocked = gls.is_some();
    out.coefficients
        .par_chunks_mut(n_coef)
        .zip(out.stdev_unscaled.par_chunks_mut(n_coef))
        .zip(out.fitted.par_chunks_mut(n_samples))
        .zip(out.sigma.par_iter_mut())
        .zip(out.df_residual.par_iter_mut())
        .enumerate()
        .for_each_init(
            || Scratch::new(n_samples, n_coef, blocked),
            |scratch, (gene, ((((coef, stdev), fit), sig), df))| {
                let start = gene * n_samples;
                fit_gene(
                    scratch,
                    &y[start..start + n_samples],
                    design,
                    n_samples,
                    n_coef,
                    weights.map(|w| w.row(gene, n_samples)),
                    gls.as_ref(),
                    design_full_rank,
                    GeneOut {
                        coef,
                        stdev,
                        fitted: fit,
                        sigma: sig,
                        df,
                    },
                );
            },
        );

    out.cov_coefficients = factor.cov;
    out.pivot = factor.pivot;
    Ok(out)
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    // References are pasted verbatim from R's 17-digit output, including digits
    // past f64 precision and a 1/sqrt(2) clippy would like as a constant, so they
    // stay checkable against the quoted `Rscript` lines.
    #![allow(clippy::excessive_precision, clippy::approx_constant)]

    use super::*;
    use approx::assert_relative_eq;

    const N_GENES: usize = 8;
    const N_SAMPLES: usize = 5;

    /// Stand-in for R's `NA` in the embedded reference values.
    const NA: f64 = f64::NAN;

    /// Response fixture, `y[i] = ((i * 7) mod 23) / 64`, row-major 8 by 5.
    ///
    /// Dyadic rationals, so R and Rust see bit-identical inputs.
    fn fixture_y() -> Vec<f64> {
        (0..N_GENES * N_SAMPLES)
            .map(|i| ((i * 7) % 23) as f64 / 64.0)
            .collect()
    }

    /// Full weight fixture, `w[i] = (((i * 5) mod 7) + 1) / 8`, row-major 8 by 5.
    fn fixture_weights() -> Vec<f64> {
        (0..N_GENES * N_SAMPLES)
            .map(|i| (((i * 5) % 7) + 1) as f64 / 8.0)
            .collect()
    }

    /// Two-column design: intercept and a two-versus-three group indicator.
    fn design_two() -> Vec<f64> {
        vec![1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
    }

    /// Three-column design whose first column is the sum of the other two.
    fn design_aliased() -> Vec<f64> {
        vec![
            1.0, 0.0, 1.0, //
            1.0, 0.0, 1.0, //
            1.0, 1.0, 0.0, //
            1.0, 1.0, 0.0, //
            1.0, 1.0, 0.0,
        ]
    }

    /// Block labels used by every correlated fixture.
    fn blocks() -> Vec<usize> {
        vec![1, 1, 2, 2, 3]
    }

    /// Compares against R, treating `NaN` and `NA` as equal.
    fn assert_close(got: &[f64], want: &[f64]) {
        assert_eq!(got.len(), want.len(), "length");
        for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() {
            if w.is_nan() {
                assert!(g.is_nan(), "index {i}: expected NA, got {g}");
            } else {
                assert!(g.is_finite(), "index {i}: expected {w}, got {g}");
                assert_relative_eq!(*g, *w, max_relative = 1e-13, epsilon = 0.0);
            }
        }
    }

    //////////////////////
    // Reference values //
    //////////////////////

    // Literals come from limma 3.66 under R 4.5. Fixtures are built in R by the
    // same closed forms as `fixture_y` and `fixture_weights`:
    //
    //   library(limma)
    //   Y <- matrix(((0:39)*7) %% 23 / 64, 8, 5, byrow=TRUE)
    //   Wfull <- matrix((((0:39)*5) %% 7 + 1)/8, 8, 5, byrow=TRUE)
    //   X2 <- cbind(1, c(0,0,1,1,1))
    //   X3 <- cbind(1, c(0,0,1,1,1), c(1,1,0,0,0))
    //   wsamp <- c(1, 2, 0.5, 4, 0.25); b <- c(1,1,2,2,3)
    //   Ym <- Y; Ym[4,3] <- NA; Ym[6,] <- NA
    //   f <- lmFit(...)
    //   cat(sprintf("%.17g", as.vector(t(f$coefficients))), sep=", ")
    //   cat(sprintf("%.17g", as.vector(t(f$stdev.unscaled))), sep=", ")
    //   cat(sprintf("%.17g", f$sigma), sep=", ")
    //   cat(sprintf("%.17g", f$df.residual), sep=", ")
    //   cat(sprintf("%.17g", as.vector(t(f$coefficients %*% t(X)))), sep=", ")
    //
    // Where limma reports `sigma = NA` (zero df.residual), the literal is `0.0`:
    // this crate's one deliberate departure.

    /// `lmFit(Y, X2)`.
    const A_COEF: [f64; 16] = [
        0.054687499999999972,
        0.15364583333333334,
        0.24218749999999997,
        -0.085937499999999986,
        0.070312499999999972,
        0.15364583333333334,
        0.25781250000000011,
        -0.085937500000000042,
        0.085937499999999972,
        0.033854166666666678,
        0.27343749999999994,
        -0.085937499999999972,
        0.10156249999999997,
        0.033854166666666678,
        0.2890625,
        -0.085937499999999972,
    ];
    const A_SDU: [f64; 16] = [
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
    ];
    const A_SIGMA: [f64; 8] = [
        0.11164557869909383,
        0.099845257878545909,
        0.11164557869909383,
        0.099845257878545896,
        0.11164557869909383,
        0.099845257878545896,
        0.11164557869909383,
        0.099845257878545896,
    ];
    const A_FITTED: [f64; 40] = [
        0.054687499999999972,
        0.054687499999999972,
        0.20833333333333331,
        0.20833333333333331,
        0.20833333333333331,
        0.24218749999999997,
        0.24218749999999997,
        0.15625,
        0.15625,
        0.15625,
        0.070312499999999972,
        0.070312499999999972,
        0.22395833333333331,
        0.22395833333333331,
        0.22395833333333331,
        0.25781250000000011,
        0.25781250000000011,
        0.17187500000000006,
        0.17187500000000006,
        0.17187500000000006,
        0.085937499999999972,
        0.085937499999999972,
        0.11979166666666666,
        0.11979166666666666,
        0.11979166666666666,
        0.27343749999999994,
        0.27343749999999994,
        0.18749999999999997,
        0.18749999999999997,
        0.18749999999999997,
        0.10156249999999997,
        0.10156249999999997,
        0.13541666666666666,
        0.13541666666666666,
        0.13541666666666666,
        0.2890625,
        0.2890625,
        0.20312500000000003,
        0.20312500000000003,
        0.20312500000000003,
    ];

    /// `lmFit(Y, X2, weights = c(1, 2, 0.5, 4, 0.25))`.
    const B_COEF: [f64; 16] = [
        0.072916666666666671,
        0.23053728070175439,
        0.26041666666666663,
        -0.10992324561403506,
        0.088541666666666741,
        0.23053728070175433,
        0.27604166666666669,
        -0.10992324561403508,
        0.10416666666666669,
        -0.07209429824561403,
        0.29166666666666669,
        -0.10992324561403508,
        0.11979166666666664,
        -0.072094298245614016,
        0.30729166666666669,
        -0.10992324561403506,
    ];
    const B_SDU: [f64; 2] = [0.57735026918962584, 0.73746840550820003];
    const B_SIGMA: [f64; 8] = [
        0.094323144471374784,
        0.074810921506209174,
        0.094323144471374742,
        0.074810921506209146,
        0.11154962751509892,
        0.074810921506209146,
        0.11154962751509892,
        0.07481092150620916,
    ];

    /// `lmFit(Y, X2, weights = Wfull)`.
    const C_COEF: [f64; 16] = [
        0.093749999999999972,
        0.066105769230769246,
        0.22851562500000003,
        -0.042436079545454544,
        0.10069444444444443,
        0.15451388888888887,
        0.24687500000000015,
        -0.05156250000000008,
        0.058593750000000035,
        0.084635416666666616,
        0.26432291666666669,
        -0.044010416666666691,
        0.083333333333333343,
        0.070833333333333331,
        0.32812500000000006,
        -0.099759615384615405,
    ];
    const C_SDU: [f64; 16] = [
        1.0690449676496978,
        1.325987088263592,
        1.0,
        1.3142574813455419,
        0.94280904158206336,
        1.3333333333333333,
        0.89442719099991608,
        1.1710800875382399,
        1.4142135623730956,
        1.6329931618554525,
        0.81649658092772603,
        1.2110601416389968,
        1.1547005383792515,
        1.3662601021279464,
        1.0690449676496978,
        1.325987088263592,
    ];
    const C_SIGMA: [f64; 8] = [
        0.07289162230162545,
        0.054945826898183045,
        0.054147465881078799,
        0.066068380670931875,
        0.082088067294221856,
        0.07410330019231505,
        0.091613384135604214,
        0.074599961439633961,
    ];

    /// `lmFit(Ym, X2)`, with `Ym[4, 3]` and all of `Ym[6, ]` set to `NA`.
    const D_COEF: [f64; 16] = [
        0.054687499999999972,
        0.15364583333333334,
        0.24218749999999997,
        -0.085937499999999986,
        0.070312499999999972,
        0.15364583333333334,
        0.2578125,
        -0.031249999999999965,
        0.085937499999999972,
        0.033854166666666678,
        NA,
        NA,
        0.10156249999999997,
        0.033854166666666678,
        0.2890625,
        -0.085937499999999972,
    ];
    const D_SDU: [f64; 16] = [
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654757,
        1.0,
        0.70710678118654746,
        0.9128709291752769,
        NA,
        NA,
        0.70710678118654746,
        0.9128709291752769,
        0.70710678118654746,
        0.9128709291752769,
    ];
    const D_SIGMA: [f64; 8] = [
        0.11164557869909383,
        0.099845257878545909,
        0.11164557869909383,
        0.077339804192278636,
        0.11164557869909383,
        0.0,
        0.11164557869909383,
        0.099845257878545896,
    ];
    const D_DF: [f64; 8] = [3.0, 3.0, 3.0, 2.0, 3.0, 0.0, 3.0, 3.0];
    const D_FITTED: [f64; 40] = [
        0.054687499999999972,
        0.054687499999999972,
        0.20833333333333331,
        0.20833333333333331,
        0.20833333333333331,
        0.24218749999999997,
        0.24218749999999997,
        0.15625,
        0.15625,
        0.15625,
        0.070312499999999972,
        0.070312499999999972,
        0.22395833333333331,
        0.22395833333333331,
        0.22395833333333331,
        0.2578125,
        0.2578125,
        0.22656250000000003,
        0.22656250000000003,
        0.22656250000000003,
        0.085937499999999972,
        0.085937499999999972,
        0.11979166666666666,
        0.11979166666666666,
        0.11979166666666666,
        NA,
        NA,
        NA,
        NA,
        NA,
        0.10156249999999997,
        0.10156249999999997,
        0.13541666666666666,
        0.13541666666666666,
        0.13541666666666666,
        0.2890625,
        0.2890625,
        0.20312500000000003,
        0.20312500000000003,
        0.20312500000000003,
    ];

    /// `lmFit(Ym, X2, weights = Wfull)`.
    const E_COEF: [f64; 16] = [
        0.093749999999999972,
        0.066105769230769246,
        0.22851562500000003,
        -0.042436079545454544,
        0.10069444444444443,
        0.15451388888888887,
        0.24687499999999996,
        -0.029427083333333336,
        0.058593750000000035,
        0.084635416666666616,
        NA,
        NA,
        0.083333333333333343,
        0.070833333333333331,
        0.32812500000000006,
        -0.099759615384615405,
    ];
    const E_SDU: [f64; 16] = [
        1.0690449676496978,
        1.325987088263592,
        1.0,
        1.3142574813455419,
        0.94280904158206336,
        1.3333333333333333,
        0.89442719099991586,
        1.2110601416389966,
        1.4142135623730956,
        1.6329931618554525,
        NA,
        NA,
        1.1547005383792515,
        1.3662601021279464,
        1.0690449676496978,
        1.325987088263592,
    ];
    const E_SIGMA: [f64; 8] = [
        0.07289162230162545,
        0.054945826898183045,
        0.054147465881078799,
        0.063048940228462927,
        0.082088067294221856,
        0.0,
        0.091613384135604214,
        0.074599961439633961,
    ];

    /// `lmFit(Y, X2, block = b, correlation = 0.25)`.
    const F_COEF: [f64; 16] = [
        0.0546875,
        0.14362980769230768,
        0.24218749999999997,
        -0.077524038461538436,
        0.070312500000000042,
        0.14362980769230765,
        0.25781249999999989,
        -0.077524038461538367,
        0.085937499999999958,
        0.033052884615384658,
        0.27343749999999994,
        -0.077524038461538394,
        0.10156249999999996,
        0.033052884615384651,
        0.28906249999999983,
        -0.077524038461538325,
    ];
    const F_SDU: [f64; 2] = [0.79056941504209466, 1.0047961905856253];
    const F_SIGMA: [f64; 8] = [
        0.1146379531823081,
        0.10410655316738308,
        0.11463795318230807,
        0.10410655316738308,
        0.12883085482904671,
        0.10410655316738308,
        0.12883085482904669,
        0.10410655316738311,
    ];

    /// `lmFit(Y, X2, block = b, correlation = 0.25, weights = Wfull)`.
    const G_COEF: [f64; 16] = [
        0.1020338931741292,
        0.046620906898569577,
        0.22414926304156485,
        -0.029290377485346103,
        0.10866723833352257,
        0.1386183553847985,
        0.24332672403599037,
        -0.037918062449936007,
        0.051037727636791902,
        0.097377781866368834,
        0.2613411183871478,
        -0.035655125769476104,
        0.077711629165160057,
        0.080295409798648854,
        0.3364088931741292,
        -0.10120602688561266,
    ];
    const G_SDU: [f64; 16] = [
        1.139580592613717,
        1.4040630530325249,
        1.112163511876666,
        1.4227748106584122,
        1.0256781561509392,
        1.4571513092032287,
        0.99664956319551301,
        1.2759337202972107,
        1.5469735769439759,
        1.7857323811461194,
        0.91075033227467828,
        1.2983569545962401,
        1.2788621058650516,
        1.5060416834761272,
        1.139580592613717,
        1.4040630530325249,
    ];
    const G_SIGMA: [f64; 8] = [
        0.070104095566521502,
        0.055924276262054526,
        0.057844680310825951,
        0.067449568692666786,
        0.094410200626087654,
        0.077120443875363126,
        0.10549134470829188,
        0.07499280377972585,
    ];
    const G_FITTED_GENE0: [f64; 5] = [
        0.1020338931741292,
        0.1020338931741292,
        0.14865480007269877,
        0.14865480007269877,
        0.14865480007269877,
    ];

    /// `lmFit(Ym, X2, block = b, correlation = 0.25, weights = Wfull)`.
    const H_COEF: [f64; 16] = [
        0.1020338931741292,
        0.046620906898569577,
        0.22414926304156485,
        -0.029290377485346103,
        0.10866723833352257,
        0.1386183553847985,
        0.24332672403599037,
        -0.025878807369323673,
        0.051037727636791902,
        0.097377781866368834,
        NA,
        NA,
        0.077711629165160057,
        0.080295409798648854,
        0.3364088931741292,
        -0.10120602688561266,
    ];
    const H_SDU: [f64; 16] = [
        1.139580592613717,
        1.4040630530325249,
        1.112163511876666,
        1.4227748106584122,
        1.0256781561509392,
        1.4571513092032287,
        0.99664956319551323,
        1.288400954083966,
        1.5469735769439759,
        1.7857323811461194,
        NA,
        NA,
        1.2788621058650516,
        1.5060416834761272,
        1.139580592613717,
        1.4040630530325249,
    ];
    const H_SIGMA: [f64; 8] = [
        0.070104095566521502,
        0.055924276262054526,
        0.057844680310825951,
        0.067507835762062693,
        0.094410200626087654,
        0.0,
        0.10549134470829188,
        0.07499280377972585,
    ];

    /// `lmFit(Y, X2, weights = wsamp, block = b, correlation = 0.25)`.
    const K_COEF: [f64; 16] = [
        0.07853837083483986,
        0.23106008862949068,
        0.26603837083483994,
        -0.107400358117879,
        0.094163370834839874,
        0.23106008862949068,
        0.28166337083484,
        -0.10740035811787897,
        0.10978837083483996,
        -0.094332132363834628,
        0.29728837083484005,
        -0.107400358117879,
        0.12541337083483997,
        -0.094332132363834642,
        0.31291337083483994,
        -0.10740035811787904,
    ];
    const K_SDU: [f64; 2] = [0.63943105293252589, 0.80103691702704827];
    const K_SIGMA: [f64; 8] = [
        0.10159645906688854,
        0.081201785694141715,
        0.10159645906688854,
        0.081201785694141715,
        0.12346737328937207,
        0.081201785694141715,
        0.12346737328937207,
        0.081201785694141729,
    ];

    //////////
    // Fits //
    //////////

    #[test]
    fn test_unweighted_matches_limma() {
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            None,
            None,
            None,
        )
        .unwrap();
        assert_close(&fit.coefficients, &A_COEF);
        assert_close(&fit.stdev_unscaled, &A_SDU);
        assert_close(&fit.sigma, &A_SIGMA);
        assert_close(&fit.df_residual, &[3.0; 8]);
        assert_close(&fit.fitted, &A_FITTED);
    }

    /// `Rscript -e 'X <- matrix(c(1,0,1,0,1,1,1,1,1,1),5,2,byrow=TRUE);
    /// Q <- qr(X); chol2inv(Q$qr, size=Q$rank)'`
    const A_COV: [f64; 4] = [0.49999999999999994, -0.5, -0.5, 0.83333333333333337];

    /// Same design under `lm.wfit(X, y, c(1,2,3,4,5)/8)`.
    const A_COV_WEIGHTED: [f64; 4] = [
        2.6666666666666656,
        -2.6666666666666661,
        -2.6666666666666661,
        3.333333333333333,
    ];

    #[test]
    fn test_cov_coefficients_match_the_shared_design_qr() {
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            None,
            None,
            None,
        )
        .unwrap();
        assert_eq!(fit.rank, 2);
        assert_eq!(fit.pivot, vec![0, 1]);
        assert_close(&fit.cov_coefficients, &A_COV);
        // Unscaled sds are the square roots of its diagonal when no gene lost an
        // observation.
        for gene in 0..N_GENES {
            assert_relative_eq!(
                fit.stdev_unscaled[gene * 2],
                A_COV[0].sqrt(),
                epsilon = 1e-12
            );
        }
    }

    #[test]
    fn test_cov_coefficients_fold_in_a_shared_weight_row() {
        let w = Recycled::by_sample(vec![0.125, 0.25, 0.375, 0.5, 0.625]);
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            None,
            None,
        )
        .unwrap();
        assert_eq!(fit.rank, 2);
        assert_close(&fit.cov_coefficients, &A_COV_WEIGHTED);
    }

    #[test]
    fn test_cov_coefficients_ignore_per_gene_weights() {
        // limma takes `cov.coefficients` from the bare design once probe weights
        // are in play, so this must match the unweighted answer.
        let w = Recycled::full(fixture_weights(), N_GENES, N_SAMPLES).unwrap();
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            None,
            None,
        )
        .unwrap();
        assert_close(&fit.cov_coefficients, &A_COV);
    }

    #[test]
    fn test_aliased_design_reports_rank_and_pivot() {
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_aliased(),
            3,
            None,
            None,
            None,
        )
        .unwrap();
        // R rejects the third column: `qr(X)$rank` is 2 and `$pivot` is 1 2 3.
        assert_eq!(fit.rank, 2);
        assert_eq!(fit.pivot, vec![0, 1, 2]);
        assert_eq!(fit.cov_coefficients.len(), 4);
        assert_close(&fit.cov_coefficients, &A_COV);
    }

    #[test]
    fn test_sample_weights_match_limma() {
        let w = Recycled::by_sample(vec![1.0, 2.0, 0.5, 4.0, 0.25]);
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            None,
            None,
        )
        .unwrap();
        assert_close(&fit.coefficients, &B_COEF);
        assert_close(&fit.sigma, &B_SIGMA);
        // Per-sample weights are shared by every gene, so is the design metric.
        for row in fit.stdev_unscaled.chunks_exact(2) {
            assert_close(row, &B_SDU);
        }
    }

    #[test]
    fn test_full_weights_match_limma() {
        let w = Recycled::full(fixture_weights(), N_GENES, N_SAMPLES).unwrap();
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            None,
            None,
        )
        .unwrap();
        assert_close(&fit.coefficients, &C_COEF);
        assert_close(&fit.stdev_unscaled, &C_SDU);
        assert_close(&fit.sigma, &C_SIGMA);
        assert_close(&fit.df_residual, &[3.0; 8]);
    }

    #[test]
    fn test_missing_value_refits_row() {
        let mut y = fixture_y();
        y[3 * N_SAMPLES + 2] = f64::NAN;
        for s in 0..N_SAMPLES {
            y[5 * N_SAMPLES + s] = f64::NAN;
        }
        let fit = lm_fit(&y, N_GENES, N_SAMPLES, &design_two(), 2, None, None, None).unwrap();
        assert_close(&fit.coefficients, &D_COEF);
        assert_close(&fit.stdev_unscaled, &D_SDU);
        assert_close(&fit.sigma, &D_SIGMA);
        assert_close(&fit.df_residual, &D_DF);
        assert_close(&fit.fitted, &D_FITTED);
    }

    #[test]
    fn test_missing_value_with_full_weights() {
        let mut y = fixture_y();
        y[3 * N_SAMPLES + 2] = f64::NEG_INFINITY;
        for s in 0..N_SAMPLES {
            y[5 * N_SAMPLES + s] = f64::NAN;
        }
        let w = Recycled::full(fixture_weights(), N_GENES, N_SAMPLES).unwrap();
        let fit = lm_fit(
            &y,
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            None,
            None,
        )
        .unwrap();
        assert_close(&fit.coefficients, &E_COEF);
        assert_close(&fit.stdev_unscaled, &E_SDU);
        assert_close(&fit.sigma, &E_SIGMA);
        assert_close(&fit.df_residual, &D_DF);
    }

    #[test]
    fn test_zero_weight_drops_the_observation() {
        // limma sets the response to NA where the weight vanishes, so this must
        // reproduce the missing-value fit.
        let mut weights = fixture_weights();
        weights[3 * N_SAMPLES + 2] = 0.0;
        for s in 0..N_SAMPLES {
            weights[5 * N_SAMPLES + s] = 0.0;
        }
        let w = Recycled::full(weights, N_GENES, N_SAMPLES).unwrap();
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            None,
            None,
        )
        .unwrap();
        assert_close(&fit.coefficients, &E_COEF);
        assert_close(&fit.stdev_unscaled, &E_SDU);
        assert_close(&fit.sigma, &E_SIGMA);
        assert_close(&fit.df_residual, &D_DF);
    }

    #[test]
    fn test_block_correlation_matches_limma() {
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            None,
            Some(&blocks()),
            Some(0.25),
        )
        .unwrap();
        assert_close(&fit.coefficients, &F_COEF);
        assert_close(&fit.sigma, &F_SIGMA);
        assert_close(&fit.df_residual, &[3.0; 8]);
        for row in fit.stdev_unscaled.chunks_exact(2) {
            assert_close(row, &F_SDU);
        }
    }

    #[test]
    fn test_block_correlation_with_full_weights() {
        let w = Recycled::full(fixture_weights(), N_GENES, N_SAMPLES).unwrap();
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            Some(&blocks()),
            Some(0.25),
        )
        .unwrap();
        assert_close(&fit.coefficients, &G_COEF);
        assert_close(&fit.stdev_unscaled, &G_SDU);
        assert_close(&fit.sigma, &G_SIGMA);
        assert_close(&fit.fitted[..N_SAMPLES], &G_FITTED_GENE0);
    }

    #[test]
    fn test_block_correlation_with_sample_weights() {
        let w = Recycled::by_sample(vec![1.0, 2.0, 0.5, 4.0, 0.25]);
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            Some(&blocks()),
            Some(0.25),
        )
        .unwrap();
        assert_close(&fit.coefficients, &K_COEF);
        assert_close(&fit.sigma, &K_SIGMA);
        for row in fit.stdev_unscaled.chunks_exact(2) {
            assert_close(row, &K_SDU);
        }
    }

    #[test]
    fn test_block_correlation_with_missing_refits_submatrix() {
        let mut y = fixture_y();
        y[3 * N_SAMPLES + 2] = f64::NAN;
        for s in 0..N_SAMPLES {
            y[5 * N_SAMPLES + s] = f64::NAN;
        }
        let w = Recycled::full(fixture_weights(), N_GENES, N_SAMPLES).unwrap();
        let fit = lm_fit(
            &y,
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            Some(&w),
            Some(&blocks()),
            Some(0.25),
        )
        .unwrap();
        assert_close(&fit.coefficients, &H_COEF);
        assert_close(&fit.stdev_unscaled, &H_SDU);
        assert_close(&fit.sigma, &H_SIGMA);
        assert_close(&fit.df_residual, &D_DF);
    }

    #[test]
    fn test_zero_correlation_reduces_to_least_squares() {
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            None,
            Some(&blocks()),
            Some(0.0),
        )
        .unwrap();
        assert_close(&fit.coefficients, &A_COEF);
        assert_close(&fit.sigma, &A_SIGMA);
    }

    #[test]
    fn test_correlation_without_block_is_ignored() {
        // `lmFit(Y, X2, correlation = 0.25)` dispatches to `lm.series`, which
        // never sees the correlation.
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_two(),
            2,
            None,
            None,
            Some(0.25),
        )
        .unwrap();
        assert_close(&fit.coefficients, &A_COEF);
        assert_close(&fit.sigma, &A_SIGMA);
    }

    #[test]
    fn test_rank_deficient_design_aliases_the_last_column() {
        // R's `lm.fit` accepts columns left to right, so the third is dropped
        // (`lmFit` prints "Coefficients not estimable: 3").
        let fit = lm_fit(
            &fixture_y(),
            N_GENES,
            N_SAMPLES,
            &design_aliased(),
            3,
            None,
            None,
            None,
        )
        .unwrap();
        for gene in 0..N_GENES {
            let coef = &fit.coefficients[gene * 3..(gene + 1) * 3];
            let sdu = &fit.stdev_unscaled[gene * 3..(gene + 1) * 3];
            assert_close(&coef[..2], &A_COEF[gene * 2..(gene + 1) * 2]);
            assert_close(&sdu[..2], &A_SDU[gene * 2..(gene + 1) * 2]);
            assert!(coef[2].is_nan());
            assert!(sdu[2].is_nan());
        }
        assert_close(&fit.sigma, &A_SIGMA);
        assert_close(&fit.df_residual, &[3.0; 8]);
        // limma's `fitted()` would propagate the NA coefficient across the row;
        // here the aliased column contributes nothing.
        assert_close(&fit.fitted, &A_FITTED);
    }

    #[test]
    fn test_single_gene() {
        let y: Vec<f64> = fixture_y()[..N_SAMPLES].to_vec();
        let fit = lm_fit(&y, 1, N_SAMPLES, &design_two(), 2, None, None, None).unwrap();
        assert_close(&fit.coefficients, &A_COEF[..2]);
        assert_close(&fit.stdev_unscaled, &A_SDU[..2]);
        assert_close(&fit.sigma, &A_SIGMA[..1]);
        assert_close(&fit.df_residual, &[3.0]);
        assert_close(&fit.fitted, &A_FITTED[..N_SAMPLES]);
    }

    #[test]
    fn test_saturated_design_gives_zero_sigma() {
        // Five coefficients on five samples: limma would report NA, this crate
        // reports zero.
        let mut design = vec![0.0; N_SAMPLES * N_SAMPLES];
        for s in 0..N_SAMPLES {
            design[s * N_SAMPLES + s] = 1.0;
        }
        let y = fixture_y();
        let fit = lm_fit(&y, N_GENES, N_SAMPLES, &design, N_SAMPLES, None, None, None).unwrap();
        assert_close(&fit.df_residual, &[0.0; 8]);
        assert_close(&fit.sigma, &[0.0; 8]);
        assert_close(&fit.coefficients, &y);
        assert_close(&fit.fitted, &y);
    }

    ////////////////////
    // Error branches //
    ////////////////////

    #[test]
    fn test_rejects_empty_input() {
        assert!(matches!(
            lm_fit(&[], 0, N_SAMPLES, &design_two(), 2, None, None, None),
            Err(EdgeErrors::EmptyCounts { n_genes: 0, .. })
        ));
        assert!(matches!(
            lm_fit(&[], N_GENES, 0, &design_two(), 2, None, None, None),
            Err(EdgeErrors::EmptyCounts { n_samples: 0, .. })
        ));
    }

    #[test]
    fn test_rejects_zero_coefficients() {
        assert!(matches!(
            lm_fit(&fixture_y(), N_GENES, N_SAMPLES, &[], 0, None, None, None),
            Err(EdgeErrors::MustBePositive(_))
        ));
    }

    #[test]
    fn test_rejects_wrong_response_length() {
        assert!(matches!(
            lm_fit(
                &[0.0; 7],
                N_GENES,
                N_SAMPLES,
                &design_two(),
                2,
                None,
                None,
                None
            ),
            Err(EdgeErrors::LengthMismatch {
                name: "y",
                expected: 40,
                got: 7
            })
        ));
    }

    #[test]
    fn test_rejects_wrong_design_length() {
        assert!(matches!(
            lm_fit(
                &fixture_y(),
                N_GENES,
                N_SAMPLES,
                &[1.0; 8],
                2,
                None,
                None,
                None
            ),
            Err(EdgeErrors::LengthMismatch {
                name: "design",
                expected: 10,
                got: 8
            })
        ));
    }

    #[test]
    fn test_rejects_mismatched_weights() {
        let w = Recycled::by_sample(vec![1.0, 2.0]);
        assert!(matches!(
            lm_fit(
                &fixture_y(),
                N_GENES,
                N_SAMPLES,
                &design_two(),
                2,
                Some(&w),
                None,
                None
            ),
            Err(EdgeErrors::LengthMismatch {
                name: "by_sample",
                ..
            })
        ));
    }

    #[test]
    fn test_rejects_wrong_block_length() {
        assert!(matches!(
            lm_fit(
                &fixture_y(),
                N_GENES,
                N_SAMPLES,
                &design_two(),
                2,
                None,
                Some(&[1, 1, 2]),
                Some(0.25)
            ),
            Err(EdgeErrors::LengthMismatch {
                name: "block",
                expected: 5,
                got: 3
            })
        ));
    }

    #[test]
    fn test_rejects_block_without_correlation() {
        assert!(matches!(
            lm_fit(
                &fixture_y(),
                N_GENES,
                N_SAMPLES,
                &design_two(),
                2,
                None,
                Some(&blocks()),
                None
            ),
            Err(EdgeErrors::InvalidArgument(_))
        ));
    }

    #[test]
    fn test_rejects_out_of_range_correlation() {
        for rho in [1.0, -1.0, 2.5, f64::NAN] {
            assert!(matches!(
                lm_fit(
                    &fixture_y(),
                    N_GENES,
                    N_SAMPLES,
                    &design_two(),
                    2,
                    None,
                    Some(&blocks()),
                    Some(rho)
                ),
                Err(EdgeErrors::InvalidArgument(_))
            ));
        }
    }

    #[test]
    fn test_rejects_indefinite_block_covariance() {
        // One block of five needs a correlation above -1/4; limma's `chol` fails
        // at -0.5 too.
        assert!(matches!(
            lm_fit(
                &fixture_y(),
                N_GENES,
                N_SAMPLES,
                &design_two(),
                2,
                None,
                Some(&[1, 1, 1, 1, 1]),
                Some(-0.5)
            ),
            Err(EdgeErrors::CholeskyFailed(_))
        ));
    }

    /// Fits once through the shared path and once forced per gene, and
    /// compares every output.
    ///
    /// ### Params
    ///
    /// * `design` - Row-major design, `n_samples * n_coef`
    /// * `n_samples` - Number of samples
    /// * `n_coef` - Number of coefficients
    /// * `weights` - Optional weights
    fn assert_shared_matches_per_gene(
        design: &[f64],
        n_samples: usize,
        n_coef: usize,
        weights: Option<&Recycled<f64>>,
    ) {
        use rand::prelude::*;
        use rand_distr::Normal;

        let n_genes = 200;
        let mut rng = StdRng::seed_from_u64(7);
        let noise = Normal::new(0.0, 0.7).unwrap();
        // A spread of means, so some genes fit tightly at high expression.
        let y: Vec<f64> = (0..n_genes * n_samples)
            .map(|i| (i / n_samples) as f64 * 0.1 + rng.sample(noise))
            .collect();

        let fit = |allow_shared| {
            lm_fit_impl(
                &y,
                n_genes,
                n_samples,
                design,
                n_coef,
                weights,
                None,
                None,
                allow_shared,
            )
            .unwrap()
        };
        let (fast, slow) = (fit(true), fit(false));

        // Absolute or relative. Measured worst absolute difference over the
        // three cases is 2.1e-14 (coefficients, fitted); relative goes to
        // 2.3e-11 only for coefficients near zero.
        let tol = 5e-14;
        for (name, a, b) in [
            ("coefficients", &fast.coefficients, &slow.coefficients),
            ("stdev_unscaled", &fast.stdev_unscaled, &slow.stdev_unscaled),
            ("sigma", &fast.sigma, &slow.sigma),
            ("df_residual", &fast.df_residual, &slow.df_residual),
            ("fitted", &fast.fitted, &slow.fitted),
        ] {
            assert_eq!(a.len(), b.len(), "{name}");
            for (x, z) in a.iter().zip(b) {
                if z.is_nan() {
                    assert!(x.is_nan(), "{name}: {x} vs NaN");
                } else {
                    assert_relative_eq!(*x, *z, max_relative = tol, epsilon = tol);
                }
            }
        }
        assert_eq!(fast.rank, slow.rank);
        assert_eq!(fast.pivot, slow.pivot);
        assert_eq!(fast.cov_coefficients, slow.cov_coefficients);
    }

    /// Intercept, a three-level factor and a continuous covariate over 30
    /// samples.
    fn design_factor_covariate() -> Vec<f64> {
        (0..30)
            .flat_map(|s| {
                let level = s % 3;
                [
                    1.0,
                    (level == 1) as u8 as f64,
                    (level == 2) as u8 as f64,
                    (s as f64 * 0.37).sin(),
                ]
            })
            .collect()
    }

    #[test]
    fn test_shared_path_matches_per_gene_path_unweighted() {
        assert_shared_matches_per_gene(&design_factor_covariate(), 30, 4, None);
    }

    #[test]
    fn test_shared_path_matches_per_gene_path_with_sample_weights() {
        let w: Vec<f64> = (0..30).map(|s| 0.5 + (s % 7) as f64 / 4.0).collect();
        assert_shared_matches_per_gene(
            &design_factor_covariate(),
            30,
            4,
            Some(&Recycled::BySample(w)),
        );
    }

    #[test]
    fn test_shared_path_matches_per_gene_path_rank_deficient() {
        // Column 4 is the sum of the two factor columns, so it is aliased.
        let design: Vec<f64> = design_factor_covariate()
            .chunks_exact(4)
            .flat_map(|r| [r[0], r[1], r[2], r[3], r[1] + r[2]])
            .collect();
        assert_shared_matches_per_gene(&design, 30, 5, None);
    }
}