## metabric_data.R --assemble the analysis-ready METABRIC objects.
##
## Inputs (supplied locally, never committed -- see data/README.md):
##   data/metabric_raw/brca_metabric_clinical_data.tsv            cBioPortal export, 2509 samples
##   data/metabric_raw/IntClustFeatures_discoverydataset997.txt   Curtis 2012 Table S43
##   data/metabric_raw/IntClustFeatures_validationdataset995.txt  Curtis 2012 Table S43
##
## Output: data/metabric.rds -- a list with x/y/surv for discovery and
## validation, plus the feature names.
##

suppressPackageStartupMessages(library(survival))

CLINICAL   <- "data/metabric_raw/brca_metabric_clinical_data.tsv"
DISCOVERY  <- "data/metabric_raw/IntClustFeatures_discoverydataset997.txt"
VALIDATION <- "data/metabric_raw/IntClustFeatures_validationdataset995.txt"
OUT        <- "data/metabric.rds"

## Keep a validation sample whose IntClust label is known but whose vital status
## is not (MB-5130). The paper reports 983 validation samples, which is the count
## you get under the "has an IntClust label" rule; requiring a resolvable
## disease-specific event as well would give 982. Its DSS event is treated as
## censored below. Set FALSE to drop it instead.
KEEP_UNRESOLVED_VITAL_STATUS <- TRUE

# ---------------------------------------------------------------- clinical ---

read_clinical <- function(path) {
  clin <- read.delim(path, check.names = FALSE, stringsAsFactors = FALSE,
                     na.strings = c("NA", ""))
  stopifnot(nrow(clin) == 2509L)

  ## cBioPortal splits IntClust 4 by ER status ("4ER+", "4ER-"). Curtis 2012 and
  ## the paper we are replicating use 10 clusters, so fold the two back together.
  ic <- clin[["Integrative Cluster"]]
  stopifnot(setdiff(na.omit(unique(ic)),
                    c(as.character(1:10), "4ER+", "4ER-")) |> length() == 0L)
  ic <- sub("^4ER[+-]$", "4", ic)

  ## Disease-specific survival. The export has no DSS column; it is derived from
  ## the three-level vital status, which is the only disease-attributed outcome
  ## available. Death from other causes is a censoring event, NOT a failure --
  ## using overall survival instead would answer a different question.
  vital <- clin[["Patient's Vital Status"]]
  stopifnot(setdiff(na.omit(unique(vital)),
                    c("Living", "Died of Disease", "Died of Other Causes")) |>
              length() == 0L)
  dss_event <- ifelse(is.na(vital), 0L, as.integer(vital == "Died of Disease"))

  ## Overall Survival Status is carried alongside vital status because the two
  ## disagree on exactly one patient (MB-5130: OS status DECEASED, vital status
  ## NA). For an all-cause endpoint that patient is a known death, so downstream
  ## code must use os_status, not vital.
  os_status <- as.integer(sub(":.*$", "", clin[["Overall Survival Status"]]))

  data.frame(
    sample_id = clin[["Sample ID"]],
    cluster   = ic,
    time      = clin[["Overall Survival (Months)"]],
    event     = dss_event,
    vital     = vital,
    os_status = os_status,
    stringsAsFactors = FALSE
  )
}

# ------------------------------------------------------------------ genomic ---

## Table S43 ships features-in-rows, samples-in-columns, with the header row one
## field short (R's write.table row-name convention). Returns samples x features,
## which is what glmnet and pamr both want (pamr transposes back internally).
read_features <- function(path) {
  ids <- scan(path, what = "", nlines = 1, sep = "\t", quiet = TRUE)
  m <- as.matrix(read.delim(path, header = FALSE, skip = 1, row.names = 1,
                            check.names = FALSE))
  ## Duplicate feature names would already have errored in read.delim(row.names
  ## = 1), but duplicate SAMPLE ids would pass silently and corrupt every join
  ## downstream, so check explicitly.
  stopifnot(ncol(m) == length(ids), nrow(m) == 754L, !anyNA(m),
            !anyDuplicated(ids), !anyDuplicated(rownames(m)))
  colnames(m) <- ids
  t(m)
}

# -------------------------------------------------------------------- build ---

clinical <- read_clinical(CLINICAL)
stopifnot(sum(!is.na(clinical$cluster)) == 1980L,
          nlevels(factor(clinical$cluster)) == 10L)

x_disc <- read_features(DISCOVERY)
x_val  <- read_features(VALIDATION)
stopifnot(nrow(x_disc) == 997L, nrow(x_val) == 995L,
          identical(colnames(x_disc), colnames(x_val)),
          length(intersect(rownames(x_disc), rownames(x_val))) == 0L)

## 39 copy-number features, 715 expression features.
features <- colnames(x_disc)
stopifnot(sum(startsWith(features, "CN-")) == 39L,
          sum(startsWith(features, "Exp-")) == 715L)

assemble <- function(x, clinical) {
  keep_ids <- clinical$sample_id[!is.na(clinical$cluster)]
  if (!KEEP_UNRESOLVED_VITAL_STATUS) {
    keep_ids <- intersect(keep_ids, clinical$sample_id[!is.na(clinical$vital)])
  }
  ids <- rownames(x)[rownames(x) %in% keep_ids]
  cl  <- clinical[match(ids, clinical$sample_id), ]
  ## One patient (MB-0284) has time exactly 0, censored -- harmless for survfit
  ## and survdiff, but assert non-negativity so a future refactor cannot slip a
  ## negative through.
  stopifnot(identical(cl$sample_id, ids), !anyNA(cl$time), all(cl$time >= 0))
  list(
    x    = x[ids, , drop = FALSE],
    y    = factor(cl$cluster, levels = as.character(1:10)),
    surv = data.frame(sample_id = ids, time = cl$time, event = cl$event,
                      stringsAsFactors = FALSE)
  )
}

discovery  <- assemble(x_disc, clinical)
validation <- assemble(x_val,  clinical)

## The paper's headline counts: 997 discovery, 983 validation, 1980 total.
stopifnot(nrow(discovery$x) == 997L, nrow(validation$x) == 983L,
          nrow(discovery$x) + nrow(validation$x) == 1980L,
          !anyNA(discovery$y), !anyNA(validation$y))

dir.create("data", showWarnings = FALSE)
saveRDS(list(discovery = discovery, validation = validation,
             features = features, clinical = clinical), OUT)

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

cat("discovery :", nrow(discovery$x), "x", ncol(discovery$x), "\n")
cat("validation:", nrow(validation$x), "x", ncol(validation$x), "\n")
cat("features  :", length(features), "=",
    sum(startsWith(features, "CN-")), "CN +",
    sum(startsWith(features, "Exp-")), "Exp\n\n")

cat("cluster sizes\n")
print(rbind(discovery  = table(discovery$y),
            validation = table(validation$y)))

cat("\ndisease-specific events (deaths from disease)\n")
print(rbind(discovery  = c(events = sum(discovery$surv$event),
                           n = nrow(discovery$surv)),
            validation = c(events = sum(validation$surv$event),
                           n = nrow(validation$surv))))

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