# load the libraries
library(tidyverse)
library(stringr)
library(edgeR)
library(limma)
library(psych)
library(scales)

# ================== TMM normalization from DGEList object =====================
apply_tmm_factors <- function(y, color = NULL, plot = TRUE) {
    # computes the tmm normalized data from the DGEList object
        # y - DGEList object
        # returns a dataframe with normalized intensities
    
    # compute and print "Sample loading" normalization factors
    lib_facs <- mean(y$samples$lib.size) / y$samples$lib.size
    cat("\nLibrary size factors:\n", 
        sprintf("%-5s -> %f\n", colnames(y$counts), lib_facs))
    
    # compute and print TMM normalization factors
    tmm_facs <- 1/y$samples$norm.factors
    cat("\nTrimmed mean of M-values (TMM) factors:\n", 
        sprintf("%-5s -> %f\n", colnames(y$counts), tmm_facs))
    
    # compute and print the final correction factors
    norm_facs <- lib_facs * tmm_facs
    cat("\nCombined (lib size and TMM) normalization factors:\n", 
        sprintf("%-5s -> %f\n", colnames(y$counts), norm_facs))

    # compute the normalized data as a new data frame
    tmt_tmm <- as.data.frame(sweep(y$counts, 2, norm_facs, FUN = "*"))
    colnames(tmt_tmm) <- str_c(colnames(y$counts), "_tmm")
    
    # visualize results and return data frame
    if(plot == TRUE) {
        boxplot(log10(tmt_tmm), col = color, notch = TRUE, main = "TMM Normalized data")
    }
    tmt_tmm
}

# ============== CV function ===================================================
CV <- function(df) {
    # Computes CVs of data frame rows
        # df - data frame, 
        # returns vector of CVs (%)
    ave <- rowMeans(df)    # compute averages
    sd <- apply(df, 1, sd) # compute standard deviations
    cv <- 100 * sd / ave   # compute CVs in percent (last thing gets returned)
}

# ================= reformat edgeR test results ================================
collect_results <- function(df, tt, x, xlab, y, ylab) {
    # Computes new columns and extracts some columns to make results frame
        # df - data in data.frame
        # tt - top tags table from edgeR test
        # x - columns for first condition
        # xlab - label for x
        # y - columns for second condition
        # ylab - label for y
        # returns a new dataframe
    
    # condition average vectors
    ave_x <- rowMeans(df[x])
    ave_y <- rowMeans(df[y])
    
    # FC, direction, candidates
    fc <- ifelse(ave_y > ave_x, (ave_y / ave_x), (-1 * ave_x / ave_y))
    direction <- ifelse(ave_y > ave_x, "up", "down")
    candidate <- cut(tt$FDR, breaks = c(-Inf, 0.01, 0.05, 0.10, 1.0), 
                     labels = c("high", "med", "low", "no"))
    
    # make data frame
    temp <- cbind(df[c(x, y)], data.frame(logFC = tt$logFC, FC = fc, 
                                          PValue = tt$PValue, FDR = tt$FDR, 
                                          ave_x = ave_x, ave_y = ave_y, 
                                          direction = direction, candidate = candidate, 
                                          Acc = tt$genes)) 
    
    # fix column headers for averages
    names(temp)[names(temp) %in% c("ave_x", "ave_y")]  <- str_c("ave_", c(xlab, ylab))    
    
    temp # return the data frame
}

# ========== Setup for MA and volcano plots ====================================
transform <- function(results, x, y) {
    # Make data frame with some transformed columns
        # results - results data frame
        # x - columns for x condition
        # y - columns for y condition
        # return new data frame
    df <- data.frame(log10((results[x] + results[y])/2), 
                     log2(results[y] / results[x]), 
                     results$candidate,
                     -log10(results$FDR))
    colnames(df) <- c("A", "M", "candidate", "P")
    
    df # return the data frame
}

# ========== MA plots using ggplot =============================================
MA_plots <- function(results, x, y, title) {
    # makes MA-plot DE candidate ggplots
        # results - data frame with edgeR results and some condition average columns
        # x - string for x-axis column
        # y - string for y-axis column
        # title - title string to use in plots
        # returns a list of plots 
    
    # uses transformed data
    temp <- transform(results, x, y)
    
    # 2-fold change lines
    ma_lines <- list(geom_hline(yintercept = 0.0, color = "black"),
                     geom_hline(yintercept = 1.0, color = "black", linetype = "dotted"),
                     geom_hline(yintercept = -1.0, color = "black", linetype = "dotted"))

    # make main MA plot
    ma <- ggplot(temp, aes(x = A, y = M)) +
        geom_point(aes(color = candidate, shape = candidate)) +
        scale_y_continuous(paste0("logFC (", y, "/", x, ")")) +
        scale_x_continuous("Ave_intensity") +
        ggtitle(title) + 
        ma_lines
    
    # make separate MA plots
    ma_facet <- ggplot(temp, aes(x = A, y = M)) +
        geom_point(aes(color = candidate, shape = candidate)) +
        scale_y_continuous(paste0("log2 FC (", y, "/", x, ")")) +
        scale_x_continuous("log10 Ave_intensity") +
        ma_lines +
        facet_wrap(~ candidate) +
        ggtitle(str_c(title, " (separated)"))

    # make the plots visible
    print(ma)
    print(ma_facet)
}    

# ========== Scatter plots using ggplot ========================================
scatter_plots <- function(results, x, y, title) {
    # makes scatter-plot DE candidate ggplots
        # results - data frame with edgeR results and some condition average columns
        # x - string for x-axis column
        # y - string for y-axis column
        # title - title string to use in plots
        # returns a list of plots
    
    # 2-fold change lines
    scatter_lines <- list(geom_abline(intercept = 0.0, slope = 1.0, color = "black"),
                          geom_abline(intercept = 0.301, slope = 1.0, color = "black", linetype = "dotted"),
                          geom_abline(intercept = -0.301, slope = 1.0, color = "black", linetype = "dotted"),
                          scale_y_log10(),
                          scale_x_log10())

    # make main scatter plot
    scatter <- ggplot(results, aes_string(x, y)) +
        geom_point(aes(color = candidate, shape = candidate)) +
        ggtitle(title) + 
        scatter_lines

    # make separate scatter plots
    scatter_facet <- ggplot(results, aes_string(x, y)) +
        geom_point(aes(color = candidate, shape = candidate)) +
        scatter_lines +
        facet_wrap(~ candidate) +
        ggtitle(str_c(title, " (separated)")) 

    # make the plots visible
    print(scatter)
    print(scatter_facet)
}

# ========== Volcano plots using ggplot ========================================
volcano_plot <- function(results, x, y, title) {
    # makes a volcano plot
        # results - a data frame with edgeR results
        # x - string for the x-axis column
        # y - string for y-axis column
        # title - plot title string
    
    # uses transformed data
    temp <- transform(results, x, y)
    
    # build the plot
    ggplot(temp, aes(x = M, y = P)) +
        geom_point(aes(color = candidate, shape = candidate)) +
        xlab("log2 FC") +
        ylab("-log10 FDR") +
        ggtitle(str_c(title, " Volcano Plot"))
}

# ============== individual protein expression plots ===========================
# function to extract the identifier part of the accesssion
get_identifier <- function(accession) {
    identifier <- str_split(accession, ";", simplify = TRUE)
    identifier[,1]
}

set_plot_dimensions <- function(width_choice, height_choice) {
    options(repr.plot.width=width_choice, repr.plot.height=height_choice)
}

plot_top_tags <- function(results, nleft, nright, top_tags) {
    # results should have data first, then test results (two condition summary table)
    # nleft, nright are number of data points in each condition
    # top_tags is number of up and number of down top DE candidates to plot
    # get top ipregulated
    up <- results %>% 
        filter(logFC >= 0) %>%
        arrange(FDR)
    up <- up[1:top_tags, ]
    
    # get top down regulated
    down <- results %>% 
        filter(logFC < 0) %>%
        arrange(FDR)
    down <- down[1:top_tags, ]
    
    # pack them
    proteins <- rbind(up, down)
        
    color = c(rep("red", nleft), rep("blue", nright))
    for (row_num in 1:nrow(proteins)) {
        row <- proteins[row_num, ]
        vec <- as.vector(unlist(row[1:(nleft + nright)]))
        names(vec) <- colnames(row[1:(nleft + nright)])
        title <- str_c(get_identifier(row$Acc), ", int: ", scientific(mean(vec), 2), 
                       ", FDR: ", scientific(row$FDR, digits = 3), 
                       ", FC: ", round(row$FC, digits = 1))
        barplot(vec, col = color, main = title,
                cex.main = 1.0, cex.names = 0.7, cex.lab = 0.7)
    }    
}

# read in the data export file
data_all_pep <- read_tsv("Cond-2_TiO2.txt")

# save accessions and remove from data
accessions <- data_all_pep$Accession
counts_pep <- data_all_pep %>% select(-Accession)


# get the data arranged for the testing (resting first, then activated)
resting <- select(counts_pep, contains("R"))
resting <- resting[, order(colnames(resting))]
activated <- select(counts_pep, contains("A"))
activated <- activated[, order(colnames(activated))]
counts_pep <- cbind(resting, activated)

# define some column indexes
R  <- 1:5 # resting
A <- 6:10 # activated

# set color vector to match sample groups
color = c(rep("red", length(R)), rep("blue", length(A)))

head(counts_pep)
length(accessions)

boxplot(log10(counts_pep), col = color,
        xlab = 'TMT samples', ylab = 'log10 Intensity', 
        main = 'Condition 2 TiO2 Starting data', notch = TRUE)

# load into edgeR DGEList object and normalize the data
group = c(rep("R", 5), rep("A", 5))
ye <- DGEList(counts = counts_pep, group = group, gene = accessions)
ye <- calcNormFactors(ye)

counts_pep_tmm <- apply_tmm_factors(ye, color)

# check sample-to-sample similarity
pairs.panels(log10(counts_pep_tmm[R]), lm = TRUE, main = "Condition 2 Phosphopeptides: resting")
pairs.panels(log10(counts_pep_tmm[A]), lm = TRUE, main = "Condition 2 Phosphopeptides: activated")

# check the clustering (set colors by condition)
plotMDS(ye, col = color, main = "Condition 2: Resting (red) and Activated (blue)")

# compute dispersions and plot BCV
ye <- estimateDisp(ye)
plotBCV(ye, main = "Phosphopeptide Variance Trend")

# the exact test object has columns like fold-change, CPM, and p-values
# we have already loaded the data into "y" above, normalized, and est. dispersion
ete <- exactTest(ye, pair = c("R", "A"))

# this counts up, down, and unchanged genes (here it is proteins)
summary(decideTestsDGE(ete, p.value = 0.10))
topTags(ete)$table

# the topTags function adds the BH FDR values to an exactTest data frame 
# make sure we do not change the row order!
tte <- topTags(ete, n = Inf, sort.by = "none")$table

results_e <- collect_results(counts_pep_tmm, tte, R, "R", A, "A")

# make an MD plot
plotMD(ete, p.value = 0.10 , main = "Phosphopeptide Exact Test: Resting vs Activated")
abline(h = c(-1, 1), col = "black")

# check the p-value distribution
ggplot(tte, aes(PValue)) + 
  geom_histogram(bins = 100, fill = "white", color = "black") + 
  geom_hline(yintercept = mean(hist(ete$table$PValue, breaks = 100, 
                                    plot = FALSE)$counts[26:100])) +
  ggtitle("Phosphopeptide Exact Test: Resting vs Activated")

# see how many candidates by category
results_e %>% count(candidate)

# plot log2 fold-changes by category
ggplot(results_e, aes(x = logFC, fill = candidate)) +
  geom_histogram(binwidth=0.1, color = "black") +
  facet_wrap(~candidate) +
  coord_cartesian(xlim = c(-5, 5)) +
  ggtitle("Log2 R by category")

### Dotted lines are 2-fold and solid line is 1-to-1

# make the DE plots
MA_plots(results_e, "ave_R", "ave_A", "Resting vs Activated - Exact Test")
scatter_plots(results_e, "ave_R", "ave_A", "Resting vs Activated - Exact Test")
volcano_plot(results_e, "ave_R", "ave_A", "Resting vs Activated - Exact Test")

# plot the top N up down proteins
top_N <- 20

set_plot_dimensions(6, 3.5)
plot_top_tags(results_e, 5, 5, top_N)
set_plot_dimensions(7, 7)

# load into edgeR DGEList object and normalize the data
group = c(rep("R", 5), rep("A", 5))
yp <- DGEList(counts = counts_pep, group = group, gene = accessions)
yp <- calcNormFactors(yp)

# create the experimental design matrix
subject <- factor(rep(c("F1", "F3", "M1", "M2", "M4"), 2))
state <- factor(c(rep("R", 5), rep("A", 5)))

# Example 4.1 in edgeR user's guide
design <- model.matrix(~subject+state)
rownames(design) <- colnames(yp)
design

# extimate the dispersion parameters and check
yp <- estimateDisp(yp, design, robust = TRUE)
yp$common.dispersion

# fit statistical models (design matrix already in y$design)
fit <- glmQLFit(yp, design, robust = TRUE)
plotQLDisp(fit, main = "Condition 2 - Paired testing")

# if we do not specify a contrast, the default is the last column
# of the design matrix - a Resting versus Activated comparison
paired <- glmQLFTest(fit) # default comparison

# check test results
topTags(paired)$table
ttp <- topTags(paired, n = Inf, sort.by = "none")$table
summary(decideTests(paired, p.value = 0.10))

# make basic MA plot
plotMD(paired, p.value = 0.10, main = "Condition 2 TiO2 - Paired test")
print(abline(h = c(-1, 1), col = "black")) # 2-fold lines
    
# check the p-value distribution
ggplot(ttp, aes(PValue)) + 
    geom_histogram(bins = 100, fill = "white", color = "black") + 
    geom_hline(yintercept = mean(hist(ttp$PValue, breaks = 100, 
                                 plot = FALSE)$counts[30:100])) +
    ggtitle("p-value distribution: paired test")

# collect and reformat the test results
results_p <- collect_results(counts_pep_tmm, ttp, R, "R", A, "A")

# correct the changed reference order for the glm testing
results_p$logFC <- -1 * results_p$logFC

# see how many candidates by category
results_p %>% count(candidate)

# plot log2 fold-changes by category
ggplot(results_p, aes(x = logFC, fill = candidate)) +
  geom_histogram(binwidth=0.1, color = "black") +
  facet_wrap(~candidate) +
  coord_cartesian(xlim = c(-5, 5)) +
  ggtitle("Log2 R by paired test category")

# make the DE plots
MA_plots(results_p, "ave_R", "ave_A", "ASLA-515 TiO2 paired: Resting vs Activated")
scatter_plots(results_p, "ave_R", "ave_A", "ASLA-515 TiO2 paired: Resting vs Activated")
volcano_plot(results_p, "ave_R", "ave_A", "ASLA-515 TiO2 paired: Resting vs Activated")

# plot the top N up down proteins
top_N <- 20

set_plot_dimensions(6, 3.5)
plot_top_tags(results_p, 5, 5, top_N)
set_plot_dimensions(7, 7)

# phosphopeptides
cv_pep <- data.frame(R_pep = CV(counts_pep_tmm[R]), A_pep = CV(counts_pep_tmm[A]))
medians <- apply(cv_pep, 2, FUN = median)
print("Phosphopeptide median CVs (%)")
round(medians, 1)

# make the boxplots
boxplot(cv_pep, notch = TRUE, main = "Phosphopeptide CV distributions", 
        ylim = c(0, 150), ylab = "CV (%)")

# save the results data frame
results <- cbind(results_e, results_p)
write.table(results, "Cond-2_TiO2_results.txt", sep = "\t", row.names = FALSE)

# log the session
sessionInfo()