## metabric_figures.R --the remaining figures from the paper.
##
## Everything is written under figures/ with names that do NOT collide with the
## original PNGs in the project root, so the two sets can be compared
## side by side. The originals are never touched; assert_safe() below refuses to
## write anywhere except figures/.
##
##   figures/nsc_agreement.png       <-> intclust_test_accuracy_plot.png (Fig 4)
##   figures/features_per_cluster.png<-> features_plot.png               (Fig 5)
##   figures/cv_error_vs_logrank.png <-> err_z_scatter.png
##   figures/km_all_clusters.png     <-> all_clusters_km_panel.png
##
## Input: data/models.rds, data/survival.rds, data/metabric.rds

suppressPackageStartupMessages({
  library(survival)
  library(ggplot2)
})

dat    <- readRDS("data/metabric.rds")
models <- readRDS("data/models.rds")
surv   <- readRDS("data/survival.rds")

clusters <- models$clusters
val      <- dat$validation
PALETTE  <- surv$palette
METHODS  <- surv$methods

dir.create("images/metabric", showWarnings = FALSE)

## Guard: the paper's own PNGs live in the project root and must survive.
assert_safe <- function(path) {
  stopifnot(startsWith(path, "figures/"), !file.exists(file.path(".", basename(path))))
  path
}
save_fig <- function(p, file, width, height) {
  ggsave(assert_safe(file), p, width = width, height = height, dpi = 200)
  invisible(file)
}

base_theme <- theme_bw(base_size = 11) +
  theme(panel.grid.minor = element_blank(),
        legend.position = "bottom")

# ------------------------------------------- Fig 4: agreement with NSC labels ---
#
# "the proportion of patients assigned to it by both glmnet and pamr out of the
# patients assigned to it by pamr only" -- i.e. recall against the NSC labels.
# pamr scores 1.0 everywhere by construction, since its predictions ARE the
# published validation labels; it is the definition of the yardstick, not a
# competitor measured by it.

agreement <- do.call(rbind, lapply(names(METHODS), function(m) {
  p <- models$pred[[m]]
  data.frame(method = METHODS[[m]], cluster = as.integer(clusters),
             score = vapply(clusters, function(k) {
               denom <- sum(val$y == k)          # patients pamr assigns to k
               if (denom == 0L) NA_real_ else sum(p == k & val$y == k) / denom
             }, numeric(1)))
}))

p_agree <- ggplot(agreement, aes(cluster, score, colour = method)) +
  geom_line(linewidth = 0.7) + geom_point(size = 2.2) +
  scale_x_continuous(breaks = 1:10) +
  ## coord_cartesian below, not scale limits: limits DROP out-of-range points
  ## silently rather than zooming, which at another seed would delete a cluster.
  scale_y_continuous(breaks = seq(0.5, 1, 0.1)) +
  coord_cartesian(ylim = c(0.5, 1)) +
  scale_colour_manual(values = PALETTE, name = NULL) +
  labs(x = "Cluster", y = "Similarity to NSC-predicted Cluster Labels",
       title = "Agreement with the published validation labels",
       subtitle = paste("pamr is pinned at 1.0 by construction: the published",
                        "labels are NSC output, so this\nmeasures agreement",
                        "with pamr rather than accuracy.")) +
  base_theme
save_fig(p_agree, "figures/nsc_agreement.png", 9, 6)

# --------------------------------------- Fig 5: features selected per cluster ---

features <- do.call(rbind, lapply(names(METHODS), function(m) {
  data.frame(method = METHODS[[m]], cluster = as.integer(clusters),
             n_features = as.integer(models[[m]]$nz))
}))

p_feat <- ggplot(features, aes(cluster, n_features, colour = method)) +
  geom_line(linewidth = 0.7) + geom_point(size = 2.2) +
  scale_x_continuous(breaks = 1:10) +
  scale_colour_manual(values = PALETTE, name = NULL) +
  labs(x = "Cluster", y = "Number of Features Selected",
       title = "Features selected per cluster",
       subtitle = paste0("Grouped is flat by construction (one shared support). ",
                         "pamr uses all 754 genes at its\nchosen threshold; the ",
                         "counts shown are genes with a non-zero shrunken ",
                         "centroid for that class.")) +
  base_theme
save_fig(p_feat, "figures/features_per_cluster.png", 9, 6)

# ------------------------------------------- CV error vs log-rank statistic ---
#
# One point per cluster per method. The x axis uses the LIKE-FOR-LIKE per-cluster
# CV error (fraction of patients truly in cluster k that held-out CV sent
# elsewhere), not the paper's per-binomial OVR rate, so the four methods are on
# one scale.

scatter <- do.call(rbind, lapply(names(surv$endpoints), function(e) {
  do.call(rbind, lapply(names(METHODS), function(m) {
    s <- surv$res[[e]]$stats
    data.frame(endpoint = surv$endpoints[[e]], method = METHODS[[m]],
               cluster = as.integer(clusters),
               cv_error = as.numeric(models$cv_by_cluster[m, ]),
               z = s$z[s$key == m])
  }))
}))
scatter$endpoint <- factor(scatter$endpoint, levels = unname(surv$endpoints))

p_scatter <- ggplot(scatter, aes(cv_error, z, colour = method)) +
  geom_point(size = 2.4) +
  geom_text(aes(label = cluster), size = 2.7, vjust = -0.9,
            show.legend = FALSE) +
  facet_wrap(~endpoint, ncol = 1) +
  scale_colour_manual(values = PALETTE, name = NULL) +
  labs(x = "CV Error for that cluster", y = "Log-rank Test Statistic (|z|)",
       title = "CV error against survival agreement, per cluster",
       subtitle = "Points labeled by cluster.") +
  base_theme
save_fig(p_scatter, "figures/cv_error_vs_logrank.png", 9, 9)

## Correlation of per-cluster CV error with per-cluster |z|, reported as both
## coefficients with confidence intervals and Holm-adjusted p-values.
corr <- do.call(rbind, lapply(
  split(scatter, list(scatter$endpoint, scatter$method), drop = TRUE),
  function(d) {
    sp <- suppressWarnings(cor.test(d$cv_error, d$z, method = "spearman"))
    pe <- cor.test(d$cv_error, d$z)
    data.frame(endpoint = as.character(d$endpoint[1]), method = d$method[1],
               spearman = round(unname(sp$estimate), 3),
               sp_p = round(sp$p.value, 3),
               pearson = round(unname(pe$estimate), 3),
               pe_p = round(pe$p.value, 3),
               pe_lo = round(pe$conf.int[1], 3), pe_hi = round(pe$conf.int[2], 3))
  }))
corr$pe_p_holm <- round(p.adjust(corr$pe_p, "holm"), 3)
rownames(corr) <- NULL

## The positive, well-powered result the correlations cannot give: how much of
## the variation in |z| is cluster and how much is method. Run per endpoint.
z_anova <- lapply(names(surv$endpoints), function(e) {
  a <- aov(z ~ factor(cluster) + factor(key), data = surv$res[[e]]$stats)
  ss <- summary(a)[[1]][, "Sum Sq"]
  list(fit = a, pct_cluster = 100 * ss[1] / sum(ss),
       pct_method = 100 * ss[2] / sum(ss))
})
names(z_anova) <- names(surv$endpoints)

# ------------------------------------------ all-clusters KM panel ------------
#
# The paper's all_clusters_km_panel.png annotates each panel with a single
# k-sample log-rank p across all five curves. That test is INVALID here and is
# deliberately not reproduced.
#
# The five "groups" are the discovery truth plus four sets of predictions drawn
# from the SAME 983 validation patients, so a patient predicted into cluster k
# by three methods contributes three rows. At cluster 5 the four predicted
# groups supply 436 rows from only 127 distinct patients (3.4x), and pamr's and
# OVR's groups alone share 91. survdiff assumes independent observations;
# duplicating patients understates the variance and biases p toward 1 --
# precisely the direction that would appear to support "the curves are
# indistinguishable". Failure to reject would not show equivalence even if the
# test were valid.
#
# Annotated instead with the two-group |z| against the reference, which uses
# disjoint patient sets and is already computed in metabric_survival.R.

dup_factor <- vapply(clusters, function(k) {
  sets <- lapply(names(METHODS), function(m) which(models$pred[[m]] == k))
  sum(lengths(sets)) / length(unique(unlist(sets)))
}, numeric(1))

for (e in names(surv$endpoints)) {
  a <- surv$res[[e]]
  lab_df <- data.frame(
    cluster = as.integer(clusters),
    label = sprintf("|z| vs truth: %.2f", a$stats$z[a$stats$key == "pamr"]))
  curves <- a$curves
  curves$facet <- factor(paste("Cluster", curves$cluster),
                         levels = paste("Cluster", 1:10))
  lab_df$facet <- factor(paste("Cluster", lab_df$cluster),
                         levels = paste("Cluster", 1:10))

  p_all <- ggplot(curves, aes(time, surv, colour = label)) +
    geom_step(linewidth = 0.5) +
    geom_text(data = lab_df, aes(x = 5, y = 0.08, label = label),
              inherit.aes = FALSE, hjust = 0, size = 3.1) +
    scale_colour_manual(values = PALETTE, name = NULL) +
    coord_cartesian(xlim = c(0, surv$cap), ylim = c(0, 1)) +
    facet_wrap(~facet, ncol = 2) +
    labs(x = "Time (months)", y = surv$endpoints[[e]],
         title = surv$endpoints[[e]]) +
    base_theme +
    theme(strip.background = element_rect(fill = "grey92", colour = NA))
  save_fig(p_all, sprintf("figures/km_all_clusters_%s.png", e), 9, 16)
}

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

cat("\n=== agreement with published validation labels (Fig 4) ===\n")
w <- reshape(agreement, idvar = "method", timevar = "cluster", direction = "wide")
colnames(w) <- sub("^score\\.", "", colnames(w))
print(format(w, digits = 3), row.names = FALSE)

cat("\n=== features per cluster (Fig 5) ===\n")
fw <- do.call(rbind, lapply(names(METHODS), function(m) as.integer(models[[m]]$nz)))
dimnames(fw) <- list(names(METHODS), clusters)
print(fw)

cat("\n=== does per-cluster CV error predict per-cluster |z|? ===\n")
print(corr, row.names = FALSE)

cat("\n=== what DOES drive |z|: cluster or method? ===\n")
for (e in names(z_anova)) {
  cat("\n--", surv$endpoints[[e]], "\n")
  print(summary(z_anova[[e]]$fit))
  cat(sprintf("cluster explains %.1f%% of the variance in |z|, method %.1f%%.\n",
              z_anova[[e]]$pct_cluster, z_anova[[e]]$pct_method))
}
cat("\n|z| is a property of which cluster you are looking at, not of which\n")
cat("classifier produced the labels.\n")

cat("\n=== why the paper's k-sample panel p is not reproduced ===\n")
cat("Its five groups share patients. Rows-per-distinct-patient among the four\n")
cat("predicted groups, by cluster:\n")
print(round(setNames(dup_factor, clusters), 2))
cat("survdiff assumes independence; this biases p toward 1. Panels are\n")
cat("annotated with the disjoint two-group |z| against the truth instead.\n")

cat("\n=== originals untouched ===\n")
chk <- system2("shasum", c("-a", "256", "-c", "data/original_figures.sha"),
               stdout = TRUE, stderr = TRUE)
cat(paste(chk, collapse = "\n"), "\n")

cat("\nwrote five figures under figures/\n")