outCompNorm <- snakemake@output[["bkgrd"]]
outEffNorm <- snakemake@output[["hiAbund"]]

# cores for parallel processing
bp_cores <- snakemake@threads-1 

# read in mitochondrial chromosome name
mito_chr <- snakemake@params[["mito_chr"]]

# blacklist BED file
blacklist_file <- snakemake@params[["blacklist"]]

# file with the standard chromosome names
# file should contain only 1 line, which is space delimited
std_chroms_file <- snakemake@input[["std_chroms"]]

# bam files to read in
bam_samps <- snakemake@params[["bam_samp_names"]]
bam.files <- snakemake@input[["bams"]]

# before naming the bam files by sample name, we ensure that the sample names and the file names match
stopifnot(identical(paste0("analysis/filt_bams/", bam_samps, "_filt_alns.bam"), bam.files))
names(bam.files) <- snakemake@params[["bam_samp_names"]]

# output RDS objects
binned_rds <- paste0(snakemake@params[["out_pref"]], "_binned.rds")
small_wins_rds <- paste0(snakemake@params[["out_pref"]], "_small_wins.rds")
filt_small_wins_rds <- paste0(snakemake@params[["out_pref"]], "_filt_small_wins.rds")

# csaw params
bkgd_bin_width <- 10000
hi_abund_win_width <- snakemake@params[["window_width"]]
win_filter_fold <- 3 # fold change from background to keep as high abundance windows
minq <- 20
dedup=TRUE

##########
##########
##########

# load packages
suppressPackageStartupMessages(library(csaw))
suppressPackageStartupMessages(library(rtracklayer))
suppressPackageStartupMessages(library(magrittr))
suppressPackageStartupMessages(library(BiocParallel))
suppressPackageStartupMessages(library(tibble))

# read in standard chromosomes file as a vector, and remove mitochondrial chromosome name
chroms_keep_vec <- strsplit(readLines(std_chroms_file)[1], " ")[[1]] %>% 
    grep(pattern=paste0("^", mito_chr, "$"), x=., perl=TRUE, value=TRUE, invert=TRUE)

message("Keeping only these chromosomes: ", paste(chroms_keep_vec, collapse=", "))

# import blacklist as genomicranges
blacklist_gr <- import(blacklist_file)

# Set params for reading in BAM files
param <- readParam(minq=minq, dedup=dedup, pe="both", restrict=chroms_keep_vec, discard=blacklist_gr)

# Eliminate composition biases
# TMM on binned counts

binned <- windowCounts(bam.files, bin=TRUE, width=bkgd_bin_width, param=param, BPPARAM=MulticoreParam(bp_cores)) # note that windows with less than 'filter' number of reads summed across libraries are removed
message("For TMM on background bins, ", length(binned), " bins were used.")

binned <- normFactors(binned, se.out=TRUE)
saveRDS(binned, binned_rds) # for faster debugging

# code adapted from https://www.biostars.org/p/394434/
# 'se' is csaw generated ranged summarized experiment with norm.factors and totals in the coldata
# 'outfile' is the output file containing the colData
calc_and_write_final_norm_facs <- function(se, outfile){
    se$final.factors <- ((se$norm.factors * colData(se)$totals) / 1000000)^-1 # Recall that the norm factors from edgeR/csaw must be multiplied by the library size to also correct for library size; this is different from and not need with the DESeq2 norm factors.
    write.table(x = as.data.frame(colData(se)) %>% tibble::rownames_to_column("sample"),
                file = outfile, sep="\t", quote = FALSE,
                col.names = TRUE, row.names = FALSE)

}
## turn off exponential notation to avoid rounding errors
options(scipen=999) 

calc_and_write_final_norm_facs(se=binned, outfile=outCompNorm)

# Eliminate efficiency biases
# TMM on high abundance regions

small_wins <- windowCounts(bam.files, width=hi_abund_win_width, param=param, BPPARAM=MulticoreParam(bp_cores)) # note that windows with less than 'filter' number of reads summed across libraries are removed
saveRDS(small_wins, small_wins_rds) # for faster debugging
#binned <- windowCounts(bam.files, bin=TRUE, width=10000, param=param) 

filter_wins <- filterWindowsGlobal(small_wins, binned)
keep <- filter_wins$filter > log2(win_filter_fold) # filter for greater than 'win_filter_fold' fold difference compared to background

keep_num <- sum(keep)
tot_wins <- length(keep)
message("For TMM on high abundance normalization, ", keep_num, " out of ", tot_wins, " (", round(keep_num/tot_wins, 2), ")", " windows were kept.")
filtered.wins <- small_wins[keep, ]

filtered.wins <- normFactors(filtered.wins, se.out=TRUE) 
saveRDS(filtered.wins, filt_small_wins_rds) # for faster debugging

calc_and_write_final_norm_facs(se=filtered.wins, outfile=outEffNorm)

options(warn=1) # print all warnings as they happen

library(DiffBind)
library(readr)
library(dplyr)
library(stringr)
library(Rsamtools)

# variables
samplesheet <- snakemake@input[['samplesheet']] 
db_config <- list(cores=snakemake@threads)
out_dir <- snakemake@params[["outdir"]]
out_samplesheet <- snakemake@output[['samplesheet']]
DB_summits <- snakemake@params[['DB_summits']]
macs2_type <- snakemake@params[['macs2_type']]
subtract_controls <- as.logical(snakemake@params[['subtract_controls']])
curr_enriched_factor <- snakemake@params[['enriched_factor']]
is_atac <- as.logical(snakemake@params[['is_atac']])

# set up Diffbind samplesheet
samples <- read_tsv(samplesheet) %>%
        dplyr::mutate(SampleID=sample, 
                      bamReads=str_glue("analysis/filt_bams/{sample}_filt_alns.bam"), 
                      bamControl=ifelse(!is.na(control) & !is_atac, str_glue("analysis/filt_bams/{control}_filt_alns.bam"), NA), # typically there is no control for ATAC-seq datasets
                      Peaks=str_glue("analysis/{macs2_type}/rm_blacklist/{sample}_macs2_narrow_peaks.rm_blacklist.narrowPeak"),
                      PeakCaller="bed",
                      Factor=sample_group,
                      out_pref=ifelse(is.na(enriched_factor), "peaks", enriched_factor)) %>%
        dplyr::select(SampleID, bamReads, bamControl, Peaks, PeakCaller, Factor, out_pref) %>%
        unique() # some samples represented twice if they had more than 1 set of fastq files

if(any(!is.na(samples$bamControl))){
    samples <- samples %>% dplyr::filter(!is.na(bamControl)) # if project has controls, then we remove all the rows for the controls
} else{
    samples <- samples %>% dplyr::select(-bamControl)
}

# filter for just the samples related to the current enriched factor
stopifnot(curr_enriched_factor %in% samples$out_pref)
samples <- samples %>% dplyr::filter(out_pref == curr_enriched_factor)

# subtract controls?
if (!subtract_controls){
    samples$bamControl <- NA
}

write_tsv(samples, out_samplesheet)

print(samples)

message(str_glue("Making DBA for {curr_enriched_factor}"))

# make DBA object
diffbind <- dba(sampleSheet = samples %>% dplyr::select(-out_pref))

# set num cores and how many reads to store in memory
diffbind$config$cores <- db_config$cores-1
diffbind$config$yieldSize <- 20000000

# get counts
diffbind$config$scanbamparam <- ScanBamParam(flag = scanBamFlag(isDuplicate=FALSE, isSecondaryAlignment=FALSE, isSupplementaryAlignment=FALSE, isPaired=TRUE, isProperPair=TRUE, isNotPassingQualityControls=FALSE, isUnmappedQuery=FALSE, hasUnmappedMate=FALSE, isMinusStrand=NA, isMateMinusStrand=NA))

if (is_atac){
    diffbind$config$singleEnd <- TRUE # if ATAC, each read represents a cut site
}

diffbind <- dba.count(diffbind, summits=DB_summits, bUseSummarizeOverlaps = TRUE, bParallel = TRUE)
print(diffbind)

# normalization
diffbind <- dba.normalize(diffbind, normalize = DBA_NORM_NATIVE, background = TRUE)

# print the normalization methods
dba.normalize(diffbind, bRetrieve=TRUE)

dir.create(out_dir, recursive=TRUE)
# write the DBA to an RDS file
saveRDS(diffbind, str_glue("{out_dir}/{curr_enriched_factor}.rds"))


# session info
sessionInfo()

suppressPackageStartupMessages(library(GenomeInfoDb))
suppressPackageStartupMessages(library(magrittr))
suppressPackageStartupMessages(library(Rsamtools))

# params from workflow
ref_fasta <- snakemake@params[["ref_fasta"]]#"/varidata/research/projects/bbc/versioned_references/2022-03-08_14.47.50_v9/data/mm10_gencode_plus_viruses_and_cfmedips/sequence/mm10_gencode_plus_viruses_and_cfmedips.fa"
outfile <- snakemake@output[[1]] #"output.txt"
print(outfile)
# make GenomicRanges from the genome
ref_gr <- as(seqinfo(FaFile(ref_fasta)), "GRanges")

# output the standard chromosome names as a space-separated text file
standardChromosomes(ref_gr) %>% paste(., collapse=" ") %>% writeLines(., outfile)

library(dplyr)
library(stringr)
library(rtracklayer)
library(ChIPpeakAnno)
library(viridis)

# a list containing vectors as elements, where each vector is a different group of peak files.
peak_files <- snakemake@input
# extract only the named elements
peak_files <- peak_files[grep("^$", names(peak_files), value = TRUE, invert = TRUE)]
sample_names <- snakemake@params[["sample_names"]]
#save.image(file=paste0(outdir, "/Snakemake_vars.Rimage"))
peaks <- lapply(rlang::set_names(names(peak_files), names(peak_files)), function(peak_group_name){
    peak_group <- peak_files[[peak_group_name]]
    stopifnot(length(sample_names) == length(peak_group))
    lapply(rlang::set_names(peak_group, sample_names), 
            function(peak_file){
                unique(rtracklayer::import(peak_file)) # use unique() to remove duplicate peaks resulting from using --call-summits in macs2
            })
})

shell:
    """
    fastqc --outdir {params.outdir} {input}
    """

shell:
    """
    trim_galore --paired {input} --output_dir analysis/trim_galore/ --cores {threads} --fastqc
    """

shell:
    """
    bwa mem \
    -t {threads} \
    {params.bwa_idx} \
    {input} | \
    samblaster {params.samblaster_params} 2>{output.samblaster_err} | \
    samtools sort \
    -m 6G \
    -@ {threads} \
    -O "BAM" \
    -o {output.outbam} \
    -

    echo "END bwamem"
    echo "END bwamem" 1>&2

    samtools index -@ {threads} {output.outbam}

    echo "END indexing"
    echo "END indexing" 1>&2

    samtools idxstats {output.outbam} > {output.idxstat}

    echo "END idxstats"
    echo "END idxstats" 1>&2

    """

shell:
    """
    export TMPDIR={params.temp}

    bamtools filter  -insertSize "<={params.maxFragmentSize}" -in {input} -out {output.bam}

    samtools index {output.bam}
    samtools idxstats {output.bam} > {output.idxstat}

    """

shell:
    """
    export TMPDIR={params.temp}

    bamCoverage \
    -p {threads} \
    --binSize {params.binsize} \
    --bam {input.bam} \
    {params.extend_reads} \
    --blackListFileName {params.blacklist} \
    -o {output.bigwig_rmdups} \
    --normalizeUsing {params.norm_method} \
    --samFlagExclude {params.sam_exclude} \
    {params.sam_keep}

    """

shell:
    """
    export TMPDIR={params.temp}

    plotPCA --transpose -in {input} -o {output.pdf} --log2 --ntop 5000 
    """

shell:
    """
    # name sort BAMs
    samtools sort -@ {threads} -n -o {output.bam} {input}

    # below adapted from https://github.com/ENCODE-DCC/atac-seq-pipeline/blob/master/src/encode_task_bam2ta.py (Latest commit 867cfe5)
    bedtools bamtobed -bedpe -mate1 -i {output.bam} | gzip -nc > {output.bedpe}

    zcat {output.bedpe} | awk 'BEGIN{{OFS="\\t"}}{{printf "%s\\t%s\\t%s\\tN\\t1000\\t%s\\n%s\\t%s\\t%s\\tN\\t1000\\t%s\\n",$1,$2,$3,$9,$4,$5,$6,$10}}' | gzip -nc > {output.bed}


    """

shell:
    """
    # ENCODE uses -p 0.01 but we use -q 0.05 here to be consistent with the rest of this workflow

    macs2 \
    callpeak \
    -t {input.trt} \
    -f 'BED' \
    --outdir {params.outdir} \
    -n {params.broad_name} \
    -g {params.species} \
    --shift -75 --extsize 150 \
    --nomodel \
    -q {params.q_cutoff} \
    --keep-dup all \
    --broad \
    --tempdir {params.temp_dir}

    echo "END broad peak calling" 1>&2
    echo "END broad peak calling"

    macs2 \
    callpeak \
    -t {input.trt} \
    -f 'BED' \
    --outdir {params.outdir} \
    -n {params.narrow_name} \
    -g {params.species} \
    --shift -75 --extsize 150 \
    --nomodel --call-summits \
    -q {params.q_cutoff} \
    --keep-dup all \
    -B \
    --SPMR \
    --tempdir {params.temp_dir}

    echo "END narrow peak calling" 1>&2
    echo "END narrow peak calling"

    # remove (-truncate leads to intervals where the end occurs before start) the intervals that extend past the end of a chromosome
    bedClip -verbose=2 {output.narrow_bedgraph} {input.chr_sizes} {output.narrow_bedgraph_clip}

    bedGraphToBigWig {output.narrow_bedgraph_clip} {input.chr_sizes} {output.narrow_bigwig} 

    echo "END narrow peak calling" 1>&2
    echo "END narrow peak calling"

    """

shell:
    """
    bedops --merge {input} 1> {output}
    """

shell:
    """
    # preseq doesn't process supplmentary alignments properly. Remove these using samtools.

    preseq c_curve \
    -v \
    {params.paired} \
    -bam \
    -o {output.ccurve} \
    {input}

    echo "Finished c_curve." >&1
    echo "Finished c_curve." >&2

    preseq lc_extrap \
    -v \
    {params.paired} \
    -bam \
    -o {output.lcextrap} \
    {input}

    echo "Finished lc_extrap." >&1
    echo "Finished lc_extrap." >&2

    """

shell:
    """
    qualimap bamqc -bam {input} --java-mem-size={resources.mem_gb}G --paint-chromosome-limits -outdir analysis/qualimap/{wildcards.sample} -nt {threads}

    """

shell:
    """
    multiqc \
    --force \
    --config bin/multiqc_config.yaml \
    --outdir {params.workdir} \
    --filename {params.outfile} \
    {params.dirs} {params.PE_dirs} analysis/macs2/*narrow*

    """