require('RUVSeq')
library("edgeR")
library(RColorBrewer)
library(gplots)

main <- function(){
    args <- commandArgs(trailingOnly = TRUE)
    inputFile <- args[1]
    outputFolder <- args[2]
    inputTable = read.table(
        inputFile, sep="\t", header=TRUE, row.names=10, stringsAsFactors=F)
    countsTable = inputTable[,10:ncol(inputTable)]
    timePoints <- c("04", "16")
    mutantStrain <- c("Deltaomr", "DeltaryhB", "Deltayrl")
    for (i in seq(length(timePoints))){
        for (j in seq(length(mutantStrain))){
            # Compare mutant vs WT
            selectedCountsTable <- countsTable[,
                 c(sprintf("HeLa.S3.WT.%s_h.R1", timePoints[i]),
                   sprintf("HeLa.S3.WT.%s_h.R2", timePoints[i]),
                   sprintf("HeLa.S3.WT.%s_h.R3", timePoints[i]),
                   sprintf("HeLa.S3.%s.%s_h.R1", mutantStrain[j], timePoints[i]),
                   sprintf("HeLa.S3.%s.%s_h.R2", mutantStrain[j], timePoints[i]),
                   sprintf("HeLa.S3.%s.%s_h.R3", mutantStrain[j], timePoints[i]))]
            colNames <- c("WT_1", "WT_2", "WT_3",
                          sprintf("%s_1", mutantStrain[j]),
                          sprintf("%s_2", mutantStrain[j]),
                          sprintf("%s_3", mutantStrain[j]))
            prefix <- sprintf("%s_vs_WT", mutantStrain[j])
            refLevel <- ("WT")
            runRUVedgeR(selectedCountsTable, timePoints[i], colNames, prefix,
                        outputFolder, refLevel)

            # Compare mutant vs mock
            selectedCountsTable <- countsTable[,
                 c(sprintf("HeLa.S3.mock.%s_h.R1", timePoints[i]),
                   sprintf("HeLa.S3.mock.%s_h.R2", timePoints[i]),
                   sprintf("HeLa.S3.mock.%s_h.R3", timePoints[i]),
                   sprintf("HeLa.S3.%s.%s_h.R1", mutantStrain[j], timePoints[i]),
                   sprintf("HeLa.S3.%s.%s_h.R2", mutantStrain[j], timePoints[i]),
                   sprintf("HeLa.S3.%s.%s_h.R3", mutantStrain[j], timePoints[i]))]
            colNames <- c("mock_1", "mock_2", "mock_3",
                          sprintf("%s_1", mutantStrain[j]),
                          sprintf("%s_2", mutantStrain[j]),
                          sprintf("%s_3", mutantStrain[j]))
            prefix <- sprintf("%s_vs_mock", mutantStrain[j])
            refLevel <- ("mock")
            runRUVedgeR(selectedCountsTable, timePoints[i], colNames, prefix,
                        outputFolder, refLevel)
        }
        # Compare WT vs mock
        selectedCountsTable <- countsTable[,
              c(sprintf("HeLa.S3.mock.%s_h.R1", timePoints[i]),
                sprintf("HeLa.S3.mock.%s_h.R2", timePoints[i]),
                sprintf("HeLa.S3.mock.%s_h.R3", timePoints[i]),
                sprintf("HeLa.S3.WT.%s_h.R1", timePoints[i]),
                sprintf("HeLa.S3.WT.%s_h.R2", timePoints[i]),
                sprintf("HeLa.S3.WT.%s_h.R3", timePoints[i]))]
        colNames <- c("mock_1", "mock_2", "mock_3", "WT_1", "WT_2", "WT_3")
        prefix <- "WT_vs_mock"
        refLevel <- ("mock")
        runRUVedgeR(selectedCountsTable, timePoints[i], colNames, prefix,
                    outputFolder, refLevel)
    }
}

runRUVedgeR <- function(mat, timePoint, colNames, prefix, outputFolder,
                        refLevel){
    filter <- apply(mat, 1, function(x) length(x[x>10])>=3)
    filtered <- mat[filter,]
    differences <- matrix(data=c(1:6), byrow=T, nrow=2)
    genes <- rownames(filtered)
    colors <- brewer.pal(3, "Set1")
    colorSelection <- colors[c(1, 1, 1, 2, 2, 2)]
    set <- as.matrix(filtered)
    colnames(set) <- colNames
    # Set level explicitly as they are otherwise alphabetically sorted
    cond <- factor(c(gsub("_.*", "", colNames)))
    cond <- relevel(cond, refLevel)


    pdf(file=sprintf("%s/edgeR-plots-%s.pdf", outputFolder, prefix))    
    plotRLE(set, ylim=c(-1,1), main="RLE, pre-norm", col=colorSelection)
    plotPCA(set, main="PCA, pre-norm", col=colorSelection)
    design <- model.matrix(~cond)

    # Run without RUV correction
    y <- DGEList(counts=set, group=cond)
    y <- calcNormFactors(y, method="upperquartile")
    y <- estimateGLMCommonDisp(y, design)
    y <- estimateGLMTagwiseDisp(y, design)
    fit <- glmFit(y, design)
    lrt <- glmLRT(fit, coef=2)
    top <- topTags(lrt, n=nrow(set))$table

    plotRLE(cpm(y, normalized.lib.sizes=T), ylim=c(-1,1), main="RLE, UQ norm", col=colorSelection)
    plotPCA(cpm(y, normalized.lib.sizes=T), main="PCA, UQ norm", col=colorSelection)

    hist(top$PValue, breaks=100, main="P-value distribution, UQ norm")

    setRUV <- RUVs(set, genes, k=3, differences)
    plotRLE(
        setRUV$normalizedCounts,
        ylim=c(-1,1),
        main="RLE, post-norm, k=3",
        col=colorSelection)
    plotPCA(
        setRUV$normalizedCounts,
        main="PCA, post-norm, k=3",
        col=colorSelection)
    design <- model.matrix(~cond + setRUV$W)
    y <- DGEList(counts=set, group=cond)
    y <- calcNormFactors(y, method="upperquartile")
    y <- estimateGLMCommonDisp(y, design)
    y <- estimateGLMTagwiseDisp(y, design)
    fit <- glmFit(y, design)
    lrt <- glmLRT(fit, coef=2)
    topRUV <- topTags(lrt, n=nrow(set))$table
    
    hist(topRUV$PValue, breaks=100,
         main="P-value distribution, post-norm, k=3")
    write.table(topRUV,
                file=sprintf("%s/edgeR-comparison-%s.csv",
                             outputFolder, prefix),
                sep="\t", quote=FALSE, row.names=TRUE, col.names=NA)    
    dev.off()
}

main()
