R/nmf.gmm.R

In nmfkc: Non-Negative Matrix Factorization with Kernel Covariates

# =====================================================================
#  R/nmf.gmm.R -- NMF-GMM: a Gaussian-mixture latent-class extension of
#                 NMF with covariates (Satoh 2026).
#
#  Model:  b_n | (z_n = k) ~ N_Q(C a_n + mu_k, Sigma_k),  y_n = X b_n + eps_n,
#          X >= 0, colSums(X) = 1, eps ~ N(0, sigma2 I_P), sum_k xi_k mu_k = 0.
#  C (Q x R) is the covariate-coefficient matrix (Theta in the paper); mu
#  (Q x K) are the class means; the mixture on the scores turns NMF-RE's
#  arg-max hard clustering into model-based soft clustering via the posterior
#  responsibilities. K = 1 (tied) reduces exactly to NMF-RE (nmfre()).
#  Optionally (class.effects) chosen rows S of A get a coefficient that
#  differs by class: C a_n + Del_k a_n[S], with sum_k xi_k Del_k = 0.
#
#  House-style API (optimization / inference split):
#    nmf.gmm           EM fit (optimization only)
#    nmf.gmm.inference given-X Wald inference on C (sandwich / outer-product
#                      information + wild bootstrap), tied covariance only
#    nmf.gmm.select    choose K by BIC / ICL (optionally ARI vs known labels)
#
#  The internal .nmfgmm.* engine below is the verbatim research EM core
#  (estep / mstep / inner ANCOVA solve / multistart driver); the public
#  wrappers only rename arguments/returns to the package conventions
#  (Theta -> C, x0 -> X.init, ...) and add BIC/ICL, labels and S3 support.
#  Reference: Satoh K. (2026), NMF-GMM.
# =====================================================================

.nmfgmm.WOODBURY_P <- 100L   # switch to the Woodbury E-step at P >= this

.nmfgmm.pos <- function(M) pmax(M, 0)
.nmfgmm.neg <- function(M) pmax(-M, 0)

## k-means++ seeding (rows of Xr are points)
.nmfgmm.kmpp <- function(Xr, K, seed) {
  set.seed(seed); n <- nrow(Xr); ce <- matrix(NA, K, ncol(Xr))
  ce[1, ] <- Xr[sample(n, 1), ]
  d2 <- colSums((t(Xr) - ce[1, ])^2)
  for (k in 2:K) {
    ## Guard the degenerate case the shared .kmeanspp.seed() already handles:
    ## when every remaining point coincides with a chosen centre, sum(d2) is 0
    ## and d2/sum(d2) is all NaN, which makes sample() error.  Reachable here,
    ## since the input is residual scores that collapse when the covariate
    ## design saturates.
    i <- if (sum(d2) > 0) sample(n, 1, prob = d2) else sample(n, 1)
    ce[k, ] <- Xr[i, ]
    d2 <- pmin(d2, colSums((t(Xr) - ce[k, ])^2))
  }
  ce
}

## ---------------------------------------------------------------------
## E-step
## ---------------------------------------------------------------------
.nmfgmm.estep <- function(par, Y, A, cov = "tied", S = NULL) {
  X <- par$X; Theta <- par$Theta; mu <- par$mu; tau2 <- par$tau2
  sigma2 <- par$sigma2; xi <- par$xi
  P <- nrow(Y); N <- ncol(Y); Q <- ncol(X); K <- length(xi)
  XtX <- crossprod(X); M <- Theta %*% A; XtY <- crossprod(X, Y)
  AS <- if (length(S)) A[S, , drop = FALSE] else NULL   # class-specific rows

  if (cov != "free") {                             # "tied" or "scalar": shared diagonal
    tau2 <- as.numeric(tau2)                       # length-Q vector (all equal if scalar)
    Dinv <- diag(1 / tau2, Q)
    Omega <- solve(Dinv + XtX / sigma2)            # Q x Q, both paths
    use_wood <- (P >= .nmfgmm.WOODBURY_P)
    if (use_wood) {
      W <- Omega / sigma2
      logdetV <- P * log(sigma2) +
        as.numeric(determinant(diag(Q) + XtX %*% diag(tau2, Q) / sigma2,
                               logarithm = TRUE)$modulus)
      YtY <- colSums(Y^2)
    } else {
      V <- X %*% (tau2 * t(X)) + sigma2 * diag(P)
      L <- chol(V); logdetV <- 2 * sum(log(diag(L))); Vinv <- chol2inv(L)
    }
    logcomp <- matrix(0, N, K); bhat <- vector("list", K)
    for (k in 1:K) {
      prior <- M + matrix(mu[, k], Q, N)
      if (length(S)) prior <- prior + par$Del[[k]] %*% AS
      if (use_wood) {
        Xprior <- X %*% prior; Xtr <- XtY - XtX %*% prior
        r_sq <- YtY - 2 * colSums(Y * Xprior) + colSums(Xprior^2)
        quad <- (r_sq - colSums(Xtr * (W %*% Xtr))) / sigma2
      } else {
        Resid <- Y - X %*% prior; quad <- colSums(Resid * (Vinv %*% Resid))
      }
      logcomp[, k] <- log(xi[k]) - 0.5 * (P * log(2 * pi) + logdetV + quad)
      bhat[[k]] <- Omega %*% (Dinv %*% prior + XtY / sigma2)
    }
    Omega_out <- Omega                             # single, class-free
  } else {                                         # cov == "free": per-class
    logcomp <- matrix(0, N, K); bhat <- vector("list", K)
    Omega_out <- vector("list", K)
    for (k in 1:K) {
      Dinv <- diag(1 / tau2[, k], Q)
      Ok <- solve(Dinv + XtX / sigma2)
      Vk <- X %*% (tau2[, k] * t(X)) + sigma2 * diag(P)
      Lk <- chol(Vk); logdetV <- 2 * sum(log(diag(Lk))); Vinv <- chol2inv(Lk)
      prior <- M + matrix(mu[, k], Q, N)
      if (length(S)) prior <- prior + par$Del[[k]] %*% AS
      Resid <- Y - X %*% prior; quad <- colSums(Resid * (Vinv %*% Resid))
      logcomp[, k] <- log(xi[k]) - 0.5 * (P * log(2 * pi) + logdetV + quad)
      bhat[[k]] <- Ok %*% (Dinv %*% prior + XtY / sigma2)
      Omega_out[[k]] <- Ok
    }
  }

  ## Vectorized softmax over rows of logcomp (bit-identical to the per-row loop:
  ## m/s broadcast down rows, gamma keeps the exp(lc - s) form, and Reduce("+")
  ## reproduces the loop's sequential double-precision loglik accumulation).
  m <- apply(logcomp, 1, max)
  s <- m + log(rowSums(exp(logcomp - m)))
  gamma <- exp(logcomp - s)
  ll <- Reduce("+", s)
  list(gamma = gamma, bhat = bhat, Omega = Omega_out, loglik = ll,
       Nj = colSums(gamma))
}

## ---------------------------------------------------------------------
## Inner Theta/mu alternation: the ANCOVA normal equations, solved to
## convergence (rel. change < inner_tol) with a cap (inner_max).  Returns
## the pair (Theta, mu) satisfying  E A' = 0  and  sum_n gamma_nk e_nk = 0
## to machine precision (paper rem:ancova).  Used by both .nmfgmm.mstep and the
## final polish in .nmfgmm.fit.
## ---------------------------------------------------------------------
.nmfgmm.solve_theta_mu <- function(Theta, mu, bhat, gamma, Nj, A, intercept,
                           inner_tol = 1e-12, inner_max = 50,
                           cov = "tied", tau2 = NULL) {
  Q <- nrow(mu); N <- ncol(A); K <- ncol(mu); R <- nrow(A)
  xi <- Nj / sum(Nj); AAt_inv <- solve(tcrossprod(A))
  gls <- (cov == "free" && !is.null(tau2))
  if (gls) {
    ## weight matrix does not depend on Theta or mu: build it once
    W <- matrix(0, Q * R, Q * R)
    for (k in 1:K) {
      Ak <- A %*% (gamma[, k] * t(A))            # R x R : sum_n gamma_nk a_n a_n'
      W <- W + kronecker(Ak, diag(1 / tau2[, k], Q))
    }
    Wc <- chol(W)
  }
  prev <- c(Theta, mu)
  for (it in 1:inner_max) {
    if (gls) {
      rhs <- matrix(0, Q, R)
      for (k in 1:K)
        rhs <- rhs + (1 / tau2[, k]) *
          (sweep(bhat[[k]], 1, mu[, k], "-") %*% (gamma[, k] * t(A)))
      Theta <- matrix(backsolve(Wc, backsolve(Wc, as.numeric(rhs),
                                              transpose = TRUE)), Q, R)
    } else {
      bbar <- matrix(0, Q, N)
      for (k in 1:K) bbar <- bbar + sweep(sweep(bhat[[k]], 1, mu[, k], "-"), 2, gamma[, k], "*")
      Theta <- (bbar %*% t(A)) %*% AAt_inv
    }
    M <- Theta %*% A
    for (k in 1:K) mu[, k] <- rowSums(sweep(bhat[[k]] - M, 2, gamma[, k], "*")) / Nj[k]
    mubar <- as.numeric(mu %*% xi); mu <- mu - mubar
    Theta[, intercept] <- Theta[, intercept] + mubar
    cur <- c(Theta, mu)
    if (sqrt(sum((cur - prev)^2)) / (sqrt(sum(cur^2)) + 1e-12) < inner_tol) break
    prev <- cur
  }
  list(Theta = Theta, mu = mu)
}

## ---------------------------------------------------------------------
## Class-specific covariate effects (class.effects).  The rows S of A get a
## coefficient that differs by class,
##   b_n | z_n = k  ~  N(Theta a_n + Del_k a_n[S] + mu_k, Sigma),
##   sum_k xi_k mu_k = 0,   sum_k xi_k Del_k = 0,
## so Theta[, S] is the xi-weighted average of the class-specific coefficients
## Theta[, S] + Del_k.  The update of (Theta, Del, mu) is still least squares,
## part by part because Sigma is diagonal, and is solved exactly here instead
## of by the alternation above.  The freed coefficients are parametrized as a
## common slope plus deviations for classes 2..K; a relative ridge (1e-8) on
## the deviation block keeps the system solvable when a class has no
## within-class variation in a freed covariate, and is numerically invisible
## otherwise.  With cov = "free" each part is weighted by 1 / tau2[q, k].
## An empty S never reaches this function: the fit then keeps the alternation.
## ---------------------------------------------------------------------
.nmfgmm.solve_classwise <- function(bhat, gamma, Nj, A, intercept, S,
                                    cov = "tied", tau2 = NULL, ridge = 1e-8) {
  K <- ncol(gamma); N <- ncol(A); R <- nrow(A); Q <- nrow(bhat[[1]])
  xi <- Nj / sum(Nj)
  cm <- setdiff(seq_len(R), c(intercept, S))     # shared covariates
  nc <- length(cm); nf <- length(S)
  p <- nc + nf + (K - 1) * nf + K
  idev <- nc + nf + seq_len((K - 1) * nf)        # deviation block
  icls <- nc + nf + (K - 1) * nf + seq_len(K)    # class intercepts
  Dl <- lapply(seq_len(K), function(k) {          # design of class k: N x p
    Dk <- matrix(0, N, p)
    if (nc) Dk[, seq_len(nc)] <- t(A[cm, , drop = FALSE])
    Dk[, nc + seq_len(nf)] <- t(A[S, , drop = FALSE])
    if (k > 1) Dk[, nc + nf + (k - 2) * nf + seq_len(nf)] <- t(A[S, , drop = FALSE])
    Dk[, icls[k]] <- 1
    Dk
  })
  fitq <- function(w, qs) {                       # w[[k]]: weights, length N
    G <- matrix(0, p, p); H <- matrix(0, p, length(qs))
    for (k in seq_len(K)) {
      G <- G + crossprod(Dl[[k]], w[[k]] * Dl[[k]])
      H <- H + crossprod(Dl[[k]], w[[k]] * t(bhat[[k]][qs, , drop = FALSE]))
    }
    sc <- mean(diag(G)); diag(G) <- diag(G) + 1e-12 * sc
    diag(G)[idev] <- diag(G)[idev] + ridge * sc
    solve(G, H)
  }
  if (cov == "free" && !is.null(tau2)) {
    B <- matrix(vapply(seq_len(Q), function(q)
      fitq(lapply(seq_len(K), function(k) gamma[, k] / tau2[q, k]), q),
      numeric(p)), p, Q)
  } else {
    B <- fitq(lapply(seq_len(K), function(k) gamma[, k]), seq_len(Q))
  }
  Theta <- matrix(0, Q, R)
  if (nc) Theta[, cm] <- t(B[seq_len(nc), , drop = FALSE])
  ci <- t(B[icls, , drop = FALSE])                # Q x K class intercepts
  Theta[, intercept] <- as.numeric(ci %*% xi)
  mu <- ci - Theta[, intercept]
  com <- t(B[nc + seq_len(nf), , drop = FALSE])   # Q x nf
  Gam <- lapply(seq_len(K), function(k)
    if (k == 1) com else com + t(B[nc + nf + (k - 2) * nf + seq_len(nf), , drop = FALSE]))
  gb <- Reduce(`+`, Map(`*`, Gam, xi))
  Theta[, S] <- gb
  list(Theta = Theta, mu = mu, Del = lapply(Gam, function(g) g - gb))
}

## ---------------------------------------------------------------------
## M-step
## ---------------------------------------------------------------------
.nmfgmm.mstep <- function(par, es, Y, A, intercept = 1, cov = "tied",
                  inner_tol = 1e-12, inner_max = 50, fixX = FALSE,
                  vfloor = NULL, S = NULL) {
  ## vfloor: lower bound on the variance parameters.  NULL (default) makes it
  ## SCALE-RELATIVE, 1e-12 * mean(Y^2), so the fit is equivariant under
  ## Y -> cY.  The old fixed 1e-6 bound was absolute and could bind on
  ## small-scale data (it did on the nir absorbances, whose entries are
  ## O(0.02)); pass vfloor = 1e-6 to reproduce the old behaviour.
  X <- par$X; tau2 <- par$tau2; sigma2 <- par$sigma2
  P <- nrow(Y); N <- ncol(Y); Q <- ncol(X); K <- ncol(par$mu)
  gamma <- es$gamma; bhat <- es$bhat; Omega <- es$Omega; Nj <- es$Nj
  if (is.null(vfloor)) vfloor <- 1e-12 * mean(Y^2)
  xi <- Nj / N

  ## --- inner Theta/mu alternation (ANCOVA normal equations); with
  ##     class-specific rows S, the exact joint least-squares solve instead ---
  if (length(S)) {
    tm <- .nmfgmm.solve_classwise(bhat, gamma, Nj, A, intercept, S,
                                  cov = cov, tau2 = par$tau2)
    Del <- tm$Del; AS <- A[S, , drop = FALSE]
  } else {
    tm <- .nmfgmm.solve_theta_mu(par$Theta, par$mu, bhat, gamma, Nj, A, intercept,
                         inner_tol, inner_max, cov = cov, tau2 = par$tau2)
  }
  Theta <- tm$Theta; mu <- tm$mu

  ## --- variances ---
  M <- Theta %*% A
  if (cov != "free") {                     # "tied" or "scalar": shared diagonal
    num <- rep(0, Q)
    for (k in 1:K) {
      resid <- sweep(bhat[[k]] - M, 1, mu[, k], "-")
      if (length(S)) resid <- resid - Del[[k]] %*% AS
      num <- num + rowSums(sweep(resid^2, 2, gamma[, k], "*"))
    }
    tau2 <- pmax((num + N * diag(Omega)) / N, vfloor)
    if (cov == "scalar") tau2 <- rep(mean(tau2), Q)   # isotropic: pool
  } else {
    for (k in 1:K) {
      resid <- sweep(bhat[[k]] - M, 1, mu[, k], "-")
      if (length(S)) resid <- resid - Del[[k]] %*% AS
      tau2[, k] <- pmax((rowSums(sweep(resid^2, 2, gamma[, k], "*")) +
                         Nj[k] * diag(Omega[[k]])) / Nj[k], vfloor)
    }
  }

  ## --- basis: responsibility-weighted semi-NMF update ---
  Cstat <- matrix(0, P, Q); Gstat <- matrix(0, Q, Q)
  for (k in 1:K) {
    Cstat <- Cstat + sweep(Y, 2, gamma[, k], "*") %*% t(bhat[[k]])
    Gstat <- Gstat + sweep(bhat[[k]], 2, gamma[, k], "*") %*% t(bhat[[k]])
    if (cov == "free") Gstat <- Gstat + Nj[k] * Omega[[k]]
  }
  if (cov != "free") Gstat <- Gstat + N * Omega
  ## fixX = TRUE (2026-08-28): hold the basis at its current value -- used for
  ## the fixed-basis two-stage-vs-joint comparison, where BOTH routes run on the
  ## SAME X (they still fit different data: the two-stage route fits the
  ## reconstituted, shifted residuals).
  Xn <- if (fixX) X else
        X * sqrt((.nmfgmm.pos(Cstat) + X %*% .nmfgmm.neg(Gstat)) /
                 (.nmfgmm.neg(Cstat) + X %*% .nmfgmm.pos(Gstat) + 1e-12))

  ## --- noise variance ---
  ss <- 0; tr <- 0
  for (k in 1:K) {
    rY <- Y - Xn %*% bhat[[k]]; ss <- ss + sum(colSums(rY^2) * gamma[, k])
    if (cov == "free") tr <- tr + Nj[k] * sum((Xn %*% Omega[[k]]) * Xn)
  }
  if (cov != "free") tr <- N * sum((Xn %*% Omega) * Xn)
  sigma2 <- max((ss + tr) / (P * N), vfloor)

  ## --- column-normalize X and rescale ---
  D <- colSums(Xn); Xn <- sweep(Xn, 2, D, "/")
  Theta <- sweep(Theta, 1, D, "*"); mu <- sweep(mu, 1, D, "*")
  if (cov != "free") tau2 <- tau2 * D^2 else tau2 <- sweep(tau2, 1, D^2, "*")
  if (cov == "scalar") tau2 <- rep(mean(tau2), Q)     # rescaling breaks isotropy

  out <- list(X = Xn, Theta = Theta, mu = mu, tau2 = tau2, sigma2 = sigma2, xi = xi)
  if (length(S)) out$Del <- lapply(Del, function(d) sweep(d, 1, D, "*"))
  out
}

## ---------------------------------------------------------------------
## Initialization  (NMF-RE fit for K=1, then k-means++ on residual scores)
## ---------------------------------------------------------------------
.nmfgmm.init_par <- function(Y, A, X0, K, sigma2_0, intercept = 1, seed = 1, cov = "tied") {
  Q <- ncol(X0); N <- ncol(Y)
  bls <- solve(crossprod(X0), crossprod(X0, Y))
  Theta <- (bls %*% t(A)) %*% solve(tcrossprod(A))
  r <- bls - Theta %*% A
  ## initialization floors are scale-relative too (see fit_nmfgmm's vfloor)
  ifloor <- 1e-12 * mean(bls^2)
  ms <- function(M) pmax(apply(M, 1, function(v) mean(v^2)), ifloor)
  if (K == 1) {
    tau2 <- if (cov == "free") matrix(ms(r), Q, 1)
            else if (cov == "scalar") rep(mean(ms(r)), Q) else ms(r)
    return(list(X = X0, Theta = Theta, mu = matrix(0, Q, 1),
                tau2 = tau2, sigma2 = sigma2_0, xi = 1))
  }
  rt <- t(r)
  km <- stats::kmeans(rt, centers = .nmfgmm.kmpp(rt, K, seed), iter.max = 50)
  xi <- as.numeric(table(factor(km$cluster, levels = 1:K))) / N
  mu <- matrix(t(km$centers), Q, K)
  if (cov != "free") {                    # "tied" or "scalar": shared diagonal
    resid <- r - mu[, km$cluster]; tau2 <- pmax(rowMeans(resid^2), ifloor)
    if (cov == "scalar") tau2 <- rep(mean(tau2), Q)
  } else {
    tau2 <- sapply(1:K, function(k) { idx <- km$cluster == k
      if (sum(idx) > 1) ms(r[, idx, drop = FALSE] - mu[, k]) else ms(r) })
    tau2 <- matrix(tau2, Q, K)
  }
  mubar <- as.numeric(mu %*% xi); mu <- mu - mubar
  Theta[, intercept] <- Theta[, intercept] + mubar
  list(X = X0, Theta = Theta, mu = mu, tau2 = tau2, sigma2 = sigma2_0, xi = xi)
}

## ---------------------------------------------------------------------
## Multistart EM driver
## ---------------------------------------------------------------------
.nmfgmm.fit <- function(Y, A, X0, K, cov = "tied", intercept = 1,
                       maxit = 200000, tol = 1e-7, ptol = 1e-9,
                       nstart = if (K == 1) 1 else 8,
                       inner_tol = 1e-12, inner_max = 50, cores = 1L, fixX = FALSE,
                       vfloor = NULL, S = NULL) {
  ## S     : rows of A whose coefficients are class-specific (class.effects);
  ##         NULL / empty is the common-effect model, computed exactly as before.
  ## tol   : relative change of the marginal log-likelihood (the usual EM rule).
  ## ptol  : relative change of the REPORTED parameters (X, Theta, mu); NULL
  ##         keeps the objective-only rule.  The objective rule cannot fire when
  ##         a nuisance variance drifts to its boundary (tau^2 -> 0 on a sample
  ##         with no random effect), even though the reported quantities have
  ##         settled; $stop_by records which rule fired.
  ## vfloor: variance lower bound, NULL = scale-relative 1e-12 * mean(Y^2), so
  ##         the fit is equivariant under Y -> cY.
  bls <- solve(crossprod(X0), crossprod(X0, Y))
  s2_0 <- max(mean((Y - X0 %*% bls)^2), 1e-12 * mean(Y^2))
  ## One EM restart. Each start re-seeds from `s` (via .nmfgmm.init_par) and the
  ## EM loop is RNG-free, so run_start(s) is a deterministic function of s alone
  ## -- independent of the other starts and of execution order.
  run_start <- function(s) {
    par <- .nmfgmm.init_par(Y, A, X0, K, s2_0, intercept, seed = s, cov = cov)
    if (length(S)) par$Del <- lapply(seq_len(K), function(k) matrix(0, ncol(X0), length(S)))
    ll_old <- -Inf; hist <- numeric(0)
    rep_old <- c(par$X, par$Theta, par$mu, unlist(par$Del)); stop_by <- "maxit"
    for (it in 1:maxit) {
      es <- .nmfgmm.estep(par, Y, A, cov = cov, S = S); hist <- c(hist, es$loglik)
      if (abs(es$loglik - ll_old) / (abs(es$loglik) + 1) < tol) {
        stop_by <- "loglik"; break
      }
      ll_old <- es$loglik
      par <- .nmfgmm.mstep(par, es, Y, A, intercept, cov = cov,
                   inner_tol = inner_tol, inner_max = inner_max, fixX = fixX,
                   vfloor = vfloor, S = S)
      if (!is.null(ptol)) {
        rep_new <- c(par$X, par$Theta, par$mu, unlist(par$Del))
        if (sqrt(sum((rep_new - rep_old)^2)) /
            (sqrt(sum(rep_new^2)) + 1e-12) < ptol) {
          ## The E-step just after the loop recomputes this on the same `par`,
          ## so doing it here as well costs one full E-step for nothing.
          ## Record only the flag; `hist` gains its last entry from that call.
          stop_by <- "parameters"
          break
        }
        rep_old <- rep_new
      }
    }
    es <- .nmfgmm.estep(par, Y, A, cov = cov, S = S)
    if (stop_by == "parameters") hist <- c(hist, es$loglik)
    ## polish: re-solve Theta/mu at the final responsibilities so the returned
    ## pair satisfies the ANCOVA normal equations (rem:ancova) to machine
    ## precision, removing the one-step outer-loop lag.  Holds X, variances,
    ## and the loglik fixed to outer-tol; .nmfgmm.ARI/BIC unaffected.
    if (length(S)) {
      tm <- .nmfgmm.solve_classwise(es$bhat, es$gamma, es$Nj, A, intercept, S,
                                    cov = cov, tau2 = par$tau2)
      par$Del <- tm$Del
    } else {
      tm <- .nmfgmm.solve_theta_mu(par$Theta, par$mu, es$bhat, es$gamma, es$Nj, A,
                           intercept, inner_tol, inner_max,
                           cov = cov, tau2 = par$tau2)
    }
    par$Theta <- tm$Theta; par$mu <- tm$mu
    list(par = par, es = es, loglik = es$loglik, iter = it, hist = hist,
         stop_by = stop_by)
  }
  ## .nmfkc.parlapply preserves input order, so which.max() picks the same
  ## first-best restart as the sequential `> best$loglik` scan -- the returned
  ## fit is identical for any `cores`.
  starts <- .nmfkc.parlapply(seq_len(nstart), run_start,
                             cores = if (nstart > 1L) cores else 1L)
  starts[[which.max(vapply(starts, function(r) r$loglik, numeric(1)))]]
}

## ---------------------------------------------------------------------
## Utilities
## ---------------------------------------------------------------------
## free (per-class) parameter count uses Q*K variances; tied uses Q.
## nS class-specific covariate rows add (K - 1) * Q * nS deviations.
.nmfgmm.pcount <- function(Q, P, R, K, cov = "tied", nS = 0L) {
  vars <- switch(cov, scalar = 1L, tied = Q, free = Q * K)
  Q * (P - 1) + Q * R + Q * (K - 1) + vars + (K - 1) + 1 + (K - 1) * Q * nS
}

## Resolve class.effects to row indices S of A.  Accepts row indices, a
## logical vector over the rows, or names; a name matches a row name of A or,
## when A was built from a formula, every row of that formula term, so
## class.effects = "diet" frees all indicator rows of a factor `diet`.
.nmfgmm.classrows <- function(spec, A, A.term, intercept) {
  if (is.null(spec) || length(spec) == 0) return(integer(0))
  R <- nrow(A)
  if (is.logical(spec)) {
    if (length(spec) != R) stop("a logical class.effects needs one entry per row of A.")
    S <- which(spec)
  } else if (is.numeric(spec)) {
    S <- as.integer(spec)
    if (any(S < 1L | S > R)) stop("class.effects indexes a row outside A.")
  } else {
    rn <- rownames(A); S <- integer(0)
    for (s in as.character(spec)) {
      hit <- integer(0)
      if (!is.null(rn)) hit <- which(rn == s)
      if (!is.null(A.term)) hit <- union(hit, which(A.term == s))
      if (!length(hit))
        stop(sprintf("class.effects: no row of A, and no formula term, is named '%s'.", s))
      S <- c(S, hit)
    }
  }
  S <- sort(unique(S))
  if (intercept %in% S)
    stop("class.effects must not include the intercept row: the class means mu ",
         "already make the intercept class-specific.")
  S
}

.nmfgmm.ARI <- function(a, b) {
  tab <- table(a, b); n <- sum(tab); cc <- function(x) sum(choose(x, 2))
  ai <- rowSums(tab); bj <- colSums(tab)
  idx <- cc(as.vector(tab)); ei <- cc(ai) * cc(bj) / choose(n, 2)
  mi <- (cc(ai) + cc(bj)) / 2
  if (mi - ei == 0) return(0); (idx - ei) / (mi - ei)
}


# =====================================================================
#  Public API
# =====================================================================

## Build a numeric R x N covariate matrix from a one-sided formula + data:
## model.matrix, drop its intercept, optionally center+scale, prepend the
## package intercept row.  Returns A with attributes A.center / A.scale.
#' @noRd
.nmfgmm.buildA <- function(formula, data, N, standardize = TRUE) {
  if (is.null(data)) stop("A is a formula; supply `data` (one row per column of Y).")
  mm <- stats::model.matrix(formula, data = data)
  ## the formula term behind each column, so class.effects can name a term
  tl <- attr(stats::terms(formula, data = data), "term.labels")
  keep <- colnames(mm) != "(Intercept)"
  term <- tl[attr(mm, "assign")[keep]]
  mm <- mm[, keep, drop = FALSE]
  if (nrow(mm) != N) stop("`data` must have one row per column of Y.")
  ctr <- scl <- NULL
  if (standardize && ncol(mm) > 0) {
    mm <- scale(mm)
    ctr <- attr(mm, "scaled:center"); scl <- attr(mm, "scaled:scale")
  }
  A <- rbind(Intercept = rep(1, N), t(mm))
  attr(A, "A.center") <- ctr; attr(A, "A.scale") <- scl
  attr(A, "A.term") <- c("(Intercept)", term)
  A
}

## Initial non-negative, column-normalized basis, shared by nmf.gmm() and
## nmf.gmm.twostage() so that both routes start from the same X0.  X.init is a
## P x Q matrix or a method name:
##   "nmf"      the basis of an ordinary NMF, nmfkc(Y, rank = Q, seed = seed)
##              with no covariates, i.e. the basis a user gets by calling
##              nmfkc() directly with the same seed (nmfkc()'s own default
##              seed is 123).  Entries are floored at 1e-8 before normalizing,
##              because the multiplicative X-update of the EM cannot move an
##              exact zero.  Needs a non-negative Y.
##   otherwise  passed to the shared initializer .init_X_method() ("nndsvd",
##              the default, "kmeans++", "kmeans", ...).
#' @noRd
.nmfgmm.initX <- function(X.init, Y, Q, seed, nstart) {
  if (is.matrix(X.init) || (is.numeric(X.init) && length(X.init) > 1)) {
    X0 <- as.matrix(X.init)
  } else {
    method <- if (is.character(X.init)) X.init else "nndsvd"
    if (identical(method, "nmf")) {
      X0 <- pmax(unname(as.matrix(nmfkc(Y, rank = Q, seed = seed, verbose = FALSE)$X)), 1e-8)
    } else if (identical(method, "spa")) {
      ## nmfkc()'s X.init = "spa": the anchor rows SPA finds in Y, the
      ## non-negative least-squares fit to them, and the anchor zeros filled
      ## -- the EM's multiplicative X-update would otherwise keep them, turning
      ## a starting value into a constraint (X.anchor is that constraint).
      X0 <- .nmfkc_anchor_start(Y, .nmfkc_spa_rows(Y, Q, "nmf.gmm"), 1e-10)
      z <- which(X0 == 0)
      if (length(z)) {
        if (!is.null(seed)) set.seed(seed)
        X0[z] <- stats::runif(length(z)) * mean(Y) / 100
      }
    } else {
      X0 <- .init_X_method(method, Y, Q, seed = seed, nstart = max(nstart, 1L))
    }
  }
  X0 / rep(pmax(colSums(X0), 1e-12), each = nrow(X0))
}

#' @title Fit NMF-GMM: a Gaussian-mixture latent-class extension of NMF with covariates
#' @description
#' This function is \strong{experimental and still under development}. The
#' interface may change in future versions: argument names, defaults and the
#' contents of the returned object are not yet stable. Code written against it
#' today may need adjusting after an update. The rest of the package does not
#' carry this caveat.
#'
#' \code{nmf.gmm} fits the model
#' \deqn{\bm b_n\mid(z_n=k)\sim N_Q(C\bm a_n+\bm\mu_k,\Sigma_k),\qquad
#'       \bm y_n = X\bm b_n+\bm\varepsilon_n,\quad X\ge 0,}
#' by a generalized EM algorithm: a \eqn{K}-component Gaussian mixture is placed
#' on the latent NMF scores \eqn{\bm b_n}, whose class-conditional mean combines
#' a covariate regression \eqn{C\bm a_n} (\eqn{C=\Theta} in the paper, the
#' \eqn{Q\times R} coefficient matrix) with a class shift \eqn{\bm\mu_k}. The
#' shared basis \eqn{X} stays non-negative and column-normalized. Clustering is
#' model based, through the posterior responsibilities. \eqn{K=1} (tied
#' covariance) reduces exactly to NMF with random effects (\code{\link{nmfre}}).
#'
#' Like the rest of the package, \code{nmf.gmm} performs \strong{optimization
#' only}; use \code{\link{nmf.gmm.inference}} for standard errors and tests of
#' \eqn{C}, and \code{\link{nmf.gmm.select}} to choose \eqn{K}.
#'
#' @param Y Data matrix \eqn{Y} (P x N).
#' @param A Covariate matrix \eqn{A} (R x N). Default \code{NULL} uses an
#'   intercept-only \eqn{1\times N} matrix (a plain Gaussian mixture on the
#'   scores, with no covariate adjustment). A row of ones should be present for
#'   the identifiability constraint \eqn{\sum_k \xi_k\bm\mu_k = 0} (see
#'   \code{intercept}).
#' @param rank Integer rank \eqn{Q} of the basis.
#' @param K Integer number of mixture components. Default 1 (= NMF-RE).
#' @param ... Additional arguments:
#'   \itemize{
#'     \item \code{cov}: score-covariance structure. \code{"tied"} (default,
#'       shared diagonal \eqn{\Sigma_k=\mathrm{diag}(\tau^2_1,\dots,\tau^2_Q)},
#'       \eqn{Q} variances); \code{"free"} (per-class diagonal, \eqn{Q\times K});
#'       or \code{"scalar"} (isotropic \eqn{\Sigma_k=\tau^2 I_Q}, a single
#'       variance --- the most parsimonious variant; at \eqn{K=1} it is the
#'       \code{\link{nmfre}} model, and \code{tau2} is returned as one number).
#'     \item \code{X.init}: initial basis. A \eqn{P\times Q} non-negative matrix,
#'       or an initialization method name: \code{"nndsvd"} (default),
#'       \code{"kmeans++"}, \code{"kmeans"}, ... (passed to the shared
#'       initializer), or \code{"nmf"}, the basis of an ordinary NMF fitted by
#'       \code{nmfkc(Y, rank = rank, seed = seed)} without covariates
#'       (requires a non-negative \code{Y}; \code{\link{nmfkc}}'s own default
#'       seed is 123). The EM restarts (\code{nstart}) all start from this one
#'       basis and differ only in the seeding of the mixture, so the fit can
#'       depend on the choice of \code{X.init}; the basis actually used is
#'       returned as \code{X0}.  \code{"spa"} starts from the anchor rows the
#'       successive projection algorithm finds in \code{Y} (see
#'       \code{\link{nmfkc}}; requires a non-negative \code{Y}).
#'     \item \code{X.anchor}: anchor rows of \eqn{X}, imposed as in
#'       \code{\link{nmfkc}}: \code{NULL} (default), \code{"spa"}, or one row
#'       of \code{Y} per basis (requires a non-negative \code{Y}).  The EM's
#'       multiplicative basis update keeps the zeros.  There is no
#'       \code{rank <= nrow(A)} limit here: the mixture gives the scores full
#'       rank.  \code{nmf.gmm.twostage()} accepts it too, finding the anchors
#'       on \code{Y}.  The rows used are returned as \code{X.anchor}.
#'     \item \code{intercept}: index of the intercept row of \code{A} (default 1);
#'       the class-mean average is absorbed into that column of \eqn{C}.
#'     \item \code{nstart}: EM restarts (default 1 for \code{K=1}, else 8).
#'     \item \code{cores}: run the \code{nstart} restarts in parallel (default
#'       \code{getOption("mc.cores", 1L)}; PSOCK cluster on Windows, forking
#'       elsewhere). Each restart is an independent self-seeded EM and results
#'       are combined in order, so the returned fit is identical for any
#'       \code{cores}. (\code{\link{nmf.gmm.select}} parallelizes over \code{K}
#'       instead and runs each inner fit with \code{cores = 1}.)
#'     \item \code{maxit}, \code{tol}: outer EM cap / tolerance (2e5, 1e-7).
#'       \code{tol} is the relative change of the marginal log-likelihood, the
#'       usual EM rule.
#'     \item \code{ptol}: a second stopping rule, on the relative change of the
#'       \emph{reported} parameters \eqn{X}, \eqn{C} and \eqn{\bm\mu}
#'       (default 1e-9; \code{NULL} disables it). The log-likelihood rule alone
#'       can fail to fire when a nuisance variance drifts to its boundary --
#'       \eqn{\tau^2\to 0} on a sample that carries no random effect -- so the
#'       objective keeps creeping up after the partition, the basis and the
#'       covariate effect have settled. The returned \code{stop_by} says which
#'       rule fired (\code{"loglik"}, \code{"parameters"} or \code{"maxit"}).
#'     \item \code{vfloor}: lower bound on the variance parameters. The default
#'       (\code{NULL}) makes it scale-relative, \eqn{10^{-12}\,\mathrm{mean}(Y^2)},
#'       so the fit is equivariant under \eqn{Y\to cY}; a fixed bound would bind
#'       on small-scale data.
#'     \item \code{fixX}: hold the basis at its initial value instead of
#'       updating it (default \code{FALSE}). Used to compare the joint fit with
#'       \code{\link{nmf.gmm.twostage}} on the \emph{same} \eqn{X}, so that the
#'       two routes share the basis. They still fit different data: the
#'       two-stage route fits the reconstituted, shifted residuals.
#'     \item \code{class.effects}: covariates whose coefficients differ by
#'       class (default \code{NULL}: every covariate has one coefficient shared
#'       by all classes). Row indices of \code{A}, row names, or --- when
#'       \code{A} is a formula --- term names, so \code{class.effects = "diet"}
#'       frees every indicator row of a factor \code{diet}. The class-\eqn{k}
#'       mean becomes \eqn{C\bm a_n+\Delta_k\bm a_{n,S}+\bm\mu_k} with
#'       \eqn{\sum_k\xi_k\Delta_k=0}, so the columns \eqn{S} of \code{C} hold
#'       the \eqn{\xi}-weighted average of the class-specific coefficients,
#'       which are returned in \code{C.class}. The intercept row cannot be
#'       freed: \eqn{\bm\mu_k} already makes it class-specific. With a
#'       continuous covariate, the class-specific coefficients are identified
#'       when the covariate varies within every class. With a categorical
#'       covariate, freeing all of its indicator rows gives every level its own
#'       class means, so the likelihood no longer ties the classes of one level
#'       to those of another: which class of one level corresponds to which
#'       class of another is then carried by the starting values rather than by
#'       the data. Not supported by \code{\link{nmf.gmm.inference}} or
#'       \code{\link{nmf.gmm.twostage}}.
#'     \item \code{seed}: RNG seed (default 1). \code{prefix}: basis-name prefix
#'       (default \code{"Basis"}).
#'     \item \code{data}: a data frame with one row per column of \code{Y},
#'       required when \code{A} is a formula (see below).
#'     \item \code{standardize}: when \code{A} is a formula, center and scale
#'       the constructed covariate columns (default \code{TRUE}). Factors are
#'       expanded to treatment indicators before standardization, so a
#'       \eqn{5}-level factor becomes four centered, scaled columns.
#'   }
#'
#'   \code{A} may also be a one-sided \strong{formula} (e.g. \code{~ size} or
#'   \code{~ diet}), evaluated in \code{data}: the design matrix is built with
#'   \code{\link[stats]{model.matrix}}, its intercept column is replaced by the
#'   package's own intercept row, and the remaining columns are standardized by
#'   default. The constructed numeric \eqn{A} is returned in the fit (fields
#'   \code{A}, \code{A.formula}, \code{A.center}, \code{A.scale}), so
#'   \code{\link{nmf.gmm.inference}} works unchanged.
#'
#' @return An object of class \code{"nmf.gmm"}: a list with \code{X}
#'   (basis), \code{X0} (the column-normalized initial basis the EM started
#'   from), \code{C} (\eqn{=\Theta}, Q x R), \code{mu} (class means, Q x K),
#'   \code{tau2}, \code{sigma2}, \code{xi} (mixing proportions), \code{gamma}
#'   (responsibilities, N x K), \code{cluster} (hard labels), \code{loglik},
#'   \code{BIC}, \code{ICL}, \code{n.params}, \code{entropy}, \code{Yhat},
#'   \code{objfunc(.iter)}, \code{iter}, \code{converged}, \code{stop_by},
#'   \code{K},
#'   \code{rank}, \code{cov}, \code{dims} and \code{runtime}. With
#'   \code{class.effects}, also \code{class.effects} (the freed rows of
#'   \code{A}, by name), \code{class.rows} (their indices) and \code{C.class}
#'   (a Q x |S| x K array of class-specific coefficients; the class-specific
#'   effect in observation space is \code{X \%*\% C.class[, , k]}).
#'
#' @seealso \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}},
#'   \code{\link{nmfre}} (the \eqn{K=1} special case).
#' @examples
#' \donttest{
#' set.seed(1)
#' Y <- matrix(abs(rnorm(20 * 60)) + 1, 20, 60)
#' A <- rbind(1, rnorm(60))                 # intercept + one covariate
#' fit <- nmf.gmm(Y, A, rank = 2, K = 2)
#' fit$BIC
#' table(fit$cluster)
#' }
#' @export
nmf.gmm <- function(Y, A = NULL, rank, K = 1, ...) {
  extra <- base::list(...)
  getopt <- function(nm, default) if (!is.null(extra[[nm]])) extra[[nm]] else default
  cov       <- getopt("cov", "tied")
  X.init    <- extra$X.init
  intercept <- getopt("intercept", 1L)
  maxit     <- getopt("maxit", 200000L)
  ## `epsilon` is the house name for the convergence tolerance in every other
  ## fitter, and this function already RETURNS it as `epsilon`.  Accepting only
  ## `tol` meant nmf.gmm(..., epsilon = 1e-9) was silently a no-op.  `tol` stays
  ## as the legacy alias.
  tol       <- getopt("epsilon", getopt("tol", 1e-7))
  ## Likewise `wild.seed` is the convention for the bootstrap; here `seed` is
  ## the fit seed, so it keeps its name.
  seed      <- getopt("seed", 1L)
  ## Keep the self-seeding of this fit out of the caller's random stream.
  .rng <- .nmfkc.rng.save(seed)
  on.exit(.nmfkc.rng.restore(.rng), add = TRUE)
  prefix    <- getopt("prefix", "Basis")
  inner_tol <- getopt("inner_tol", 1e-12)
  inner_max <- getopt("inner_max", 50L)
  ptol      <- getopt("ptol", 1e-9)
  vfloor    <- getopt("vfloor", NULL)
  fixX      <- getopt("fixX", FALSE)
  cov <- match.arg(cov, c("tied", "free", "scalar"))

  t0 <- proc.time()
  Y <- as.matrix(Y); P <- nrow(Y); N <- ncol(Y); Q <- as.integer(rank)
  A.formula <- NULL; A.center <- NULL; A.scale <- NULL; A.term <- NULL
  if (inherits(A, "formula")) {
    A.formula <- A
    A <- .nmfgmm.buildA(A, extra$data, N, getopt("standardize", TRUE))
    A.center <- attr(A, "A.center"); A.scale <- attr(A, "A.scale")
    A.term <- attr(A, "A.term"); attr(A, "A.term") <- NULL
  }
  if (is.null(A)) A <- matrix(1, 1, N)
  A <- as.matrix(A); if (ncol(A) != N && nrow(A) == N) A <- t(A)
  if (ncol(A) != N) stop("A must have N columns (R x N).")
  R <- nrow(A)
  S <- .nmfgmm.classrows(extra$class.effects, A, A.term, intercept)
  nstart <- getopt("nstart", if (K == 1) 1L else 8L)
  cores  <- getopt("cores", getOption("mc.cores", 1L))

  ## --- initial non-negative, column-normalized basis X0 ---
  ## X.anchor: anchor rows of X, kept as exact zeros for the whole fit (the
  ## EM's X-update is multiplicative and the column scaling divides, so a
  ## zero stays a zero).  No rank check as in nmfkc(): the scores carry the
  ## mixture means and a Gaussian part, so B has full rank anyway.
  X.anchor <- extra$X.anchor
  anchor.rows <- NULL
  if (!is.null(X.anchor)) {
    an <- .nmfkc_anchor_setup(X.anchor, Y, Q, D_A = NULL, fun = "nmf.gmm",
                              seed = seed)
    anchor.rows <- an$rows
    X0 <- an$X / rep(pmax(colSums(an$X), 1e-12), each = P)
    if (!is.null(X.init) && !identical(X.init, "spa"))
      message("nmf.gmm(): X.anchor builds its own starting basis, so X.init was not used.")
  } else {
    X0 <- .nmfgmm.initX(X.init, Y, Q, seed, nstart)
  }

  fit <- .nmfgmm.fit(Y, A, X0, K, cov = cov, intercept = intercept,
                     maxit = maxit, tol = tol, ptol = ptol, nstart = nstart,
                     inner_tol = inner_tol, inner_max = inner_max, cores = cores,
                     fixX = fixX, vfloor = vfloor, S = S)
  par <- fit$par; es <- fit$es

  ## --- labels (house style) ---
  X <- par$X; C <- par$Theta; mu <- par$mu
  blab <- paste0(prefix, seq_len(Q)); clab <- paste0("Class", seq_len(K))
  alab <- rownames(A); if (is.null(alab)) alab <- paste0("Cov", seq_len(R))
  dimnames(X) <- list(rownames(Y), blab)
  dimnames(X0) <- list(rownames(Y), blab)
  dimnames(C) <- list(blab, alab)
  dimnames(mu) <- list(blab, clab)
  gamma <- es$gamma; colnames(gamma) <- clab
  cluster <- max.col(gamma, ties.method = "first")
  tau2 <- par$tau2
  if (cov == "scalar") tau2 <- as.numeric(tau2)[1]   # single isotropic variance
  else if (cov == "tied") { tau2 <- as.numeric(tau2); names(tau2) <- blab }
  else dimnames(tau2) <- list(blab, clab)

  ## --- BIC / ICL ---
  n.params <- .nmfgmm.pcount(Q, P, R, K, cov = cov, nS = length(S))
  bic <- -2 * fit$loglik + n.params * log(N)
  ent <- -sum(ifelse(gamma > 1e-300, gamma * log(gamma), 0))
  icl <- bic + 2 * ent

  ## --- least-squares scores (Q x N): bls = (X'X)^{-1} X'Y ---
  ## The geometric projection of Y onto the estimated basis X, used by
  ## plot(type = "scores"). This matches the paper's score-space figure: the
  ## covariate mean C %*% A is removed at plot time to show the size-adjusted
  ## residual scores. (Deliberately NOT the responsibility-averaged posterior
  ## mean sum_k gamma_nk E[b_n | z=k], which shrinks points toward their class
  ## centroid and would make the clusters look artificially separated.)
  scores <- solve(crossprod(X), crossprod(X, Y))
  dimnames(scores) <- list(blab, colnames(Y))

  ## --- responsibility-averaged reconstruction Yhat = X (C A + mu gamma') ---
  ##     (plus sum_k gamma_nk Del_k a_n[S] with class-specific rows)
  Bhat <- C %*% A + mu %*% t(gamma)
  if (length(S)) for (k in 1:K)
    Bhat <- Bhat + sweep(par$Del[[k]] %*% A[S, , drop = FALSE], 2, gamma[, k], "*")
  Yhat <- X %*% Bhat
  dimnames(Yhat) <- dimnames(Y)

  dims <- sprintf("Y(%d,%d) ~ X(%d,%d) [K=%d, %s]", P, N, P, Q, K, cov)
  runtime <- as.numeric((proc.time() - t0)["elapsed"])

  out <- structure(list(
    call = match.call(), dims = dims, runtime = runtime,
    rank = Q, K = K, cov = cov, intercept = intercept,
    X = X, X0 = X0, C = C, mu = mu, tau2 = tau2, sigma2 = par$sigma2, xi = par$xi,
    gamma = gamma, cluster = cluster, scores = scores, Yhat = Yhat,
    loglik = fit$loglik, BIC = bic, ICL = icl, entropy = ent, n.params = n.params,
    objfunc = -fit$loglik, objfunc.iter = -fit$hist, iter = fit$iter,
    ## Standard convergence names shared with the other fitters; tol is this
    ## model's epsilon.
    maxit = maxit, epsilon = tol,
    ## stop_by records WHICH rule fired ("loglik", "parameters" or "maxit").
    ## iter < maxit alone cannot tell a ptol stop from a loglik stop, and it
    ## calls a run that used its last iteration productively "not converged".
    converged = fit$stop_by != "maxit", stop_by = fit$stop_by, A = A,
    A.formula = A.formula, A.center = A.center, A.scale = A.scale
  ), class = "nmf.gmm")
  if (length(S)) {
    ## class-specific coefficients C[, S] + Del_k; C[, S] is their xi-average
    C.class <- array(0, c(Q, length(S), K), dimnames = list(blab, alab[S], clab))
    for (k in 1:K) C.class[, , k] <- C[, S, drop = FALSE] + par$Del[[k]]
    out$class.effects <- alab[S]; out$class.rows <- S; out$C.class <- C.class
  }
  ## Added only when anchors were used, so every other fit keeps its fields.
  if (!is.null(anchor.rows)) out$X.anchor <- stats::setNames(anchor.rows, blab)
  out
}


#' @title Statistical inference for an NMF-GMM fit (given basis)
#' @description
#' This function is \strong{experimental and still under development}. The
#' interface may change in future versions: argument names, defaults and the
#' contents of the returned object are not yet stable. Code written against it
#' today may need adjusting after an update. The rest of the package does not
#' carry this caveat.
#'
#' \code{nmf.gmm.inference} adds Wald inference on the covariate-coefficient
#' matrix \eqn{C} (\eqn{=\Theta}) of a fitted \code{\link{nmf.gmm}} object,
#' conditional on the estimated basis \eqn{\hat X} and mixture. It reports the
#' analytic outer-product (mixture-information) standard error --- the
#' coverage-calibrated choice at \eqn{K>1} --- together with a
#' \strong{wild-bootstrap} standard error and percentile confidence interval,
#' mirroring \code{\link{nmfre.inference}} / \code{\link{nmfkc.inference}}.
#' Supported for the tied-covariance variant.
#'
#' @param object A fitted \code{"nmf.gmm"} object (with \code{cov = "tied"}).
#' @param Y The data matrix used in the fit.
#' @param A The covariate matrix used in the fit. Default \code{object$A}.
#' @param ... Additional arguments, named as in the other wild-bootstrap
#'   inference functions of the package: \code{wild.B} (bootstrap replicates,
#'   default 500), \code{wild.level} (confidence level, default 0.95) and
#'   \code{seed} (default 123).
#'
#' @return \code{object} with class \code{c("nmf.gmm.inference", "nmf.gmm")}
#'   and added components: \code{coefficients} (data frame with \code{Basis},
#'   \code{Covariate}, \code{Estimate}, \code{SE}, \code{BSE}, \code{z_value},
#'   \code{p_value}, \code{CI_low}, \code{CI_high}), \code{C.se.outer},
#'   \code{C.se.sandwich}, \code{C.se.boot} (Q x R matrices) and \code{Xi}
#'   (\eqn{= X C}).
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.select}}
#' @examples
#' \donttest{
#' set.seed(1)
#' Y <- matrix(abs(rnorm(20 * 60)) + 1, 20, 60); A <- rbind(1, rnorm(60))
#' fit <- nmf.gmm(Y, A, rank = 2, K = 2)
#' fit <- nmf.gmm.inference(fit, Y, A, wild.B = 300)
#' head(fit$coefficients)
#' }
#' @export
nmf.gmm.inference <- function(object, Y, A = object$A, ...) {
  if (!inherits(object, "nmf.gmm")) stop("object must be of class 'nmf.gmm'.")
  if (!identical(object$cov, "tied"))
    stop("nmf.gmm.inference currently supports cov = 'tied' only.")
  if (length(object$class.rows))
    stop("nmf.gmm.inference does not yet support class-specific effects (class.effects).")
  extra <- base::list(...)
  B     <- if (!is.null(extra$wild.B))     extra$wild.B     else 500L
  ## CONVENTIONS.md 1: the bootstrap controls carry the wild. prefix, and every
  ## sibling .inference honours wild.seed.  Reading only `seed` made
  ## nmf.gmm.inference(..., wild.seed = 7) a silent no-op.  `seed` stays as the
  ## alias so existing calls keep working.
  seed  <- if (!is.null(extra$wild.seed)) extra$wild.seed
           else if (!is.null(extra$seed)) extra$seed else 123L
  ## Keep our own seeding out of the caller's random stream.
  .rng <- .nmfkc.rng.save(seed)
  on.exit(.nmfkc.rng.restore(.rng), add = TRUE)
  level <- if (!is.null(extra$wild.level)) extra$wild.level else 0.95

  X <- object$X; C <- object$C; mu <- object$mu
  tau2 <- as.numeric(object$tau2); s2 <- object$sigma2; gamma <- object$gamma
  Y <- as.matrix(Y); A <- as.matrix(A)
  P <- nrow(Y); N <- ncol(Y); Q <- ncol(X); R <- nrow(A)

  V     <- X %*% (tau2 * t(X)) + s2 * diag(P)
  Vinv  <- chol2inv(chol(V))
  XtVX  <- crossprod(X, Vinv %*% X)
  Info  <- kronecker(tcrossprod(A), XtVX)               # vec(C) column-stacked
  Hinv  <- solve(Info)

  mubar <- mu %*% t(gamma)                              # Q x N
  Rm    <- Y - X %*% (C %*% A + mubar)
  G     <- crossprod(X, Vinv %*% Rm)                    # Q x N
  score <- matrix(0, Q * R, N)
  for (n in 1:N) score[, n] <- as.vector(tcrossprod(G[, n], A[, n]))

  J      <- tcrossprod(score) * (N / (N - 1))
  V_sand <- Hinv %*% J %*% Hinv
  se_sand <- matrix(sqrt(pmax(diag(V_sand), 0)), Q, R)
  SS     <- tcrossprod(score)
  V_J    <- tryCatch(solve(SS),
                     error = function(e) solve(SS + diag(1e-8 * mean(diag(SS)), nrow(SS))))
  se_J   <- matrix(sqrt(pmax(diag(V_J), 0)), Q, R)

  ## wild bootstrap on vec(C)
  set.seed(seed)
  th <- as.vector(C); boot <- matrix(NA_real_, Q * R, B)
  for (b in 1:B) {
    w <- stats::rexp(N, rate = 1) - 1
    boot[, b] <- th + as.vector(Hinv %*% (score %*% w))
  }
  a <- 1 - level
  lo  <- matrix(apply(boot, 1, stats::quantile, probs = a/2,     names = FALSE), Q, R)
  hi  <- matrix(apply(boot, 1, stats::quantile, probs = 1 - a/2, names = FALSE), Q, R)
  bsd <- matrix(apply(boot, 1, stats::sd), Q, R)

  z <- C / se_J
  coefficients <- data.frame(
    Basis     = rep(rownames(C), times = R),
    Covariate = rep(colnames(C), each = Q),
    Estimate  = as.vector(C),
    SE        = as.vector(se_J),
    BSE       = as.vector(bsd),
    z_value   = as.vector(z),
    p_value   = as.vector(2 * stats::pnorm(-abs(z))),
    CI_low    = as.vector(lo),
    CI_high   = as.vector(hi),
    stringsAsFactors = FALSE, row.names = NULL)

  Xi <- X %*% C; dimnames(Xi) <- list(rownames(X), colnames(C))
  object$coefficients <- coefficients
  object$C.se.outer <- se_J; object$C.se.sandwich <- se_sand; object$C.se.boot <- bsd
  dimnames(object$C.se.outer) <- dimnames(object$C.se.sandwich) <-
    dimnames(object$C.se.boot) <- dimnames(C)
  object$Xi <- Xi; object$C.p.side <- "two.sided"
  if (!inherits(object, "nmf.gmm.inference"))
    class(object) <- c("nmf.gmm.inference", class(object))
  object
}


#' @title Choose the number of mixture components K for NMF-GMM
#' @description
#' This function is \strong{experimental and still under development}. The
#' interface may change in future versions: argument names, defaults and the
#' contents of the returned object are not yet stable. Code written against it
#' today may need adjusting after an update. The rest of the package does not
#' carry this caveat.
#'
#' \code{nmf.gmm.select} fits \code{\link{nmf.gmm}} over a vector of \eqn{K}
#' values and reports the log-likelihood, BIC and ICL for each, selecting
#' \code{K.best} by BIC (\code{K.best.icl} by ICL). If a vector of known labels
#' is supplied via \code{truth}, the adjusted Rand index of the hard clustering
#' is added.
#'
#' @param Y,A,rank As in \code{\link{nmf.gmm}}.
#' @param K Integer vector of candidate component counts. Default \code{1:5}.
#' @param ... Additional arguments passed to \code{\link{nmf.gmm}} (e.g.
#'   \code{cov}, \code{X.init}, \code{nstart}, \code{maxit}, \code{seed}). Also
#'   accepts \code{truth} (a length-N vector of known class labels for the ARI
#'   column), \code{verbose} (logical, print the table; default \code{TRUE}), and
#'   \code{cores} (evaluate the \code{K} candidates in parallel; default
#'   \code{getOption("mc.cores", 1L)}). Parallelism uses a PSOCK cluster on
#'   Windows and forking elsewhere; because each \code{K} is an independent
#'   self-seeded fit and results are returned in order, the table and the
#'   selected \code{K} are identical for any \code{cores}.
#'
#' @return An object of class \code{"nmf.gmm.select"}: a list with the
#'   \code{table} (data frame over \code{K}), \code{K.best} (min BIC),
#'   \code{K.best.icl} (min ICL) and \code{fits} (the fitted objects).
#' @seealso \code{\link{nmf.gmm}}
#' @examples
#' \donttest{
#' set.seed(1)
#' Y <- matrix(abs(rnorm(20 * 80)) + 1, 20, 80); A <- rbind(1, rnorm(80))
#' sel <- nmf.gmm.select(Y, A, rank = 2, K = 1:4, verbose = FALSE)
#' sel$K.best
#' }
#' @export
nmf.gmm.select <- function(Y, A = NULL, rank, K = 1:5, ...) {
  extra <- base::list(...)
  truth   <- extra$truth
  verbose <- if (!is.null(extra$verbose)) extra$verbose else TRUE
  cores   <- if (!is.null(extra$cores)) extra$cores else getOption("mc.cores", 1L)
  fit_args <- extra[!names(extra) %in% c("truth", "verbose", "cores")]

  ## Each K is an independent, self-seeded nmf.gmm() fit; .nmfkc.parlapply
  ## preserves input order so the table and selection are identical to the
  ## sequential loop for any `cores`. The inner fit runs its nstart restarts
  ## sequentially (cores = 1) so the K-level parallelism here does not nest.
  fit_one <- function(k)
    do.call(nmf.gmm, c(list(Y = Y, A = A, rank = rank, K = k, cores = 1L), fit_args))
  fits <- .nmfkc.parlapply(K, fit_one, cores = cores)
  rows <- lapply(seq_along(K), function(i) {
    f <- fits[[i]]
    ari <- if (!is.null(truth) && K[i] > 1) .nmfgmm.ARI(f$cluster, truth) else NA_real_
    data.frame(K = K[i], logLik = f$loglik, n.params = f$n.params,
               BIC = f$BIC, ICL = f$ICL, ARI = ari)
  })
  tab <- do.call(rbind, rows)
  K.best     <- tab$K[which.min(tab$BIC)]
  K.best.icl <- tab$K[which.min(tab$ICL)]
  if (isTRUE(verbose)) {
    show <- tab; show$logLik <- round(show$logLik, 1)
    show$BIC <- round(show$BIC, 1); show$ICL <- round(show$ICL, 1)
    show$ARI <- round(show$ARI, 3)
    if (all(is.na(show$ARI))) show$ARI <- NULL
    print(show, row.names = FALSE)
    cat(sprintf("BIC selects K=%d; ICL selects K=%d\n", K.best, K.best.icl))
  }
  structure(list(table = tab, K.best = K.best, K.best.icl = K.best.icl,
                 fits = fits), class = "nmf.gmm.select")
}


#' @title Two-stage (adjust-then-cluster) baseline for NMF-GMM
#' @description
#' This function is \strong{experimental and still under development}. The
#' interface may change in future versions: argument names, defaults and the
#' contents of the returned object are not yet stable. Code written against it
#' today may need adjusting after an update. The rest of the package does not
#' carry this caveat.
#'
#' \code{nmf.gmm.twostage} runs the \emph{two-stage} route that
#' \code{\link{nmf.gmm}} is designed to improve on, as a matched baseline:
#' (1) estimate least-squares scores on the initial basis, (2) regress the
#' scores on the covariates \emph{blind to the class} and keep the residuals,
#' (3) reconstitute the residuals in observation space, shift them to
#' non-negativity, and (4) refit an intercept-only \code{nmf.gmm} from the
#' \emph{same} basis initialization. Both routes start from the same basis and
#' use the same mixture family and settings. They differ in the order of
#' adjustment and clustering and, unless \code{fixX = TRUE}, in the basis each
#' refits to its own data: the two-stage route refits it to the
#' reconstituted, shifted residuals. When the covariate is associated with the
#' class, the class-blind regression of stage (2) displaces the class means
#' (the oracle proposition of Satoh 2026); when the covariate is
#' (near-)mean-independent of the class, that displacement vanishes in the
#' population, so both routes target the same class means, although the
#' fitted partitions need not coincide in a finite sample.
#'
#' @param Y Data matrix \eqn{Y} (P x N).
#' @param A Covariate matrix \eqn{A} (R x N) including an intercept row, or a
#'   one-sided formula evaluated in \code{data} (as in \code{\link{nmf.gmm}}).
#'   An intercept-only \code{A} is an error: there is nothing to adjust for.
#' @param rank Integer rank \eqn{Q} of the basis.
#' @param K Integer number of mixture components.
#' @param ... Additional arguments as in \code{\link{nmf.gmm}} (\code{cov},
#'   \code{X.init}, \code{nstart}, \code{maxit}, \code{seed}, \code{data},
#'   \code{standardize}, ...); they are applied to both stages.
#'
#' @return An object of class \code{c("nmf.gmm.twostage", "nmf.gmm")}: the
#'   stage-2 fit (all \code{\link{nmf.gmm}} fields and S3 methods apply),
#'   plus a \code{twostage} list with the non-negativity \code{shift}, the
#'   covariate matrix \code{A} that was removed, and the shared basis
#'   initialization \code{X0}.
#'
#' @seealso \code{\link{nmf.gmm}} (the joint route this baselines).
#' @examples
#' \donttest{
#' set.seed(1)
#' Y <- matrix(abs(rnorm(12 * 40)) + 1, 12, 40)
#' A <- rbind(1, rnorm(40))
#' ts <- nmf.gmm.twostage(Y, A, rank = 2, K = 2, nstart = 2, maxit = 100)
#' table(ts$cluster)
#' }
#' @export
nmf.gmm.twostage <- function(Y, A = NULL, rank, K = 1, ...) {
  extra <- base::list(...)
  getopt <- function(nm, default) if (!is.null(extra[[nm]])) extra[[nm]] else default
  intercept <- getopt("intercept", 1L)
  seed      <- getopt("seed", 1L)
  if (!is.null(extra$class.effects))
    stop("class-specific effects (class.effects) cannot be removed before clustering: ",
         "the two-stage route has no class to attach them to. Use nmf.gmm().")
  .rng <- .nmfkc.rng.save(seed)
  on.exit(.nmfkc.rng.restore(.rng), add = TRUE)
  Y <- as.matrix(Y); N <- ncol(Y); Q <- as.integer(rank)
  A.formula <- NULL
  if (inherits(A, "formula")) {
    A.formula <- A
    A <- .nmfgmm.buildA(A, extra$data, N, getopt("standardize", TRUE))
  }
  if (is.null(A)) stop("nmf.gmm.twostage() needs a covariate: supply A or a formula with `data`.")
  A <- as.matrix(A); if (ncol(A) != N && nrow(A) == N) A <- t(A)
  if (ncol(A) != N) stop("A must have N columns (R x N).")
  if (nrow(A) < 2) stop("A holds only an intercept; there is nothing to adjust for.")

  ## stage 0: the same initial basis nmf.gmm would use
  nstart <- getopt("nstart", if (K == 1) 1L else 8L)
  anchor.rows <- NULL
  if (!is.null(extra$X.anchor)) {
    ## anchors are found on Y, not on the shifted residuals, and reach the
    ## refit as zeros in X.init = X0, which its multiplicative X-update keeps
    an <- .nmfkc_anchor_setup(extra$X.anchor, Y, Q, D_A = NULL,
                              fun = "nmf.gmm.twostage", seed = seed)
    anchor.rows <- an$rows
    X0 <- an$X / rep(pmax(colSums(an$X), 1e-12), each = nrow(Y))
  } else {
    X0 <- .nmfgmm.initX(extra$X.init, Y, Q, seed, nstart)
  }

  ## stage 1: least-squares scores; remove the covariates blind to the class
  B <- solve(crossprod(X0), crossprod(X0, Y))
  Acov <- t(A[-intercept, , drop = FALSE])
  Bres <- t(stats::resid(stats::lm(t(B) ~ Acov)))

  ## stage 2: reconstitute, shift to non-negativity, refit from the same basis
  Yres <- X0 %*% Bres
  shift <- max(0, -min(Yres))
  Yres <- Yres + shift
  dimnames(Yres) <- dimnames(Y)
  pass <- extra
  pass$data <- NULL; pass$standardize <- NULL; pass$X.init <- NULL
  pass$X.anchor <- NULL
  fit <- do.call(nmf.gmm, c(list(Y = Yres, A = NULL, rank = rank, K = K,
                                 X.init = X0), pass))
  if (!is.null(anchor.rows))
    fit$X.anchor <- stats::setNames(anchor.rows, colnames(fit$X))
  fit$twostage <- list(shift = shift, A = A, A.formula = A.formula, X0 = X0)
  fit$call <- match.call()
  class(fit) <- c("nmf.gmm.twostage", class(fit))
  fit
}


# =====================================================================
#  S3 methods
# =====================================================================

#' @title Extract the covariate-coefficient matrix from an NMF-GMM fit
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @param object An object of class \code{"nmf.gmm"}.
#' @param ... Ignored.
#' @return The \eqn{Q\times R} coefficient matrix \eqn{C} (\eqn{=\Theta}).
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
coef.nmf.gmm <- function(object, ...) {
  ## Follow coef.nmf's precedence: the coefficients table once inference has
  ## been run, the raw matrix otherwise.  Returning object$C unconditionally
  ## meant coef() gave the same matrix before and after nmf.gmm.inference(),
  ## so the table it builds was unreachable through the generic.
  if (!is.null(object$coefficients)) object$coefficients else object$C
}

#' @title Fitted (responsibility-averaged) reconstruction of an NMF-GMM fit
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @param object An object of class \code{"nmf.gmm"}.
#' @param ... Ignored.
#' @return The \eqn{P\times N} fitted matrix \eqn{\hat Y = X(CA + \mu\gamma^\top)}.
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
fitted.nmf.gmm <- function(object, ...) object$Yhat


#' @title Residuals from an NMF-GMM fit
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @description
#' \eqn{Y - \hat Y}, matching \code{\link{residuals.nmf}}.  Without this method
#' \code{residuals()} fell through to \code{stats::residuals.default} and
#' returned \code{NULL}, even though \code{fitted()} worked.
#' @param object An object of class \code{"nmf.gmm"}.
#' @param Y The observation matrix the model was fitted on.
#' @param ... Ignored.
#' @return The residual matrix.
#' @seealso \code{\link{nmf.gmm}}, \code{\link{fitted.nmf.gmm}}
#' @export
residuals.nmf.gmm <- function(object, Y, ...) Y - stats::fitted(object)

#' @title Class assignments / responsibilities from an NMF-GMM fit
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @param object An object of class \code{"nmf.gmm"}.
#' @param type \code{"class"} (default) returns the hard cluster labels;
#'   \code{"responsibility"} returns the N x K posterior-probability matrix.
#' @param ... Ignored.
#' @return A vector of labels or the responsibility matrix.
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
predict.nmf.gmm <- function(object, type = c("class", "responsibility"), ...) {
  type <- match.arg(type)
  if (type == "class") object$cluster else object$gamma
}

#' @title Print an NMF-GMM fit
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @param x An object of class \code{"nmf.gmm"}.
#' @param ... Ignored.
#' @return \code{x}, invisibly.
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
print.nmf.gmm <- function(x, ...) {
  cat("NMF-GMM: Gaussian-mixture latent-class NMF with covariates\n")
  cat(x$dims, "\n")
  cat(sprintf("K=%d components, rank Q=%d, covariance=%s\n", x$K, x$rank, x$cov))
  if (length(x$class.effects))
    cat("Class-specific covariate effects:", paste(x$class.effects, collapse = ", "), "\n")
  cat(sprintf("logLik=%.2f, BIC=%.2f, ICL=%.2f (params=%d)\n",
              x$loglik, x$BIC, x$ICL, x$n.params))
  ## Delegate so this reads exactly like every other fitter's convergence line:
  ## the iteration count against the cap, so a run that merely ran out of
  ## iterations is not mistaken for a fitted one.
  .print.convergence(x, label = "Convergence: ")
  cat("Mixing proportions (xi):", paste(round(x$xi, 3), collapse = " "), "\n")
  cat("Cluster sizes:\n"); print(table(cluster = x$cluster))
  if (!is.null(x$coefficients))
    cat("\n(coefficients table for C available in $coefficients)\n")
  invisible(x)
}

#' @title Summary of an NMF-GMM fit
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @param object An object of class \code{"nmf.gmm"}.
#' @param ... Ignored.
#' @return An object of class \code{"summary.nmf.gmm"}.
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
summary.nmf.gmm <- function(object, ...) {
  ans <- list(dims = object$dims, K = object$K, rank = object$rank,
              cov = object$cov, loglik = object$loglik, BIC = object$BIC,
              ICL = object$ICL, n.params = object$n.params, xi = object$xi,
              sizes = table(object$cluster), iter = object$iter,
              maxit = object$maxit, epsilon = object$epsilon,
              converged = object$converged,
              runtime = object$runtime, coefficients = object$coefficients,
              C.p.side = object$C.p.side)
  class(ans) <- "summary.nmf.gmm"
  ans
}

#' @title Print method for summary.nmf.gmm objects
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @param x A \code{"summary.nmf.gmm"} object.
#' @param digits Minimum significant digits.
#' @param ... Ignored.
#' @return \code{x}, invisibly.
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
print.summary.nmf.gmm <- function(x, digits = max(3L, getOption("digits") - 3L), ...) {
  cat("NMF-GMM fit\n"); cat(x$dims, "\n")
  cat(sprintf("K=%d, rank=%d, cov=%s | logLik=%.2f, BIC=%.2f, ICL=%.2f, params=%d\n",
              x$K, x$rank, x$cov, x$loglik, x$BIC, x$ICL, x$n.params))
  .print.convergence(x, label = "Convergence: ")
  cat("Mixing proportions:", paste(round(x$xi, 3), collapse = " "), "\n")
  cat("Cluster sizes:\n"); print(x$sizes)
  if (!is.null(x$coefficients)) {
    cf <- x$coefficients
    stars <- ifelse(!is.finite(cf$p_value), " ",
             ifelse(cf$p_value < 0.001, "***",
             ifelse(cf$p_value < 0.01, "**",
             ifelse(cf$p_value < 0.05, "*",
             ifelse(cf$p_value < 0.1, ".", " ")))))
    cat("\nCoefficients C (conditional on the estimated basis X):\n")
    out <- data.frame(cf[, c("Basis", "Covariate")],
                      Estimate = round(cf$Estimate, 4), SE = round(cf$SE, 4),
                      BSE = round(cf$BSE, 4), z = round(cf$z_value, 2),
                      p_value = format.pval(cf$p_value, digits = 3, eps = 1e-4),
                      sig = stars)
    print(out, row.names = FALSE)
    cat("---\nSignif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1\n")
    cat("(SE = outer-product mixture-information; BSE = wild bootstrap.)\n")
  }
  invisible(x)
}

#' @title Plot method for nmf.gmm objects
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @description
#' \code{type = "convergence"} (default) plots the EM objective
#' (\eqn{-\log L}) over iterations.  \code{type = "adjusted.scores"} draws the
#' cluster scatterplot, coloured by the hard cluster assignment (or by
#' \code{group}).
#'
#' @details
#' \strong{The scatterplot shows the adjusted scores, not the scores.}  What is
#' drawn is \eqn{\bm b_n - C\bm a_n}, the least-squares scores with the
#' covariate effect removed --- which is the space the mixture actually acts in,
#' since the model puts the covariate effect in the component means.  The
#' distinction matters: on the \code{crabs} example the \emph{raw} scores
#' separate the four true groups no better with a size covariate than without
#' (0.635 versus 0.648, between/within on the first two principal components),
#' while the adjusted space separates them at 4.628.  The visible separation is
#' produced by the adjustment, so a panel labelled "scores" would claim
#' something the picture does not show.  \code{type = "scores"} is accepted as
#' an alias.
#'
#' The projection depends on the rank: a rank-2 fit is drawn in its own two
#' coordinates (rotating it through \code{prcomp} would trade two interpretable
#' part axes for arbitrary ones), a higher rank is projected onto its first two
#' principal components, and a rank-1 fit becomes a strip plot with
#' deterministic stacking.
#' @param x An object of class \code{"nmf.gmm"}.
#' @param type \code{"convergence"} or \code{"adjusted.scores"}
#'   (\code{"scores"} is an alias for the latter).
#' @param group Optional length-N factor/vector to colour the scatterplot by
#'   (e.g. known labels). Default \code{NULL} colours by \code{x$cluster}.
#' @param ... Additional graphical parameters passed to \code{plot}.
#' @return Invisible \code{NULL}.
#' @examples
#' \donttest{
#' set.seed(1)
#' Y <- matrix(abs(rnorm(20 * 80)) + 1, 20, 80); A <- rbind(1, rnorm(80))
#' fit <- nmf.gmm(Y, A, rank = 2, K = 3)
#' plot(fit, type = "adjusted.scores")
#' }
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
plot.nmf.gmm <- function(x, type = c("convergence", "adjusted.scores", "scores"),
                         group = NULL, ...) {
  type <- match.arg(type)
  if (type == "scores") type <- "adjusted.scores"   # kept as an alias
  extra <- list(...)
  if (type == "convergence") {
    args <- list(x = seq_along(x$objfunc.iter), y = x$objfunc.iter)
    if (is.null(extra$type)) args$type <- "l"
    if (is.null(extra$xlab)) args$xlab <- "EM iteration"
    if (is.null(extra$ylab)) args$ylab <- "-loglik"
    if (is.null(extra$main)) args$main <- "nmf.gmm convergence"
    do.call("plot", c(args, extra))
    return(invisible(NULL))
  }
  ## type == "adjusted.scores".  What is drawn is NOT b_n but b_n - C a_n, the
  ## space the mixture actually acts in: the model puts the covariate effect in
  ## the component means, so the clusters live in the residual after that effect
  ## is removed.  On the crabs example the raw scores separate the four groups
  ## no better with a size covariate than without (0.635 vs 0.648 by
  ## between/within on the first two PCs), while the adjusted space separates
  ## them at 4.628 -- the separation is created by the adjustment, and calling
  ## the panel "scores" would claim otherwise.
  if (is.null(x$scores)) stop("no $scores stored; refit with nmf.gmm().")
  adj <- x$scores - x$C %*% x$A                      # remove the covariate mean
  if (length(x$class.rows)) {                        # class-specific rows: remove
    S <- x$class.rows                                # the effect of each unit's class
    for (k in seq_len(x$K)) {
      idx <- x$cluster == k
      if (any(idx)) {
        Dk <- matrix(x$C.class[, , k], nrow(adj)) - x$C[, S, drop = FALSE]
        adj[, idx] <- adj[, idx, drop = FALSE] - Dk %*% x$A[S, idx, drop = FALSE]
      }
    }
  }
  g   <- if (!is.null(group)) as.factor(group) else as.factor(x$cluster)
  pal <- grDevices::hcl.colors(max(nlevels(g), 2), "Dark 3")
  col <- pal[as.integer(g)]
  lab <- rownames(adj)
  if (is.null(lab)) lab <- paste0("basis", seq_len(nrow(adj)))
  if (nrow(adj) == 2L) {
    ## A rank-2 fit is already 2-D.  Rotating it through prcomp would replace
    ## two interpretable part axes with arbitrary ones for no gain.
    pc <- t(adj)
    xl <- paste("adjusted", lab[1]); yl <- paste("adjusted", lab[2])
  } else if (nrow(adj) > 2L) {
    pc <- stats::prcomp(t(adj))$x[, 1:2, drop = FALSE]
    xl <- "adjusted score PC1"; yl <- "adjusted score PC2"
  } else {                                            # rank 1: strip plot
    ## Deterministic stacking rather than runif(): a plot method must not
    ## advance the caller's random stream.
    v   <- as.numeric(adj)
    nb  <- max(2L, min(50L, ncol(adj) %/% 2L))
    off <- stats::ave(seq_along(v), cut(v, breaks = nb, labels = FALSE),
                      FUN = function(i) seq_along(i) - (length(i) + 1) / 2)
    pc <- cbind(v, off)
    xl <- "adjusted score"; yl <- ""
  }
  args <- list(x = pc[, 1], y = pc[, 2], col = col)
  if (is.null(extra$pch))  args$pch  <- 19
  if (is.null(extra$xlab)) args$xlab <- xl
  if (is.null(extra$ylab)) args$ylab <- yl
  if (is.null(extra$main)) args$main <- sprintf("nmf.gmm adjusted scores (K=%d)", x$K)
  do.call("plot", c(args, extra))
  graphics::legend("topright", legend = levels(g), col = pal[seq_len(nlevels(g))],
                   pch = 19, bty = "n", cex = 0.8)
  invisible(NULL)
}

#' @title Print method for nmf.gmm.select objects
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @param x An object of class \code{"nmf.gmm.select"}.
#' @param ... Ignored.
#' @return \code{x}, invisibly.
#' @seealso \code{\link{nmf.gmm}}, \code{\link{nmf.gmm.inference}}, \code{\link{nmf.gmm.select}}
#' @export
print.nmf.gmm.select <- function(x, ...) {
  tab <- x$table; tab$logLik <- round(tab$logLik, 1)
  tab$BIC <- round(tab$BIC, 1); tab$ICL <- round(tab$ICL, 1); tab$ARI <- round(tab$ARI, 3)
  if (all(is.na(tab$ARI))) tab$ARI <- NULL
  print(tab, row.names = FALSE)
  cat(sprintf("BIC selects K=%d; ICL selects K=%d\n", x$K.best, x$K.best.icl))
  invisible(x)
}


#' @title Plot method for nmf.gmm.select objects
#' @section Development status:
#' Part of the \strong{experimental} NMF-GMM family: see \code{\link{nmf.gmm}}.
#' The interface may change in future versions --- argument names, defaults
#' and the contents of the returned object are not yet stable.
#' @description
#' Draws BIC and ICL against \eqn{K}, with the selected \eqn{K} marked.  Every
#' other selector in the package (\code{\link{nmfkc.rank}},
#' \code{\link{nmfkc.ard}}, \code{\link{nmfkc.consensus}},
#' \code{\link{nmfkc.bicv}}) has a plot method; without this one
#' \code{plot(nmf.gmm.select(...))} fell through to \code{plot.default} and
#' errored with "'x' is a list, but does not have components 'x' and 'y'".
#' @param x An object of class \code{"nmf.gmm.select"}.
#' @param ... Passed to the underlying plot call.
#' @return \code{x}, invisibly.
#' @seealso \code{\link{nmf.gmm.select}}, \code{\link{nmf.gmm}}
#' @export
plot.nmf.gmm.select <- function(x, ...) {
  tab <- x$table
  extra <- list(...)
  dflt <- function(...) {
    d <- list(...)
    c(d[setdiff(names(d), names(extra))], extra)
  }
  ylim <- range(c(tab$BIC, tab$ICL), finite = TRUE)
  do.call("plot", dflt(x = tab$K, y = tab$BIC, type = "b", pch = 19,
                       col = "steelblue4", ylim = ylim, xlab = "K",
                       ylab = "information criterion",
                       main = "nmf.gmm.select: BIC / ICL"))
  graphics::lines(tab$K, tab$ICL, type = "b", pch = 17, lty = 2,
                  col = "firebrick")
  graphics::abline(v = x$K.best, col = "steelblue4", lty = 3)
  if (!identical(x$K.best, x$K.best.icl))
    graphics::abline(v = x$K.best.icl, col = "firebrick", lty = 3)
  graphics::legend("topright",
                   legend = c(sprintf("BIC (best K = %d)", x$K.best),
                              sprintf("ICL (best K = %d)", x$K.best.icl)),
                   col = c("steelblue4", "firebrick"), pch = c(19, 17),
                   lty = c(1, 2), bty = "n", cex = 0.85)
  invisible(x)
}