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=F)
    # 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]

    generateRUVnormalizeTable(countsTable, outputFolder,
                              "RUV-Normalized-WT_and_DeltapinT_00_h-16_h")
    quit()
    
    timePoints <- c("00", "02", "04", "08", "16", "24")
    for (i in seq(length(timePoints))){
        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.DeltapinT.%s_h.R1", timePoints[i]),
                                 sprintf("HeLa.S3.DeltapinT.%s_h.R2", timePoints[i]),
                                 sprintf("HeLa.S3.DeltapinT.%s_h.R3", timePoints[i]))]
        colNames <- c("WT_1", "WT_2", "WT_3",
                      "DeltapinT_1", "DeltapinT_2", "DeltapinT_3")
        prefix <- sprintf("DeltapinT_%s_h_vs_WT_%s_h", timePoints[i], timePoints[i])
        refLevel <- c("WT")
        runRUVedgeR(selectedCountsTable, colNames,
                    prefix, outputFolder, refLevel)
        }
}

runRUVedgeR <- function(mat, 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
    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, prior.count=1.0)
    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)
    
    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, prior.count=1.0)
    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()
}

generateRUVnormalizeTable <- function(countsTable, outputFolder, outputPrefix){

    countsTable = countsTable[, !(names(countsTable) %in% c(
        "HeLa.S3.DeltapinT.24_h.R1", "HeLa.S3.DeltapinT.24_h.R2", "HeLa.S3.DeltapinT.24_h.R3",
        "HeLa.S3.WT.24_h.R1", "HeLa.S3.WT.24_h.R2", "HeLa.S3.WT.24_h.R3",
        "HeLa.S3.mock.02_h.R1", "HeLa.S3.mock.02_h.R2", "HeLa.S3.mock.02_h.R3",
        "HeLa.S3.mock.04_h.R1", "HeLa.S3.mock.04_h.R2", "HeLa.S3.mock.04_h.R3",
        "HeLa.S3.mock.08_h.R1", "HeLa.S3.mock.08_h.R2", "HeLa.S3.mock.08_h.R3",
        "HeLa.S3.mock.16_h.R1", "HeLa.S3.mock.16_h.R2", "HeLa.S3.mock.16_h.R3",
        "HeLa.S3.mock.24_h.R1", "HeLa.S3.mock.24_h.R2", "HeLa.S3.mock.24_h.R3"))]
    
    filter <- apply(countsTable, 1, function(x) length(x[x>10])>=3)
    filtered <- countsTable[filter,]
    genes <- rownames(filtered)
    set <- as.matrix(filtered)
    colnames(set) <- colnames(countsTable)

    differences <- matrix(
        data=c(1:length(colnames(countsTable))),
        byrow=T,
        nrow=length(colnames(countsTable))/3)
    
    setRUV <- RUVs(set, genes, k=3, differences)
    pdf(file=sprintf("%s/%s_plot.pdf", outputFolder, outputPrefix))
    plotRLE(
        setRUV$normalizedCounts,
        ylim=c(-1,1),
        main="RLE, post-norm, k=3")
    dev.off()
    write.table(setRUV$normalizedCounts,
                file=sprintf("%s/%s.csv", outputFolder, outputPrefix),
                sep="\t", quote=FALSE)
}

main()
