# Module 4: what kind of noise do counts have?  Poisson vs negative binomial.  Video: T6.
# Working directory: the "workshop" folder.
suppressPackageStartupMessages(library(DESeq2))
set.seed(42)

# ---- A. simulate: Poisson vs negative binomial with the same mean -----------------------------
mu <- 100
pois <- rpois(10000, mu)
alpha <- 0.1
nb <- rnbinom(10000, mu = mu, size = 1 / alpha)      # R's "size" is 1 / dispersion
c(mean_pois = mean(pois), var_pois = var(pois), mean_nb = mean(nb), var_nb = var(nb),
  expected_var_nb = mu + alpha * mu^2)
hist(nb, breaks = 60, col = "#c2185b80", border = NA, main = "same mean, different spread", xlab = "count")
hist(pois, breaks = 30, col = "#3f7cac80", border = NA, add = TRUE)
legend("topright", c("Poisson", "NB, alpha = 0.1"), fill = c("#3f7cac80", "#c2185b80"))
# EXERCISE A1: for alpha = 0.1, what is the coefficient of variation (sd / mean) when mu = 10, 100, 10000?
#              Compare with sqrt(alpha). When does the Poisson part stop mattering?

# ---- B. real data: variance between replicates vs mean ---------------------------------------------
meta <- read.delim("data/metadata.txt", row.names = 1)
count <- round(as.matrix(read.delim("data/count_gene.txt", row.names = 1, check.names = FALSE)[, rownames(meta)]))
dds <- DESeqDataSetFromMatrix(count, data.frame(row.names = rownames(meta), conds = factor(meta$condition)), ~conds)
dds <- estimateSizeFactors(dds)
nc <- counts(dds, normalized = TRUE)
grp <- meta$condition
m <- sapply(unique(grp), function(g) rowMeans(nc[, grp == g]))
v <- sapply(unique(grp), function(g) apply(nc[, grp == g], 1, var))
mv <- data.frame(mean = rowMeans(m), var = rowMeans(v))
mv <- mv[mv$mean > 1, ]
plot(mv$mean, pmax(mv$var, 0.1), log = "xy", pch = 16, cex = 0.3, col = "#3f7cac60",
     xlab = "mean of normalised counts", ylab = "variance between replicates")
abline(0, 1, lty = 2)                                   # variance = mean (Poisson)
cat("share of genes with variance > mean:", mean(mv$var > mv$mean), "\n")

# ---- C. DESeq2's dispersion estimates ----------------------------------------------------------------
dds <- estimateDispersions(dds)
plotDispEsts(dds)       # black: gene-wise, red: fitted trend, blue: final (shrunk)
# EXERCISE C1: what is the fitted dispersion for highly expressed genes (baseMean > 1000)?
mc <- mcols(dds)
median(mc$dispFit[mc$baseMean > 1000], na.rm = TRUE)
# EXERCISE C2: pick the gene with the largest gene-wise dispersion among genes with baseMean > 1000.
#              Plot its counts with plotCounts(). Why might its replicates disagree so much?