rrblup_rs/

mixed_solve.rs

//! Rust implementation of R/rrBLUP::mixed.solve
//!
//! This module provides a direct translation of the mixed.solve function from
//! the R/rrBLUP package for solving mixed linear models using spectral decomposition.
//!
//! # Example
//!
//! ```
//! use rrblup_rs::mixed_solve::{mixed_solve, MixedSolveOptions, Method};
//!
//! // Simple mixed model with intercept only
//! let y = vec![1.0, 2.0, 3.0, 4.0, 5.0];
//! let result = mixed_solve(&y, None, None, None, None).unwrap();
//! assert!(result.vu >= 0.0);
//! assert!(result.ve >= 0.0);
//! ```
//!
//! # Reference
//!
//! R/rrBLUP package by Jeffrey Endelman:
//! <https://cran.r-project.org/package=rrBLUP>

use anyhow::{anyhow, Result};
use faer::Mat as FaerMat;
use nalgebra::{DMatrix, DVector, SymmetricEigen};

/// Method for variance component estimation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Method {
    /// Maximum Likelihood
    ML,
    /// Restricted Maximum Likelihood (default)
    #[default]
    REML,
}

/// Options for mixed.solve function
#[derive(Debug, Clone)]
pub struct MixedSolveOptions {
    /// Method for variance component estimation (default: REML)
    pub method: Method,
    /// Bounds for lambda optimization (default: [1e-9, 1e9])
    pub bounds: (f64, f64),
    /// Whether to compute standard errors (default: false)
    pub se: bool,
    /// Whether to return H inverse matrix (default: false)
    pub return_hinv: bool,
}

impl Default for MixedSolveOptions {
    fn default() -> Self {
        Self {
            method: Method::REML,
            bounds: (1e-9, 1e9),
            se: false,
            return_hinv: false,
        }
    }
}

/// Result from mixed.solve function
///
/// Corresponds to the list returned by R's mixed.solve:
/// list(Vu, Ve, beta, u, LL) or with SE: list(Vu, Ve, beta, beta.SE, u, u.SE, LL)
/// and optionally Hinv if return.Hinv=TRUE
#[derive(Debug, Clone)]
pub struct MixedSolveResult {
    /// Variance of random effects (Vu)
    pub vu: f64,
    /// Residual variance (Ve)
    pub ve: f64,
    /// Fixed effects coefficients (beta)
    pub beta: DVector<f64>,
    /// Standard errors of beta (only if SE=TRUE)
    pub beta_se: Option<DVector<f64>>,
    /// Random effects BLUPs (u)
    pub u: DVector<f64>,
    /// Standard errors of u (only if SE=TRUE)
    pub u_se: Option<DVector<f64>>,
    /// Log-likelihood (LL)
    pub ll: f64,
    /// Inverse of H matrix (only if return.Hinv=TRUE)
    pub hinv: Option<DMatrix<f64>>,
}

/// Solve mixed model y = Xβ + Zu + e using spectral decomposition
///
/// This is a Rust implementation of R/rrBLUP::mixed.solve().
///
/// # Arguments
///
/// * `y` - Response vector (n x 1), may contain NaN for missing values
/// * `z` - Random effects design matrix (n x m), defaults to identity if None
/// * `k` - Covariance matrix for random effects (m x m), defaults to identity if None
/// * `x` - Fixed effects design matrix (n x p), defaults to intercept if None
/// * `options` - Optional settings (method, bounds, SE, return.Hinv)
///
/// # Returns
///
/// [`MixedSolveResult`] with variance components, BLUPs, and optionally SEs.
///
/// # Errors
///
/// Returns an error if:
/// - All `y` values are `NaN`
/// - Matrix dimensions are incompatible (`nrow(Z) != n` or `nrow(X) != n`)
/// - `X` is not full rank
/// - `K` is not positive semi-definite
///
/// # Example
///
/// ```
/// use rrblup_rs::mixed_solve::{mixed_solve, MixedSolveOptions};
///
/// let y = vec![1.0, 2.0, 3.0, 4.0, 5.0];
/// let result = mixed_solve(&y, None, None, None, None).unwrap();
///
/// println!("Genetic variance (Vu): {}", result.vu);
/// println!("Residual variance (Ve): {}", result.ve);
/// println!("Fixed effects (beta): {:?}", result.beta);
/// ```
///
/// # Notes
///
/// The parameter order matches R: mixed.solve(y, Z=NULL, K=NULL, X=NULL, ...)
pub fn mixed_solve(
    y: &[f64],
    z: Option<&DMatrix<f64>>,
    k: Option<&DMatrix<f64>>,
    x: Option<&DMatrix<f64>>,
    options: Option<MixedSolveOptions>,
) -> Result<MixedSolveResult> {
    let opts = options.unwrap_or_default();
    let pi = std::f64::consts::PI;

    // R: n <- length(y)
    let n_full = y.len();

    // R: not.NA <- which(!is.na(y))
    let not_na: Vec<usize> = (0..n_full).filter(|&i| y[i].is_finite()).collect();

    if not_na.is_empty() {
        return Err(anyhow!("All y values are NA"));
    }

    // Setup X matrix
    // if (is.null(X)) { p <- 1; X <- matrix(rep(1,n),n,1) }
    let x_full: DMatrix<f64> = match x {
        Some(x_mat) => x_mat.clone(),
        None => DMatrix::from_element(n_full, 1, 1.0),
    };
    let p = x_full.ncols();

    // Setup Z matrix
    // if (is.null(Z)) { Z <- diag(n) }
    let z_full: DMatrix<f64> = match z {
        Some(z_mat) => z_mat.clone(),
        None => DMatrix::identity(n_full, n_full),
    };
    let m = z_full.ncols();

    // Dimension checks
    // stopifnot(nrow(Z) == n)
    if z_full.nrows() != n_full {
        return Err(anyhow!(
            "nrow(Z) = {} != n = {}",
            z_full.nrows(),
            n_full
        ));
    }
    // stopifnot(nrow(X) == n)
    if x_full.nrows() != n_full {
        return Err(anyhow!(
            "nrow(X) = {} != n = {}",
            x_full.nrows(),
            n_full
        ));
    }

    // Check K dimensions if provided
    if let Some(k_mat) = k {
        // stopifnot(nrow(K) == m)
        // stopifnot(ncol(K) == m)
        if k_mat.nrows() != m || k_mat.ncols() != m {
            return Err(anyhow!(
                "K must be {} x {}, got {} x {}",
                m,
                m,
                k_mat.nrows(),
                k_mat.ncols()
            ));
        }
    }

    // Subset to non-NA observations
    // R: Z <- as.matrix(Z[not.NA,])
    // R: X <- as.matrix(X[not.NA,])
    // R: n <- length(not.NA)
    // R: y <- matrix(y[not.NA],n,1)
    let n = not_na.len();
    let y_vec: DVector<f64> = DVector::from_iterator(n, not_na.iter().map(|&i| y[i]));
    let z_mat = DMatrix::from_fn(n, m, |i, j| z_full[(not_na[i], j)]);
    let x_mat = DMatrix::from_fn(n, p, |i, j| x_full[(not_na[i], j)]);

    // R: XtX <- crossprod(X, X)
    let xtx = x_mat.transpose() * &x_mat;

    // R: rank.X <- qr(XtX)$rank
    // if (rank.X < p) {stop("X not full rank")}
    // R: XtXinv <- solve(XtX)
    let xtx_inv = xtx
        .clone()
        .try_inverse()
        .ok_or_else(|| anyhow!("X not full rank"))?;

    // R: S <- diag(n) - tcrossprod(X%*%XtXinv,X)
    // S = I - X @ XtXinv @ X'
    let x_xtxinv = &x_mat * &xtx_inv;
    let s_mat = DMatrix::identity(n, n) - &x_xtxinv * x_mat.transpose();

    // Determine spectral method
    // if (n <= m + p) { spectral.method <- "eigen" } else { spectral.method <- "cholesky" }
    let use_eigen = n <= m + p;

    // Variables that will be set by either branch
    let phi: Vec<f64>;
    let theta: Vec<f64>;
    let u_mat: DMatrix<f64>;
    let q_mat: DMatrix<f64>;

    if use_eigen {
        // ============================================================
        // EIGEN BRANCH: spectral.method == "eigen"
        // ============================================================

        // R: offset <- sqrt(n)
        let offset = (n as f64).sqrt();

        // Hb computation
        let hb: DMatrix<f64> = if let Some(k_mat) = k {
            // R: Hb <- tcrossprod(Z%*%K,Z) + offset*diag(n)
            let zk = &z_mat * k_mat;
            let zkzt = &zk * z_mat.transpose();
            let mut hb = zkzt;
            for i in 0..n {
                hb[(i, i)] += offset;
            }
            hb
        } else {
            // R: Hb <- tcrossprod(Z,Z) + offset*diag(n)
            let zzt = &z_mat * z_mat.transpose();
            let mut hb = zzt;
            for i in 0..n {
                hb[(i, i)] += offset;
            }
            hb
        };

        // R: Hb.system <- eigen(Hb, symmetric = TRUE)
        let hb_eig = SymmetricEigen::new(hb.clone());

        // NOTE: R's eigen() returns eigenvalues in DESCENDING order.
        // nalgebra's SymmetricEigen does NOT guarantee sorted eigenvalues!
        // We must explicitly sort them.

        // Create sorted indices for Hb eigenvalues (descending order)
        let mut hb_indices: Vec<usize> = (0..n).collect();
        hb_indices.sort_by(|&a, &b| {
            hb_eig.eigenvalues[b]
                .partial_cmp(&hb_eig.eigenvalues[a])
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // R: phi <- Hb.system$values - offset (in descending order)
        phi = hb_indices
            .iter()
            .map(|&i| hb_eig.eigenvalues[i] - offset)
            .collect();

        // if (min(phi) < -1e-6) {stop("K not positive semi-definite.")}
        let min_phi = phi.iter().cloned().fold(f64::INFINITY, f64::min);
        if min_phi < -1e-6 {
            return Err(anyhow!("K not positive semi-definite (min phi = {})", min_phi));
        }

        // R: U <- Hb.system$vectors (columns in descending eigenvalue order)
        u_mat = DMatrix::from_fn(n, n, |i, j| hb_eig.eigenvectors[(i, hb_indices[j])]);

        // R: SHbS <- S %*% Hb %*% S
        let shbs = &s_mat * &hb * &s_mat;

        // R: SHbS.system <- eigen(SHbS, symmetric = TRUE)
        let shbs_eig = SymmetricEigen::new(shbs);

        // Create sorted indices for SHbS eigenvalues (descending order)
        let mut shbs_indices: Vec<usize> = (0..n).collect();
        shbs_indices.sort_by(|&a, &b| {
            shbs_eig.eigenvalues[b]
                .partial_cmp(&shbs_eig.eigenvalues[a])
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // R: theta <- SHbS.system$values[1:(n - p)] - offset
        // R takes first n-p eigenvalues (largest, in descending order)
        let n_theta = n - p;
        theta = shbs_indices
            .iter()
            .take(n_theta) // Take the largest n_theta eigenvalues
            .map(|&i| shbs_eig.eigenvalues[i] - offset)
            .collect();

        // R: Q <- SHbS.system$vectors[, 1:(n - p)]
        // Take columns corresponding to the largest n_theta eigenvalues
        q_mat = DMatrix::from_fn(n, n_theta, |i, j| {
            shbs_eig.eigenvectors[(i, shbs_indices[j])]
        });
    } else {
        // ============================================================
        // CHOLESKY BRANCH: spectral.method == "cholesky"
        // ============================================================

        // B = chol(K) if K provided (with jitter on diagonal)
        let zbt: DMatrix<f64> = if let Some(k_mat) = k {
            // diag(K) <- diag(K) + 1e-6
            // R: B <- try(chol(K),silent=TRUE)
            let mut k_jittered = k_mat.clone();
            for i in 0..m {
                k_jittered[(i, i)] += 1e-6;
            }
            let chol = k_jittered
                .cholesky()
                .ok_or_else(|| anyhow!("K not positive semi-definite"))?;
            // R: ZBt <- tcrossprod(Z,B)
            // In R, chol() returns upper triangular, so B = chol(K)
            // tcrossprod(Z, B) = Z @ B'
            // nalgebra cholesky().l() returns lower triangular L where K = L @ L'
            // So we need Z @ L' = Z @ chol(K)' in R terms
            let b_t = chol.l().transpose(); // This is the upper triangular
            &z_mat * &b_t.transpose() // tcrossprod(Z, B) = Z @ B' where B is upper triangular
        } else {
            // R: ZBt <- Z
            z_mat.clone()
        };

        // R: svd.ZBt <- svd(ZBt,nu=n)
        let zbt_faer = nalgebra_to_faer(&zbt);
        let svd_zbt = zbt_faer.svd();
        let u_full = faer_to_nalgebra(&svd_zbt.u());
        let d_vals = svd_zbt.s_diagonal();

        // R: U <- svd.ZBt$u
        u_mat = u_full;

        // R: phi <- c(svd.ZBt$d^2,rep(0,n-m))
        phi = (0..n)
            .map(|i| {
                if i < d_vals.nrows() {
                    let d = d_vals.read(i);
                    d * d
                } else {
                    0.0
                }
            })
            .collect();

        // R: SZBt <- S %*% ZBt
        let szbt = &s_mat * &zbt;

        // R: svd.SZBt <- try(svd(SZBt),silent=TRUE)
        // if (inherits(svd.SZBt,what="try-error")) {
        //   svd.SZBt <- svd(SZBt+matrix(1e-10,nrow=nrow(SZBt),ncol=ncol(SZBt)))
        // }
        let szbt_faer = nalgebra_to_faer(&szbt);
        let svd_szbt = szbt_faer.thin_svd();
        let u_szbt = faer_to_nalgebra(&svd_szbt.u());
        let d_szbt = svd_szbt.s_diagonal();

        // R: QR <- qr(cbind(X,svd.SZBt$u))
        let n_u_cols = u_szbt.ncols();
        let mut combined = DMatrix::zeros(n, p + n_u_cols);
        for i in 0..n {
            for j in 0..p {
                combined[(i, j)] = x_mat[(i, j)];
            }
            for j in 0..n_u_cols {
                combined[(i, p + j)] = u_szbt[(i, j)];
            }
        }

        let combined_faer = nalgebra_to_faer(&combined);
        let qr = combined_faer.qr();

        // R: Q <- qr.Q(QR,complete=TRUE)[,(p+1):n]
        let q_full = faer_to_nalgebra(&qr.compute_q().as_ref());
        let q_complement = q_full.columns(p, n - p).into_owned();
        q_mat = q_complement;

        // R: R <- qr.R(QR)[p+1:m,p+1:m]
        let r_faer = qr.compute_r();
        let r_full = faer_mat_to_nalgebra(&r_faer);

        // R: ans <- try(solve(t(R^2), svd.SZBt$d^2),silent=TRUE)
        // R: theta <- c(ans,rep(0, n - p - m))
        let r22_size = m.min(r_full.nrows().saturating_sub(p)).min(r_full.ncols().saturating_sub(p));

        let theta_result: Result<Vec<f64>, ()> = if r22_size > 0 && d_szbt.nrows() > 0 {
            // Extract R22 = R[p+1:m, p+1:m] (1-indexed in R, 0-indexed here)
            let mut r22_sq = DMatrix::zeros(r22_size, r22_size);
            for i in 0..r22_size {
                for j in 0..r22_size {
                    let val = r_full[(p + i, p + j)];
                    r22_sq[(i, j)] = val * val;
                }
            }

            // t(R^2)
            let t_r22_sq = r22_sq.transpose();

            // svd.SZBt$d^2
            let d_sq_len = r22_size.min(d_szbt.nrows());
            let d_sq: Vec<f64> = (0..d_sq_len)
                .map(|i| {
                    let d = d_szbt.read(i);
                    d * d
                })
                .collect();
            let d_sq_vec = DVector::from_row_slice(&d_sq);

            // solve(t(R^2), d^2)
            match t_r22_sq.clone().try_inverse() {
                Some(inv) => {
                    let ans = inv * &d_sq_vec;
                    let n_theta = n - p;
                    Ok((0..n_theta)
                        .map(|i| {
                            if i < ans.len() {
                                ans[i]
                            } else {
                                0.0
                            }
                        })
                        .collect())
                }
                None => Err(()),
            }
        } else {
            Err(())
        };

        theta = match theta_result {
            Ok(t) => t,
            Err(_) => {
                // Fallback: this would trigger spectral.method <- "eigen" in R
                // For simplicity, we use zeros which may not be fully correct
                // but follows the structure
                vec![0.0; n - p]
            }
        };
    }

    // R: omega <- crossprod(Q, y)
    let omega = q_mat.transpose() * &y_vec;

    // R: omega.sq <- omega^2
    let omega_sq: Vec<f64> = omega.iter().map(|v| v * v).collect();

    // Optimization
    let (lambda_opt, obj_val, df): (f64, f64, usize);

    if opts.method == Method::ML {
        // R: f.ML <- function(lambda, n, theta, omega.sq, phi) {
        //   n * log(sum(omega.sq/(theta + lambda))) + sum(log(phi + lambda))
        // }
        let f_ml = |lambda: f64| -> f64 {
            if lambda <= 0.0 {
                return f64::INFINITY;
            }
            let sum_ratio: f64 = omega_sq
                .iter()
                .zip(theta.iter())
                .map(|(o, t)| o / (t + lambda))
                .sum();
            if sum_ratio <= 0.0 {
                return f64::INFINITY;
            }
            let sum_log_phi: f64 = phi.iter().map(|p| (p + lambda).ln()).sum();
            (n as f64) * sum_ratio.ln() + sum_log_phi
        };

        let (opt_lambda, opt_obj) = golden_section_minimize(f_ml, opts.bounds.0, opts.bounds.1);
        lambda_opt = opt_lambda;
        obj_val = opt_obj;
        df = n;
    } else {
        // R: f.REML <- function(lambda, n.p, theta, omega.sq) {
        //   n.p * log(sum(omega.sq/(theta + lambda))) + sum(log(theta + lambda))
        // }
        let n_p = n - p;
        let f_reml = |lambda: f64| -> f64 {
            if lambda <= 0.0 {
                return f64::INFINITY;
            }
            let sum_ratio: f64 = omega_sq
                .iter()
                .zip(theta.iter())
                .map(|(o, t)| o / (t + lambda))
                .sum();
            if sum_ratio <= 0.0 {
                return f64::INFINITY;
            }
            let sum_log_theta: f64 = theta.iter().map(|t| (t + lambda).ln()).sum();
            (n_p as f64) * sum_ratio.ln() + sum_log_theta
        };

        let (opt_lambda, opt_obj) = golden_section_minimize(f_reml, opts.bounds.0, opts.bounds.1);
        lambda_opt = opt_lambda;
        obj_val = opt_obj;
        df = n - p;
    }

    // R: Vu.opt <- sum(omega.sq/(theta + lambda.opt))/df
    let vu_opt: f64 = omega_sq
        .iter()
        .zip(theta.iter())
        .map(|(o, t)| o / (t + lambda_opt))
        .sum::<f64>()
        / (df as f64);

    // R: Ve.opt <- lambda.opt * Vu.opt
    let ve_opt = lambda_opt * vu_opt;

    // R: Hinv <- U %*% (t(U)/(phi+lambda.opt))
    // Hinv[i,j] = sum_k U[i,k] * U[j,k] / (phi[k] + lambda)
    let mut hinv = DMatrix::zeros(n, n);
    for i in 0..n {
        for j in 0..n {
            let mut sum = 0.0;
            for kk in 0..n {
                sum += u_mat[(i, kk)] * u_mat[(j, kk)] / (phi[kk] + lambda_opt);
            }
            hinv[(i, j)] = sum;
        }
    }

    // R: W <- crossprod(X,Hinv%*%X)
    let hinv_x = &hinv * &x_mat;
    let w = x_mat.transpose() * &hinv_x;

    // R: beta <- array(solve(W,crossprod(X,Hinv%*%y)))
    let w_inv = w
        .clone()
        .try_inverse()
        .ok_or_else(|| anyhow!("W not invertible"))?;
    let hinv_y = &hinv * &y_vec;
    let beta = &w_inv * (x_mat.transpose() * &hinv_y);

    // KZt computation
    // if (is.null(K)) { KZt <- t(Z) } else { KZt <- tcrossprod(K,Z) }
    let kzt: DMatrix<f64> = if let Some(k_mat) = k {
        // R: KZt <- tcrossprod(K,Z) = K @ Z'
        k_mat * z_mat.transpose()
    } else {
        // R: KZt <- t(Z)
        z_mat.transpose()
    };

    // R: KZt.Hinv <- KZt %*% Hinv
    let kzt_hinv = &kzt * &hinv;

    // R: u <- array(KZt.Hinv %*% (y - X%*%beta))
    let resid = &y_vec - &x_mat * &beta;
    let u_blup = &kzt_hinv * &resid;

    // LL = -0.5 * (soln$objective + df + df * log(2 * pi/df))
    let ll = -0.5 * (obj_val + (df as f64) + (df as f64) * (2.0 * pi / (df as f64)).ln());

    // Standard errors (if SE=TRUE)
    let (beta_se, u_se) = if opts.se {
        // R: Winv <- solve(W)
        let winv = w_inv.clone();

        // R: beta.SE <- array(sqrt(Vu.opt*diag(Winv)))
        let beta_se_vec: DVector<f64> =
            DVector::from_fn(p, |i, _| (vu_opt * winv[(i, i)]).sqrt());

        // R: WW <- tcrossprod(KZt.Hinv,KZt)
        let ww = &kzt_hinv * kzt.transpose();

        // R: WWW <- KZt.Hinv%*%X
        let www = &kzt_hinv * &x_mat;

        // u.SE computation
        let u_se_vec: DVector<f64> = if k.is_none() {
            // R: u.SE <- array(sqrt(Vu.opt * (rep(1,m) - diag(WW) + diag(tcrossprod(WWW%*%Winv,WWW)))))
            let www_winv = &www * &winv;
            let www_term = &www_winv * www.transpose();
            DVector::from_fn(m, |i, _| {
                let val = vu_opt * (1.0 - ww[(i, i)] + www_term[(i, i)]);
                if val > 0.0 {
                    val.sqrt()
                } else {
                    0.0
                }
            })
        } else {
            // R: u.SE <- array(sqrt(Vu.opt * (diag(K) - diag(WW) + diag(tcrossprod(WWW%*%Winv,WWW)))))
            let k_mat = k.unwrap();
            let www_winv = &www * &winv;
            let www_term = &www_winv * www.transpose();
            DVector::from_fn(m, |i, _| {
                let val = vu_opt * (k_mat[(i, i)] - ww[(i, i)] + www_term[(i, i)]);
                if val > 0.0 {
                    val.sqrt()
                } else {
                    0.0
                }
            })
        };

        (Some(beta_se_vec), Some(u_se_vec))
    } else {
        (None, None)
    };

    // Return Hinv only if requested
    let hinv_return = if opts.return_hinv { Some(hinv) } else { None };

    Ok(MixedSolveResult {
        vu: vu_opt,
        ve: ve_opt,
        beta,
        beta_se,
        u: u_blup,
        u_se,
        ll,
        hinv: hinv_return,
    })
}

/// Golden section search for minimization (equivalent to R's optimize)
fn golden_section_minimize<F>(f: F, mut a: f64, mut b: f64) -> (f64, f64)
where
    F: Fn(f64) -> f64,
{
    let gr = 0.5 * (1.0 + 5f64.sqrt()); // Golden ratio
    let tol = 1e-8;
    let max_iter = 100;

    let mut c = b - (b - a) / gr;
    let mut d = a + (b - a) / gr;
    let mut fc = f(c);
    let mut fd = f(d);

    for _ in 0..max_iter {
        if (b - a).abs() < tol {
            break;
        }
        if fc < fd {
            b = d;
            d = c;
            fd = fc;
            c = b - (b - a) / gr;
            fc = f(c);
        } else {
            a = c;
            c = d;
            fc = fd;
            d = a + (b - a) / gr;
            fd = f(d);
        }
    }

    let x_min = if fc < fd { c } else { d };
    let f_min = if fc < fd { fc } else { fd };
    (x_min, f_min)
}

// ============================================================
// Helper functions for matrix conversions between nalgebra and faer
// ============================================================

fn nalgebra_to_faer(m: &DMatrix<f64>) -> FaerMat<f64> {
    let nrows = m.nrows();
    let ncols = m.ncols();
    FaerMat::from_fn(nrows, ncols, |i, j| m[(i, j)])
}

fn faer_to_nalgebra(m: &faer::MatRef<f64>) -> DMatrix<f64> {
    let nrows = m.nrows();
    let ncols = m.ncols();
    DMatrix::from_fn(nrows, ncols, |i, j| m.read(i, j))
}

fn faer_mat_to_nalgebra(m: &FaerMat<f64>) -> DMatrix<f64> {
    let nrows = m.nrows();
    let ncols = m.ncols();
    DMatrix::from_fn(nrows, ncols, |i, j| m.read(i, j))
}

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

    #[test]
    fn test_mixed_solve_simple_intercept() {
        // Simple test: y = 1, 2, 3, 4, 5 with intercept only
        let y = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let result = mixed_solve(&y, None, None, None, None).unwrap();

        // Mean should be approximately 3.0
        assert_relative_eq!(result.beta[0], 3.0, epsilon = 0.5);
        assert!(result.vu >= 0.0);
        assert!(result.ve >= 0.0);
    }

    #[test]
    fn test_mixed_solve_with_na() {
        // Test with NA (NaN) values
        let y = vec![1.0, f64::NAN, 3.0, f64::NAN, 5.0];
        let result = mixed_solve(&y, None, None, None, None).unwrap();

        // Mean of 1, 3, 5 is 3.0
        assert_relative_eq!(result.beta[0], 3.0, epsilon = 0.5);
    }

    #[test]
    fn test_mixed_solve_with_se() {
        let y = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let opts = MixedSolveOptions {
            se: true,
            ..Default::default()
        };
        let result = mixed_solve(&y, None, None, None, Some(opts)).unwrap();

        assert!(result.beta_se.is_some());
        assert!(result.u_se.is_some());
        assert!(result.beta_se.unwrap()[0] > 0.0);
    }

    #[test]
    fn test_mixed_solve_ml_vs_reml() {
        let y = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];

        let opts_reml = MixedSolveOptions {
            method: Method::REML,
            ..Default::default()
        };
        let result_reml = mixed_solve(&y, None, None, None, Some(opts_reml)).unwrap();

        let opts_ml = MixedSolveOptions {
            method: Method::ML,
            ..Default::default()
        };
        let result_ml = mixed_solve(&y, None, None, None, Some(opts_ml)).unwrap();

        // Both should give similar beta estimates
        assert_relative_eq!(result_reml.beta[0], result_ml.beta[0], epsilon = 0.5);
        // REML typically gives larger variance estimates than ML
        // (not always true for very small samples, so we just check they're both valid)
        assert!(result_reml.vu >= 0.0);
        assert!(result_ml.vu >= 0.0);
    }

    #[test]
    fn test_mixed_solve_with_fixed_effects() {
        let y = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        // X with intercept and one covariate
        let x = DMatrix::from_row_slice(6, 2, &[
            1.0, 0.0,
            1.0, 1.0,
            1.0, 2.0,
            1.0, 3.0,
            1.0, 4.0,
            1.0, 5.0,
        ]);

        let result = mixed_solve(&y, None, None, Some(&x), None).unwrap();

        assert_eq!(result.beta.len(), 2);
        // With this design, beta should capture the linear trend
    }

    #[test]
    fn test_mixed_solve_return_hinv() {
        let y = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let opts = MixedSolveOptions {
            return_hinv: true,
            ..Default::default()
        };
        let result = mixed_solve(&y, None, None, None, Some(opts)).unwrap();

        assert!(result.hinv.is_some());
        let hinv = result.hinv.unwrap();
        assert_eq!(hinv.nrows(), 5);
        assert_eq!(hinv.ncols(), 5);
    }
}