## Download TCGA-ESCA RNA-seq counts, then compare what several normalization
## methods do to the per-sample distributions. Figures are ComplexHeatmap
## density heatmaps, one column per sample.
##
## Usage: Rscript normalize_tcga.R <work_dir> <figure_dir>
## The work dir holds the GDC download cache; nothing from it is committed.

args     <- commandArgs(trailingOnly = TRUE)
WORK     <- if (length(args) >= 1) args[1] else "."
FIGDIR   <- if (length(args) >= 2) args[2] else "."
PROJECT  <- "TCGA-ESCA"
CACHE    <- file.path(WORK, "esca_se.rds")

dir.create(FIGDIR, showWarnings = FALSE, recursive = TRUE)
setwd(WORK)

suppressPackageStartupMessages({
  library(TCGAbiolinks); library(SummarizedExperiment); library(edgeR)
  library(DESeq2); library(limma); library(ComplexHeatmap); library(circlize)
})

## ---- 1. data ---------------------------------------------------------------
if (file.exists(CACHE)) {
  se <- readRDS(CACHE)
} else {
  q <- GDCquery(project = PROJECT, data.category = "Transcriptome Profiling",
                data.type = "Gene Expression Quantification",
                workflow.type = "STAR - Counts")
  GDCdownload(q)
  se <- GDCprepare(q)
  saveRDS(se, CACHE)
}

cts  <- assay(se, "unstranded")
info <- as.data.frame(colData(se))
grp  <- ifelse(info$sample_type == "Solid Tissue Normal", "Normal", "Tumor")

## protein-coding only, then drop genes with too little signal to inform anything
pc   <- rowData(se)$gene_type == "protein_coding"
cts  <- cts[pc, ]
dge  <- DGEList(counts = cts, group = factor(grp))
keep <- filterByExpr(dge, group = factor(grp))
dge  <- dge[keep, , keep.lib.sizes = FALSE]
cts  <- dge$counts

libsize <- colSums(cts)
cat(sprintf("genes %d  samples %d  (%d tumor / %d normal)\n",
            nrow(cts), ncol(cts), sum(grp == "Tumor"), sum(grp == "Normal")))
cat(sprintf("sample types: %s\n",
            paste(sprintf("%s=%d", names(table(info$sample_type)),
                          as.integer(table(info$sample_type))), collapse = ", ")))
cat(sprintf("library size: min %.1fM  median %.1fM  max %.1fM  (%.1fx spread)\n",
            min(libsize)/1e6, median(libsize)/1e6, max(libsize)/1e6,
            max(libsize)/min(libsize)))

## ---- 2. normalizations -----------------------------------------------------
## Each returns a log2 matrix so the panels are directly comparable.
L <- list()
L[["Raw counts"]] <- log2(cts + 1)

## library size only
L[["CPM"]] <- cpm(dge, log = TRUE, prior.count = 1)

## upper quartile: scale each sample by its own 75th percentile
d_uq <- calcNormFactors(dge, method = "upperquartile")
L[["Upper quartile"]] <- cpm(d_uq, log = TRUE, prior.count = 1)

## TMM: trimmed mean of M-values against a reference sample
d_tmm <- calcNormFactors(dge, method = "TMM")
L[["TMM"]] <- cpm(d_tmm, log = TRUE, prior.count = 1)

## DESeq2 median of ratios
dds <- DESeqDataSetFromMatrix(cts, data.frame(grp = factor(grp)), ~ grp)
dds <- estimateSizeFactors(dds)
L[["Median of ratios"]] <- log2(counts(dds, normalized = TRUE) + 1)

## quantile: force every sample onto one common distribution
L[["Quantile"]] <- normalizeQuantiles(log2(cts + 1))

## ---- 3. how far apart are the samples, per method? -------------------------
spread <- do.call(rbind, lapply(names(L), function(n) {
  m   <- L[[n]]
  med <- apply(m, 2, median)
  q1  <- apply(m, 2, quantile, 0.25)
  q3  <- apply(m, 2, quantile, 0.75)
  data.frame(method = n,
             median_range = diff(range(med)),
             median_sd    = sd(med),
             iqr_range    = diff(range(q3 - q1)),
             q3_range     = diff(range(q3)),
             cor_med_lib  = cor(med, libsize))
}))
print(spread, row.names = FALSE, digits = 3)
write.csv(spread, file.path(FIGDIR, "normalization_spread.csv"), row.names = FALSE)

nf <- data.frame(sample = colnames(cts), group = grp, lib_size = libsize,
                 uq_factor  = d_uq$samples$norm.factors,
                 tmm_factor = d_tmm$samples$norm.factors,
                 deseq_size_factor = sizeFactors(dds))
write.csv(nf, file.path(FIGDIR, "normalization_factors.csv"), row.names = FALSE)
cat("\nnormalization factor ranges:\n")
cat(sprintf("  upper quartile %.2f to %.2f\n", min(nf$uq_factor), max(nf$uq_factor)))
cat(sprintf("  TMM            %.2f to %.2f\n", min(nf$tmm_factor), max(nf$tmm_factor)))
cat(sprintf("  DESeq2 size    %.2f to %.2f\n", min(nf$deseq_size_factor), max(nf$deseq_size_factor)))
cat(sprintf("  correlation TMM vs UQ: %.3f\n", cor(nf$tmm_factor, nf$uq_factor)))

## how concentrated is each library in a handful of genes?
topn  <- 100L
share <- apply(cts, 2, function(x) sum(sort(x, decreasing = TRUE)[seq_len(topn)]) / sum(x))
cat(sprintf("\ntop %d genes hold %.1f%% to %.1f%% of each library (median %.1f%%)\n",
            topn, 100*min(share), 100*max(share), 100*median(share)))

## ---- 4. figures ------------------------------------------------------------
ylim <- c(-2, 16)
ann  <- HeatmapAnnotation(
  Type = grp,
  `Library size (M)` = anno_barplot(libsize / 1e6, gp = gpar(fill = "#9aa7b8", col = NA),
                                    height = unit(11, "mm")),
  col = list(Type = c(Tumor = "#3E6E8E", Normal = "#C9A227")),
  annotation_name_gp = gpar(fontsize = 8), show_legend = TRUE)

panel <- function(mat, title, file, annotate = TRUE, w = 1500, h = 950) {
  png(file.path(FIGDIR, file), width = w, height = h, res = 150)
  ht <- densityHeatmap(mat, ylim = ylim, ylab = "log2 expression",
                       title = title, title_gp = gpar(fontsize = 13, fontface = "bold"),
                       cluster_columns = FALSE, show_column_names = FALSE,
                       top_annotation = if (annotate) ann else NULL)
  draw(ht, merge_legend = TRUE)
  dev.off()
  cat("wrote", file, "\n")
}

panel(L[["Raw counts"]], "Raw counts, log2(count + 1)", "01-raw.png")
files <- c("CPM" = "02-cpm.png", "Upper quartile" = "03-upperquartile.png",
           "TMM" = "04-tmm.png", "Median of ratios" = "05-medianratios.png",
           "Quantile" = "06-quantile.png")
for (n in names(files)) panel(L[[n]], n, files[[n]])

## boxplot of per-sample medians across methods
png(file.path(FIGDIR, "07-medians.png"), width = 1500, height = 800, res = 150)
op <- par(mar = c(8, 4.5, 2, 1))
meds <- sapply(L, function(m) apply(m, 2, median))
boxplot(meds, ylab = "per-sample median log2 expression", las = 2,
        col = "#dbe3ea", border = "#4a5b6b", outline = TRUE, pch = 16, cex = 0.6)
title("Spread of per-sample medians (lower is more consistent)", cex.main = 1)
par(op); dev.off()
cat("wrote 07-medians.png\n")

## does the per-sample median track sequencing depth?
png(file.path(FIGDIR, "08-median-vs-libsize.png"), width = 1500, height = 700, res = 150)
op <- par(mfrow = c(1, 2), mar = c(4.4, 4.6, 3.2, 1))
cols <- ifelse(grp == "Tumor", "#3E6E8E", "#C9A227")
## CPM sits at a different absolute level to raw counts, so a shared y range would
## make its scatter look tighter than it is. Give both panels the same y *span*
## centred on their own median instead, which makes the spread comparable by eye.
span <- 1.15 * max(diff(range(apply(L[["Raw counts"]], 2, median))),
                   diff(range(apply(L[["CPM"]], 2, median))))
for (n in c("Raw counts", "CPM")) {
  md <- apply(L[[n]], 2, median); x <- libsize / 1e6
  yl <- median(md) + c(-1, 1) * span / 2
  plot(x, md, pch = 21, bg = cols, col = "white", cex = 1.15, ylim = yl,
       xlab = "library size (millions of reads)",
       ylab = "per-sample median log2 expression",
       main = sprintf("%s   (r = %+.2f)", n, cor(md, libsize)), cex.main = 1.05)
  abline(lm(md ~ x), col = "#b02020", lwd = 2)
}
par(op); dev.off()
cat("wrote 08-median-vs-libsize.png\n")

## how much of each library sits in its most-expressed genes
png(file.path(FIGDIR, "09-library-concentration.png"), width = 1400, height = 760, res = 150)
op <- par(mar = c(4.4, 4.6, 3.2, 1))
ks <- unique(round(exp(seq(log(1), log(nrow(cts)), length.out = 140))))
cum <- apply(cts, 2, function(x) cumsum(sort(x, decreasing = TRUE))[ks] / sum(x))
plot(NA, xlim = range(ks), ylim = c(0, 1), log = "x",
     xlab = "number of most-expressed genes", ylab = "cumulative share of the library",
     main = "A library total is dominated by a few genes", cex.main = 1.05)
for (j in seq_len(ncol(cum))) lines(ks, cum[, j], col = rgb(0.24, 0.35, 0.47, 0.16), lwd = 1)
lines(ks, apply(cum, 1, median), col = "#b02020", lwd = 2.5)
abline(v = topn, lty = 2, col = "gray35")
text(topn, 0.06, sprintf("  top %d genes: %.0f%% to %.0f%%", topn,
                         100*min(share), 100*max(share)), adj = 0, cex = 0.85)
legend("bottomright", c("one sample", "median across samples"), bty = "n", cex = 0.85,
       col = c(rgb(0.24,0.35,0.47,0.7), "#b02020"), lwd = c(1, 2.5))
par(op); dev.off()
cat("wrote 09-library-concentration.png\n")

saveRDS(list(spread = spread, nf = nf, libsize = libsize, grp = grp, share = share,
             n_genes = nrow(cts), n_samples = ncol(cts),
             sample_types = table(info$sample_type)),
        file.path(FIGDIR, "summary.rds"))
cat("done\n")