shell:
    "fastp -i {input.r1} -I {input.r2} -o {output.r1} -O {output.r2} --detect_adapter_for_pe --thread {threads} -j {log} -h /dev/null"

shell:
    "fastqc -t {threads} --outdir data/fastqc {input} > {log} 2>&1"

shell:
    "fastq_screen --aligner bowtie2 --threads {threads} --outdir data/fastq_screen "
    "--conf {config[FASTQ_SCREEN]} --force {input} > {log} 2>&1"

shell:
    "bowtie2 --local --very-sensitive-local "
    "--no-unal --no-mixed --threads {threads} "
    "--no-discordant --phred33 "
    "-I 10 -X 700 -x {config[GENOME]} "
    "-1 {input.r1} -2 {input.r2} 2>{log.err} | samtools view -@ {threads} -Sbh - > {output}"

shell:
    "sambamba sort --tmpdir=data/aligned -t {threads} -o {output} {input} > {log} 2>&1"

shell:
    "bamCoverage -b {input[0]} -o {output} --binSize 10 --smoothLength 50 --normalizeUsing CPM -p {threads} "

shell:
    "cat {input} | sort -k1,1 -k2,2n | "
    "bedtools merge | "
    "bedtools intersect -a - -b {input} -c | "
    "awk -v OFS='\t' '$4>=2 {{print}}' > {output}"

run:
    pd.options.plotting.backend = "plotly"
    dfs = []
    for i in input:
        cond_marker = [os.path.basename(i).split(".")[0]]
        temp_df = pd.read_csv(i, sep = "\t", index_col = 0, names = cond_marker)
        dfs.append(temp_df)
    df = pd.concat(dfs, axis = 1)
    fraglen = df.plot()
    fraglen.update_layout( 
        title='Fragment Length Distribution', 
        xaxis_title='Fragment Length (bp)', 
        yaxis_title='Counts', 
        legend_title_text='Samples')
    fraglen.write_html(str(output))

shell:
    "preseq c_curve -B {resources.defect_mode} -l 1000000000 -o {output} {input} > {log} 2>&1"

script:
    "src/deseq2.R"

import pandas as pd

# map samples to fastqs
def get_samples():
    """
    return list of samples from samplesheet.tsv
    """
    return list(st.index)

def get_marks():
    """
    return list of marks from samplesheet.tsv
    """
    return list(st['mark'])

def get_mark_conditions():
    """
    return list of samples by condition
    """
    st['mark_condition']=st['mark'].astype(str)+"_"+st['condition']
    return st['mark_condition'].unique().tolist()

def get_tracks_by_mark_condition(wildcards):
    """
    return list of tracks by mark_condition
    """
    st['mark_condition']=st['mark'].astype(str)+"_"+st['condition']
    samps=st.groupby(["mark_condition"])['sample'].apply(list)[wildcards.mark_condition]
    return [f"data/tracks/{s}.bw" for s in samps]

def get_peaks_by_mark_condition(wildcards):
    """
    return list of peaks by mark_condition
    """
    st['mark_condition']=st['mark'].astype(str)+"_"+st['condition']
    samps=st.groupby(["mark_condition"])['sample'].apply(list)[wildcards.mark_condition]
    return [f"data/callpeaks/{s}_peaks.bed" for s in samps]

def get_bowtie2_input(wildcards):
    """
    returns reads associated with a sample
    """
    r1=st.loc[wildcards.sample]['R1']
    r2=st.loc[wildcards.sample]['R2']
    return r1,r2

def get_reads():
    """
    get list of all reads
    """
    rlist=list(st['R1'])+list(st['R2'])
    rlist=[os.path.basename(f).split('.')[0] for f in rlist]
    return rlist

def get_igg(wildcards):
    """
    Returns the igg file for the sample unless
    the sample is IgG then no control file is used.
    """ 
    if config['USEIGG']:
        igg=st.loc[wildcards.sample]['igg']
        iggbam=f'data/markd/{igg}.sorted.markd.bam'
        isigg=config['IGG'] in wildcards.sample
        if not isigg:
            return f'-c {iggbam}'
        else:
            return ""
    else:
        return ""

def get_callpeaks(wildcards):
    """
    Returns the callpeaks input files
    """
    bam=f"data/markd/{wildcards.sample}.sorted.markd.bam"
    bai=f"data/markd/{wildcards.sample}.sorted.markd.bam.bai"
    if config["USEIGG"]:
        igg=st.loc[wildcards.sample]['igg']
        iggbam=f'data/markd/{igg}.sorted.markd.bam'
        iggbam=f'data/markd/{igg}.sorted.markd.bam.bai'
        isigg=config['IGG'] in wildcards.sample
        if not isigg:
            return [bam,bai,iggbam]
        else:
            return [bam,bai]
    else:
        return [bam,bai]

def callpeaks_params(wildcards):
    """
    Returns callpeaks parameters specified by the user in the samplesheet
    """
    params = st.loc[wildcards.sample]['gopeaks']
    if params == "-":
        return ""
    else:
        return params

def defect_mode(wildcards, attempt):
    if attempt == 1:
        return ""
    elif attempt > 1:
        return "-D"

if ! command -v bedtools &> /dev/null
then
    echo "bedtools could not be found"
    exit
fi

# check args
if [[ $# -eq 0 ]]
then
    echo "no arguments supplied"
    exit 1
fi

MARK=$1
N_INTERSECTS=$2
OUTFILE=$3

# array of replicates per mark
declare -a mpks=(data/callpeaks/*${MARK}*_peaks.bed)

all_samples=$(echo ${mpks[@]} | tr ' ' '\n' | cut -d/ -f3 | cut -d_ -f1-3)
all_conditions=$(echo ${mpks[@]} | tr ' ' '\n' | cut -d/ -f3 | cut -d_ -f1)

for condition in $all_conditions; do

    # file I/O
    tmp_output="data/counts/$MARK.$condition.tmp.bed"

    # list all replicates in one condition
    all_replicates=$(find data/callpeaks/ -name "*$condition*$MARK*.bed" | sort | tr '\n' ' ')

    # find widest peak that appear in at least n replicates + export to tmp file.
    cat $all_replicates | cut -f1-3 | sort -k1,1 -k2,2n | bedtools merge | \
    bedtools intersect -c -a stdin -b $all_replicates | \
    awk -v n=$N_INTERSECTS '$4 >= n' | awk -v OFS='\t' '{print $1,$2,$3}' > $tmp_output

done

# merge intervals for all tmp files and export as consensus peak.
all_temp_files=$(find data/counts -name "*$MARK*.tmp.bed" | sort | tr '\n' ' ')
cat $all_temp_files | sort -k1,1 -k2,2n | bedtools merge > data/counts/${MARK}_consensus.bed
rm $all_temp_files

# count the number of reads per union peak
bedtools multicov -bams data/markd/*${MARK}*.bam -bed data/counts/${MARK}_consensus.bed -D > ${OUTFILE}_tmp

# label the counts table
ls data/markd/*${MARK}*.bam  | sed 's!.*/!!' | cut -d. -f1 | xargs |  tr ' ' '\t' | awk '{print "chrom\tstart\tend\t" $0}' | cat - ${OUTFILE}_tmp > ${OUTFILE}

# remove tmp 
rm ${OUTFILE}_tmp

library("tidyverse")
library("vsn")
library("pheatmap")
library("ggplot2")
library("RColorBrewer")
library("DESeq2")
library("PoiClaClu")
library("gplots")
library("plyranges")
library("tidyr")
library("dendextend")
library("ggrepel")
library("tools")


message("setting snakemake parameters")

exp_prefix = snakemake@params[["mark"]]
outdir = snakemake@params[["outdir"]]
numclusters = snakemake@params[["numclusters"]]
message(paste0("using numclusters = ", numclusters))
stopifnot(is.numeric(numclusters))

input = snakemake@input[["counts"]]
meta = snakemake@input[["meta"]]
genes = snakemake@input[["genes"]]
pcaPlot = snakemake@output[["pcaPlot"]]
pcaPlotVsd = snakemake@output[["pcaPlotVsd"]]
normCount = snakemake@output[["normCounts"]]
lnormCount = snakemake@output[["lnormCounts"]]
sdMeanRld = snakemake@output[["sdMeanRld"]]
sdMeanVsd = snakemake@output[["sdMeanVsd"]]
sampleDistPlotVsd = snakemake@output[["sampleDistVsd"]]
sampleDistPlotRld = snakemake@output[["sampleDistRld"]]
rds = snakemake@output[["rds"]]

if (!dir.exists(outdir)) {
    dir.create(outdir)
}

# read in genes file
genestab = read_tsv(genes, col_names=c("seqnames", "start", "end", "name", "score", "strand")) %>% GRanges()

# counts join
peaks = read_tsv(input) %>% 
    mutate(row=row_number()) %>%
    mutate(name=paste0("peak", as.character(row)))


# read in counts table
counts = read.delim(input, header=T, stringsAsFactors = F, check.names=F) %>%
    mutate(row=row_number()) %>%
    mutate(peak=paste0("peak", as.character(row)))
rownames(counts) = counts$peak
counts = counts[,4:ncol(counts)]

# read in metadata
meta <- read.csv(meta, header = T, stringsAsFactors = F)

# make the meta reflect the available sample names
meta = meta[meta$sample %in% colnames(counts),]

# make sure meta rownames and counts colnames in the same order.
rownames(meta) <- meta$sample
counts <- counts[, rownames(meta)]
stopifnot(rownames(meta) == colnames(counts))

message('calculating deseq...')
print(head(counts))
print(head(meta))
# make deseqdataset
ddsMat <- DESeqDataSetFromMatrix(countData = counts, colData = meta, design = ~condition)
dds <- DESeq(ddsMat)

# write normalized counts tables
normCounts <- counts(dds, normalized=TRUE)
normCounts <- 0.00001+(normCounts) # ensure nonzero 
lnormCounts <- log2(normCounts)
lnormCounts <- as.data.frame(lnormCounts)
lnormCounts$ID <- counts$peak
print(head(lnormCounts))
write.csv(normCounts, normCount)
write.csv(lnormCounts, lnormCount) 

# calculate rlog transform
rld <- rlog(dds, blind=FALSE)
vsd <- varianceStabilizingTransformation(dds, blind = TRUE,
  fitType = "parametric")
ntd <- normTransform(dds)

message('plotting count transforms...')
save.image()
# plot the transform for the first two samples
pdf(sdMeanRld)
rldCounts <- meanSdPlot(assay(rld))
rldCounts$gg + labs(title="meanSdPlot - rlog transformed")
dev.off()

pdf(sdMeanVsd)
vsdCounts <- meanSdPlot(assay(vsd))
vsdCounts$gg + labs(title="meanSdPlot - vsd transformed")
dev.off()

message('plotting poisson distance sample cross correlation...')
# plot sample distances 
pdf(sampleDistPlotVsd)
sampleDists <- dist(t(assay(vsd)))
sampleDistMatrix <- as.matrix(sampleDists)
rownames(sampleDistMatrix) <- paste(meta$sample, meta$condition, sep="-")
colnames(sampleDistMatrix) <- NULL
colors <- colorRampPalette( rev(brewer.pal(9, "Blues")) )(255)
pheatmap(sampleDistMatrix,
         clustering_distance_rows=sampleDists,
         clustering_distance_cols=sampleDists,
         col=colors,
         main=paste(exp_prefix, "vsd sample distance matrix"))
dev.off()

pdf(sampleDistPlotRld)
sampleDists <- dist(t(assay(rld)))
sampleDistMatrix <- as.matrix(sampleDists)
rownames(sampleDistMatrix) <- paste(meta$sample, meta$condition, sep="-")
colnames(sampleDistMatrix) <- NULL
colors <- colorRampPalette( rev(brewer.pal(9, "Blues")) )(255)
pheatmap(sampleDistMatrix,
         clustering_distance_rows=sampleDists,
         clustering_distance_cols=sampleDists,
         col=colors,
         main=paste(exp_prefix, "rld sample distance matrix"))
dev.off()

# plot PCA rld
pdf(pcaPlot)
plotPCA(rld, intgroup = c("condition")) + labs(title=paste0(exp_prefix,"-rld")) + geom_text_repel(aes(label=name))
dev.off()

# plot PCA vsd
pdf(pcaPlotVsd)
plotPCA(rld, intgroup = c("condition")) + labs(title=paste0(exp_prefix,"-vsd")) + geom_text_repel(aes(label=name))
dev.off()

# make contrasts
c = combn(unique(meta$condition), 2)
lsc=split(c, rep(1:ncol(c), each = nrow(c)))
message('writing contrast results...')
# write differential results for each contrast
for (k in 1:length(lsc)) {
    cl=lsc[[k]]
    rname=paste0(cl[1],"vs",cl[2])
    res=results(dds, 
            independentFiltering = TRUE,
            cooksCutoff = FALSE,
            contrast = c("condition", cl[1], cl[2]),
            alpha=0.05)

    # estimate shrinkage 
    resLFC <- lfcShrink(dds, coef=2, type="apeglm")

    #plotMA
    maplot=paste0(outdir,"/",exp_prefix,"/",exp_prefix,"-",rname,"plotMA.pdf")
    pdf(maplot)
    par(mfrow=c(1,2), mar=c(4,4,2,1))
    xlim <- c(1,1e5); ylim <- c(-2,2)
    plotMA(resLFC, xlim=xlim, ylim=ylim, alpha=0.05, main=paste(rname, "apeglm"))
    plotMA(res, xlim=xlim, ylim=ylim, alpha=0.05, main=paste(rname, "normal"))
    dev.off()

    tableName=paste0(outdir,"/",exp_prefix,"/",cl[1],'-',cl[2],'-',exp_prefix,'-','peaks-annot.tsv')

    # annottate peaks
    diffexp=res %>% 
        as_tibble(rownames='name') %>%
        left_join(peaks) %>%
        GRanges() %>%
        join_nearest(genestab, suffix=c(".peak", ".gene")) %>%
        as_tibble() 
    write_tsv(diffexp, path=tableName)

    # differential peak files lfc > 0 | lfc > 0  padj < 0.05 | padj < 0.01
    deup05=paste0(outdir,"/",exp_prefix,"/",cl[1],'-',cl[2],'-',exp_prefix,'-','differential-up-05.bed')
    dedown05=paste0(outdir,"/",exp_prefix,"/",cl[1],'-',cl[2],'-',exp_prefix,'-','differential-down-05.bed')
    deup01=paste0(outdir,"/",exp_prefix,"/",cl[1],'-',cl[2],'-',exp_prefix,'-','differential-up-01.bed')
    dedown01=paste0(outdir,"/",exp_prefix,"/",cl[1],'-',cl[2],'-',exp_prefix,'-','differential-down-01.bed')
    # sig up peaks
    diffexp %>%
        filter(log2FoldChange > 0, padj < 0.05) %>%
        select(seqnames, start, end, name.peak,name.gene,score,strand) %>%
        write_tsv(deup05, col_names=FALSE)

    # sig down peaks
    diffexp %>%
        filter(log2FoldChange < 0, padj < 0.05) %>%
        select(seqnames, start, end, name.peak,name.gene,score,strand) %>%
        write_tsv(dedown05, col_names=FALSE)

    # sig up peaks
    diffexp %>%
        filter(log2FoldChange > 0, padj < 0.01) %>%
        select(seqnames, start, end, name.peak,name.gene,score,strand) %>%
        write_tsv(deup01, col_names=FALSE)

    # sig down peaks
    diffexp %>%
        filter(log2FoldChange < 0, padj < 0.01) %>%
        select(seqnames, start, end, name.peak,name.gene,score,strand) %>%
        write_tsv(dedown01, col_names=FALSE)

    summaryName=paste0(outdir,"/",exp_prefix,"/",cl[1],"-",cl[2],"-",exp_prefix,"-diffexp-summary.txt")
    diffexp %>%
        summarise(Condition=paste0(exp_prefix,"-",rname),
                  DP05=nrow(filter(.,padj<0.05)),
                  DP01=nrow(filter(.,padj<0.01))) %>%
        write_tsv(summaryName)

    message("creating heatmap...")

    sig = diffexp %>% filter(padj < 0.05)

    # calculate heatmap if there are more than 10 sig peaks
    if (nrow(sig) > 10) {
        annot_filename=gsub(".tsv", "-05-clust.tsv", basename(tableName))

        # make direction annotation df 
        direction <- sig %>%
            select(name.peak,log2FoldChange) %>%
            mutate(direction=ifelse(log2FoldChange>0,"up","down")) %>%
            data.frame()

        rownames(direction) = direction[,1]
        direction[,1]=NULL

        # subset counts by sig peaks
        print(head(lnormCounts))
        counts <- rownames_to_column(lnormCounts, var="peak")  %>% as_tibble()
        print(head(counts))
        sigcounts <- counts %>% filter(peak %in% sig$name.peak)

        # organize samples by condition                               
        meta_conditions <- meta %>% group_by(condition) %>% 
            summarise(sample=paste(sample, collapse = ","))

        # rowMean counts for each condition
        for (c in meta_conditions$condition) {
            samples = filter(meta_conditions, condition==c)$sample
            s=unlist(strsplit(samples, ","))
            sigcounts <- sigcounts %>% mutate(!!c := rowMeans(sigcounts[s]))
        }

        # retain mean counts per condition
        sigcounts <- sigcounts[append("peak", meta_conditions$condition)]

        # scale the counts
        scaled <- t(apply(as.matrix(sigcounts[,-1]), 1, scale))
        colnames(scaled) = colnames(sigcounts[-1])

        rownames(scaled) <- sigcounts$peak
        # generate clusters
        clust=hclust(dist(scaled), method="ward.D2")
        clustdf=data.frame(cluster=cutree(as.dendrogram(clust), k = numclusters))
        annots=merge.data.frame(clustdf, direction, by=0, all = T) %>% select(-log2FoldChange)
        rownames(annots)=annots[,1]
        annots[,1]=NULL

        # get a numebr of unique colors
        cols=RColorBrewer::brewer.pal(numclusters, name="Set1")

        # name the colors by cluster
        names(cols) = unique(clustdf$cluster)

        # set cluster colors
        colors = list(
        cluster=cols
        )	

        heattitle=paste0(exp_prefix,"-",rname)
        # create heatmap
        heat <- pheatmap(scaled,
            main=heattitle,
            show_rownames = F,
            show_colnames = T,
            cluster_rows = T,
            cluster_method = "ward.D2",
            annotation_row = annots,
            annotation_colors = colors)

        save_pheatmap <- function(x, filename, width=7, height=7) {
            if (tolower(tools::file_ext(filename)) == "png") {
                png(filename = filename, width = width, height = height)
            }
            else {
                pdf(filename, width = width, height = height)
            }
            grid::grid.newpage()
            grid::grid.draw(x$gtable)
            dev.off()
        }

        heatmap_filename=paste0(outdir,"/",exp_prefix,"/",cl[1],"-",cl[2],"-",exp_prefix,"-heatmap.pdf")
        save_pheatmap(heat, heatmap_filename)

        annot_filename=gsub(".tsv", "-05-clust.tsv", tableName)
        # annotate cluster and direction
        annots[["name.peak"]] = rownames(annots)
        sig %>% left_join(annots) %>% write_delim(annot_filename, delim="\t")
    }
}

# save RDS of deseq object
saveRDS(dds, rds)

usage(){
	echo "Usage: " 
        echo "  $0 -c chrom_sizes_file -o output_file rep1.bw rep2.bw rep3.bw"
	exit 1
}

file_exits(){
	local f="$1"
	[[ -f "$f" ]] && return 0 || return 1
}

cleanup() {
    rm ${out_name}.tmp.bdg &> /dev/null
}

[[ $# -eq 0 ]] && usage

script_args=()
while [ $OPTIND -le "$#" ]
do
    if getopts c:o: option
    then
        case $option
        in
            c) chrom_size="$OPTARG";;
            o) out_name="$OPTARG";;
        esac
    else
        script_args+=("${!OPTIND}")
        ((OPTIND++))
    fi
done


if [ -z "$chrom_size" ] 
then 
    echo "Please supply the chrom.sizes file." 
    usage
else 
    if ( ! file_exits "$chrom_size" )
    then 
        echo "error: file $chrom_size not found"
        usage 
    fi
fi


if [ -z "$out_name" ] 
then 
    echo "Please supply an output file name."
    usage
fi

echo "Using chrom.sizes file: $chrom_size" 
echo "Output file name will be: $out_name" 

# require bigWigMerge
if ! command -v wiggletools &> /dev/null
then
    echo "wiggletools not found"
    exit
fi

# require bigWigToBedGraph
if ! command -v bigWigToBedGraph &> /dev/null
then 
    echo "bigWigToBedGraph no found"
    exit
fi


echo "merging the following files..."
echo ${script_args[@]}

cmd="wiggletools mean ${script_args[@]} | grep -vE '_|EBV|GL|JH' | awk 'NF==4{print}' | sort -S 8G -k1,1 -k2,2n > ${out_name}.tmp.bdg"
echo $cmd
eval $cmd

if [ $? -eq 0 ]; then
    echo "..."
else
    echo "Command wiggletools failed to complete..."
    echo "quitting.."
    cleanup
    exit 1
fi

# convert to bigwig
bedGraphToBigWig ${out_name}.tmp.bdg $chrom_size $out_name

if [ $? -eq 0 ]; then
    echo "..."
else
    echo "Command bedGraphToBigWig failed..."
    echo "quitting.."
    cleanup
    exit 1
fi

cleanup
echo "done"