## metabric_survival.R --do the patients a model routes into cluster k die like real
## cluster-k patients?
##
## This is the only comparison in the paper where all four methods are on equal
## footing. CV error and label agreement both grade against NSC-derived labels,
## which pamr reproduces by construction; the survival data was never seen by any
## classifier during training.
##
## For each cluster k and each method m:
##   reference  = DISCOVERY patients whose true label is k
##   comparison = VALIDATION patients that m assigned to k
## and a two-sample log-rank test between them. For two groups the log-rank
## statistic is z^2 = chisq, so |z| = sqrt(chisq).
##
## The endpoint is overall survival censored at 180 months, for the reason set
## out below.
##
## Input:  data/metabric.rds, data/models.rds
## Output: data/survival.rds and four figures under figures/

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

## ------------------------------------------------------------- endpoints ---
##
## The paper labels its KM axis "Disease-specific Survival". Its numbers are not
## disease-specific survival.
##
## The pamr row of the log-rank table is a pure data quantity -- pamr's
## predictions ARE the published validation labels, so that row compares
## discovery cluster k against validation cluster k with no model in it. It
## therefore identifies the endpoint, and a scan over event definitions and
## censoring times identifies it uniquely: matching all ten of the paper's
## per-cluster values (not just their median, which dozens of combinations
## satisfy) leaves exactly one cell, in a censoring window one month wide.
##
##     endpoint                pamr median |z|    max dev vs paper's ten
##     DSS, full follow-up               1.173                     1.474
##     DSS, capped at 180 mo             1.175                     1.540
##     OS,  full follow-up               0.959                     0.678
##     OS,  capped at 179 mo             0.581                     0.086
##     OS,  capped at 180 mo             0.581                     0.006  <- unique
##     OS,  capped at 181 mo             0.625                     0.101
##
## So the published analysis used OVERALL survival censored at 180 months: death
## from any cause is an event, and 496 of 1980 patients (25%) who died of other
## causes are events here but censorings under the endpoint the paper describes.
##
## 180 months is also the window the paper plots.
ENDPOINTS <- c(os180 = "Overall survival (any death), censored at 180 months")
CAP       <- 180

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

clusters <- models$clusters
disc     <- dat$discovery
val      <- dat$validation
clin     <- dat$clinical

## Overall survival is derived here rather than in metabric_data.R so this script
## stays self-contained. Read from Overall Survival Status, NOT vital status:
## MB-5130 has vital status NA but OS status 1:DECEASED, so the vital-status
## route scores a known death as alive.
overall_event <- function(ids) {
  s <- clin$os_status[match(ids, clin$sample_id)]
  ifelse(is.na(s), 0L, as.integer(s))
}

## Administrative censoring at CAP.
apply_endpoint <- function(surv) {
  event <- overall_event(surv$sample_id)
  time  <- surv$time
  event <- ifelse(time > CAP, 0L, event)
  data.frame(sample_id = surv$sample_id, time = pmin(time, CAP), event = event)
}

METHODS <- c(ungrouped = "Ungrouped glmnet",
             grouped   = "Grouped glmnet",
             ovr       = "OVR glmnet",
             pamr      = "pamr")

PALETTE <- c("True (Discovery)"  = "#35B7CD",
             "pamr"              = "#5FBB46",
             "Ungrouped glmnet"  = "#D48FC8",
             "Grouped glmnet"    = "#C9256A",
             "OVR glmnet"        = "#26356E")

## pamr's predictions must equal the published validation labels. This holds at
## SEED = 42, where CV selects the threshold 0.582 that the paper reports; the
## choice is knife-edge and another fold draw can flip it. Guard, not argument.
stopifnot(identical(as.character(models$pred$pamr), as.character(val$y)))

# ------------------------------------------------------------- log-rank |z| ---

## TIE HANDLING. survdiff has no `ties` argument -- only `rho` and `timefix` --
## so there is no Breslow/Efron default to get wrong here, and the paper states
## it used survdiff. But the log-rank test IS asymptotically the score test from
## a Cox model with BRESLOW ties, so the analysis is implicitly Breslow-like and
## it is fair to ask whether Efron would move anything. Checked directly, all 40
## comparisons:
##
##   max |difference| from survdiff
##     Cox score test, Breslow           4.1e-04   (i.e. same test)
##     Cox score test, Efron             1.4e-03
##
## Medians move in the fourth decimal (OVR 0.38026 -> 0.37965) and every ranking
## is unchanged. The endpoint identification is unchanged too: max deviation
## from the paper's ten pamr values is 0.0060 under survdiff and 0.0050 under
## Efron. Re-checked at seed 42 by tmp/efron.R.
##
## The reason is that ties are sparse: 965 events at 869 distinct times, largest
## tie block 3, and only ~19% of events share a time with any other. Breslow and
## Efron differ only in how they handle simultaneous events, so with blocks this
## small there is almost nothing for them to disagree about.
logrank_z <- function(ref, cmp) {
  if (nrow(cmp) == 0L || sum(cmp$event) + sum(ref$event) == 0L) {
    return(c(z = NA_real_, n = nrow(cmp), p = NA_real_))
  }
  d <- rbind(
    data.frame(time = ref$time, event = ref$event, grp = "reference"),
    data.frame(time = cmp$time, event = cmp$event, grp = "comparison")
  )
  sd <- survdiff(Surv(time, event) ~ grp, data = d)
  c(z = sqrt(sd$chisq), n = nrow(cmp),
    p = pchisq(sd$chisq, df = 1, lower.tail = FALSE))
}

## survfit0() prepends the (0, 1) origin so the step starts at 1.
km_df <- function(rows, label) {
  if (nrow(rows) == 0L) return(NULL)
  f <- survfit0(survfit(Surv(time, event) ~ 1, data = rows))
  data.frame(time = f$time, surv = f$surv, n.censor = f$n.censor, label = label)
}

# ----------------------------------------------- everything, one endpoint ---
#
# WHAT THE pamr ROW IS, AND IS NOT. It carries no classifier error, so it is the
# reference the other three are measured against, not a fourth competitor:
# nothing can beat the ground truth at recovering the ground truth.
#
# WHETHER IT IS A DRIFT FLOOR OR A NOISE FLOOR is settled by permuting cohort
# membership within cluster (permutation_floor): observed 0.581 against a null
# mean of 0.681, p = 0.656, so under this endpoint the pamr row is noise.
#
# And "excess near zero" carries almost no information, for two reasons that the
# report quantifies: (a) under H0, |z| = sqrt(chisq_1) is half-normal with median
# 0.674, NOT zero, so near-zero excess is what noise predicts; (b) the compared
# groups largely ARE the reference group (Jaccard 0.55-0.89), so each excess is a
# difference of two strongly correlated statistics and roughly half the clusters
# fall below the floor regardless of signal. A sign test on that was previously
# reported here and was near-vacuous; it has been removed.

PERM_REPS <- 500
HALF_NORMAL_MEDIAN <- qnorm(0.75)

analyze <- function(endpoint) {
  d <- apply_endpoint(disc$surv)
  v <- apply_endpoint(val$surv)
  dy <- as.character(disc$y)

  stats <- do.call(rbind, lapply(names(METHODS), function(m) {
    do.call(rbind, lapply(clusters, function(k) {
      r <- logrank_z(d[dy == k, , drop = FALSE],
                     v[models$pred[[m]] == k, , drop = FALSE])
      data.frame(method = METHODS[[m]], key = m, cluster = as.integer(k),
                 z = r[["z"]], n = r[["n"]], p = r[["p"]])
    }))
  }))
  rownames(stats) <- NULL
  medians <- tapply(stats$z, stats$method, median, na.rm = TRUE)

  curves <- do.call(rbind, lapply(clusters, function(k) {
    parts <- c(
      list(km_df(d[dy == k, , drop = FALSE], "True (Discovery)")),
      lapply(names(METHODS), function(m) {
        km_df(v[models$pred[[m]] == k, , drop = FALSE], METHODS[[m]])
      })
    )
    out <- do.call(rbind, Filter(Negate(is.null), parts))
    out$cluster <- as.integer(k)
    out
  }))
  curves$label <- factor(curves$label, levels = names(PALETTE))

  ## Excess over the pamr reference.
  pz <- stats$z[stats$key == "pamr"]
  names(pz) <- stats$cluster[stats$key == "pamr"]
  baseline <- do.call(rbind, lapply(c("ungrouped", "grouped", "ovr"), function(m) {
    b <- stats[stats$key == m, ]
    data.frame(method = METHODS[[m]], key = m, cluster = b$cluster, z = b$z,
               floor = as.numeric(pz[as.character(b$cluster)]),
               excess = b$z - as.numeric(pz[as.character(b$cluster)]), n = b$n)
  }))
  rownames(baseline) <- NULL

  baseline_summary <- do.call(rbind, lapply(
    split(baseline, factor(baseline$key, levels = c("ungrouped", "grouped", "ovr"))),
    function(b) data.frame(
      method = b$method[1],
      below_floor = sum(b$excess < 0), n_clusters = nrow(b),
      median_excess = median(b$excess), worst_excess = max(b$excess),
      jaccard_med = median(vapply(clusters, function(k) {
        a <- which(models$pred[[b$key[1]]] == k); z <- which(val$y == k)
        length(intersect(a, z)) / length(union(a, z))
      }, numeric(1))),
      cor_with_floor = cor(b$z, b$floor),
      ci_lo = t.test(b$excess)$conf.int[1],
      ci_hi = t.test(b$excess)$conf.int[2])))
  rownames(baseline_summary) <- NULL

  null_calibration <- do.call(rbind, lapply(split(stats, stats$key), function(b) {
    ks <- suppressWarnings(ks.test(b$z, function(q) 2 * pnorm(q) - 1))
    data.frame(method = b$method[1], median_z = median(b$z),
               ks_D = unname(ks$statistic), ks_p = ks$p.value)
  }))
  rownames(null_calibration) <- NULL

  set.seed(1)
  obs <- median(stats$z[stats$key == "pamr"])
  null <- replicate(PERM_REPS, median(vapply(clusters, function(k) {
    a <- d[dy == k, ]; b <- v[as.character(val$y) == k, ]
    ab <- rbind(a, b); g <- sample(rep(1:2, c(nrow(a), nrow(b))))
    sqrt(survdiff(Surv(time, event) ~ g, data = ab)$chisq)
  }, numeric(1))))

  list(endpoint = endpoint, stats = stats, medians = medians, curves = curves,
       baseline = baseline, baseline_summary = baseline_summary,
       null_calibration = null_calibration,
       friedman_p = friedman.test(do.call(cbind, split(stats$z, stats$key)))$p.value,
       permutation_floor = list(observed = obs, null_mean = mean(null),
                                lo = unname(quantile(null, 0.025)),
                                hi = unname(quantile(null, 0.975)),
                                p = mean(null >= obs)),
       events_per_cluster = vapply(clusters, function(k) {
         sum(d$event[dy == k]) + sum(v$event[models$pred$pamr == k])
       }, numeric(1)),
       surv_disc = d, surv_val = v)
}

res <- lapply(names(ENDPOINTS), analyze)
names(res) <- names(ENDPOINTS)

# ------------------------------------------------------------------- figures ---

dir.create("figures", showWarnings = FALSE)

km_panel <- function(a, ks, shown, file, width, height) {
  d <- a$curves[a$curves$cluster %in% ks & a$curves$label %in% shown, ]
  d$facet <- factor(paste("Cluster", d$cluster), levels = paste("Cluster", ks))
  p <- ggplot(d, aes(time, surv, colour = label)) +
    geom_step(linewidth = 0.6) +
    geom_point(data = d[d$n.censor > 0, ], shape = 3, size = 1.1,
               show.legend = FALSE) +
    scale_colour_manual(values = PALETTE[shown], limits = shown, name = NULL) +
    coord_cartesian(xlim = c(0, CAP), ylim = c(0, 1)) +
    facet_wrap(~facet, ncol = 2) +
    labs(x = "Time (months)", y = ENDPOINTS[[a$endpoint]],
         title = paste0("Clusters ", min(ks), "-", max(ks), ": ",
                        ENDPOINTS[[a$endpoint]])) +
    theme_bw(base_size = 11) +
    theme(legend.position = "bottom", panel.grid.minor = element_blank(),
          strip.background = element_rect(fill = "grey92", colour = NA))
  ggsave(file, p, width = width, height = height, dpi = 200)
}

PANEL_CURVES <- c("True (Discovery)", "pamr", "Ungrouped glmnet", "OVR glmnet")

for (e in names(ENDPOINTS)) {
  a <- res[[e]]
  km_panel(a, 1:6,  PANEL_CURVES, sprintf("figures/km_clusters_1_6_%s.png", e), 9, 10)
  km_panel(a, 7:10, PANEL_CURVES, sprintf("figures/km_clusters_7_10_%s.png", e), 9, 7)

  lab <- sprintf("%s (median |z| = %.3f)", names(a$medians), a$medians)
  names(lab) <- names(a$medians)
  s <- a$stats; s$method_lab <- lab[s$method]
  p_lr <- ggplot(s, aes(cluster, z, colour = method_lab)) +
    geom_hline(yintercept = HALF_NORMAL_MEDIAN, linetype = "dashed",
               colour = "grey45") +
    annotate("text", x = 10.4, y = HALF_NORMAL_MEDIAN, hjust = 1, vjust = -0.5,
             size = 3, colour = "grey30", label = "median under H0 (0.674)") +
    geom_line(linewidth = 0.7) + geom_point(size = 2) +
    scale_x_continuous(breaks = 1:10) +
    scale_colour_manual(values = setNames(PALETTE[names(a$medians)], lab),
                        name = "Method") +
    labs(x = "Cluster", y = "Log-rank Test Statistic (|z|)",
         title = ENDPOINTS[[e]],
         subtitle = sprintf(paste("All four are indistinguishable from the",
                   "half-normal null (KS p %.2f-%.2f);\nFriedman across methods",
                   "p = %.2f. The ordering is a ranking of noise."),
                   min(a$null_calibration$ks_p), max(a$null_calibration$ks_p),
                   a$friedman_p)) +
    theme_bw(base_size = 11) + theme(panel.grid.minor = element_blank())
  ggsave(sprintf("figures/logrank_%s.png", e), p_lr, width = 9, height = 6, dpi = 200)

  p_bl <- ggplot(a$baseline, aes(cluster, excess, colour = method)) +
    geom_hline(yintercept = 0, linewidth = 0.7, colour = "#5FBB46") +
    geom_line(linewidth = 0.7) + geom_point(size = 2) +
    scale_x_continuous(breaks = 1:10) +
    scale_colour_manual(values = PALETTE, name = "Method") +
    labs(x = "Cluster",
         y = expression("Excess  " * group("|", z, "|") * "  over pamr reference"),
         title = paste("Relative to the pamr reference --", ENDPOINTS[[e]]),
         subtitle = paste("Zero = discovery cluster k vs TRUE validation",
                          "cluster k (no classifier error). Under H0 that\n",
                          "reference is itself half-normal, so excess near zero",
                          "is what noise predicts, not agreement.\nThe predicted",
                          "groups are 55-89% identical to the reference group.")) +
    theme_bw(base_size = 11) + theme(panel.grid.minor = element_blank())
  ggsave(sprintf("figures/logrank_vs_baseline_%s.png", e), p_bl,
         width = 9, height = 6, dpi = 200)
}

saveRDS(list(res = res, endpoints = ENDPOINTS, cap = CAP,
             palette = PALETTE, methods = METHODS,
             half_normal_median = HALF_NORMAL_MEDIAN),
        "data/survival.rds")

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

nm <- unname(METHODS)

cat("\n=== median |z| ===\n")
print(data.frame(method = nm,
                 median_z = round(as.numeric(res$os180$medians[nm]), 3),
                 rank = rank(as.numeric(res$os180$medians[nm]))),
      row.names = FALSE)

for (e in names(ENDPOINTS)) {
  a <- res[[e]]
  cat("\n\n######## ", e, " -- ", ENDPOINTS[[e]], "\n", sep = "")

  cat("\n-- log-rank |z| per cluster\n")
  w <- reshape(a$stats[, c("method", "cluster", "z")], idvar = "method",
               timevar = "cluster", direction = "wide")
  colnames(w) <- sub("^z\\.", "", colnames(w))
  print(format(w, digits = 3), row.names = FALSE)

  cat("\n-- excess over the pamr reference\n")
  bs <- a$baseline_summary
  num <- vapply(bs, is.numeric, logical(1)); bs[num] <- lapply(bs[num], round, 3)
  print(bs, row.names = FALSE)
  cat("jaccard_med: how much of each predicted group IS the true cluster.\n")
  cat("ci_lo/ci_hi span ~", round(mean(bs$ci_hi - bs$ci_lo), 2),
      "|z| units around point estimates of ~0.0-0.1: this cannot\n")
  cat("exclude an excess of +/-0.35. It is not evidence of equality.\n")

  cat("\n-- distinguishable from noise?  (H0 median |z| =",
      round(HALF_NORMAL_MEDIAN, 3), ")\n")
  nc <- a$null_calibration; nc[-1] <- lapply(nc[-1], round, 3)
  print(nc, row.names = FALSE)
  cat(sprintf("Friedman across methods p = %.3f;  %d of %d tests reject at 0.05 (expect %.0f)\n",
              a$friedman_p, sum(a$stats$p < 0.05, na.rm = TRUE), nrow(a$stats),
              0.05 * nrow(a$stats)))

  pf <- a$permutation_floor
  cat(sprintf("\n-- is there any cohort difference to measure?\nobserved %.3f, permutation null mean %.3f, 95%% [%.3f, %.3f], p = %.3f\n",
              pf$observed, pf$null_mean, pf$lo, pf$hi, pf$p))

  cat("\n-- events per comparison (|z| is not comparable across clusters)\n")
  print(setNames(as.integer(a$events_per_cluster), clusters))
}

cat("\n\nwrote data/survival.rds and four figures under figures/\n")