## metabric_glmnet.R --fit the four classifiers and cross-validate them on one
## shared set of folds.
##
##   ungrouped   cv.glmnet, family = "multinomial", type.multinomial = "ungrouped"
##   grouped     cv.glmnet, family = "multinomial", type.multinomial = "grouped"
##   ovr         ten binomial cv.glmnet fits, one per cluster
##   pamr        nearest shrunken centroids
##
## We use seed 42L per LT paper
## Input:  data/metabric.rds   (from metabric_data.R)
## Output: data/models.rds

suppressPackageStartupMessages({
  library(glmnet)
  library(pamr)
})

SEED   <- 42L
NFOLDS <- 10L

CACHE  <- "data/metabric.rds"
OUT    <- "data/models.rds"

stopifnot(file.exists(CACHE))
dat <- readRDS(CACHE)

x  <- dat$discovery$x
y  <- dat$discovery$y
xv <- dat$validation$x
n  <- nrow(x)
clusters <- levels(y)

## One fold assignment, used by all four models.
## Keeping the same LT construction.
set.seed(SEED)
foldid <- sample(rep(seq_len(NFOLDS), trunc(n / NFOLDS) + 1L))[seq_len(n)]

## Multinomial fits
##
## keep = TRUE retains fit.preval, the held-out linear predictors for every
## sample at every lambda. That gives per-cluster CV error without refitting;
## it does not change the fit itself.
fit_multinomial <- function(type) {
  cv.glmnet(x, y,
            family = "multinomial", type.multinomial = type,
            type.measure = "class", foldid = foldid, keep = TRUE)
}

message("fitting ungrouped multinomial ...")
cv_ungrouped <- fit_multinomial("ungrouped")
##
## The grouped fit emits 12 convergence warnings: with 10 classes tied to a
## single support, coordinate descent does not reach tolerance at the smallest
## lambdas within maxit = 1e5, and glmnet truncates the path from 87 values to
## 54. This is benign here: refitting at maxit = 1e6 and 1e7 clears the warnings
## and restores the full 87-value path, but returns the same lambda.min, the
## same CV error, and the identical 184-feature support -- lambda.min sits well
## inside the surviving path, not against the truncated end.
##
message("fitting grouped multinomial ...")
cv_grouped   <- fit_multinomial("grouped")

## Nonzero features per class at lambda.min, dropping the intercept. For the
## grouped fit every class shares one support, so all ten counts are equal and
## the union equals any one of them.
nonzero_by_class <- function(cvfit) {
  co <- coef(cvfit, s = "lambda.min")
  vapply(co, function(m) sum(m[-1, 1] != 0), integer(1))
}

support_union <- function(cvfit) {
  co <- coef(cvfit, s = "lambda.min")
  Reduce(union, lapply(co, function(m) rownames(m)[-1][m[-1, 1] != 0]))
}

cv_error_at_min <- function(cvfit) cvfit$cvm[cvfit$index["min", ]]

## OVR fit
## One binomial lasso per cluster: the cluster is 1, everything else 0. Each
## classifier gets its own lambda.min -- that per-class choice of regularization
## is the whole point of the comparison.
message("fitting one-versus-rest (10 binomial models) ...")
ovr <- lapply(clusters, function(k) {
  yk <- as.integer(y == k)
  cv.glmnet(x, yk, family = "binomial",
            type.measure = "class", foldid = foldid, keep = TRUE)
})
names(ovr) <- clusters

ovr_err   <- vapply(ovr, cv_error_at_min, numeric(1))
ovr_nz    <- vapply(ovr, function(f) sum(coef(f, s = "lambda.min")[-1, 1] != 0),
                    integer(1))
ovr_union <- Reduce(union, lapply(ovr, function(f) {
  m <- coef(f, s = "lambda.min")
  rownames(m)[-1][m[-1, 1] != 0]
}))

## The LT weighted average: each classifier's minimum CV error times its
## cluster's size, summed, over the total. Cluster sizes sum to n, so this is a
## weighted mean with weights n_k / n.
cluster_n <- as.integer(table(y)[clusters])
ovr_weighted_err <- sum(ovr_err * cluster_n) / n

## Combining ten binomial classifiers into one label: argmax of the predicted
## probability.
ovr_predict <- function(newx) {
  ## each predict() returns an n x 1 matrix, so cbind gives n patients x 10
  ## clusters, in the order of `ovr` -- which is `clusters`, not levels(y).
  p <- do.call(cbind, lapply(ovr, predict, newx = newx, s = "lambda.min",
                             type = "response"))
  factor(clusters[max.col(p, ties.method = "first")], levels = clusters)
}

## pamr fit
## pamr wants genes in rows. Same folds as above, handed over as a list of
## held-out index vectors.
message("fitting pamr ...")
pamr_data  <- list(x = t(x), y = y, geneid = colnames(x))
pamr_fit   <- pamr.train(pamr_data)
pamr_folds <- split(seq_len(n), foldid)
## No set.seed here: pamr.cv only randomizes when `folds` is NULL (nsccv calls
## balanced.folds in that branch only), so with folds supplied it is
## deterministic. Seeding here would imply otherwise.
pamr_cv <- pamr.cv(pamr_fit, pamr_data, folds = pamr_folds)
stopifnot(sum(pamr_cv$error == min(pamr_cv$error)) == 1L)  # unique min
pamr_best_i   <- which.min(pamr_cv$error)
pamr_best_thr <- pamr_cv$threshold[pamr_best_i]
pamr_err      <- pamr_cv$error[pamr_best_i]
## $size is object$nonzero: the gene count of the FULL-DATA training fit at each
## threshold, not a cross-validated quantity.
pamr_size     <- pamr_cv$size[pamr_best_i]

## Per-class feature counts at the chosen threshold: genes whose shrunken
## centroid differs from the overall centroid for that class.
##
## Computed from the centroid matrix directly rather than by parsing
## pamr.listgenes, which rounds to 4 decimals before returning and so reports a
## non-zero-but-tiny difference as zero -- that undercounted clusters 7 and 9 by
## one gene each. pamr.listgenes also sets options(width = 500) without
## restoring it, which then reformats every print() later in this script.
pamr_nz <- local({
  cen <- pamr.predict(pamr_fit, pamr_data$x, threshold = pamr_best_thr,
                      type = "cent")
  d <- cen - pamr_fit$centroid.overall
  setNames(as.integer(colSums(d != 0)), clusters)
})

# ------------------------------------------------- per-cluster CV error ------

## Error broken down by TRUE cluster, from the held-out CV predictions: of the
## discovery patients really in cluster k, what fraction did CV put elsewhere?
## This is the quantity in the lower panel of LT's Figures 2 and 3, and the
## x-axis of its CV-error-vs-log-rank scatter.
##
## One definition, applied identically to all four methods, so the per-cluster
## numbers are comparable across them.
##
## This is a different quantity from ovr_weighted_err above, which averages the
## ten BINOMIAL error rates -- each counting errors over all 997 patients rather
## than over cluster k only. Both are kept; they answer different questions.
per_class_error <- function(pred, truth) {
  vapply(clusters, function(k) mean(pred[truth == k] != k), numeric(1))
}

## cvstats() trims lambda/cvm wherever cvsd is NA and $index points into the
## TRIMMED vector, but fit.preval is stored untrimmed. If trimming ever fires,
## indexing fit.preval by $index silently selects the wrong lambda. It does not
## fire for any fit here, but match() against the underlying path is correct by
## construction, so use that and assert the equivalence.
preval_index <- function(cvfit) {
  i <- match(cvfit$lambda.min, cvfit$glmnet.fit$lambda)
  stopifnot(!is.na(i), dim(cvfit$fit.preval)[length(dim(cvfit$fit.preval))] ==
              length(cvfit$glmnet.fit$lambda))
  i
}

heldout_multinomial <- function(cvfit) {
  ## fit.preval is on the LINK scale for multinomial; softmax is order
  ## preserving across classes, so argmax of the link is the CV class.
  lp <- cvfit$fit.preval[, , preval_index(cvfit)]
  factor(clusters[max.col(lp, ties.method = "first")], levels = clusters)
}

## Ten binomial classifiers, each at its own lambda.min, combined by argmax.
## plogis is monotone, so ranking the linear predictors ranks the probabilities.
ovr_heldout_lp <- vapply(ovr, function(f) f$fit.preval[, preval_index(f)],
                         numeric(n))
heldout_ovr <- factor(clusters[max.col(ovr_heldout_lp, ties.method = "first")],
                      levels = clusters)

heldout_pamr <- factor(pamr_cv$yhat[[pamr_best_i]], levels = clusters)
stopifnot(length(heldout_pamr) == n, !anyNA(heldout_pamr))

heldout <- list(
  ungrouped = heldout_multinomial(cv_ungrouped),
  grouped   = heldout_multinomial(cv_grouped),
  ovr       = heldout_ovr,
  pamr      = heldout_pamr
)
cv_by_cluster <- do.call(rbind, lapply(heldout, per_class_error, truth = y))

## Overall multiclass CV error under that same definition: the fraction of the
## 997 discovery patients that held-out CV assigned to the wrong cluster. For
## the two multinomial fits this reproduces cv.glmnet's own cvm exactly, and for
## pamr it reproduces pamr.cv's error exactly. For OVR it is a different number
## from the weighted binomial average the arm is scored on.
cv_overall <- vapply(heldout, function(p) mean(p != y), numeric(1))

## The decisive objection, which needs no metric argument at all: the reported
## number and the labels come from two DIFFERENT rules. type.measure = "class"
## on a binomial scores that binomial thresholded at p > 0.5, in isolation; the
## labels come from argmax over all ten probabilities (ovr_predict, above).
## Only the second is a classifier. Applied to all ten at once on the held-out
## discovery patients, the first leaves a quarter unlabeled and gives an eighth
## two or more labels -- so 0.0572 does not score the rule that produced the
## paper's validation labels or its KM curves.
ovr_label_counts <- local({
  pos <- ovr_heldout_lp > 0          # p > 0.5 on the link scale
  k <- rowSums(pos)
  c(none = sum(k == 0), one = sum(k == 1), multiple = sum(k >= 2))
})

## Predictions

pred <- list(
  ungrouped = factor(drop(predict(cv_ungrouped, newx = xv, s = "lambda.min",
                                  type = "class")), levels = clusters),
  grouped   = factor(drop(predict(cv_grouped, newx = xv, s = "lambda.min",
                                  type = "class")), levels = clusters),
  ovr       = ovr_predict(xv),
  pamr      = factor(pamr.predict(pamr_fit, t(xv), threshold = pamr_best_thr),
                     levels = clusters)
)
stopifnot(vapply(pred, function(p) all(!is.na(p)) && length(p) == nrow(xv),
                 logical(1)))

## Collect all results for the book so as to avoid recomputes

models <- list(
  seed = SEED, foldid = foldid, clusters = clusters,
  ungrouped = list(fit = cv_ungrouped, err = cv_error_at_min(cv_ungrouped),
                   nz = nonzero_by_class(cv_ungrouped),
                   total = length(support_union(cv_ungrouped)),
                   lambda_min = cv_ungrouped$lambda.min,
                   lambda_1se = cv_ungrouped$lambda.1se),
  grouped   = list(fit = cv_grouped, err = cv_error_at_min(cv_grouped),
                   nz = nonzero_by_class(cv_grouped),
                   total = length(support_union(cv_grouped)),
                   lambda_min = cv_grouped$lambda.min,
                   lambda_1se = cv_grouped$lambda.1se),
  ovr       = list(fit = ovr, err = ovr_weighted_err, err_by_class = ovr_err,
                   nz = ovr_nz, total = length(ovr_union),
                   sum_nz = sum(ovr_nz)),
  pamr      = list(fit = pamr_fit, cv = pamr_cv, err = pamr_err,
                   nz = pamr_nz, total = pamr_size, threshold = pamr_best_thr),
  pred = pred,
  heldout = heldout,
  cv_by_cluster = cv_by_cluster,
  cv_overall = cv_overall,
  ovr_label_counts = ovr_label_counts,
  cluster_n = setNames(cluster_n, clusters)
)
saveRDS(models, OUT)

# ------------------------------------------------------------------- report ---

fitted <- data.frame(
  model = c("ungrouped", "grouped", "ovr", "pamr"),
  err   = round(c(models$ungrouped$err, models$grouped$err,
                  models$ovr$err, models$pamr$err), 4),
  total = c(models$ungrouped$total, models$grouped$total,
            models$ovr$total, models$pamr$total)
)

cat("\n=== CV error and sparsity (seed =", SEED, ") ===\n")
print(fitted, row.names = FALSE)

cat("\nlambda.min  ungrouped:", signif(models$ungrouped$lambda_min, 3),
    " grouped:", signif(models$grouped$lambda_min, 3), "\n")
cat("lambda.1se  ungrouped:", signif(models$ungrouped$lambda_1se, 3), "\n")
cat("pamr best threshold  :", signif(models$pamr$threshold, 3), "\n")

cat("\n=== features selected per cluster ===\n")
print(rbind(ungrouped = models$ungrouped$nz, grouped = models$grouped$nz,
            ovr = models$ovr$nz, pamr = models$pamr$nz))
cat("\novr: union", models$ovr$total, " sum over classifiers", models$ovr$sum_nz,
    "\n")

cat("\n=== CV error by TRUE cluster (same definition for all four) ===\n")
print(round(models$cv_by_cluster, 4))

cat("\n=== overall CV error ===\n")
lfl <- data.frame(
  model      = names(cv_overall),
  as_scored  = round(c(models$ungrouped$err, models$grouped$err,
                       models$ovr$err, models$pamr$err), 4),
  multiclass = round(as.numeric(cv_overall), 4)
)
print(lfl, row.names = FALSE)
cat("\nas_scored is the size-weighted binomial average for OVR and multiclass\n")
cat("for the other three; multiclass applies one definition to all four.\n")

cat("\n=== how many labels does thresholding each binomial give? ===\n")
cat("Thresholding each binomial at p > 0.5 on the held-out discovery patients:\n")
cat(sprintf("  no label   : %3d (%.1f%%)\n  exactly one: %3d (%.1f%%)\n",
            ovr_label_counts[["none"]], 100 * ovr_label_counts[["none"]] / n,
            ovr_label_counts[["one"]], 100 * ovr_label_counts[["one"]] / n))
cat(sprintf("  two or more: %3d (%.1f%%)\n",
            ovr_label_counts[["multiple"]],
            100 * ovr_label_counts[["multiple"]] / n))
cat("So the p > 0.5 rule that scores each binomial leaves ",
    sprintf("%.1f%%", 100 * (ovr_label_counts[["none"]] +
                             ovr_label_counts[["multiple"]]) / n),
    " of the cohort\nwithout a usable label; the labels come from argmax over",
    " the ten instead.\n", sep = "")

cat("\n=== OVR per-cluster CV error (per binomial classifier) ===\n")
print(round(models$ovr$err_by_class, 4))
cat("weighted average:", round(models$ovr$err, 4), "\n")

cat("\n=== agreement with the published validation labels ===\n")
truth <- dat$validation$y
for (nm in names(pred)) {
  cat(sprintf("  %-10s %.4f\n", nm, mean(pred[[nm]] == truth)))
}

cat("\nwrote", OUT, "\n")