script:
    "../scripts/col_mod_featurecounts.py"

shell:
    "Rscript ../workflow/scripts/plot_macs_qc.R -i {params.input} -s {params.sample_control_combinations}  -o {output.plot} -p {output.summmary} 2> {log}"

shell:
    "Rscript ../workflow/scripts/plot_homer_annotatepeaks.R -i {params.input} -s {params.sample_control_combinations}  -o {output.plot} -p {output.summmary} 2> {log}"

wrapper:
    "0.64.0/bio/ucsc/bedGraphToBigWig"

wrapper:
    "0.72.0/bio/fastqc"

wrapper:
    "0.64.0/bio/multiqc"

shell:
    "../workflow/scripts/gtf2bed {input} > {output} 2> {log}"

import sys
import pandas as pd
import os.path

sys.stderr = open(snakemake.log[0], "w")

samples = pd.read_csv(snakemake.input.get("samples_file"), sep="\t")


def get_group_control_combination(bam_path):
    sample = os.path.basename(bam_path.split(".sorted.bam")[0])
    sample_row = samples.loc[samples["sample"] == sample]
    group = sample_row["group"].iloc[0]
    control = sample_row["control"].iloc[0]
    if pd.isnull(control):
        control = sample
    return "{}_{}_{}".format(group, control, sample)


def modify_header(old_header):
    return [get_group_control_combination(i) if i in snakemake.input["bam"] else i for i in old_header]


f_counts_tab = pd.read_csv(snakemake.input["featurecounts"], sep="\t", skiprows=1)
header = list(f_counts_tab.columns.values)
header_mod = modify_header(header)
f_counts_frame = pd.DataFrame(f_counts_tab.values)
f_counts_frame.columns = header_mod
f_counts_frame.to_csv(snakemake.output[0], sep='\t', index=False)

    ## - FIRST SIX COLUMNS OF FEATURECOUNTS_FILE SHOULD BE INTERVAL INFO. REMAINDER OF COLUMNS SHOULD BE SAMPLES-SPECIFIC COUNTS.
    ## - SAMPLE NAMES HAVE TO END IN "_R1" REPRESENTING REPLICATE ID. LAST 3 CHARACTERS OF SAMPLE NAME WILL BE TRIMMED TO OBTAIN GROUP ID FOR DESEQ2 COMPARISONS.
    ## - BAM_SUFFIX IS PORTION OF FILENAME AFTER SAMPLE NAME IN FEATURECOUNTS COLUMN SAMPLE NAMES E.G. ".rmDup.bam" if "DRUG_R1.rmDup.bam"
    ## - PACKAGES BELOW NEED TO BE AVAILABLE TO LOAD WHEN RUNNING R

################################################
################################################
## LOAD LIBRARIES                             ##
################################################
################################################

#library(optparse)
library(DESeq2)
library(vsn)
library(ggplot2)
library(RColorBrewer)
library(pheatmap)
library(lattice)
library(BiocParallel)

################################################
################################################
## PARSE COMMAND-LINE PARAMETERS              ##
################################################
################################################

#option_list <- list(make_option(c("-i", "--featurecount_file"), type="character", default=NULL, help="Feature count file generated by the SubRead featureCounts command.", metavar="path"),
#                    make_option(c("-b", "--bam_suffix"), type="character", default=NULL, help="Portion of filename after sample name in featurecount file header e.g. '.rmDup.bam' if 'DRUG_R1.rmDup.bam'", metavar="string"),
#                    make_option(c("-o", "--outdir"), type="character", default='./', help="Output directory", metavar="path"),
#                    make_option(c("-p", "--outprefix"), type="character", default='differential', help="Output prefix", metavar="string"),
#                    make_option(c("-s", "--outsuffix"), type="character", default='', help="Output suffix for comparison-level results", metavar="string"),
#                    make_option(c("-v", "--vst"), type="logical", default=FALSE, help="Run vst transform instead of rlog", metavar="boolean"),
#                    make_option(c("-c", "--cores"), type="integer", default=1, help="Number of cores", metavar="integer"))
#
#opt_parser <- OptionParser(option_list=option_list)
#opt <- parse_args(opt_parser)
#
#if (is.null(opt$featurecount_file)){
#    print_help(opt_parser)
#    stop("Please provide featurecount file.", call.=FALSE)
#}
#if (is.null(opt$bam_suffix)){
#    print_help(opt_parser)
#    stop("Please provide bam suffix in header of featurecount file.", call.=FALSE)
#}

################################################
################################################
## READ IN COUNTS FILE                        ##
################################################
################################################

featurecount_file <- snakemake@input[[1]]   # AVI: adapted to snakemake
# AVI: suffix and prefix are already removed in rule featurecounts_modified_colnames
#bam_suffix <- ".bam"                       # AVI
#if (snakemake@params[["singleend"]]) {    # AVI
#    bam_suffix <- ".mLb.clN.sorted.bam"    # AVI
#}

count.table <- read.delim(file=featurecount_file,header=TRUE)  # AVI: removed 'skip=1', this was already done in rule featurecounts_modified_colnames
#colnames(count.table) <- gsub(bam_suffix,"",colnames(count.table))  # AVI
#colnames(count.table) <- as.character(lapply(colnames(count.table), function (x) tail(strsplit(x,'.',fixed=TRUE)[[1]],1)))  # AVI
rownames(count.table) <- count.table$Geneid
interval.table <- count.table[,1:6]
count.table <- count.table[,7:ncol(count.table),drop=FALSE]

################################################
################################################
## RUN DESEQ2                                 ##
################################################
################################################

# AVI: this is handled by snakemake
#if (file.exists(opt$outdir) == FALSE) {
#    dir.create(opt$outdir,recursive=TRUE)
#}
#setwd(opt$outdir)

samples.vec <- sort(colnames(count.table))
groups <- sub("_[^_]+$", "", samples.vec)
print(unique(groups))
if (length(unique(groups)) == 1) {
    quit(save = "no", status = 0, runLast = FALSE)
}

DDSFile <- snakemake@output[["dds"]]  # AVI: adapted to snakemake
if (file.exists(DDSFile) == FALSE) {
    counts <- count.table[,samples.vec,drop=FALSE]
    print(head(counts))
    coldata <- data.frame(row.names=colnames(counts),condition=groups)
    print(head(coldata))

    # AVI: set threads limit to prevent the "'bplapply' receive data failed" error
    # see also https://github.com/kdkorthauer/dmrseq/issues/7
    threads <- floor(snakemake@threads[[1]] * 0.75)

    dds <- DESeqDataSetFromMatrix(countData = round(counts), colData = coldata, design = ~ condition)
    dds <- DESeq(dds, parallel=TRUE, BPPARAM=MulticoreParam(ifelse(threads>0, threads, 1)))  # AVI: set threads limit

    if (!snakemake@params[["vst"]]) {
        rld <- rlog(dds)
    } else {
        rld <- vst(dds)
    }
    save(dds,rld,file=DDSFile)
}

#################################################
#################################################
### PLOT QC                                    ##
#################################################
#################################################

PlotPCAFile <- snakemake@output[["plot_pca"]]  # AVI: adapted to snakemake
PlotHeatmapFile <- snakemake@output[["plot_heatmap"]]  # AVI: adapted to snakemake
if (file.exists(PlotPCAFile) == FALSE) {
    # pdf(file=PlotFile,onefile=TRUE,width=7,height=7)  # AVI: splitted in separate pdf files

    ## PCA
    pdf(file=PlotPCAFile,onefile=TRUE,width=7,height=7)  # AVI: added to create separate pdf files
    pca.data <- DESeq2::plotPCA(rld,intgroup=c("condition"),returnData=TRUE)
    percentVar <- round(100 * attr(pca.data, "percentVar"))
    plot <- ggplot(pca.data, aes(PC1, PC2, color=condition)) +
            geom_point(size=3) +
            xlab(paste0("PC1: ",percentVar[1],"% variance")) +
            ylab(paste0("PC2: ",percentVar[2],"% variance")) +
            theme(panel.grid.major = element_blank(),
                  panel.grid.minor = element_blank(),
                  panel.background = element_blank(),
                  panel.border = element_rect(colour = "black", fill=NA, size=1))
    print(plot)
    dev.off()  # AVI: added to create separate pdf files
}

    ## WRITE PC1 vs PC2 VALUES TO FILE
if (file.exists(PlotHeatmapFile) == FALSE) {   # AVI: added for splitted pdf files
    pca.vals <- pca.data[,1:2]
    colnames(pca.vals) <- paste(colnames(pca.vals),paste(percentVar,'% variance',sep=""), sep=": ")
    pca.vals <- cbind(sample = rownames(pca.vals), pca.vals)
    write.table(pca.vals,file=snakemake@output[["pca_data"]],row.names=FALSE,col.names=TRUE,sep="\t",quote=TRUE) # AVI: adapted to snakemake

    ## SAMPLE CORRELATION HEATMAP
    pdf(file=PlotHeatmapFile,onefile=TRUE,width=7,height=7)  # AVI: added to create separate pdf files
    sampleDists <- dist(t(assay(rld)))
    sampleDistMatrix <- as.matrix(sampleDists)
    colors <- colorRampPalette( rev(brewer.pal(9, "Blues")) )(255)
    pheatmap(sampleDistMatrix,clustering_distance_rows=sampleDists,clustering_distance_cols=sampleDists,col=colors)
    dev.off()  # AVI: added to create separate pdf files

    ## WRITE SAMPLE DISTANCES TO FILE  # AVI: adapted to snakemake
    write.table(cbind(sample = rownames(sampleDistMatrix), sampleDistMatrix),file=snakemake@output[["dist_data"]],row.names=FALSE,col.names=TRUE,sep="\t",quote=FALSE)


}

#################################################
#################################################
### SAVE SIZE FACTORS                          ##
#################################################
#################################################

#SizeFactorsDir <- "sizeFactors/"
#if (file.exists(SizeFactorsDir) == FALSE) {
#    dir.create(SizeFactorsDir,recursive=TRUE)
#}

NormFactorsFile <- snakemake@output[["size_factors_rdata"]]  # AVI: adapted to snakemake
if (file.exists(NormFactorsFile) == FALSE) {
    normFactors <- sizeFactors(dds)
    save(normFactors,file=NormFactorsFile)

    for (name in names(sizeFactors(dds))) {
        sizeFactorFile <- snakemake@output[["size_factors_res"]]  # AVI: adapted to snakemake
        if (file.exists(sizeFactorFile) == FALSE) {
            write(as.numeric(sizeFactors(dds)[name]),file=sizeFactorFile)
        }
    }
}

#################################################
#################################################
### WRITE LOG FILE                             ##
#################################################
#################################################

LogFile <- snakemake@log[[1]]  # AVI: adapted to snakemake
if (file.exists(LogFile) == FALSE) {
    cat("\nSamples =",samples.vec,"\n\n",file=LogFile,append=TRUE,sep=', ')
    cat("Groups =",groups,"\n\n",file=LogFile,append=TRUE,sep=', ')
    cat("Dimensions of count matrix =",dim(counts),"\n\n",file=LogFile,append=FALSE,sep=' ')
    cat("\n",file=LogFile,append=TRUE,sep='')
}

#################################################
#################################################
### LOOP THROUGH COMPARISONS                   ##
#################################################
#################################################

ResultsFile <- snakemake@output[["results"]] # AVI: adapted to snakemake
if (file.exists(ResultsFile) == FALSE) {

    raw.counts <- counts(dds,normalized=FALSE)
    colnames(raw.counts) <- paste(colnames(raw.counts),'raw',sep='.')
    pseudo.counts <- counts(dds,normalized=TRUE)
    colnames(pseudo.counts) <- paste(colnames(pseudo.counts),'pseudo',sep='.')

    deseq2_results_list <- list()
    comparisons <- combn(unique(groups),2)
    for (idx in 1:ncol(comparisons)) {

        control.group <- comparisons[1,idx]
        treat.group <- comparisons[2,idx]
        CompPrefix <- paste(control.group,treat.group,sep="vs")
        cat("Saving results for ",CompPrefix," ...\n",sep="")

        # AVI: this is handled by snakemake
        #CompOutDir <- paste(CompPrefix,'/',sep="")
        #if (file.exists(CompOutDir) == FALSE) {
        #    dir.create(CompOutDir,recursive=TRUE)
        #}

        control.samples <- samples.vec[which(groups == control.group)]
        treat.samples <- samples.vec[which(groups == treat.group)]
        comp.samples <- c(control.samples,treat.samples)

        comp.results <- results(dds,contrast=c("condition",c(control.group,treat.group)))
        comp.df <- as.data.frame(comp.results)
        comp.table <- cbind(interval.table, as.data.frame(comp.df), raw.counts[,paste(comp.samples,'raw',sep='.')], pseudo.counts[,paste(comp.samples,'pseudo',sep='.')])

        ## WRITE RESULTS FILE
        CompResultsFile <- snakemake@output[["results"]]
        write.table(comp.table, file=CompResultsFile, col.names=TRUE, row.names=FALSE, sep='\t', quote=FALSE)

        ## FILTER RESULTS BY FDR & LOGFC AND WRITE RESULTS FILE
        # pdf(file=snakemake@output[["deseq2_plots"]],width=10,height=8) # AVI: splitted in separate pdf files
        if (length(comp.samples) > 2) {
            for (MIN_FDR in c(0.01,0.05)) {

                ## SUBSET RESULTS BY FDR
                pass.fdr.table <- subset(comp.table, padj < MIN_FDR)
                pass.fdr.up.table <- subset(comp.table, padj < MIN_FDR & log2FoldChange > 0)
                pass.fdr.down.table <- subset(comp.table, padj < MIN_FDR & log2FoldChange < 0)

                ## SUBSET RESULTS BY FDR AND LOGFC
                pass.fdr.logFC.table <- subset(comp.table, padj < MIN_FDR & abs(log2FoldChange) >= 1)
                pass.fdr.logFC.up.table <- subset(comp.table, padj < MIN_FDR & abs(log2FoldChange) >= 1 & log2FoldChange > 0)
                pass.fdr.logFC.down.table <- subset(comp.table, padj < MIN_FDR & abs(log2FoldChange) >= 1 & log2FoldChange < 0)

                ## WRITE RESULTS FILE
                if (MIN_FDR == 0.01) {
                    CompResultsFile <- snakemake@output[["FDR_1_perc_res"]]
                    CompBEDFile <- snakemake@output[["FDR_1_perc_bed"]]
                    MAplotFile <- snakemake@output[["plot_FDR_1_perc_MA"]]  # AVI: added to create separate pdf files
                    VolcanoPlotFile <- snakemake@output[["plot_FDR_1_perc_volcano"]]  # AVI: added to create separate pdf files
                }
                if (MIN_FDR == 0.05) {
                    CompResultsFile <- snakemake@output[["FDR_5_perc_res"]]
                    CompBEDFile <- snakemake@output[["FDR_5_perc_bed"]]
                    MAplotFile <- snakemake@output[["plot_FDR_5_perc_MA"]]  # AVI: added to create separate pdf files
                    VolcanoPlotFile <- snakemake@output[["plot_FDR_5_perc_volcano"]]  # AVI: added to create separate pdf files
                }

                write.table(pass.fdr.table, file=CompResultsFile, col.names=TRUE, row.names=FALSE, sep='\t', quote=FALSE)
                write.table(pass.fdr.table[,c("Chr","Start","End","Geneid","log2FoldChange","Strand")], file=CompBEDFile, col.names=FALSE, row.names=FALSE, sep='\t', quote=FALSE)

                ## MA PLOT & VOLCANO PLOT
                pdf(file=MAplotFile,width=10,height=8)  # AVI: added to create separate pdf files
                DESeq2::plotMA(comp.results, main=paste("MA plot FDR <= ",MIN_FDR,sep=""), ylim=c(-2,2),alpha=MIN_FDR)
                dev.off()  # AVI: added to create separate pdf files

                pdf(file=VolcanoPlotFile,width=10,height=8)  # AVI: added to create separate pdf files
                plot(comp.table$log2FoldChange, -1*log10(comp.table$padj), col=ifelse(comp.table$padj<=MIN_FDR, "red", "black"), xlab="logFC", ylab="-1*log10(FDR)", main=paste("Volcano plot FDR <=",MIN_FDR,sep=" "), pch=20)
                dev.off()  # AVI: added to create separate pdf files

                ## ADD COUNTS TO LOGFILE
                cat(CompPrefix," genes with FDR <= ",MIN_FDR,": ",nrow(pass.fdr.table)," (up=",nrow(pass.fdr.up.table),", down=",nrow(pass.fdr.down.table),")","\n",file=LogFile,append=TRUE,sep="")
                cat(CompPrefix," genes with FDR <= ",MIN_FDR," & FC > 2: ",nrow(pass.fdr.logFC.table)," (up=",nrow(pass.fdr.logFC.up.table),", down=",nrow(pass.fdr.logFC.down.table),")","\n",file=LogFile,append=TRUE,sep="")

            }
            cat("\n",file=LogFile,append=TRUE,sep="")
        # AVI: creates required output files with message
        } else {
            stop("More than 2 samples treated with the same antibody are needed to calculate the FDR & LOGFC.")
        }

        ## SAMPLE CORRELATION HEATMAP
        pdf(file=snakemake@output[["plot_sample_corr_heatmap"]],width=10,height=8)  # AVI: added to create separate pdf files
        rld.subset <- assay(rld)[,comp.samples]
        sampleDists <- dist(t(rld.subset))
        sampleDistMatrix <- as.matrix(sampleDists)
        colors <- colorRampPalette( rev(brewer.pal(9, "Blues")) )(255)
        pheatmap(sampleDistMatrix,clustering_distance_rows=sampleDists,clustering_distance_cols=sampleDists,col=colors)
        dev.off()  # AVI: added to create separate pdf files

        ## SCATTER PLOT FOR RLOG COUNTS
        pdf(file=snakemake@output[["plot_scatter"]],width=10,height=8)  # AVI: added to create separate pdf files
        combs <- combn(comp.samples,2,simplify=FALSE)
        clabels <- sapply(combs,function(x){paste(x,collapse=' & ')})
        plotdat <- data.frame(x=unlist(lapply(combs, function(x){rld.subset[, x[1] ]})),y=unlist(lapply(combs, function(y){rld.subset[, y[2] ]})),comp=rep(clabels, each=nrow(rld.subset)))
        plot <- xyplot(y~x|comp,plotdat,
                       panel=function(...){
                           panel.xyplot(...)
                           panel.abline(0,1,col="red")
                       },
                       par.strip.text=list(cex=0.5))
        print(plot)
        dev.off()  # AVI: added to create separate pdf files

        colnames(comp.df) <- paste(CompPrefix,".",colnames(comp.df),sep="")
        deseq2_results_list[[idx]] <- comp.df

    }

    ## WRITE RESULTS FROM ALL COMPARISONS TO FILE
    deseq2_results_table <- cbind(interval.table,do.call(cbind, deseq2_results_list),raw.counts,pseudo.counts)
    write.table(deseq2_results_table, file=ResultsFile, col.names=TRUE, row.names=FALSE, sep='\t', quote=FALSE)

}

#################################################
#################################################
### R SESSION INFO                             ##
#################################################
#################################################

cat(unlist(sessionInfo()),file=LogFile,append=TRUE,sep='')

import os
import yaml
from smart_open import open

# download blacklist and trim it for a specific chromosome


def copy_blacklist(igenomes, blacklist_path):
    with open(igenomes) as fin:
        with open(blacklist_path, 'w') as fout:
            for line in fin:
                fout.write(line)


def get_blacklist_from_igenomes(igenomes, blacklist_path):
    with open(igenomes) as f:
        igenomes = yaml.load(f, Loader=yaml.FullLoader)
        if "blacklist" in igenomes["params"]["genomes"][build]:
            blacklist_link = igenomes["params"]["genomes"][build]["blacklist"]
            with open(blacklist_link) as fin:
                with open(blacklist_path, 'w') as fout:
                    for line in fin:
                        if line.startswith("chr"):
                            line = line.replace("chr", "", 1)
                        if chromosome:
                            if line.startswith("{}\t".format(chromosome)):
                                fout.write(line)
                        else:
                            fout.write(line)
        else:
            open(blacklist_path, 'a').close()


igenomes = snakemake.input[0]
blacklist_path = snakemake.output.get("blacklist_path", "")

build = snakemake.params.get("build", "")
chromosome = snakemake.params.get("chromosome", "")
blacklist = snakemake.params.get("blacklist", "")

if blacklist:
    copy_blacklist(igenomes, blacklist_path)
else:
    get_blacklist_from_igenomes(igenomes, blacklist_path)