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, stringsAsFactors=FALSE)

    # Generate unique row name
    row.names(inputTable) <- paste(inputTable$Locus_tag, inputTable$Gene_name,
                                   inputTable$Start, inputTable$Feature,
                                   inputTable$Attributes, sep="|")
    countsTable = inputTable[,13:ncol(inputTable)-2]
    colNames <- gsub("in_vitro.", "", colnames(countsTable))
    colnames(countsTable) <- colNames

    # DeltapinT.00_h vs WT.00_h
    selectedCountsTable <- countsTable[, c(
        "WT.00_h.R1", "WT.00_h.R2", "WT.00_h.R3",
        "DeltapinT.00_h.R1", "DeltapinT.00_h.R2", "DeltapinT.00_h.R3")]
    selColNames <- colnames(selectedCountsTable)
    refLevel <- "WT.00_h"
    prefix <- "DeltapinT.00_h_vs_WT.00_h"
    runRUVedgeR(selectedCountsTable, selColNames, prefix,
                outputFolder, refLevel)

    # DeltapinT.01_h vs WT.01_h
    selectedCountsTable <- countsTable[, c(
        "WT.01_h.R1", "WT.01_h.R2", "WT.01_h.R3",
        "DeltapinT.01_h.R1", "DeltapinT.01_h.R2", "DeltapinT.01_h.R3")]
    selColNames <- colnames(selectedCountsTable)
    refLevel <- "WT.01_h"
    prefix <- "DeltapinT.01_h_vs_WT.01_h"
    runRUVedgeR(selectedCountsTable, selColNames, prefix,
                outputFolder, refLevel)

    # ComppinT.00_h vs WT.00_h
    selectedCountsTable <- countsTable[, c(
        "WT.00_h.R1", "WT.00_h.R2", "WT.00_h.R3",
        "ComppinT.00_h.R1", "ComppinT.00_h.R2", "ComppinT.00_h.R3")]
    selColNames <- colnames(selectedCountsTable)
    refLevel <- "WT.00_h"
    prefix <- "ComppinT.00_h_vs_WT.00_h"
    runRUVedgeR(selectedCountsTable, selColNames, prefix,
                outputFolder, refLevel)

    # ComppinT.01_h vs WT.01_h
    selectedCountsTable <- countsTable[, c(
        "WT.01_h.R1", "WT.01_h.R2", "WT.01_h.R3",
        "ComppinT.01_h.R1", "ComppinT.01_h.R2", "ComppinT.01_h.R3")]
    selColNames <- colnames(selectedCountsTable)
    refLevel <- "WT.01_h"
    prefix <- "ComppinT.01_h_vs_WT.01_h"
    runRUVedgeR(selectedCountsTable, selColNames, prefix,
                outputFolder, refLevel)
}

runRUVedgeR <- function(mat, colNames, prefix, outputFolder, refLevel){
    minNumberOfReads <- 50
    filter <- apply(mat, 1, function(x){length(x[x>minNumberOfReads])>=3})
    filtered <- mat[filter,]
    differences <- matrix(data=c(1:6), byrow=T, nrow=2)
    genes <- rownames(filtered)
    colors <- brewer.pal(3, "Set1")
    set <- as.matrix(filtered)
    colnames(set) <- colNames
    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=colors[c(1,1,1,2,2,2)])
    plotPCA(set, main="PCA, pre-norm", col=colors[c(1,1,1,2,2,2)])
    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, prior.count=1)
    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=colors[c(1,1,1,2,2,2)])
    plotPCA(cpm(y, normalized.lib.sizes=T), main="PCA, UQ norm",
            col=colors[c(1,1,1,2,2,2)])

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

    setRUV <- RUVs(set, genes, k=1, differences)
    plotRLE(
        setRUV$normalizedCounts,
        ylim=c(-1,1),
        main="RLE, post-norm, k=3",
        col=colors[c(1,1,1,2,2,2)])
    plotPCA(
        setRUV$normalizedCounts,
        main="PCA, post-norm, k=3",
        col=colors[c(1,1,1,2,2,2)])
    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, prior.count=1)
    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()
