# Module 3: exploring the samples with PCA.  Video: T4.
# Working directory: the "workshop" folder.
suppressPackageStartupMessages({ library(DESeq2); library(ggplot2) })
dir.create("results", showWarnings = FALSE)

meta <- read.delim("data/metadata.txt", row.names = 1)
tpm <- read.delim("data/tpm_gene.txt", row.names = 1, check.names = FALSE)[, rownames(meta)]
count <- round(as.matrix(read.delim("data/count_gene.txt", row.names = 1, check.names = FALSE)[, rownames(meta)]))
meta$MI <- ifelse(grepl("^MI", meta$condition), "MI", "SHAM")
meta$day <- sub(".*_", "", meta$condition)

plot_pca <- function(p, title) {
  ve <- round(100 * p$sdev^2 / sum(p$sdev^2), 1)
  d <- data.frame(p$x[, 1:2], meta)
  ggplot(d, aes(PC1, PC2, colour = MI, shape = day)) + geom_point(size = 4) +
    labs(title = title, x = paste0("PC1 (", ve[1], "%)"), y = paste0("PC2 (", ve[2], "%)")) +
    scale_colour_manual(values = c(MI = "#c2185b", SHAM = "#3f7cac")) + theme_bw()
}

# ---- A. the original workshop: PCA on raw TPM --------------------------------------------------
p_raw <- prcomp(t(tpm), scale. = FALSE)
print(plot_pca(p_raw, "raw TPM"))
ggsave("results/pca_raw_tpm.png", width = 5, height = 4)
# which genes carry the most weight on PC1?
head(sort(abs(p_raw$rotation[, 1]), decreasing = TRUE), 10)
# EXERCISE A1: how much of the total variance comes from the 10 genes with the highest variance?
v <- apply(tpm, 1, var)
sum(sort(v, decreasing = TRUE)[1:10]) / sum(v)

# ---- B. log scale -------------------------------------------------------------------------------
p_log <- prcomp(t(log2(tpm[rowSums(tpm) > 0, ] + 1)))
print(plot_pca(p_log, "log2(TPM + 1)"))
ggsave("results/pca_log_tpm.png", width = 5, height = 4)

# ---- C. DESeq2's variance stabilising transformation (vst) ---------------------------------------
dds <- DESeqDataSetFromMatrix(count, data.frame(row.names = rownames(meta), conds = factor(meta$condition)), ~conds)
vsd <- vst(dds, blind = TRUE)
top <- order(apply(assay(vsd), 1, var), decreasing = TRUE)[1:500]
p_vst <- prcomp(t(assay(vsd)[top, ]))
print(plot_pca(p_vst, "vst, top 500 genes"))
ggsave("results/pca_vst.png", width = 5, height = 4)
# DESeq2 has a shortcut that does the same thing:
plotPCA(vsd, intgroup = "conds")

# ---- D. scree plot: how much variance does each PC explain? ----------------------------------------
ve <- 100 * p_vst$sdev^2 / sum(p_vst$sdev^2)
barplot(ve, names.arg = paste0("PC", seq_along(ve)), ylab = "% of variance", main = "vst, top 500 genes")
round(ve, 1)

# ---- E. sample-to-sample distances (another view of the same thing) ---------------------------------
d <- as.matrix(dist(t(assay(vsd))))
pheatmap::pheatmap(d, annotation_col = meta[, c("MI", "day")], main = "Euclidean distance, vst")